Merge branch 'dev' into dev

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:

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

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

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()

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()

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