[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.

Author: tqchen

python/tvm_ffi/cpp/extension.py

@@ -40,21 +40,50 @@
 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: