windows install rocm if not installed

installer.py

@@ -678,32 +678,42 @@ def install_rocm_zluda():
         return torch_command
     from modules import rocm
 
-    log.info('ROCm: AMD toolkit detected')
+    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')
+    except Exception as e:
+        log.warning(f'ROCm agent enumerator failed: {e}')
+
     #os.environ.setdefault('TENSORFLOW_PACKAGE', 'tensorflow')
 
     device = None
-    try:
-        amd_gpus = rocm.get_agents()
-        if len(amd_gpus) == 0:
-            log.warning('ROCm: no agent was found')
+    if len(amd_gpus) == 0:
+        log.warning('ROCm: no agent was found')
+    else:
+        log.info(f'ROCm: agents={[gpu.name for gpu in amd_gpus]}')
+        if args.device_id is None:
+            index = 0
+            for idx, gpu in enumerate(amd_gpus):
+                index = idx
+                if not gpu.is_apu:
+                    # although apu was found, there can be a dedicated card. do not break loop.
+                    # if no dedicated card was found, apu will be used.
+                    break
+            os.environ.setdefault('HIP_VISIBLE_DEVICES', str(index))
+            device = amd_gpus[index]
         else:
-            log.info(f'ROCm: agents={[gpu.name for gpu in amd_gpus]}')
-            if args.device_id is None:
-                index = 0
-                for idx, gpu in enumerate(amd_gpus):
-                    index = idx
-                    if not gpu.is_apu:
-                        # although apu was found, there can be a dedicated card. do not break loop.
-                        # if no dedicated card was found, apu will be used.
-                        break
-                os.environ.setdefault('HIP_VISIBLE_DEVICES', str(index))
-                device = amd_gpus[index]
-            else:
-                device_id = int(args.device_id)
-                if device_id < len(amd_gpus):
-                    device = amd_gpus[device_id]
-    except Exception as e:
-        log.warning(f'ROCm agent enumerator failed: {e}')
+            device_id = int(args.device_id)
+            if device_id < len(amd_gpus):
+                device = amd_gpus[device_id]
+
+    if sys.platform == "win32" and args.use_rocm and not rocm.is_installed:
+        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()
 
     msg = f'ROCm: version={rocm.version}'
     if device is not None:
@@ -714,7 +724,7 @@ def install_rocm_zluda():
         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
                 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/{rocm.get_distribution(device)}')
+                torch_command = os.environ.get('TORCH_COMMAND', f'torch torchvision --index-url https://rocm.nightlies.amd.com/v2-staging/{device.therock}')
             else:
                 check_python(supported_minors=[12], reason='ROCm Windows preview requires Python version 3.12')
                 torch_command = os.environ.get('TORCH_COMMAND', '--no-cache-dir https://repo.radeon.com/rocm/windows/rocm-rel-6.4.4/torch-2.8.0a0%2Bgitfc14c65-cp312-cp312-win_amd64.whl https://repo.radeon.com/rocm/windows/rocm-rel-6.4.4/torchvision-0.24.0a0%2Bc85f008-cp312-cp312-win_amd64.whl')

modules/rocm.py

@@ -105,6 +105,20 @@ class Agent:
         self.is_apu = (self.gfx_version & 0xFFF0 == 0x1150) or self.gfx_version in (0x801, 0x902, 0x90c, 0x1013, 0x1033, 0x1035, 0x1036, 0x1103,)
         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:
+        if (self.gfx_version & 0xFFF0) == 0x1100:
+            return "gfx110X-dgpu"
+        if self.gfx_version == 0x1151:
+            return "gfx1151"
+        if (self.gfx_version & 0xFFF0) == 0x1200:
+            return "gfx120X-all"
+        if (self.gfx_version & 0xFFF0) == 0x940:
+            return "gfx94X-dcgpu"
+        if self.gfx_version == 0x950:
+            return "gfx950-dcgpu"
+        raise Exception(f"Unsupported GPU architecture: {self.name}")
+
     def get_gfx_version(self) -> Union[str, None]:
         if self.gfx_version >= 0x1100 and self.gfx_version < 0x1200:
             return "11.0.0"
@@ -204,29 +218,29 @@ def get_flash_attention_command(agent: Agent) -> str:
     return "--no-build-isolation " + os.environ.get("FLASH_ATTENTION_PACKAGE", default)
 
 
-def get_distribution(agent: Agent) -> str:
-    if (agent.gfx_version & 0xFFF0) == 0x1100:
-        return "gfx110X-dgpu"
-    if agent.gfx_version == 0x1151:
-        return "gfx1151"
-    if (agent.gfx_version & 0xFFF0) == 0x1200:
-        return "gfx120X-all"
-    if (agent.gfx_version & 0xFFF0) == 0x940:
-        return "gfx94X-dcgpu"
-    if agent.gfx_version == 0x950:
-        return "gfx950-dcgpu"
-    raise Exception(f"Unsupported GPU architecture: {agent.name}")
-
-
 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")]
 
+    def driver_get_agents() -> List[Agent]:
+        # unsafe and experimental feature
+        from modules import windows_hip_ffi
+        hip = windows_hip_ffi.HIP()
+        count = hip.get_device_count()
+        agents = [None] * count
+        for i in range(count):
+            prop = hip.get_device_properties(i)
+            name = prop.gcnArchName.decode('utf-8').strip('\x00')
+            agents[i] = Agent(name)
+        del hip
+        return agents
+
     is_wsl: bool = False
 else:
     def get_agents() -> List[Agent]:
@@ -250,16 +264,16 @@ else:
 
     is_wsl: bool = os.environ.get('WSL_DISTRO_NAME', 'unknown' if spawn('wslpath -w /') else None) is not None
 environment = None
-err = None
-try:
-    environment = find()
-except Exception as e:
-    err = e
 blaslt_tensile_libpath = ""
 is_installed = False
 version = None
-if environment is not None:
-    if isinstance(environment, ROCmEnvironment):
-        blaslt_tensile_libpath = os.environ.get("HIPBLASLT_TENSILE_LIBPATH", os.path.join(environment.path, "bin" if sys.platform == "win32" else "lib", "hipblaslt", "library"))
-    is_installed = True
-    version = get_version()
+
+def refresh():
+    global environment, blaslt_tensile_libpath, is_installed, version
+    environment = find()
+    if environment is not None:
+        if isinstance(environment, ROCmEnvironment):
+            blaslt_tensile_libpath = os.environ.get("HIPBLASLT_TENSILE_LIBPATH", os.path.join(environment.path, "bin" if sys.platform == "win32" else "lib", "hipblaslt", "library"))
+        is_installed = True
+        version = get_version()
+refresh()

modules/windows_hip_ffi.py

@@ -0,0 +1,44 @@
+import sys
+
+if sys.platform == "win32":
+    import ctypes
+    import ctypes.wintypes
+
+    class hipDeviceProp(ctypes.Structure):
+        _fields_ = [
+            ('__front__', ctypes.c_byte * 396),
+            ('gcnArchName', ctypes.c_char * 256),
+            ('__rear__', ctypes.c_byte * 820)
+        ]
+
+    class HIP:
+        def __init__(self):
+            ctypes.windll.kernel32.LoadLibraryA.restype = ctypes.wintypes.HMODULE
+            ctypes.windll.kernel32.LoadLibraryA.argtypes = [ctypes.c_char_p]
+            # amdhip64.dll is a part of AMDGPU drivers
+            self.handle = ctypes.windll.kernel32.LoadLibraryA(b"amdhip64.dll")
+            ctypes.windll.kernel32.GetLastError.restype = ctypes.wintypes.DWORD
+            ctypes.windll.kernel32.GetLastError.argtypes = []
+            assert ctypes.windll.kernel32.GetLastError() == 0
+            ctypes.windll.kernel32.GetProcAddress.restype = ctypes.c_void_p
+            ctypes.windll.kernel32.GetProcAddress.argtypes = [ctypes.wintypes.HMODULE, ctypes.c_char_p]
+            self.hipGetDeviceCount = ctypes.CFUNCTYPE(ctypes.c_int, ctypes.POINTER(ctypes.c_int))(
+                ctypes.windll.kernel32.GetProcAddress(self.handle, b"hipGetDeviceCount"))
+            self.hipGetDeviceProperties = ctypes.CFUNCTYPE(ctypes.c_int, ctypes.POINTER(hipDeviceProp), ctypes.c_int)(
+                ctypes.windll.kernel32.GetProcAddress(self.handle, b"hipGetDeviceProperties"))
+
+        def __del__(self):
+            #ctypes.windll.kernel32.FreeLibrary.argtypes = [ctypes.wintypes.HMODULE]
+            #ctypes.windll.kernel32.FreeLibrary(self.handle)
+            # Hopefully it does not make conflicts with amdhip64_7.dll
+            pass
+
+        def get_device_count(self):
+            count = ctypes.c_int()
+            assert self.hipGetDeviceCount(ctypes.byref(count)) == 0
+            return count.value
+
+        def get_device_properties(self, device_id):
+            prop = hipDeviceProp()
+            assert self.hipGetDeviceProperties(ctypes.byref(prop), device_id) == 0
+            return prop