From cba5a63bb7cc3486ea31af3505aa1ca2a8c44d9d Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Bj=C3=B6rn=20Ellensohn?= <bjoern.ellensohn@unileoben.ac.at>
Date: Mon, 10 Aug 2026 21:34:06 +0200
Subject: [PATCH] fix: derive IS_ROCM from torch.version.hip, respect
 PYTORCH_ROCM_ARCH

Two bug-pattern fixes, same shape as bjoernellens1/gsplat's earlier
fixes for the identical problems:
- IS_ROCM was a hardcoded module-level True, meaning this always
  compiled with ROCm/HIP flags regardless of the actual installed
  PyTorch build. Now derived from torch.version.hip.
- hipcc_flags hardcoded --offload-arch=gfx942 (CDNA/Instinct only).
  Now checks PYTORCH_ROCM_ARCH first, falls back to gfx942 only if
  unset -- verified building correctly for gfx1151 (RDNA3.5) with the
  env var set, matching gsplat's release/1.5.3b2 commit 4515618's
  fix for the same issue.
---
 setup.py | 30 ++++++++++++++++++++++++++++--
 1 file changed, 28 insertions(+), 2 deletions(-)

diff --git a/setup.py b/setup.py
index aa0ade9..fe0e4f3 100644
--- a/setup.py
+++ b/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.