Implement memory estimation for AMDGPUs.

Stable.

modules/devices.py

@@ -31,7 +31,7 @@ def get_dml_device_string():
     from modules import shared
     if shared.cmd_opts.device_id is not None:
         return f"privateuseone:{shared.cmd_opts.device_id}"
-    return "privateuseone"
+    return "privateuseone:0"
 
 
 def get_optimal_device_name():

modules/dml/__init__.py

@@ -3,27 +3,27 @@ import torch_directml
 
 import modules.dml.hijack
 
-from modules.dml.optimizer.unknown import UnknownOptimizer
+from .optimizer.unknown import UnknownOptimizer
 
 class DirectML():
-    def get_optimizer(self, device: torch.device):
+    def get_optimizer(device: torch.device):
         assert(device.type == 'privateuseone')
         try:
             device_name = torch_directml.device_name(device.index)
             if 'NVIDIA' in device_name or 'GeForce' in device_name:
-                from modules.dml.optimizer.nvidia import nVidiaOptimizer as optimizer
+                from .optimizer.nvidia import nVidiaOptimizer as optimizer
             elif 'AMD' in device_name or 'Radeon' in device_name:
-                from modules.dml.optimizer.amd import AMDOptimizer as optimizer
+                from .optimizer.amd import AMDOptimizer as optimizer
             elif 'Intel' in device_name:
-                from modules.dml.optimizer.intel import IntelOptimizer as optimizer
+                from .optimizer.intel import IntelOptimizer as optimizer
             else:
                 return UnknownOptimizer
             return optimizer
         except:
             return UnknownOptimizer
 
-    def memory_stats(self, device: torch.device):
-        optimizer = self.get_optimizer(device)
+    def memory_stats(device: torch.device):
+        optimizer = DirectML.get_optimizer(device)
         return optimizer.memory_stats(device.index)
 
 # Alternative of torch.cuda for DirectML.

modules/dml/optimizer/amd/__init__.py

@@ -1,7 +1,7 @@
 from modules.dml.optimizer.optimizer import Optimizer
-from driver.atiadlxx import ATIADLxx
+from .driver.atiadlxx import ATIADLxx
 
 class AMDOptimizer(Optimizer):
     driver: ATIADLxx = ATIADLxx()
-    def memory_stats(self, index):
+    def memory_stats(index):
         return (AMDOptimizer.driver.iHyperMemorySize, AMDOptimizer.driver.get_dedicated_vram_usage(index))

modules/dml/optimizer/amd/driver/atiadlxx.py

@@ -14,8 +14,11 @@ class ATIADLxx(object):
         AdapterInfoArray = (AdapterInfo * num_adapters.value)()
         ADL2_Adapter_AdapterInfo_Get(self.context, C.cast(AdapterInfoArray, LPAdapterInfo), C.sizeof(AdapterInfoArray))
         self.devices = []
+        busNumbers = []
         for adapter in AdapterInfoArray:
-            self.devices.append(adapter)
+            if adapter.iBusNumber not in busNumbers: # filter duplicate device
+                self.devices.append(adapter)
+                busNumbers.append(adapter.iBusNumber)
         self.iHyperMemorySize = self.get_memory_info2(0).iHyperMemorySize
 
     def get_memory_info2(self, adapterIndex: int) -> ADLMemoryInfo2:

modules/dml/optimizer/amd/driver/atiadlxx_apis.py

@@ -1,50 +1,41 @@
 import ctypes as C
-import platform
 from .atiadlxx_structures import *
 
-_platform = platform.system()
+atiadlxx = C.WinDLL("atiadlxx.dll")
 
-try:
-    if _platform == "Windows":
-        atiadlxx = C.WinDLL("atiadlxx.dll")
+ADL_MAIN_MALLOC_CALLBACK = C.CFUNCTYPE(C.c_void_p, C.c_int)
+ADL_MAIN_FREE_CALLBACK = C.CFUNCTYPE(None, C.POINTER(C.c_void_p))
 
-        ADL_MAIN_MALLOC_CALLBACK = C.CFUNCTYPE(C.c_void_p, C.c_int)
-        ADL_MAIN_FREE_CALLBACK = C.CFUNCTYPE(None, C.POINTER(C.c_void_p))
+@ADL_MAIN_MALLOC_CALLBACK
+def ADL_Main_Memory_Alloc(iSize):
+    return C._malloc(iSize)
 
-        @ADL_MAIN_MALLOC_CALLBACK
-        def ADL_Main_Memory_Alloc(iSize):
-            return C._malloc(iSize)
+@ADL_MAIN_FREE_CALLBACK
+def ADL_Main_Memory_Free(lpBuffer):
+    if lpBuffer[0] is not None:
+        C._free(lpBuffer[0])
+        lpBuffer[0] = None
 
-        @ADL_MAIN_FREE_CALLBACK
-        def ADL_Main_Memory_Free(lpBuffer):
-            if lpBuffer[0] is not None:
-                C._free(lpBuffer[0])
-                lpBuffer[0] = None
+ADL2_Main_Control_Create = atiadlxx.ADL2_Main_Control_Create
+ADL2_Main_Control_Create.restype = C.c_int
+ADL2_Main_Control_Create.argtypes = [ADL_MAIN_MALLOC_CALLBACK, C.c_int, ADL_CONTEXT_HANDLE]
 
-        ADL2_Main_Control_Create = atiadlxx.ADL2_Main_Control_Create
-        ADL2_Main_Control_Create.restype = C.c_int
-        ADL2_Main_Control_Create.argtypes = [ADL_MAIN_MALLOC_CALLBACK, C.c_int, ADL_CONTEXT_HANDLE]
+ADL2_Adapter_NumberOfAdapters_Get = atiadlxx.ADL2_Adapter_NumberOfAdapters_Get
+ADL2_Adapter_NumberOfAdapters_Get.restype = C.c_int
+ADL2_Adapter_NumberOfAdapters_Get.argtypes = [ADL_CONTEXT_HANDLE, C.POINTER(C.c_int)]
 
-        ADL2_Adapter_NumberOfAdapters_Get = atiadlxx.ADL2_Adapter_NumberOfAdapters_Get
-        ADL2_Adapter_NumberOfAdapters_Get.restype = C.c_int
-        ADL2_Adapter_NumberOfAdapters_Get.argtypes = [ADL_CONTEXT_HANDLE, C.POINTER(C.c_int)]
+ADL2_Adapter_AdapterInfo_Get = atiadlxx.ADL2_Adapter_AdapterInfo_Get
+ADL2_Adapter_AdapterInfo_Get.restype = C.c_int
+ADL2_Adapter_AdapterInfo_Get.argtypes = [ADL_CONTEXT_HANDLE, LPAdapterInfo, C.c_int]
 
-        ADL2_Adapter_AdapterInfo_Get = atiadlxx.ADL2_Adapter_AdapterInfo_Get
-        ADL2_Adapter_AdapterInfo_Get.restype = C.c_int
-        ADL2_Adapter_AdapterInfo_Get.argtypes = [ADL_CONTEXT_HANDLE, LPAdapterInfo, C.c_int]
+ADL2_Adapter_MemoryInfo2_Get = atiadlxx.ADL2_Adapter_MemoryInfo2_Get
+ADL2_Adapter_MemoryInfo2_Get.restype = C.c_int
+ADL2_Adapter_MemoryInfo2_Get.argtypes = [ADL_CONTEXT_HANDLE, C.c_int, C.POINTER(ADLMemoryInfo2)]
 
-        ADL2_Adapter_MemoryInfo2_Get = atiadlxx.ADL2_Adapter_MemoryInfo2_Get
-        ADL2_Adapter_MemoryInfo2_Get.restype = C.c_int
-        ADL2_Adapter_MemoryInfo2_Get.argtypes = [ADL_CONTEXT_HANDLE, C.c_int, C.POINTER(ADLMemoryInfo2)]
+ADL2_Adapter_DedicatedVRAMUsage_Get = atiadlxx.ADL2_Adapter_DedicatedVRAMUsage_Get
+ADL2_Adapter_DedicatedVRAMUsage_Get.restype = C.c_int
+ADL2_Adapter_DedicatedVRAMUsage_Get.argtypes = [ADL_CONTEXT_HANDLE, C.c_int, C.POINTER(C.c_int)]
 
-        ADL2_Adapter_DedicatedVRAMUsage_Get = atiadlxx.ADL2_Adapter_DedicatedVRAMUsage_Get
-        ADL2_Adapter_DedicatedVRAMUsage_Get.restype = C.c_int
-        ADL2_Adapter_DedicatedVRAMUsage_Get.argtypes = [ADL_CONTEXT_HANDLE, C.c_int, C.POINTER(C.c_int)]
-
-        ADL2_Adapter_VRAMUsage_Get = atiadlxx.ADL2_Adapter_VRAMUsage_Get
-        ADL2_Adapter_VRAMUsage_Get.restype = C.c_int
-        ADL2_Adapter_VRAMUsage_Get.argtypes = [ADL_CONTEXT_HANDLE, C.c_int, C.POINTER(C.c_int)]
-    else:
-        print("Warning: experimental graphic memory optimization for AMDGPU is disabled. Because this is not Windows platform.")
-except FileNotFoundError:
-    print("Warning: memory optimization for AMDGPU is disabled. Because couldn't find 'atiadlxx.dll'. Please install GPU driver downloaded from AMD.com.")
+ADL2_Adapter_VRAMUsage_Get = atiadlxx.ADL2_Adapter_VRAMUsage_Get
+ADL2_Adapter_VRAMUsage_Get.restype = C.c_int
+ADL2_Adapter_VRAMUsage_Get.argtypes = [ADL_CONTEXT_HANDLE, C.c_int, C.POINTER(C.c_int)]

modules/dml/optimizer/intel/__init__.py

@@ -1,7 +1,7 @@
 from modules.dml.optimizer.optimizer import Optimizer
 
 class IntelOptimizer(Optimizer):
-    def memory_stats():
+    def memory_stats(index):
         raise NotImplementedError()
         # DML TODO: Implement
         return

modules/dml/optimizer/nvidia/__init__.py

@@ -1,7 +1,7 @@
 from modules.dml.optimizer.optimizer import Optimizer
 
 class nVidiaOptimizer(Optimizer):
-    def memory_stats():
+    def memory_stats(index):
         raise NotImplementedError()
         # DML TODO: Implement
         return

modules/dml/optimizer/optimizer.py

@@ -4,5 +4,5 @@ from typing import *
 class Optimizer(metaclass=ABCMeta):
     driver: Any = None
     @abstractmethod
-    def memory_stats(self, index: int) -> Tuple[int, int]:
+    def memory_stats(index: int) -> Tuple[int, int]:
         pass

modules/dml/optimizer/unknown/__init__.py

@@ -1,6 +1,6 @@
 from modules.dml.optimizer.optimizer import Optimizer
 
 class UnknownOptimizer(Optimizer):
-    def memory_stats():
+    def memory_stats(index):
         # DML TODO: Implement
         return (1073741824, 0)