zluda init

modules/rocm.py

@@ -3,7 +3,6 @@ import sys
 import glob
 import ctypes
 import shutil
-import logging
 import subprocess
 from types import ModuleType
 from typing import Union, overload, TYPE_CHECKING
@@ -282,30 +281,32 @@ if sys.platform == "win32":
             os.environ["PATH"] = ";".join(paths_no_rocm)
             return
 
-        build_targets = torch.cuda.get_arch_list()
-        agents = get_agents()
-        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
             import numpy as np
             from installer import log
             from modules.devices import get_hip_agent
+            from modules.rocm_triton_windows import apply_triton_patches
+
+            build_targets = torch.cuda.get_arch_list()
+            agents = get_agents()
+            if all(available.name not in build_targets for available in agents):
+                log.warning('ROCm: torch-rocm is installed, but none of build targets is available')
+                # use cpu instead of crashing
+                torch.cuda.is_available = lambda: False
 
             agent = get_hip_agent()
             if not agent.blaslt_supported:
-                log.log(logging.DEBUG if torch.version.hip is None else logging.WARNING, f'ROCm: hipBLASLt unavailable agent={agent}')
+                log.warning(f'ROCm: hipBLASLt unavailable agent={agent}')
             if (agent.gfx_version & 0xFFF0) == 0x1200:
                 # disable MIOpen for gfx120x
                 torch.backends.cudnn.enabled = False
                 log.debug('ROCm: disabled MIOpen')
 
+            if sys.platform == "win32":
+                apply_triton_patches()
+
             original_cholesky_ex = torch.linalg.cholesky_ex
             @wraps(original_cholesky_ex)
             def cholesky_ex(A: torch.Tensor, upper=False, check_errors=False, out=None) -> torch.return_types.linalg_cholesky_ex:
@@ -348,6 +349,7 @@ else: # sys.platform != "win32"
             try:
                 if shutil.which("conda") is not None:
                     # Preload stdc++ library. This will bypass Anaconda stdc++ library.
+                    # (hsa-runtime64 depends on stdc++)
                     load_library_global("/lib/x86_64-linux-gnu/libstdc++.so.6")
                 # Preload rocr4wsl. The user don't have to replace the library file.
                 load_library_global("/opt/rocm/lib/libhsa-runtime64.so")

modules/shared.py

@@ -862,13 +862,6 @@ mem_mon = modules.memmon.MemUsageMonitor("MemMon", devices.device)
 history = history.History()
 if devices.backend == "directml":
     directml_do_hijack()
-elif sys.platform == "win32" and (devices.backend == "zluda" or devices.backend == "rocm"):
-    from modules.rocm_triton_windows import apply_triton_patches
-    apply_triton_patches()
-
-    if devices.backend == "zluda":
-        from modules.zluda import initialize_zluda
-        initialize_zluda()
 from modules import sdnq # pylint: disable=unused-import # register to diffusers and transformers
 log.debug('Quantization: registered=SDNQ')
 

modules/zluda.py

@@ -1,17 +1,14 @@
 import sys
 from typing import Union
-import torch
-from torch._prims_common import DeviceLikeType
-from modules import shared, devices, zluda_installer
 from modules.zluda_installer import core, default_agent # pylint: disable=unused-import
-from modules.onnx_impl.execution_providers import available_execution_providers, ExecutionProvider
 
 
 PLATFORM = sys.platform
 do_nothing = lambda _: None # pylint: disable=unnecessary-lambda-assignment
 
 
-def test(device: DeviceLikeType) -> Union[Exception, None]:
+def test(device) -> Union[Exception, None]:
+    import torch
     device = torch.device(device)
     try:
         ten1 = torch.randn((2, 4,), device=device)
@@ -23,40 +20,43 @@ def test(device: DeviceLikeType) -> Union[Exception, None]:
         return e
 
 
-def initialize_zluda():
-    shared.cmd_opts.device_id = None
-    if not devices.cuda_ok or not devices.has_zluda():
-        return
-
-    torch.backends.cudnn.enabled = zluda_installer.MIOpen_enabled if shared.opts.cudnn_enabled == 'default' else shared.opts.cudnn_enabled == 'true'
-    if hasattr(torch.backends.cuda, "enable_cudnn_sdp"):
-        if not zluda_installer.MIOpen_enabled:
-            torch.backends.cuda.enable_cudnn_sdp(False)
-            torch.backends.cuda.enable_cudnn_sdp = do_nothing
-    else:
-        torch.backends.cuda.enable_cudnn_sdp = do_nothing
-    torch.backends.cuda.enable_flash_sdp(False)
-    torch.backends.cuda.enable_flash_sdp = torch.backends.cuda.enable_cudnn_sdp
-    torch.backends.cuda.enable_mem_efficient_sdp(False)
-    torch.backends.cuda.enable_mem_efficient_sdp = do_nothing
-
-    # ONNX Runtime is not supported
+def zluda_init():
     try:
-        import onnxruntime as ort
-        ort.capi._pybind_state.get_available_providers = lambda: [v for v in available_execution_providers if v != ExecutionProvider.CUDA] # pylint: disable=protected-access
-        ort.get_available_providers = ort.capi._pybind_state.get_available_providers # pylint: disable=protected-access
-        if shared.opts.onnx_execution_provider == ExecutionProvider.CUDA:
-            shared.opts.onnx_execution_provider = ExecutionProvider.CPU
-    except Exception as e:
-        shared.log.warning(f'ZLUDA ONNX runtime: {e}')
-        shared.opts.onnx_execution_provider = ExecutionProvider.CPU
+        import torch
+        from installer import log
+        from modules import devices, zluda_installer
+        from modules.shared import cmd_opts
+        from modules.rocm_triton_windows import apply_triton_patches
+        from modules.onnx_impl.execution_providers import available_execution_providers, ExecutionProvider
 
-    device = devices.get_optimal_device()
-    result = test(device)
-    if result is not None:
-        shared.log.warning(f'ZLUDA device failed to pass basic operation test: index={device.index}, device_name={torch.cuda.get_device_name(device)}')
-        shared.log.error(result)
-        torch.cuda.is_available = lambda: False
-        devices.cuda_ok = False
-        devices.backend = 'cpu'
-        devices.device = devices.cpu
+        cmd_opts.device_id = None
+
+        device = devices.get_optimal_device()
+        result = test(device)
+        if result is not None:
+            log.warning(f'ZLUDA device failed to pass basic operation test: index={device.index}, device_name={torch.cuda.get_device_name(device)}')
+            torch.cuda.is_available = lambda: False
+            devices.cuda_ok = False
+            devices.backend = 'cpu'
+            devices.device = devices.cpu
+            return False, result
+
+        if not zluda_installer.default_agent.blaslt_supported:
+            log.warning(f'ROCm: hipBLASLt unavailable agent={zluda_installer.default_agent}')
+
+        apply_triton_patches()
+
+        torch.backends.cudnn.enabled = zluda_installer.MIOpen_enabled
+        if hasattr(torch.backends.cuda, "enable_cudnn_sdp"):
+            if not zluda_installer.MIOpen_enabled:
+                torch.backends.cuda.enable_cudnn_sdp(False)
+                torch.backends.cuda.enable_cudnn_sdp = do_nothing
+        else:
+            torch.backends.cuda.enable_cudnn_sdp = do_nothing
+        torch.backends.cuda.enable_flash_sdp(False)
+        torch.backends.cuda.enable_flash_sdp = torch.backends.cuda.enable_cudnn_sdp
+        torch.backends.cuda.enable_mem_efficient_sdp(False)
+        torch.backends.cuda.enable_mem_efficient_sdp = do_nothing
+    except Exception as e:
+        return False, e
+    return True, None

modules/zluda_installer.py

@@ -161,6 +161,7 @@ def load():
     def postinstall():
         import torch
         torch.version.hip = rocm.version
+
         platform = sys.platform
         sys.platform = ""
         from torch.utils import cpp_extension
@@ -172,3 +173,6 @@ def load():
             return os.path.join(cpp_extension.ROCM_HOME, *paths)
         cpp_extension._join_rocm_home = _join_rocm_home # pylint: disable=protected-access
     rocm.postinstall = postinstall
+
+    from modules.zluda import zluda_init
+    rocm.rocm_init = zluda_init