tqchen commented on code in PR #585:
URL: https://github.com/apache/tvm-ffi/pull/585#discussion_r3229104998
##########
python/tvm_ffi/_optional_torch_c_dlpack.py:
##########
@@ -44,6 +44,15 @@
logger = logging.getLogger(__name__)
+def _torch_extension_device(torch_module: Any) -> str:
+ """Return the torch backend name used in the optional extension library
name."""
Review Comment:
if torch_module.cuda.is_available():
if getattr(torch_module.version, "cuda", None) is not None:
return "cuda"
if getattr(torch_module.version, "hip", None) is not None:
return "rocm"
return "cuda"
return "cpu"
##########
addons/torch_c_dlpack_ext/torch_c_dlpack_ext/core.py:
##########
@@ -39,6 +39,15 @@ def _create_dlpack_exchange_api_capsule(ptr_as_int: int) ->
Any:
return capsule
+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 getattr(torch_module.version, "hip", None) is not None:
Review Comment:
would be good to use the following pattern, mainly also detect cuda version
to be robust.
```python
if torch_module.cuda.is_available():
if getattr(torch_module.version, "cuda", None) is not None:
return "cuda"
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]
