From 9195116e4663c778aecef77c80b14f6cdb6ebadc Mon Sep 17 00:00:00 2001 From: Seunghoon Lee Date: Wed, 7 Jan 2026 14:59:42 +0900 Subject: [PATCH] zluda init --- modules/rocm.py | 24 ++++++------ modules/shared.py | 7 ---- modules/zluda.py | 80 +++++++++++++++++++------------------- modules/zluda_installer.py | 4 ++ 4 files changed, 57 insertions(+), 58 deletions(-) diff --git a/modules/rocm.py b/modules/rocm.py index 157ea0d4e..2395e968e 100644 --- a/modules/rocm.py +++ b/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") diff --git a/modules/shared.py b/modules/shared.py index 3a3bdab5f..369810b5b 100644 --- a/modules/shared.py +++ b/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') diff --git a/modules/zluda.py b/modules/zluda.py index 4bb263f22..5358751eb 100644 --- a/modules/zluda.py +++ b/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 diff --git a/modules/zluda_installer.py b/modules/zluda_installer.py index aa573e547..91b5d8b20 100644 --- a/modules/zluda_installer.py +++ b/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