From 0de373288472d7085e85d24a15e2c2a85d986175 Mon Sep 17 00:00:00 2001 From: Seunghoon Lee Date: Sun, 28 Sep 2025 11:31:43 +0900 Subject: [PATCH 1/8] prioritize python pacakage check over PATH --- modules/rocm.py | 17 ++++++++--------- 1 file changed, 8 insertions(+), 9 deletions(-) diff --git a/modules/rocm.py b/modules/rocm.py index 8dfcd9ec8..3ffff069a 100644 --- a/modules/rocm.py +++ b/modules/rocm.py @@ -130,17 +130,16 @@ class Agent: def find() -> Union[Environment, None]: + try: # TheRock + import _rocm_sdk_core # pylint: disable=unused-import + return PythonPackageEnvironment() + except ImportError: + pass + + # system-wide installation hip_path = shutil.which("hipconfig") if hip_path is not None: - py_path = os.path.dirname(sys.executable) - if hip_path.startswith(py_path): - try: - import _rocm_sdk_core # pylint: disable=unused-import - return PythonPackageEnvironment() - except ImportError: - pass - else: - return ROCmEnvironment(dirname(resolve_link(hip_path), 2)) + return ROCmEnvironment(dirname(resolve_link(hip_path), 2)) if sys.platform == "win32": hip_path = os.environ.get("HIP_PATH", None) From 95ee0ed29c6d02e581611d72b24d798b6c5c9acf Mon Sep 17 00:00:00 2001 From: Seunghoon Lee Date: Sun, 28 Sep 2025 17:15:12 +0900 Subject: [PATCH 2/8] load amdhip64_6.dll instead of amdhip64.dll --- modules/windows_hip_ffi.py | 15 +++++++++------ 1 file changed, 9 insertions(+), 6 deletions(-) diff --git a/modules/windows_hip_ffi.py b/modules/windows_hip_ffi.py index cc9e4a1e4..c612aa43d 100644 --- a/modules/windows_hip_ffi.py +++ b/modules/windows_hip_ffi.py @@ -1,6 +1,7 @@ import sys if sys.platform == "win32": + import os import ctypes import ctypes.wintypes @@ -15,8 +16,11 @@ if sys.platform == "win32": 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") + path = os.environ.get("windir", "C:\\Windows") + "\\System32\\amdhip64_6.dll" + if not os.path.isfile(path): + path = os.environ.get("windir", "C:\\Windows") + "\\System32\\amdhip64_7.dll" + assert os.path.isfile(path) + self.handle = ctypes.windll.kernel32.LoadLibraryA(path.encode('utf-8')) ctypes.windll.kernel32.GetLastError.restype = ctypes.wintypes.DWORD ctypes.windll.kernel32.GetLastError.argtypes = [] assert ctypes.windll.kernel32.GetLastError() == 0 @@ -28,10 +32,9 @@ if sys.platform == "win32": 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 + # Hopefully this will prevent conflicts with amdhip64_7.dll from ROCm Python packages or HIP SDK + ctypes.windll.kernel32.FreeLibrary.argtypes = [ctypes.wintypes.HMODULE] + ctypes.windll.kernel32.FreeLibrary(self.handle) def get_device_count(self): count = ctypes.c_int() From e8a507fe2a19ce984f961bd182089fadc57354fd Mon Sep 17 00:00:00 2001 From: Seunghoon Lee Date: Sun, 28 Sep 2025 17:34:47 +0900 Subject: [PATCH 3/8] skip rocm.conceal if rocm torch is installed --- installer.py | 6 ++++-- 1 file changed, 4 insertions(+), 2 deletions(-) diff --git a/installer.py b/installer.py index ea7f5b496..48f0588c9 100644 --- a/installer.py +++ b/installer.py @@ -977,8 +977,10 @@ def check_torch(): if not args.ignore: sys.exit(1) if rocm.is_installed: - if sys.platform == "win32": # CPU, DirectML, ZLUDA - rocm.conceal() + if sys.platform == "win32": + import torch + if torch.version.hip is None: # CPU, DirectML, ZLUDA + rocm.conceal() elif rocm.is_wsl: # ROCm WSL rocm.preload_hsa_runtime() if args.version: From 6e8abf8dc38a6ec2362bd572bd08f4aca641dbc0 Mon Sep 17 00:00:00 2001 From: Seunghoon Lee Date: Sun, 28 Sep 2025 18:03:02 +0900 Subject: [PATCH 4/8] hijack cholesky decomposition for therock pytorch --- installer.py | 7 +---- modules/rocm.py | 53 +++++++++++++++++++++++++++++--------- modules/zluda_installer.py | 4 +-- 3 files changed, 44 insertions(+), 20 deletions(-) diff --git a/installer.py b/installer.py index 48f0588c9..1932a20af 100644 --- a/installer.py +++ b/installer.py @@ -977,12 +977,7 @@ def check_torch(): if not args.ignore: sys.exit(1) if rocm.is_installed: - if sys.platform == "win32": - import torch - if torch.version.hip is None: # CPU, DirectML, ZLUDA - rocm.conceal() - elif rocm.is_wsl: # ROCm WSL - rocm.preload_hsa_runtime() + rocm.postinstall() if args.version: return if not args.skip_all: diff --git a/modules/rocm.py b/modules/rocm.py index 3ffff069a..c7aa6771b 100644 --- a/modules/rocm.py +++ b/modules/rocm.py @@ -5,6 +5,7 @@ import shutil import subprocess from typing import Union, List from enum import Enum +from functools import wraps def resolve_link(path_: str) -> str: @@ -28,17 +29,6 @@ def load_library_global(path_: str): ctypes.CDLL(path_, mode=ctypes.RTLD_GLOBAL) -def conceal(): - os.environ.pop("ROCM_HOME", None) - os.environ.pop("ROCM_PATH", None) - paths = os.environ["PATH"].split(";") - paths_no_rocm = [] - for path_ in paths: - if "rocm" not in path_.lower(): - paths_no_rocm.append(path_) - os.environ["PATH"] = ";".join(paths_no_rocm) - - class Environment: pass @@ -240,6 +230,43 @@ if sys.platform == "win32": del hip return agents + def postinstall(): + import torch + if torch.version.hip is None: + os.environ.pop("ROCM_HOME", None) + os.environ.pop("ROCM_PATH", None) + paths = os.environ["PATH"].split(";") + paths_no_rocm = [] + for path_ in paths: + if "rocm" not in path_.lower(): + paths_no_rocm.append(path_) + os.environ["PATH"] = ";".join(paths_no_rocm) + return + + cholesky_ex_gpu = torch.linalg.cholesky_ex + @wraps(cholesky_ex_gpu) + def cholesky_ex(A: torch.Tensor, upper=False, check_errors=False, out=None) -> torch.Tensor: + if A.device.type != 'cpu': + return cholesky_ex_gpu(A, upper=upper, check_errors=check_errors, out=out) + + assert not upper + assert not check_errors + assert out is None + + n = A.size(0) + L = torch.zeros_like(A) + + for i in range(n): + for j in range(i + 1): + s = torch.dot(L[i, :j], L[j, :j].conj()) + if i == j: + val = A[i, i] - s + L[i, j] = torch.sqrt(val.real) + else: + L[i, j] = (A[i, j] - s) / L[j, j] + return L, torch.tensor(0, dtype=torch.int32, device=A.device) + torch.linalg.cholesky_ex = cholesky_ex + is_wsl: bool = False else: def get_agents() -> List[Agent]: @@ -251,7 +278,9 @@ else: agents = [x.strip().split(" ")[-1] for x in agents if x.startswith(' Name:') and "CPU" not in x] return [Agent(x) for x in agents] - def preload_hsa_runtime(): + def postinstall(): + if not is_wsl: + return try: if shutil.which("conda") is not None: # Preload stdc++ library. This will bypass Anaconda stdc++ library. diff --git a/modules/zluda_installer.py b/modules/zluda_installer.py index 2922f196f..c6fb55c23 100644 --- a/modules/zluda_installer.py +++ b/modules/zluda_installer.py @@ -163,7 +163,7 @@ def load(): ctypes.windll.LoadLibrary(os.path.join(rocm.environment.path, 'bin', 'MIOpen.dll')) ctypes.windll.LoadLibrary(os.path.join(path, 'cudnn64_9.dll')) - def conceal(): + def postinstall(): import torch torch.version.hip = rocm.version platform = sys.platform @@ -176,4 +176,4 @@ def load(): def _join_rocm_home(*paths) -> str: return os.path.join(cpp_extension.ROCM_HOME, *paths) cpp_extension._join_rocm_home = _join_rocm_home # pylint: disable=protected-access - rocm.conceal = conceal + rocm.postinstall = postinstall From 9c19e8a9b8c9196be5cf361f6841ecbcf6fbedf9 Mon Sep 17 00:00:00 2001 From: Seunghoon Lee Date: Sun, 28 Sep 2025 18:34:10 +0900 Subject: [PATCH 5/8] return correct item --- modules/rocm.py | 4 +++- 1 file changed, 3 insertions(+), 1 deletion(-) diff --git a/modules/rocm.py b/modules/rocm.py index c7aa6771b..9728313e0 100644 --- a/modules/rocm.py +++ b/modules/rocm.py @@ -261,10 +261,12 @@ if sys.platform == "win32": s = torch.dot(L[i, :j], L[j, :j].conj()) if i == j: val = A[i, i] - s + if val.real <= 0 or (A.dtype.is_complex and val.imag != 0): + return torch.return_types.linalg_cholesky_ex((L, torch.tensor(i + 1, dtype=torch.int32, device='cpu')), {}) L[i, j] = torch.sqrt(val.real) else: L[i, j] = (A[i, j] - s) / L[j, j] - return L, torch.tensor(0, dtype=torch.int32, device=A.device) + return torch.return_types.linalg_cholesky_ex((L, torch.tensor(0, dtype=torch.int32, device='cpu')), {}) torch.linalg.cholesky_ex = cholesky_ex is_wsl: bool = False From 06b14b070b136d8ae36335aedd57b42c8cfa426e Mon Sep 17 00:00:00 2001 From: Seunghoon Lee Date: Sun, 28 Sep 2025 18:43:37 +0900 Subject: [PATCH 6/8] hijack torch.linalg.cholesky() --- modules/rocm.py | 11 ++++++++++- 1 file changed, 10 insertions(+), 1 deletion(-) diff --git a/modules/rocm.py b/modules/rocm.py index 9728313e0..0d489b72a 100644 --- a/modules/rocm.py +++ b/modules/rocm.py @@ -245,7 +245,7 @@ if sys.platform == "win32": cholesky_ex_gpu = torch.linalg.cholesky_ex @wraps(cholesky_ex_gpu) - def cholesky_ex(A: torch.Tensor, upper=False, check_errors=False, out=None) -> torch.Tensor: + def cholesky_ex(A: torch.Tensor, upper=False, check_errors=False, out=None) -> torch.return_types.linalg_cholesky_ex: if A.device.type != 'cpu': return cholesky_ex_gpu(A, upper=upper, check_errors=check_errors, out=out) @@ -269,6 +269,15 @@ if sys.platform == "win32": return torch.return_types.linalg_cholesky_ex((L, torch.tensor(0, dtype=torch.int32, device='cpu')), {}) torch.linalg.cholesky_ex = cholesky_ex + cholesky_gpu = torch.linalg.cholesky + @wraps(cholesky_gpu) + def cholesky(A: torch.Tensor, upper=False, out=None) -> torch.Tensor: + if A.device.type != 'cpu': + return cholesky_gpu(A, upper=upper, out=out) + L, _ = torch.linalg.cholesky_ex(A, upper=upper, out=out) + return L + torch.linalg.cholesky = cholesky + is_wsl: bool = False else: def get_agents() -> List[Agent]: From a47959b114b68ea8eb1f22e49183af9f70478c53 Mon Sep 17 00:00:00 2001 From: Disty0 Date: Sun, 28 Sep 2025 13:33:15 +0300 Subject: [PATCH 7/8] move ROCm Windows hijacks outside of torch install --- modules/devices.py | 2 ++ modules/rocm.py | 83 +++++++++++++++++++++++----------------------- modules/shared.py | 5 +++ 3 files changed, 49 insertions(+), 41 deletions(-) diff --git a/modules/devices.py b/modules/devices.py index abe979c17..804fdaa9b 100644 --- a/modules/devices.py +++ b/modules/devices.py @@ -48,6 +48,8 @@ def has_mps() -> bool: def has_xpu() -> bool: return bool(hasattr(torch, 'xpu') and torch.xpu.is_available()) +def has_rocm() -> bool: + return bool(torch.version.hip is not None and torch.cuda.is_available()) def has_zluda() -> bool: if not cuda_ok: diff --git a/modules/rocm.py b/modules/rocm.py index 0d489b72a..643e02692 100644 --- a/modules/rocm.py +++ b/modules/rocm.py @@ -243,40 +243,39 @@ if sys.platform == "win32": os.environ["PATH"] = ";".join(paths_no_rocm) return - cholesky_ex_gpu = torch.linalg.cholesky_ex - @wraps(cholesky_ex_gpu) - def cholesky_ex(A: torch.Tensor, upper=False, check_errors=False, out=None) -> torch.return_types.linalg_cholesky_ex: - if A.device.type != 'cpu': - return cholesky_ex_gpu(A, upper=upper, check_errors=check_errors, out=out) + def rocm_init(): + try: + import torch + import numpy as np - assert not upper - assert not check_errors - assert out is None + cholesky_ex_gpu = torch.linalg.cholesky_ex + @wraps(cholesky_ex_gpu) + def cholesky_ex(A: torch.Tensor, upper=False, check_errors=False, out=None) -> torch.return_types.linalg_cholesky_ex: + if A.device.type != 'cpu': + return cholesky_ex_gpu(A, upper=upper, check_errors=check_errors, out=out) - n = A.size(0) - L = torch.zeros_like(A) + assert not check_errors + L = torch.from_numpy(np.linalg.cholesky(A.numpy(), upper=upper)) + info = torch.tensor(0, dtype=torch.int32, device='cpu') + if out is not None: + out[0].copy_(L) + out[1].copy_(info) + return torch.return_types.linalg_cholesky_ex((L, info), {}) + torch.linalg.cholesky_ex = cholesky_ex - for i in range(n): - for j in range(i + 1): - s = torch.dot(L[i, :j], L[j, :j].conj()) - if i == j: - val = A[i, i] - s - if val.real <= 0 or (A.dtype.is_complex and val.imag != 0): - return torch.return_types.linalg_cholesky_ex((L, torch.tensor(i + 1, dtype=torch.int32, device='cpu')), {}) - L[i, j] = torch.sqrt(val.real) - else: - L[i, j] = (A[i, j] - s) / L[j, j] - return torch.return_types.linalg_cholesky_ex((L, torch.tensor(0, dtype=torch.int32, device='cpu')), {}) - torch.linalg.cholesky_ex = cholesky_ex - - cholesky_gpu = torch.linalg.cholesky - @wraps(cholesky_gpu) - def cholesky(A: torch.Tensor, upper=False, out=None) -> torch.Tensor: - if A.device.type != 'cpu': - return cholesky_gpu(A, upper=upper, out=out) - L, _ = torch.linalg.cholesky_ex(A, upper=upper, out=out) - return L - torch.linalg.cholesky = cholesky + cholesky_gpu = torch.linalg.cholesky + @wraps(cholesky_gpu) + def cholesky(A: torch.Tensor, upper=False, out=None) -> torch.Tensor: + if A.device.type != 'cpu': + return cholesky_gpu(A, upper=upper, out=out) + L = torch.from_numpy(np.linalg.cholesky(A.numpy(), upper=upper)) + if out is not None: + out.copy_(L) + return L + torch.linalg.cholesky = cholesky + except Exception as e: + return False, e + return True, None is_wsl: bool = False else: @@ -290,16 +289,18 @@ else: return [Agent(x) for x in agents] def postinstall(): - if not is_wsl: - return - try: - if shutil.which("conda") is not None: - # Preload stdc++ library. This will bypass Anaconda stdc++ library. - 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") - except OSError: - pass + if is_wsl: + try: + if shutil.which("conda") is not None: + # Preload stdc++ library. This will bypass Anaconda stdc++ library. + 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") + except OSError: + pass + + def rocm_init(): + return True, None is_wsl: bool = os.environ.get('WSL_DISTRO_NAME', 'unknown' if spawn('wslpath -w /') else None) is not None environment = None diff --git a/modules/shared.py b/modules/shared.py index 78ff5ac0c..6f54bff2c 100644 --- a/modules/shared.py +++ b/modules/shared.py @@ -83,6 +83,11 @@ elif cmd_opts.use_directml: ok, e = directml_init() if not ok: log.error(f'DirectML initialization failed: {e}') +elif cmd_opts.use_rocm or devices.has_rocm(): + from modules.rocm import rocm_init + ok, e = rocm_init() + if not ok: + log.error(f'ROCm initialization failed: {e}') devices.backend = devices.get_backend(cmd_opts) devices.device = devices.get_optimal_device() mem_stat = memory_stats() From 4a70e82b0c8e021cf48d62445bc4b5e503fa2f7a Mon Sep 17 00:00:00 2001 From: Disty0 Date: Sun, 28 Sep 2025 14:05:20 +0300 Subject: [PATCH 8/8] ROCm always use numpy on cholesky --- modules/devices.py | 2 ++ modules/rocm.py | 13 +++++-------- 2 files changed, 7 insertions(+), 8 deletions(-) diff --git a/modules/devices.py b/modules/devices.py index 804fdaa9b..a98b53d02 100644 --- a/modules/devices.py +++ b/modules/devices.py @@ -48,9 +48,11 @@ def has_mps() -> bool: def has_xpu() -> bool: return bool(hasattr(torch, 'xpu') and torch.xpu.is_available()) + def has_rocm() -> bool: return bool(torch.version.hip is not None and torch.cuda.is_available()) + def has_zluda() -> bool: if not cuda_ok: return False diff --git a/modules/rocm.py b/modules/rocm.py index 643e02692..30e90f000 100644 --- a/modules/rocm.py +++ b/modules/rocm.py @@ -251,12 +251,10 @@ if sys.platform == "win32": cholesky_ex_gpu = torch.linalg.cholesky_ex @wraps(cholesky_ex_gpu) def cholesky_ex(A: torch.Tensor, upper=False, check_errors=False, out=None) -> torch.return_types.linalg_cholesky_ex: - if A.device.type != 'cpu': - return cholesky_ex_gpu(A, upper=upper, check_errors=check_errors, out=out) - assert not check_errors - L = torch.from_numpy(np.linalg.cholesky(A.numpy(), upper=upper)) - info = torch.tensor(0, dtype=torch.int32, device='cpu') + return_device = A.device + L = torch.from_numpy(np.linalg.cholesky(A.to("cpu").numpy(), upper=upper)).to(return_device) + info = torch.tensor(0, dtype=torch.int32, device=return_device) if out is not None: out[0].copy_(L) out[1].copy_(info) @@ -266,9 +264,8 @@ if sys.platform == "win32": cholesky_gpu = torch.linalg.cholesky @wraps(cholesky_gpu) def cholesky(A: torch.Tensor, upper=False, out=None) -> torch.Tensor: - if A.device.type != 'cpu': - return cholesky_gpu(A, upper=upper, out=out) - L = torch.from_numpy(np.linalg.cholesky(A.numpy(), upper=upper)) + return_device = A.device + L = torch.from_numpy(np.linalg.cholesky(A.to("cpu").numpy(), upper=upper)).to(return_device) if out is not None: out.copy_(L) return L