handle different sizes of hipDeviceProp_t

modules/rocm.py

@@ -97,16 +97,24 @@ class Agent:
 
     @property
     def therock(self) -> Union[str, None]:
-        if (self.gfx_version & 0xFFF0) == 0x1100:
-            return "gfx110X-dgpu"
-        if self.gfx_version == 0x1151:
-            return "gfx1151"
         if (self.gfx_version & 0xFFF0) == 0x1200:
             return "gfx120X-all"
-        if (self.gfx_version & 0xFFF0) == 0x940:
-            return "gfx94X-dcgpu"
-        if self.gfx_version == 0x950:
-            return "gfx950-dcgpu"
+        if (self.gfx_version & 0xFFF0) == 0x1100:
+            return "gfx110X-all"
+        if self.gfx_version == 0x1150:
+            return "gfx1150"
+        if self.gfx_version == 0x1151:
+            return "gfx1151"
+        #if (self.gfx_version & 0xFFF0) == 0x1030:
+        #    return "gfx103X-dgpu"
+        #if (self.gfx_version & 0xFFF0) == 0x1010:
+        #    return "gfx101X-dgpu"
+        #if (self.gfx_version & 0xFFF0) == 0x900:
+        #    return "gfx90X-dcgpu"
+        #if (self.gfx_version & 0xFFF0) == 0x940:
+        #    return "gfx94X-dcgpu"
+        #if self.gfx_version == 0x950:
+        #    return "gfx950-dcgpu"
         return None
 
     def get_gfx_version(self) -> Union[str, None]:
@@ -239,15 +247,9 @@ if sys.platform == "win32":
     def driver_get_agents() -> List[Agent]:
         # unsafe and experimental feature
         from modules import windows_hip_ffi
-        hip = windows_hip_ffi.HIP()
-        count = hip.get_device_count()
-        _agents = [None] * count
-        for i in range(count):
-            prop = hip.get_device_properties(i)
-            name = prop.gcnArchName.decode('utf-8').strip('\x00')
-            _agents[i] = Agent(name)
-        del hip
-        return _agents
+        archs = windows_hip_ffi.get_archs()
+        # filter out None (is there any better way?)
+        return [Agent(x) for x in archs if x is not None]
 
     def postinstall():
         import torch

modules/windows_hip_ffi.py

@@ -7,9 +7,7 @@ if sys.platform == "win32":
 
     class hipDeviceProp(ctypes.Structure):
         _fields_ = [
-            ('__front__', ctypes.c_byte * 396),
-            ('gcnArchName', ctypes.c_char * 256),
-            ('__rear__', ctypes.c_byte * 820)
+            ('bytes', ctypes.c_byte * 1472) # 1472 in amdhip64_6.dll, shorter in amdhip64_7.dll?
         ]
 
     class HIP:
@@ -17,9 +15,11 @@ if sys.platform == "win32":
             ctypes.windll.kernel32.LoadLibraryA.restype = ctypes.wintypes.HMODULE
             ctypes.windll.kernel32.LoadLibraryA.argtypes = [ctypes.c_char_p]
             self.handle = None
-            path = os.environ.get("windir", "C:\\Windows") + "\\System32\\amdhip64_6.dll"
+            path = os.environ.get("windir", "C:\\Windows") + "\\System32\\amdhip64_7.dll"
             if not os.path.isfile(path):
-                path = os.environ.get("windir", "C:\\Windows") + "\\System32\\amdhip64_7.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.dll"
             assert os.path.isfile(path)
             self.handle = ctypes.windll.kernel32.LoadLibraryA(path.encode('utf-8'))
             ctypes.windll.kernel32.GetLastError.restype = ctypes.wintypes.DWORD
@@ -47,4 +47,41 @@ if sys.platform == "win32":
         def get_device_properties(self, device_id):
             prop = hipDeviceProp()
             assert self.hipGetDeviceProperties(ctypes.byref(prop), device_id) == 0
-            return prop
+            return prop.bytes
+
+    def get_archs():
+        hip = HIP()
+
+        count = hip.get_device_count()
+        archs = [None] * count
+        for i in range(count):
+            prop = hip.get_device_properties(i)[:]
+
+            name = ""
+            idx = 0
+            while idx < len(prop):
+                try:
+                    idx = prop.index(0x67, idx) + 1 # 'g'
+                except ValueError:
+                    break
+                if prop[idx] != 0x66: # 'f'
+                    continue
+                if prop[idx + 1] != 0x78: # 'x'
+                    continue
+
+                idx = idx + 2
+                while prop[idx] != 0x00:
+                    c = prop[idx]
+                    if (c < 0x30 or c > 0x39) and (c < 0x61 or c > 0x66): # hexadecimal
+                        name = ""
+                        continue
+                    name += chr(c)
+                    idx += 1
+                break
+
+            # if name == "", hipDeviceProp does not contain arch name
+            if name:
+                archs[i] = "gfx" + name
+
+        del hip
+        return archs