This is an automated email from the ASF dual-hosted git repository.
tqchen pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/tvm-ffi.git
The following commit(s) were added to refs/heads/main by this push:
new 95ea3a08 [FIX] Resolve JIT GPU backend by explicit signals, not stray
/opt/rocm (#652)
95ea3a08 is described below
commit 95ea3a08fcee2f6103c4dd7d829f37abd3ae9b7e
Author: Tianqi Chen <[email protected]>
AuthorDate: Sun Jul 5 08:19:13 2026 +0800
[FIX] Resolve JIT GPU backend by explicit signals, not stray /opt/rocm
(#652)
Fixes #651.
The JIT extension backend selector decided CUDA-vs-HIP by probing for a
ROCm
home, whose fallback returns `/opt/rocm` whenever that directory merely
exists.
A CUDA/NVIDIA host carrying a stray `/opt/rocm` was therefore
misclassified as
HIP even with `nvcc` present and no `hipcc`.
`_detect_gpu_backend()` now resolves the backend from explicit signals,
in
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. Detection defaults to CUDA when otherwise
inconclusive, and
the existence of `/opt/rocm` alone never forces HIP. Toolkit-home,
compiler,
include-path, flag, and architecture discovery are unchanged and run
only after
the backend is resolved, so CUDA and HIP state never mix.
---
python/tvm_ffi/cpp/extension.py | 43 ++++++++++++++++++++++++++++++++++-------
1 file changed, 36 insertions(+), 7 deletions(-)
diff --git a/python/tvm_ffi/cpp/extension.py b/python/tvm_ffi/cpp/extension.py
index 32b3b17c..f27a48fe 100644
--- a/python/tvm_ffi/cpp/extension.py
+++ b/python/tvm_ffi/cpp/extension.py
@@ -40,21 +40,50 @@ BACKEND_STR = Literal["cuda", "hip"]
logger = logging.getLogger(__name__)
+def _find_compiler(name: str, *home_vars: str) -> str | None:
+ """Locate a compiler on ``PATH`` or under an env-specified toolkit home's
``bin``."""
+ found = shutil.which(name)
+ if found is not None:
+ return found
+ for var in home_vars:
+ home = os.environ.get(var)
+ if home:
+ found = shutil.which(name, path=str(Path(home) / "bin"))
+ if found is not None:
+ return found
+ return None
+
+
@functools.lru_cache
def _detect_gpu_backend() -> BACKEND_STR:
- """Auto-detect whether to use CUDA or HIP (ROCm).
-
- Returns 'hip' if ROCm/HIP is available, 'cuda' otherwise.
+ """Auto-detect whether to use CUDA or HIP (ROCm), defaulting to CUDA.
+
+ Resolution order: the ``TVM_FFI_GPU_BACKEND`` override, then PyTorch build
+ signals (``torch.version.hip`` / ``torch.version.cuda``), then a
discoverable
+ ``hipcc`` when no ``nvcc`` can be found. Compilers are resolved from
``PATH``
+ and from the ``CUDA_HOME``/``CUDA_PATH`` and ``ROCM_HOME``/``ROCM_PATH``
+ toolkit homes, so an env-specified CUDA install is honored even when
``nvcc``
+ is off ``PATH``. 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
+
+ version = getattr(torch, "version", None)
+ if getattr(version, "hip", None):
+ return "hip"
+ if getattr(version, "cuda", None):
+ return "cuda"
+ except Exception:
+ pass
+ if _find_compiler("nvcc", "CUDA_HOME", "CUDA_PATH") is None and (
+ _find_compiler("hipcc", "ROCM_HOME", "ROCM_PATH") is not None
+ ):
return "hip"
- except RuntimeError:
- return "cuda"
+ return "cuda"
def _resolve_gpu_backend(backend: str | None) -> BACKEND_STR:
