gemini-code-assist[bot] commented on code in PR #652:
URL: https://github.com/apache/tvm-ffi/pull/652#discussion_r3524022952
##########
python/tvm_ffi/cpp/extension.py:
##########
@@ -42,19 +42,28 @@
@functools.lru_cache
def _detect_gpu_backend() -> BACKEND_STR:
- """Auto-detect whether to use CUDA or HIP (ROCm).
+ """Auto-detect whether to use CUDA or HIP (ROCm), defaulting to CUDA.
- Returns 'hip' if ROCm/HIP is available, 'cuda' otherwise.
+ Resolution order: the ``TVM_FFI_GPU_BACKEND`` override, then PyTorch build
+ signals (``torch.version.hip`` / ``torch.version.cuda``), then an available
+ ``hipcc`` when no ``nvcc`` is present. The mere existence of ``/opt/rocm``
+ is not treated as evidence of HIP.
"""
- # Check environment variable override first
backend = os.environ.get("TVM_FFI_GPU_BACKEND", "").lower()
if backend in ("cuda", "hip"):
return backend # type: ignore[return-value]
try:
- _find_rocm_home()
+ import torch # noqa: PLC0415
+
+ if getattr(torch.version, "hip", None):
+ return "hip"
+ if getattr(torch.version, "cuda", None):
+ return "cuda"
+ except Exception:
+ pass
+ if shutil.which("nvcc") is None and shutil.which("hipcc") is not None:
return "hip"
Review Comment:
If `CUDA_HOME` or `CUDA_PATH` is set but `nvcc` is not in the system `PATH`
(which is a common configuration), `shutil.which("nvcc")` will return `None`.
If `hipcc` happens to be in the `PATH`, this check will incorrectly resolve the
backend to `"hip"`, even though a valid CUDA installation is specified via
environment variables.
Since `_find_cuda_home()` and `_find_rocm_home()` check these environment
variables and look for compilers under their respective `bin` directories, we
should also check for the compilers in those locations using `shutil.which`
with the `path` argument.
```suggestion
cuda_home = os.environ.get("CUDA_HOME") or os.environ.get("CUDA_PATH")
has_nvcc = shutil.which("nvcc") is not None
if not has_nvcc and cuda_home is not None:
has_nvcc = shutil.which("nvcc", path=str(Path(cuda_home) / "bin"))
is not None
rocm_home = os.environ.get("ROCM_HOME") or os.environ.get("ROCM_PATH")
has_hipcc = shutil.which("hipcc") is not None
if not has_hipcc and rocm_home is not None:
has_hipcc = shutil.which("hipcc", path=str(Path(rocm_home) / "bin"))
is not None
if not has_nvcc and has_hipcc:
return "hip"
```
##########
python/tvm_ffi/cpp/extension.py:
##########
@@ -42,19 +42,28 @@
@functools.lru_cache
def _detect_gpu_backend() -> BACKEND_STR:
- """Auto-detect whether to use CUDA or HIP (ROCm).
+ """Auto-detect whether to use CUDA or HIP (ROCm), defaulting to CUDA.
- Returns 'hip' if ROCm/HIP is available, 'cuda' otherwise.
+ Resolution order: the ``TVM_FFI_GPU_BACKEND`` override, then PyTorch build
+ signals (``torch.version.hip`` / ``torch.version.cuda``), then an available
+ ``hipcc`` when no ``nvcc`` is present. The mere existence of ``/opt/rocm``
+ is not treated as evidence of HIP.
"""
- # Check environment variable override first
backend = os.environ.get("TVM_FFI_GPU_BACKEND", "").lower()
if backend in ("cuda", "hip"):
return backend # type: ignore[return-value]
try:
- _find_rocm_home()
+ import torch # noqa: PLC0415
+
+ if getattr(torch.version, "hip", None):
+ return "hip"
+ if getattr(torch.version, "cuda", None):
+ return "cuda"
+ except Exception:
+ pass
Review Comment:
Directly accessing `torch.version` can raise an `AttributeError` if the
`version` submodule/attribute is not present on the imported `torch` module
(e.g., in some minimal, mocked, or incomplete environments). While this is
caught by the broad `except Exception`, it will silently fail the PyTorch-based
detection.
Using `getattr(torch, "version", None)` first is safer and more robust.
```suggestion
try:
import torch # noqa: PLC0415
torch_version = getattr(torch, "version", None)
if torch_version is not None:
if getattr(torch_version, "hip", None):
return "hip"
if getattr(torch_version, "cuda", None):
return "cuda"
except Exception:
pass
```
--
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]
