From 552c22356943f5574ec13de7c9f090def7bb90d9 Mon Sep 17 00:00:00 2001
From: Seunghoon Lee <lshqqytiger@naver.com>
Date: Sat, 18 Oct 2025 01:17:30 +0900
Subject: [PATCH] use driver library, more checks for windows rocm

---
 installer.py    | 20 +++++++++++---------
 modules/rocm.py | 33 ++++++++++++++++++++++++---------
 2 files changed, 35 insertions(+), 18 deletions(-)

diff --git a/installer.py b/installer.py
index 835478690..80eb9a635 100644
--- a/installer.py
+++ b/installer.py
@@ -682,11 +682,8 @@ def install_rocm_zluda():
 
     amd_gpus = []
     try:
-        if sys.platform == "win32" and not rocm.is_installed:
-            amd_gpus = rocm.driver_get_agents()
-        else:
-            amd_gpus = rocm.get_agents()
-            log.info('ROCm: AMD toolkit detected')
+        amd_gpus = rocm.get_agents()
+        log.info('ROCm: AMD toolkit detected')
     except Exception as e:
         log.warning(f'ROCm agent enumerator failed: {e}')
 
@@ -712,10 +709,13 @@ def install_rocm_zluda():
             if device_id < len(amd_gpus):
                 device = amd_gpus[device_id]
 
-    if sys.platform == "win32" and args.use_rocm and not rocm.is_installed:
+    if sys.platform == "win32" and args.use_rocm and not rocm.is_installed and device is not None:
         check_python(supported_minors=[11, 12, 13], reason='ROCm backend requires a Python version between 3.11 and 3.13')
-        install(f"rocm rocm-sdk-core --index-url https://rocm.nightlies.amd.com/v2-staging/{device.therock}")
-        rocm.refresh()
+        if device.therock is None:
+            log.warning('No supported ROCm agent was found. Skipping ROCm package installation.')
+        else:
+            install(f"rocm rocm-sdk-core --index-url https://rocm.nightlies.amd.com/v2-staging/{device.therock}")
+            rocm.refresh()
 
     msg = f'ROCm: version={rocm.version}'
     if device is not None:
@@ -724,7 +724,9 @@ def install_rocm_zluda():
 
     if sys.platform == "win32":
         if args.use_rocm: # TODO install: switch to pytorch source when it becomes available
-            if device is not None and isinstance(rocm.environment, rocm.PythonPackageEnvironment): # TheRock
+            if device is None:
+                log.warning('No ROCm agent was found. Please make sure that graphics driver is installed and up to date.')
+            if isinstance(rocm.environment, rocm.PythonPackageEnvironment):
                 check_python(supported_minors=[11, 12, 13], reason='ROCm backend requires a Python version between 3.11 and 3.13')
                 torch_command = os.environ.get('TORCH_COMMAND', f'torch torchvision --index-url https://rocm.nightlies.amd.com/v2-staging/{device.therock}')
             else:
diff --git a/modules/rocm.py b/modules/rocm.py
index 1b03bfd32..abdc7fa32 100644
--- a/modules/rocm.py
+++ b/modules/rocm.py
@@ -96,7 +96,7 @@ class Agent:
         self.blaslt_supported = os.path.exists(os.path.join(blaslt_tensile_libpath, f"Kernels.so-000-{name}.hsaco" if sys.platform == "win32" else f"extop_{name}.co"))
 
     @property
-    def therock(self) -> str:
+    def therock(self) -> Union[str, None]:
         if (self.gfx_version & 0xFFF0) == 0x1100:
             return "gfx110X-dgpu"
         if self.gfx_version == 0x1151:
@@ -107,7 +107,7 @@ class Agent:
             return "gfx94X-dcgpu"
         if self.gfx_version == 0x950:
             return "gfx950-dcgpu"
-        raise RuntimeError(f"Unsupported GPU architecture: {self.name}")
+        return None
 
     def get_gfx_version(self) -> Union[str, None]:
         if self.gfx_version >= 0x1100 and self.gfx_version < 0x1200:
@@ -209,13 +209,16 @@ def get_flash_attention_command(agent: Agent) -> str:
 
 if sys.platform == "win32":
     def get_agents() -> List[Agent]:
-        if isinstance(environment, ROCmEnvironment):
-            out = spawn("amdgpu-arch", cwd=os.path.join(environment.path, 'bin'))
-        else:
-            # Assume that amdgpu-arch is in PATH (venv/Scripts/amdgpu-arch.exe)
-            out = spawn("amdgpu-arch")
-        out = out.strip()
-        return [Agent(x.split(' ')[-1].strip()) for x in out.split("\n")]
+        return agents
+        #if isinstance(environment, ROCmEnvironment):
+        #    out = spawn("amdgpu-arch", cwd=os.path.join(environment.path, 'bin'))
+        #else:
+        #    # Assume that amdgpu-arch is in PATH (venv/Scripts/amdgpu-arch.exe)
+        #    out = spawn("amdgpu-arch")
+        #out = out.strip()
+        #if out == "":
+        #    return []
+        #return [Agent(x.split(' ')[-1].strip()) for x in out.split("\n")]
 
     def driver_get_agents() -> List[Agent]:
         # unsafe and experimental feature
@@ -243,6 +246,14 @@ if sys.platform == "win32":
             os.environ["PATH"] = ";".join(paths_no_rocm)
             return
 
+        build_targets = torch.cuda.get_arch_list()
+        for available in agents:
+            if available.name in build_targets:
+                return
+
+        # use cpu instead of crashing
+        torch.cuda.is_available = lambda: False
+
     def rocm_init():
         try:
             import torch
@@ -275,6 +286,7 @@ if sys.platform == "win32":
         return True, None
 
     is_wsl: bool = False
+    agents: List[Agent] = [] # temp
 else:
     def get_agents() -> List[Agent]:
         try:
@@ -307,6 +319,9 @@ version = None
 
 def refresh():
     global environment, blaslt_tensile_libpath, is_installed, version # pylint: disable=global-statement
+    if sys.platform == "win32":
+        global agents
+        agents = driver_get_agents()
     environment = find()
     if environment is not None:
         if isinstance(environment, ROCmEnvironment):