From 14569aca248f6b0b2bb3c5ff63c846c5101691d5 Mon Sep 17 00:00:00 2001
From: Seunghoon Lee <lshqqytiger@naver.com>
Date: Thu, 18 Jul 2024 00:59:12 +0900
Subject: [PATCH] prevent segfault when no hip device found

---
 installer.py               |  9 +++++----
 modules/zluda_installer.py | 29 ++++++++++++++---------------
 2 files changed, 19 insertions(+), 19 deletions(-)

diff --git a/installer.py b/installer.py
index e01609ebe..80acacab9 100644
--- a/installer.py
+++ b/installer.py
@@ -484,7 +484,8 @@ 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
-    if len(hip_visible_devices) > 0:
+    hip_found_device = len(hip_visible_devices) > 0
+    if hip_found_device:
         idx, gpu, arch = hip_visible_devices[0]
         log.debug(f'ROCm agent used by default: idx={idx} gpu={gpu} arch={arch}')
         os.environ.setdefault('HIP_VISIBLE_DEVICES', str(idx))
@@ -557,17 +558,17 @@ def install_rocm_zluda(torch_command):
 
         if bool(int(os.environ.get("TORCH_BLAS_PREFER_HIPBLASLT", "1"))):
             supported_archs = []
-            hipblaslt_available = True
+            hipblaslt_available = hip_found_device
             libpath = os.environ.get("HIPBLASLT_TENSILE_LIBPATH", "/opt/rocm/lib/hipblaslt/library")
             for file in os.listdir(libpath):
                 if not file.startswith('extop_'):
                     continue
                 supported_archs.append(file[6:-3])
-            for gpu in amd_gpus:
+            for gpu in hip_visible_devices:
                 if gpu not in supported_archs:
                     hipblaslt_available = False
                     break
-            log.info(f'hipBLASLt supported_archs={supported_archs}, available={hipblaslt_available}')
+            log.debug(f'hipBLASLt supported_archs={supported_archs}, available={hipblaslt_available}')
             if hipblaslt_available:
                 import ctypes
                 # Preload hipBLASLt.
diff --git a/modules/zluda_installer.py b/modules/zluda_installer.py
index 82eada27e..e9db71c37 100644
--- a/modules/zluda_installer.py
+++ b/modules/zluda_installer.py
@@ -7,20 +7,6 @@ import urllib.request
 from typing import Tuple
 
 
-class HIPSDKVersion:
-    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}"
-
-
 class HIPSDK:
     is_installed = False
 
@@ -36,11 +22,24 @@ class HIPSDK:
         rocm_path = rf'{program_files}\AMD\ROCm'
         default_version = None
         if os.path.exists(rocm_path):
+            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}"
+
             versions = os.listdir(rocm_path)
             for s in versions:
                 version = None
                 try:
-                    version = HIPSDKVersion(s)
+                    version = Version(s)
                 except Exception:
                     continue
                 if default_version is None: