diff --git a/installer.py b/installer.py index ea7f5b496..1932a20af 100644 --- a/installer.py +++ b/installer.py @@ -977,10 +977,7 @@ def check_torch(): if not args.ignore: sys.exit(1) if rocm.is_installed: - if sys.platform == "win32": # 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/devices.py b/modules/devices.py index abe979c17..a98b53d02 100644 --- a/modules/devices.py +++ b/modules/devices.py @@ -49,6 +49,10 @@ 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 8dfcd9ec8..30e90f000 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 @@ -130,17 +120,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) @@ -241,6 +230,50 @@ 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 + + def rocm_init(): + try: + import torch + import numpy as np + + 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: + assert not check_errors + 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) + return torch.return_types.linalg_cholesky_ex((L, info), {}) + 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: + 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 + torch.linalg.cholesky = cholesky + except Exception as e: + return False, e + return True, None + is_wsl: bool = False else: def get_agents() -> List[Agent]: @@ -252,15 +285,19 @@ 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(): - 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 postinstall(): + 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() 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() 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