Fix IS_ROCM hardcoding and PYTORCH_ROCM_ARCH detection in setup.py

setup.py

@@ -2,13 +2,39 @@
 import os
 import os.path as osp
 import platform
+import re
 import sys
 
 from setuptools import find_packages, setup
 
-IS_ROCM = True
+try:
+    import torch
+
+    IS_ROCM = torch.version.hip is not None
+except ImportError:
+    IS_ROCM = False
 ROCM_HOME = "/opt/rocm"
 
+
+def get_rocm_arch():
+    """
+    Resolves the GPU architecture (gfx code) to build for.
+
+    Checks the PYTORCH_ROCM_ARCH environment variable first, since
+    PYTORCH_ROCM_ARCH may list multiple archs (comma/semicolon
+    separated); only the first one is used here. Falls back to the
+    hardcoded gfx942 default only if the env var is unset.
+
+    Returns:
+        str: The gfx code (e.g., 'gfx942', 'gfx1151').
+    """
+    env_arch = os.environ.get("PYTORCH_ROCM_ARCH", "").strip()
+    if env_arch:
+        gfx_code = re.split(r"[,;]", env_arch)[0].strip()
+        print(f"Using PYTORCH_ROCM_ARCH from environment: {gfx_code}")
+        return gfx_code
+    return "gfx942"
+
 __version__ = None
 exec(open("nerfacc/version.py", "r").read())
 
@@ -69,7 +95,7 @@ def get_extensions():
             extra_compile_args["cxx"] += ["-arch", "arm64"]
             extra_link_args += ["-arch", "arm64"]
 
-        hipcc_flags = ["-O3", "-D__HIP_PLATFORM_AMD__", "-DUSE_ROCM" , "--offload-arch=gfx942"]
+        hipcc_flags = ["-O3", "-D__HIP_PLATFORM_AMD__", "-DUSE_ROCM" , f"--offload-arch={get_rocm_arch()}"]
 
         if torch.version.hip:
             # USE_ROCM was added to later versions of PyTorch.