[Feature] Support environments built by TheRock build system

flashinfer/hip_utils.py

@@ -6,40 +6,123 @@
 FLASHINFER_SUPPORTED_ROCM_ARCHS = ["gfx942"]
 
 
-def get_system_rocm_version():
+def get_rocm_home():
     """
-    Attempt to detect the system ROCm version.
+    Get the ROCM_HOME directory from environment variables or default path.
 
     Returns:
-        str: ROCm version like "6.4" or "7.0", or None if not detectable
+        str: Path to ROCm installation (e.g., "/opt/rocm")
     """
     import os
-    import re
-    import subprocess
 
-    # Method 1: Try /opt/rocm/.info/version (most reliable)
-    rocm_path = os.environ.get("ROCM_PATH", "/opt/rocm")
-    version_file = os.path.join(rocm_path, ".info", "version")
+    return os.environ.get("ROCM_PATH") or os.environ.get("ROCM_HOME") or "/opt/rocm"
+
+
+def is_therock_build() -> bool:
+    """
+    Check if ROCm was built using TheRock build system.
+
+    Returns:
+        bool: True if TheRock build is detected, False otherwise
+    """
+    import os
+
+    # First, try checking for rocm_sdk package
+    try:
+        import rocm_sdk
+
+        if hasattr(rocm_sdk, "__version__") and rocm_sdk.__version__:
+            return True
+    except ImportError:
+        pass
+
+    # Fall back to checking for TheRock manifest file
+    rocm_home = get_rocm_home()
+    manifest_path = os.path.join(rocm_home, "share", "therock", "therock_manifest.json")
+    return os.path.isfile(manifest_path)
+
+
+def get_system_rocm_version_from_info_file():
+    """
+    Try to get ROCm version from .info/version file located in ROCM_HOME.
+
+    Returns:
+        str: ROCm version like "7.1.0" or None if not found
+    """
+    import os
+
+    rocm_home = get_rocm_home()
+    version_file = os.path.join(rocm_home, ".info", "version")
     try:
         with open(version_file, "r") as f:
             version = f.read().strip()
             return ".".join(version.split(".")[:3])
     except (FileNotFoundError, IOError):
+        return None
+
+
+def get_system_rocm_version_from_hipconfig():
+    """
+    Try to get ROCm version from hipconfig --version command.
+
+    Returns:
+        str: ROCm version like "7.1.0" or None if not found
+    """
+    import re
+    import subprocess
+
+    try:
+        result = subprocess.run(
+            ["hipconfig", "--version"],
+            capture_output=True,
+            text=True,
+            timeout=5,
+            check=False,
+        )
+        if result.returncode == 0:
+            match = re.search(r"(\d+\.\d+\.\d+)", result.stdout)
+            if match:
+                return match.group(1)
+    except (subprocess.TimeoutExpired, FileNotFoundError):
         pass
 
-    # Method 2: Try amd-smi command
+    return None
+
+
+def get_system_rocm_version_from_amd_smi():
+    """
+    Try to get ROCm version from amd-smi command.
+
+    Returns:
+        str: ROCm version like "7.1.0" or None if not found
+    """
+    import re
+    import subprocess
+
     try:
         result = subprocess.run(
             ["amd-smi"], capture_output=True, text=True, timeout=5, check=False
         )
         if result.returncode == 0:
-            match = re.search(r"ROCm version:\s*(\d+\.\d+\.\d)", result.stdout)
+            match = re.search(r"ROCm version:\s*(\d+\.\d+\.\d+)", result.stdout)
             if match:
                 return match.group(1)
     except (subprocess.TimeoutExpired, FileNotFoundError):
         pass
 
-    # Method 3: Try dpkg (Ubuntu/Debian)
+    return None
+
+
+def get_system_rocm_version_from_dpkg():
+    """
+    Try to get ROCm version from dpkg (Ubuntu/Debian package manager).
+
+    Returns:
+        str: ROCm version like "7.1.0" or None if not found
+    """
+    import re
+    import subprocess
+
     try:
         result = subprocess.run(
             ["dpkg", "-l", "rocm-core"],
@@ -49,7 +132,7 @@ def get_system_rocm_version():
             check=False,
         )
         if result.returncode == 0:
-            match = re.search(r"rocm-core\s+(\d+\.\d+\.\d)", result.stdout)
+            match = re.search(r"rocm-core\s+(\d+\.\d+\.\d+)", result.stdout)
             if match:
                 return match.group(1)
     except (subprocess.TimeoutExpired, FileNotFoundError):
@@ -58,6 +141,37 @@ def get_system_rocm_version():
     return None
 
 
+def get_system_rocm_version():
+    """
+    Attempt to detect the system ROCm version.
+
+    For standard builds, tries methods in order of reliability.
+    For TheRock builds, prioritizes hipconfig as it's more reliable.
+
+    Returns:
+        str: ROCm version like "7.1.0" or None if not detectable
+    """
+    # For TheRock builds, prioritize hipconfig
+    if is_therock_build():
+        return get_system_rocm_version_from_hipconfig()
+
+    # Try standard detection methods in order of reliability
+    detection_methods = [
+        get_system_rocm_version_from_info_file,
+        get_system_rocm_version_from_amd_smi,
+        get_system_rocm_version_from_dpkg,
+        get_system_rocm_version_from_hipconfig,
+    ]
+
+    for method in detection_methods:
+        version = method()
+        if version:
+            return version
+        print(f"ROCm version not found using {method.__name__}. Trying next method...")
+
+    return None
+
+
 def validate_rocm_arch(arch_list: str = None, verbose: bool = False) -> str:
     """
     Validate ROCm architecture against system ROCm version.
@@ -77,42 +191,33 @@ def validate_rocm_arch(arch_list: str = None, verbose: bool = False) -> str:
 
     # ROCm compatibility matrix: version -> supported gfx architectures
     # Refer: https://rocm.docs.amd.com/en/latest/compatibility/compatibility-matrix.html
+    #        https://github.com/ROCm/TheRock/blob/main/SUPPORTED_GPUS.md#rocm-on-linux
+    # Update lists for adding or removing a version or arch
+    # Add new tuple for adding a new version group
+    _ROCM_ARCH_GROUPS = [
+        (
+            ["7.3", "7.2", "7.1", "7.0"],
+            [
+                "gfx950",
+                "gfx1201",
+                "gfx1200",
+                "gfx1101",
+                "gfx1100",
+                "gfx1030",
+                "gfx942",
+                "gfx90a",
+                "gfx908",
+            ],
+        ),
+        (
+            ["6.4", "6.3"],
+            ["gfx1100", "gfx1030", "gfx942", "gfx90a", "gfx908"],
+        ),
+    ]
+
+    # Build the compatibility matrix
     ROCM_COMPAT_MATRIX = {
-        "7.2": [
-            "gfx950",
-            "gfx1201",
-            "gfx1200",
-            "gfx1101",
-            "gfx1100",
-            "gfx1030",
-            "gfx942",
-            "gfx90a",
-            "gfx908",
-        ],
-        "7.1": [
-            "gfx950",
-            "gfx1201",
-            "gfx1200",
-            "gfx1101",
-            "gfx1100",
-            "gfx1030",
-            "gfx942",
-            "gfx90a",
-            "gfx908",
-        ],
-        "7.0": [
-            "gfx950",
-            "gfx1201",
-            "gfx1200",
-            "gfx1101",
-            "gfx1100",
-            "gfx1030",
-            "gfx942",
-            "gfx90a",
-            "gfx908",
-        ],
-        "6.4": ["gfx1100", "gfx1030", "gfx942", "gfx90a", "gfx908"],
-        "6.3": ["gfx1100", "gfx1030", "gfx942", "gfx90a", "gfx908"],
+        version: archs for versions, archs in _ROCM_ARCH_GROUPS for version in versions
     }
 
     # Get architecture list from parameter, env var, or default
@@ -124,7 +229,7 @@ def validate_rocm_arch(arch_list: str = None, verbose: bool = False) -> str:
     if system_rocm_version is None:
         raise RuntimeError(
             "Could not detect ROCm installation. Please ensure ROCm is installed and "
-            "accessible (check ROCM_PATH or /opt/rocm)."
+            "accessible (check ROCM_PATH, ROCM_HOME or /opt/rocm)."
         )
 
     # Parse version to major.minor for compatibility check

flashinfer/jit/cpp_ext_hip.py

@@ -19,10 +19,12 @@
     _get_pybind11_abi_build_flags,
 )
 
-from torch.utils.cpp_extension import ROCM_HOME
+from flashinfer.hip_utils import get_rocm_home
 
 from . import env as jit_env
 
+ROCM_HOME = get_rocm_home()
+
 
 def _get_glibcxx_abi_build_flags() -> List[str]:
     glibcxx_abi_cflags = [
@@ -112,11 +114,7 @@ def generate_ninja_build_for_op(
         ldflags += extra_ldflags
 
     cxx = os.environ.get("CXX", "c++")
-    rocm_home = ROCM_HOME or "/opt/rocm"
-    amdclang = os.environ.get("PYTORCH_AMDCLANG", "$rocm_home/bin/amdclang++")
-
-    cxx = os.environ.get("CXX", "c++")
-    rocm_home = ROCM_HOME or "/opt/rocm"
+    rocm_home = ROCM_HOME
     amdclang = os.environ.get("PYTORCH_AMDCLANG", "$rocm_home/bin/amdclang++")
 
     lines = [
@@ -164,11 +162,7 @@ def generate_ninja_build_for_op(
     for source in sources:
         is_hip = source.suffix == ".cu"
         object_suffix = ".cuda.o" if is_hip else ".o"
-        cmd = ""
-        if is_hip:
-            cmd = "hip_compile"
-        else:
-            cmd = "compile"
+        cmd = "hip_compile" if is_hip else "compile"
         obj_name = source.with_suffix(object_suffix).name
         obj = f"$name/{obj_name}"
         objects.append(obj)