refactor rocm & zluda

2026-09-18 16:54:33 +02:00 · 2024-07-22 15:29:49 +09:00
parent d3a7095683
commit 0284c77ae5
3 changed files with 161 additions and 125 deletions
@@ -452,28 +452,23 @@ def check_onnx():


 def install_rocm_zluda(torch_command):
+    from modules import rocm
+
+    if not rocm.is_installed:
+        log.warning('Could not find ROCm toolkit installed.')
+        log.info('Using CPU-only torch')
+        return os.environ.get('TORCH_COMMAND', 'torch torchvision')
+
    check_python(supported_minors=[10, 11], reason='ROCm or ZLUDA backends require Python 3.10 or 3.11')
-    is_windows = platform.system() == 'Windows'
    log.info('AMD ROCm toolkit detected')
    os.environ.setdefault('PYTORCH_HIP_ALLOC_CONF', 'garbage_collection_threshold:0.8,max_split_size_mb:512')
    # if not is_windows:
    #    os.environ.setdefault('TENSORFLOW_PACKAGE', 'tensorflow-rocm')
    try:
-        if is_windows:
-            command = subprocess.run('hipinfo', shell=True, check=False, stdout=subprocess.PIPE, stderr=subprocess.PIPE)
-            amd_gpus = command.stdout.decode(encoding="utf8", errors="ignore").split('\n')
-            amd_gpus = [x.split(' ')[-1].strip() for x in amd_gpus if x.startswith('gcnArchName:')]
-        elif os.environ.get('WSL_DISTRO_NAME', None) is not None: # WSL does not have 'rocm_agent_enumerator'
-            command = subprocess.run('rocminfo', shell=True, check=False, stdout=subprocess.PIPE, stderr=subprocess.PIPE)
-            amd_gpus = command.stdout.decode(encoding="utf8", errors="ignore").split('\n')
-            amd_gpus = [x.strip().split(" ")[-1] for x in amd_gpus if x.startswith('  Name:') and "CPU" not in x]
-        else:
-            command = subprocess.run('rocm_agent_enumerator', shell=True, check=False, stdout=subprocess.PIPE, stderr=subprocess.PIPE)
-            amd_gpus = command.stdout.decode(encoding="utf8", errors="ignore").split('\n')
-            amd_gpus = [x for x in amd_gpus if x and x != 'gfx000']
-        log.debug(f'ROCm agents detected: {amd_gpus}')
+        amd_gpus = rocm.get_agents()
+        log.info(f'ROCm agents detected: {amd_gpus}')
    except Exception as e:
-        log.debug(f'ROCm agent enumerator failed: {e}')
+        log.warning(f'ROCm agent enumerator failed: {e}')
        amd_gpus = []

    hip_visible_devices = [] # use the first available amd gpu by default
@@ -484,6 +479,7 @@ def install_rocm_zluda(torch_command):
        if gpu in ['gfx1030', 'gfx1031', 'gfx1032', 'gfx1034']: # experimental navi 2x support
            hip_visible_devices.append((idx, gpu, 'navi2x'))
            break
+
    hip_found_device = len(hip_visible_devices) > 0
    if hip_found_device:
        idx, gpu, arch = hip_visible_devices[0]
@@ -498,65 +494,59 @@ def install_rocm_zluda(torch_command):
        else:
            log.debug(f'HSA_OVERRIDE_GFX_VERSION auto config is skipped for {gpu}')

-    try:
-        command = subprocess.run('hipconfig --version', shell=True, check=False, stdout=subprocess.PIPE, stderr=subprocess.PIPE)
-        arr = command.stdout.decode(encoding="utf8", errors="ignore").split('.')
-        rocm_ver = f'{arr[0]}.{arr[1]}' if len(arr) >= 2 else None
-        log.debug(f'ROCm version detected: {rocm_ver}')
-    except Exception as e:
-        log.debug(f'ROCm hipconfig failed: {e}')
-        rocm_ver = None
+    log.info(f'ROCm version detected: {rocm.version}')

-    if args.use_zluda:
-        log.warning("ZLUDA support: experimental")
-        error = None
-        from modules import zluda_installer
-        try:
-            if args.reinstall_zluda:
-                zluda_installer.uninstall()
-            zluda_path = zluda_installer.get_path()
-            zluda_installer.install(zluda_path)
-            zluda_installer.make_copy(zluda_path)
-        except Exception as e:
-            error = e
-            log.warning(f'Failed to install ZLUDA: {e}')
-        if error is None:
+    if sys.platform == "win32":
+        if args.use_zluda:
+            log.warning("ZLUDA support: experimental")
+            error = None
+            from modules import zluda_installer
            try:
-                zluda_installer.load(zluda_path)
-                torch_command = os.environ.get('TORCH_COMMAND', 'torch==2.3.0 torchvision --index-url https://download.pytorch.org/whl/cu118')
-                log.info(f'Using ZLUDA in {zluda_path}')
+                if args.reinstall_zluda:
+                    zluda_installer.uninstall()
+                zluda_path = zluda_installer.get_path()
+                zluda_installer.install(zluda_path)
+                zluda_installer.make_copy(zluda_path)
            except Exception as e:
                error = e
-                log.warning(f'Failed to load ZLUDA: {e}')
-        if error is not None:
+                log.warning(f'Failed to install ZLUDA: {e}')
+            if error is None:
+                try:
+                    zluda_installer.load(zluda_path)
+                    torch_command = os.environ.get('TORCH_COMMAND', 'torch==2.3.0 torchvision --index-url https://download.pytorch.org/whl/cu118')
+                    log.info(f'Using ZLUDA in {zluda_path}')
+                except Exception as e:
+                    error = e
+                    log.warning(f'Failed to load ZLUDA: {e}')
+            if error is not None:
+                log.info('Using CPU-only torch')
+                torch_command = os.environ.get('TORCH_COMMAND', 'torch torchvision')
+        else: # TODO TBD after ROCm for Windows is released
+            log.warning("HIP SDK is detected, but no Torch release for Windows available")
+            log.info("For ZLUDA support specify '--use-zluda'")
            log.info('Using CPU-only torch')
            torch_command = os.environ.get('TORCH_COMMAND', 'torch torchvision')
-    elif is_windows: # TODO TBD after ROCm for Windows is released
-        log.warning("HIP SDK is detected, but no Torch release for Windows available")
-        log.info("For ZLUDA support specify '--use-zluda'")
-        log.info('Using CPU-only torch')
-        torch_command = os.environ.get('TORCH_COMMAND', 'torch torchvision')

-        # conceal ROCm installed
-        conceal_rocm()
+            # conceal ROCm installed
+            conceal_rocm()
    else:
-        if rocm_ver is None: # assume the latest if version check fails
+        if rocm.version is None: # assume the latest if version check fails
            torch_command = os.environ.get('TORCH_COMMAND', 'torch torchvision --index-url https://download.pytorch.org/whl/rocm6.0')
-        elif rocm_ver == "6.1": # need nightlies
+        elif rocm.version == "6.1": # need nightlies
            if args.experimental:
                torch_command = os.environ.get('TORCH_COMMAND', 'torch torchvision --pre --index-url https://download.pytorch.org/whl/nightly/rocm6.1')
            else:
                torch_command = os.environ.get('TORCH_COMMAND', 'torch torchvision --index-url https://download.pytorch.org/whl/rocm6.0')
-        elif float(rocm_ver) < 5.5: # oldest supported version is 5.5
-            log.warning(f"Unsupported ROCm version detected: {rocm_ver}")
+        elif float(rocm.version) < 5.5: # oldest supported version is 5.5
+            log.warning(f"Unsupported ROCm version detected: {rocm.version}")
            log.warning("Minimum supported ROCm version is 5.5")
            torch_command = os.environ.get('TORCH_COMMAND', 'torch torchvision --index-url https://download.pytorch.org/whl/rocm5.5')
        else:
-            torch_command = os.environ.get('TORCH_COMMAND', f'torch torchvision --index-url https://download.pytorch.org/whl/rocm{rocm_ver}')
-        if rocm_ver is not None:
-            ort_version = os.environ.get('ONNXRUNTIME_VERSION', None)
-            ort_package = os.environ.get('ONNXRUNTIME_PACKAGE', f"--pre onnxruntime-training{'' if ort_version is None else ('==' + ort_version)} --index-url https://pypi.lsh.sh/{rocm_ver[0]}{rocm_ver[2]} --extra-index-url https://pypi.org/simple")
-            install(ort_package, 'onnxruntime-training')
+            torch_command = os.environ.get('TORCH_COMMAND', f'torch torchvision --index-url https://download.pytorch.org/whl/rocm{rocm.version}')
+
+        ort_version = os.environ.get('ONNXRUNTIME_VERSION', None)
+        ort_package = os.environ.get('ONNXRUNTIME_PACKAGE', f"--pre onnxruntime-training{'' if ort_version is None else ('==' + ort_version)} --index-url https://pypi.lsh.sh/{rocm.version[0]}{rocm.version[2]} --extra-index-url https://pypi.org/simple")
+        install(ort_package, 'onnxruntime-training')

        if bool(int(os.environ.get("TORCH_BLAS_PREFER_HIPBLASLT", "1"))):
            supported_archs = []
@@ -655,8 +645,8 @@ def is_rocm_available(allow_rocm):
        log.debug('DirectML installation is detected. Skipping HIP SDK check.')
        return False
    if platform.system() == 'Windows':
-        from modules.zluda_installer import HIPSDK
-        return HIPSDK.is_installed
+        from modules.rocm import is_installed
+        return is_installed
    else:
        return shutil.which('rocminfo') is not None or os.path.exists('/opt/rocm/bin/rocminfo') or os.path.exists('/dev/kfd')

@@ -0,0 +1,103 @@
+import os
+import sys
+import shutil
+import subprocess
+from typing import Union, List
+
+
+def resolve_link(path_: str) -> str:
+    if not os.path.islink(path_):
+        return path_
+    return resolve_link(os.readlink(path_))
+
+
+def dirname(path_: str, r: int = 1) -> str:
+    for _ in range(0, r):
+        path_ = os.path.dirname(path_)
+    return path_
+
+
+def spawn(command: str) -> str:
+    process = subprocess.run(command, shell=True, check=False, stdout=subprocess.PIPE, stderr=subprocess.PIPE)
+    return process.stdout.decode(encoding="utf8", errors="ignore")
+
+
+if sys.platform == "win32":
+    def find() -> Union[str, None]:
+        hip_path = shutil.which("hipconfig")
+        if hip_path is not None:
+            return dirname(resolve_link(hip_path), 2)
+
+        hip_path = os.environ.get("HIP_PATH", None)
+        if hip_path is not None:
+            return hip_path
+
+        program_files = os.environ.get('ProgramFiles', r'C:\Program Files')
+        hip_path = rf'{program_files}\AMD\ROCm'
+        if not os.path.exists(hip_path):
+            return None
+
+        class Version:
+            major: int
+            minor: int
+
+            def __init__(self, string: str):
+                self.major, self.minor = [int(v) for v in string.strip().split(".")]
+
+            def __gt__(self, other):
+                return self.major * 10 + other.minor > other.major * 10 + other.minor
+
+            def __str__(self):
+                return f"{self.major}.{self.minor}"
+
+        latest = None
+        versions = os.listdir(hip_path)
+        for s in versions:
+            item = None
+            try:
+                item = Version(s)
+            except Exception:
+                continue
+            if latest is None:
+                latest = item
+                continue
+            if item > latest:
+                latest = item
+
+        if latest is None:
+            return None
+
+        return os.path.join(hip_path, str(latest))
+
+    def get_version() -> str: # cannot just run hipconfig as it requires Perl installed on Windows.
+        return os.path.basename(path)
+
+    def get_agents() -> List[str]:
+        return [x.split(' ')[-1].strip() for x in spawn("hipinfo").split("\n") if x.startswith('gcnArchName:')]
+
+    is_wsl: bool = False
+else:
+    def find() -> Union[str, None]:
+        rocm_path = shutil.which("hipconfig")
+        if rocm_path is not None:
+            return dirname(resolve_link(rocm_path), 2)
+        if not os.path.exists("/opt/rocm"):
+            return None
+        return resolve_link("/opt/rocm")
+
+    def get_version() -> str:
+        arr = spawn(f"{os.path.join(path, 'hipconfig')} --version").split(".")
+        return f'{arr[0]}.{arr[1]}' if len(arr) >= 2 else None
+
+    def get_agents() -> List[str]:
+        if is_wsl: # WSL does not have 'rocm_agent_enumerator'
+            agents = spawn("rocminfo").split("\n")
+            return [x.strip().split(" ")[-1] for x in agents if x.startswith('  Name:') and "CPU" not in x]
+        else:
+            agents = spawn("rocm_agent_enumerator").split("\n")
+            return [x for x in agents if x and x != 'gfx000']
+
+    is_wsl: bool = os.environ.get('WSL_DISTRO_NAME', None) is not None
+path = find()
+is_installed = path is not None
+version = get_version()
@@ -2,68 +2,8 @@ import os
 import ctypes
 import shutil
 import zipfile
-import platform
 import urllib.request
-from typing import Tuple
-
-
-class HIPSDK:
-    is_installed = False
-
-    version: str
-    path: str
-    targets: Tuple[str]
-
-    class Version:
-        major: int
-        minor: int
-
-        def __init__(self, version: str):
-            self.major, self.minor = [int(v) for v in version.strip().split(".")]
-
-        def __gt__(self, other):
-            return self.major * 10 + other.minor > other.major * 10 + other.minor
-
-        def __str__(self):
-            return f"{self.major}.{self.minor}"
-
-    def __init__(self):
-        if platform.system() != 'Windows':
-            raise RuntimeError('ZLUDA cannot be automatically installed on Linux. Please select --use-cuda for ZLUDA or --use-rocm for ROCm.')
-
-        program_files = os.environ.get('ProgramFiles', r'C:\Program Files')
-        rocm_path = rf'{program_files}\AMD\ROCm'
-        default_version = None
-        if os.path.exists(rocm_path):
-            versions = os.listdir(rocm_path)
-            for s in versions:
-                version = None
-                try:
-                    version = HIPSDK.Version(s)
-                except Exception:
-                    continue
-                if default_version is None:
-                    default_version = version
-                    continue
-                if version > default_version:
-                    default_version = version
-
-        self.path = os.environ.get('HIP_PATH', None)
-        if self.path is None:
-            if os.environ.get("HIP_PATH_61", None) is not None:
-                self.version = "6.1"
-            elif os.environ.get("HIP_PATH_57", None) is not None:
-                self.version = "5.7"
-            elif default_version is None:
-                raise RuntimeError('Could not find AMD HIP SDK, please install it from https://www.amd.com/en/developer/resources/rocm-hub/hip-sdk.html')
-            else:
-                self.version = str(default_version)
-        else:
-            self.version = os.path.basename(self.path) or os.path.basename(os.path.dirname(self.path))
-
-        self.targets = ['rocblas.dll', 'rocsolver.dll', f'hiprtc{"".join([v.zfill(2) for v in self.version.split(".")])}.dll']
-HIPSDK = HIPSDK()
-HIPSDK.is_installed = True
+from modules import rocm


 DLL_MAPPING = {
@@ -71,6 +11,7 @@ DLL_MAPPING = {
    'cusparse.dll': 'cusparse64_11.dll',
    'nvrtc.dll': 'nvrtc64_112_0.dll',
 }
+HIPSDK_TARGETS = ['rocblas.dll', 'rocsolver.dll', f'hiprtc{"".join([v.zfill(2) for v in rocm.version.split(".")])}.dll']
 ZLUDA_TARGETS = ('nvcuda.dll', 'nvml.dll',)


@@ -83,10 +24,12 @@ def install(zluda_path: os.PathLike) -> None:
        return

    default_hash = None
-    if HIPSDK.version == "6.1":
+    if rocm.version == "6.1":
        default_hash = 'd7714d84c0c13bbf816eaaac32693e4e75e58a87'
-    elif HIPSDK.version == "5.7":
+    elif rocm.version == "5.7":
        default_hash = '11cc5844514f93161e0e74387f04e2c537705a82'
+    else:
+        raise RuntimeError(f'Unsupported HIP SDK version: {rocm.version}')
    urllib.request.urlretrieve(f'https://github.com/lshqqytiger/ZLUDA/releases/download/rel.{os.environ.get("ZLUDA_HASH", default_hash)}/ZLUDA-windows-amd64.zip', '_zluda')
    with zipfile.ZipFile('_zluda', 'r') as archive:
        infos = archive.infolist()
@@ -112,8 +55,8 @@ def make_copy(zluda_path: os.PathLike) -> None:


 def load(zluda_path: os.PathLike) -> None:
-    for v in HIPSDK.targets:
-        ctypes.windll.LoadLibrary(os.path.join(HIPSDK.path, 'bin', v))
+    for v in HIPSDK_TARGETS:
+        ctypes.windll.LoadLibrary(os.path.join(rocm.path, 'bin', v))
    for v in ZLUDA_TARGETS:
        ctypes.windll.LoadLibrary(os.path.join(zluda_path, v))
    for v in DLL_MAPPING.values():