diff --git a/modules/devices.py b/modules/devices.py index 5f5e18768..e317d91f4 100644 --- a/modules/devices.py +++ b/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(): diff --git a/modules/dml/__init__.py b/modules/dml/__init__.py index 90beb96d7..b77db18d0 100644 --- a/modules/dml/__init__.py +++ b/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. diff --git a/modules/dml/optimizer/amd/__init__.py b/modules/dml/optimizer/amd/__init__.py index 37ea3ea91..23c6e57e3 100644 --- a/modules/dml/optimizer/amd/__init__.py +++ b/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)) diff --git a/modules/dml/optimizer/amd/driver/atiadlxx.py b/modules/dml/optimizer/amd/driver/atiadlxx.py index abb5eca96..81a578a5c 100644 --- a/modules/dml/optimizer/amd/driver/atiadlxx.py +++ b/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: diff --git a/modules/dml/optimizer/amd/driver/atiadlxx_apis.py b/modules/dml/optimizer/amd/driver/atiadlxx_apis.py index 23b7da3b1..fef70b543 100644 --- a/modules/dml/optimizer/amd/driver/atiadlxx_apis.py +++ b/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.") \ No newline at end of file +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)] diff --git a/modules/dml/optimizer/amd/memory.py b/modules/dml/optimizer/amd/memory.py deleted file mode 100644 index e69de29bb..000000000 diff --git a/modules/dml/optimizer/intel/__init__.py b/modules/dml/optimizer/intel/__init__.py index d17cd59d2..bffda69f1 100644 --- a/modules/dml/optimizer/intel/__init__.py +++ b/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 diff --git a/modules/dml/optimizer/nvidia/__init__.py b/modules/dml/optimizer/nvidia/__init__.py index d467e8ae5..e5fda97b8 100644 --- a/modules/dml/optimizer/nvidia/__init__.py +++ b/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 diff --git a/modules/dml/optimizer/optimizer.py b/modules/dml/optimizer/optimizer.py index 2b2d9ed64..db88c6579 100644 --- a/modules/dml/optimizer/optimizer.py +++ b/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 diff --git a/modules/dml/optimizer/unknown/__init__.py b/modules/dml/optimizer/unknown/__init__.py index f2fbc2ae2..40ff476cd 100644 --- a/modules/dml/optimizer/unknown/__init__.py +++ b/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)