gemini-code-assist[bot] commented on code in PR #288:
URL: https://github.com/apache/tvm-ffi/pull/288#discussion_r2566636817
##########
tests/python/test_cubin_launcher.py:
##########
@@ -43,6 +43,20 @@ def _is_cuda_available() -> bool:
return torch.cuda.is_available()
+def _is_cuda_version_greater_than_13() -> bool:
+ """Check if CUDA version is greater than 13.0."""
+ if torch is None or not torch.cuda.is_available():
+ return False
+ if torch.version.cuda is None:
+ return False
+ try:
+ # Parse version string into tuple of integers (e.g., "12.1" -> (12, 1))
+ version_parts = tuple(int(x) for x in torch.version.cuda.split("."))
+ return version_parts > (13, 0)
+ except (ValueError, TypeError, AttributeError):
+ return False
Review Comment:
This new function `_is_cuda_version_greater_than_13` and its usage to skip
tests seem unrelated to the main goal of this pull request, which is to upgrade
the DLPack exchange API. It's generally better to keep pull requests focused on
a single concern. If these test skips are necessary to fix unrelated test
failures, they should ideally be in a separate PR with a clear explanation.
This change reduces test coverage for CUDA versions <= 13.0.
##########
python/tvm_ffi/_optional_torch_c_dlpack.py:
##########
@@ -61,6 +87,11 @@ def load_torch_c_dlpack_extension() -> Any: # noqa:
PLR0912, PLR0915
import torch_c_dlpack_ext # type: ignore # noqa: PLC0415, F401
if hasattr(torch.Tensor, "__c_dlpack_exchange_api__"):
+ if isinstance(torch.Tensor.__c_dlpack_exchange_api__, int):
+ # Brings up to speed with the new PyCapsule behavior
+ torch.Tensor.__c_dlpack_exchange_api__ =
_create_dlpack_exchange_api_capsule(
+ torch.Tensor.__c_dlpack_exchange_api__
+ )
return None
Review Comment:
The attribute `torch.Tensor.__c_dlpack_exchange_api__` is accessed multiple
times. For better readability and to avoid repeated lookups, it's good practice
to store it in a local variable.
```suggestion
if hasattr(torch.Tensor, "__c_dlpack_exchange_api__"):
api_attr = torch.Tensor.__c_dlpack_exchange_api__
if isinstance(api_attr, int):
# Brings up to speed with the new PyCapsule behavior
torch.Tensor.__c_dlpack_exchange_api__ =
_create_dlpack_exchange_api_capsule(
api_attr
)
return None
```
##########
tests/python/test_dlpack_exchange_api.py:
##########
@@ -46,9 +47,22 @@ def test_dlpack_exchange_api() -> None:
assert torch is not None
assert hasattr(torch.Tensor, "__c_dlpack_exchange_api__")
- api_ptr = torch.Tensor.__c_dlpack_exchange_api__
- assert isinstance(api_ptr, int), "API pointer should be an integer"
- assert api_ptr != 0, "API pointer should not be NULL"
+ api_attr = torch.Tensor.__c_dlpack_exchange_api__
+
+ # Handle both PyCapsule and integer types
+ if isinstance(api_attr, int):
+ # Direct integer pointer
+ api_ptr = api_attr
+ assert api_ptr != 0, "API pointer should not be NULL"
+ else:
+ # PyCapsule - extract the pointer as integer
+ pythonapi = ctypes.pythonapi
+ # Set restype to c_size_t to get integer directly (avoids c_void_p
quirks)
+ pythonapi.PyCapsule_GetPointer.restype = ctypes.c_size_t
+ pythonapi.PyCapsule_GetPointer.argtypes = [ctypes.py_object,
ctypes.c_char_p]
+ capsule_name = b"dlpack_exchange_api"
+ api_ptr = pythonapi.PyCapsule_GetPointer(api_attr, capsule_name)
+ assert api_ptr != 0, "API pointer from PyCapsule should not be NULL"
Review Comment:
Modifying `ctypes.pythonapi` function signatures (`restype`, `argtypes`) can
have unintended side effects on other parts of the code or other tests that
might use the same C API functions. It's safer to save and restore the original
signatures in a `try...finally` block. Also, the test could be made more robust
by explicitly checking for the `PyCapsule` type and failing for any other
unexpected types.
```python
# Handle both PyCapsule and integer types
if isinstance(api_attr, int):
# Direct integer pointer
api_ptr = api_attr
assert api_ptr != 0, "API pointer should not be NULL"
elif type(api_attr).__name__ == "PyCapsule":
# PyCapsule - extract the pointer as integer
pythonapi = ctypes.pythonapi
# Save original argtypes/restype to avoid side-effects
original_restype = getattr(pythonapi.PyCapsule_GetPointer,
"restype", None)
original_argtypes = getattr(pythonapi.PyCapsule_GetPointer,
"argtypes", None)
try:
# Set restype to c_size_t to get integer directly (avoids
c_void_p quirks)
pythonapi.PyCapsule_GetPointer.restype = ctypes.c_size_t
pythonapi.PyCapsule_GetPointer.argtypes = [ctypes.py_object,
ctypes.c_char_p]
capsule_name = b"dlpack_exchange_api"
api_ptr = pythonapi.PyCapsule_GetPointer(api_attr, capsule_name)
assert api_ptr != 0, "API pointer from PyCapsule should not be
NULL"
finally:
# Restore original argtypes/restype
if original_restype is not None:
pythonapi.PyCapsule_GetPointer.restype = original_restype
if original_argtypes is not None:
pythonapi.PyCapsule_GetPointer.argtypes = original_argtypes
else:
pytest.fail(f"Unexpected type for __c_dlpack_exchange_api__:
{type(api_attr)}")
```
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
