diff --git a/installer.py b/installer.py index 98ffef561..1bbb03719 100644 --- a/installer.py +++ b/installer.py @@ -1015,11 +1015,8 @@ def check_torch(): if rocm.is_installed: if sys.platform == "win32": # CPU, DirectML, ZLUDA rocm.conceal() - elif rocm.is_wsl: # WSL ROCm - try: - rocm.load_hsa_runtime() - except OSError: - log.error("ROCm: failed to preload HSA runtime") + else: # ROCm + rocm.load_libraries() if args.version: return if not args.skip_all: diff --git a/modules/rocm.py b/modules/rocm.py index 0222218c8..365502413 100644 --- a/modules/rocm.py +++ b/modules/rocm.py @@ -73,7 +73,11 @@ class ROCmEnvironment(Environment): class PythonPackageEnvironment(Environment): def __init__(self): import _rocm_sdk_core - super().__init__(os.path.join(_rocm_sdk_core.__path__[0], "bin", "amdhip64_7.dll")) + if sys.platform == "win32": + path = os.path.join(_rocm_sdk_core.__path__[0], "bin", "amdhip64_7.dll") + else: + raise NotImplementedError + super().__init__(path) class MicroArchitecture(Enum): @@ -216,13 +220,13 @@ if sys.platform == "win32": return [Agent(x.split(' ')[-1].strip()) for x in out.split("\n")] def get_distribution(agent: Agent) -> str: - if agent.gfx_version >= 0x1100 and agent.gfx_version < 0x1110: + if (agent.gfx_version & 0xFFF0) == 0x1100: return "gfx110X-dgpu" if agent.gfx_version == 0x1151: return "gfx1151" - if agent.gfx_version >= 0x1200 and agent.gfx_version < 0x1210: + if (agent.gfx_version & 0xFFF0) == 0x1200: return "gfx120X-all" - if agent.gfx_version >= 0x940 and agent.gfx_version < 0x950: + if (agent.gfx_version & 0xFFF0) == 0x940: return "gfx94X-dcgpu" if agent.gfx_version == 0x950: return "gfx950-dcgpu" @@ -240,16 +244,17 @@ 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 load_hsa_runtime() -> None: + def load_libraries() -> None: try: - # Preload stdc++ library. This will ignore Anaconda stdc++ library. + # Preload stdc++ library. This will bypass Anaconda stdc++ library. load_library_global("/lib/x86_64-linux-gnu/libstdc++.so.6") # Use tcmalloc if possible. load_library_global("/usr/lib/x86_64-linux-gnu/libtcmalloc_minimal.so.4") + if is_wsl: + # Preload rocr4wsl. + load_library_global("/opt/rocm/lib/libhsa-runtime64.so") except OSError: pass - # Preload HSA Runtime library. - load_library_global("/opt/rocm/lib/libhsa-runtime64.so") def set_blaslt_enabled(enabled: bool) -> None: if enabled: