tqchen commented on code in PR #585:
URL: https://github.com/apache/tvm-ffi/pull/585#discussion_r3225935915
##########
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"
Review Comment:
use the same code as
```python
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"
```
--
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]
