gemini-code-assist[bot] commented on code in PR #585:
URL: https://github.com/apache/tvm-ffi/pull/585#discussion_r3223513541
##########
addons/torch_c_dlpack_ext/torch_c_dlpack_ext/core.py:
##########
@@ -39,6 +39,13 @@ def _create_dlpack_exchange_api_capsule(ptr_as_int: int) ->
Any:
return capsule
+def _torch_extension_suffix() -> str:
+ """Return the backend suffix used by the prebuilt extension library."""
+ if torch.cuda.is_available():
+ return "rocm" if torch.version.hip is not None else "cuda"
+ return "cpu"
Review Comment:
Accessing `torch.version.hip` directly will raise an `AttributeError` on
PyTorch builds that do not include ROCm support (e.g., standard CUDA or CPU
builds). This will cause the extension loading to fail entirely. It is safer to
use `getattr` to check for the existence of the backend version attributes.
```suggestion
if torch.cuda.is_available():
if getattr(torch.version, "hip", None) is not None:
return "rocm"
return "cuda"
return "cpu"
```
##########
python/tvm_ffi/_optional_torch_c_dlpack.py:
##########
@@ -44,6 +44,17 @@
logger = logging.getLogger(__name__)
+def _torch_extension_device(torch_module: Any) -> str:
+ """Return the torch backend name used in the optional extension library
name."""
+ if torch_module.cuda.is_available():
+ if torch_module.version.cuda is not None:
+ return "cuda"
+ if torch_module.version.hip is not None:
+ return "rocm"
+ raise ValueError("Cannot determine whether to build with CUDA or
ROCm.")
+ return "cpu"
Review Comment:
Direct access to `torch_module.version.cuda` and `torch_module.version.hip`
can raise `AttributeError` if the respective backend support is not compiled
into the PyTorch binary. Additionally, raising a `ValueError` here can crash
the import of `tvm_ffi` because the call at line 108 is not wrapped in a
catch-all exception block. It's better to use `getattr` and provide a sensible
default (like "cuda" if `is_available()` is true) to ensure the optional loader
fails gracefully.
```python
if torch_module.cuda.is_available():
if getattr(torch_module.version, "hip", None) is not None:
return "rocm"
return "cuda"
return "cpu"
```
--
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]
