From 578f206bf7805fe75d310aaef36b87ad8d4a9033 Mon Sep 17 00:00:00 2001 From: Seunghoon Lee Date: Sat, 15 Jul 2023 16:19:38 +0900 Subject: [PATCH] DirectML rework. --- modules/dml/__init__.py | 42 ++------- modules/dml/optimizer/amd/__init__.py | 7 -- modules/dml/optimizer/amd/driver/atiadlxx.py | 46 ---------- .../dml/optimizer/amd/driver/atiadlxx_apis.py | 45 ---------- .../optimizer/amd/driver/atiadlxx_defines.py | 1 - .../amd/driver/atiadlxx_structures.py | 87 ------------------- modules/dml/optimizer/intel/__init__.py | 6 -- modules/dml/optimizer/nvidia/__init__.py | 6 -- modules/dml/optimizer/optimizer.py | 8 -- modules/dml/optimizer/unknown/__init__.py | 5 -- modules/shared.py | 5 +- 11 files changed, 9 insertions(+), 249 deletions(-) delete mode 100644 modules/dml/optimizer/amd/__init__.py delete mode 100644 modules/dml/optimizer/amd/driver/atiadlxx.py delete mode 100644 modules/dml/optimizer/amd/driver/atiadlxx_apis.py delete mode 100644 modules/dml/optimizer/amd/driver/atiadlxx_defines.py delete mode 100644 modules/dml/optimizer/amd/driver/atiadlxx_structures.py delete mode 100644 modules/dml/optimizer/intel/__init__.py delete mode 100644 modules/dml/optimizer/nvidia/__init__.py delete mode 100644 modules/dml/optimizer/optimizer.py delete mode 100644 modules/dml/optimizer/unknown/__init__.py diff --git a/modules/dml/__init__.py b/modules/dml/__init__.py index 26639c428..a747392ce 100644 --- a/modules/dml/__init__.py +++ b/modules/dml/__init__.py @@ -1,39 +1,9 @@ -# pylint: disable=no-member,no-self-argument,no-method-argument import torch -import torch_directml # pylint: disable=import-error -import modules.dml.hijack -import modules.dml.amp as amp -from modules.dml.opts import override_opts -from .optimizer.unknown import UnknownOptimizer +def directml_init(): + from modules.dml.backend import DirectML # pylint: disable=ungrouped-imports + from modules.dml.opts import override_opts # pylint: disable=ungrouped-imports + # Alternative of torch.cuda for DirectML. + torch.dml = DirectML -class DirectML(): - is_autocast_enabled = False - autocast_gpu_dtype = torch.float16 - - 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 .optimizer.nvidia import nVidiaOptimizer as optimizer - elif 'AMD' in device_name or 'Radeon' in device_name: - from .optimizer.amd import AMDOptimizer as optimizer - elif 'Intel' in device_name: - from .optimizer.intel import IntelOptimizer as optimizer - else: - return UnknownOptimizer - return optimizer - except Exception: - return UnknownOptimizer - - def memory_stats(device: torch.device): - optimizer = DirectML.get_optimizer(device) - return optimizer.memory_stats(device.index) - - -DirectML.amp = amp -# Alternative of torch.cuda for DirectML. -torch.dml = DirectML - -override_opts() + override_opts() diff --git a/modules/dml/optimizer/amd/__init__.py b/modules/dml/optimizer/amd/__init__.py deleted file mode 100644 index 23c6e57e3..000000000 --- a/modules/dml/optimizer/amd/__init__.py +++ /dev/null @@ -1,7 +0,0 @@ -from modules.dml.optimizer.optimizer import Optimizer -from .driver.atiadlxx import ATIADLxx - -class AMDOptimizer(Optimizer): - driver: ATIADLxx = ATIADLxx() - 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 deleted file mode 100644 index 81a578a5c..000000000 --- a/modules/dml/optimizer/amd/driver/atiadlxx.py +++ /dev/null @@ -1,46 +0,0 @@ -import ctypes as C -from .atiadlxx_apis import * -from .atiadlxx_structures import * -from .atiadlxx_defines import * - -class ATIADLxx(object): - iHyperMemorySize = 0 - - def __init__(self): - self.context = ADL_CONTEXT_HANDLE() - ADL2_Main_Control_Create(ADL_Main_Memory_Alloc, 1, C.byref(self.context)) - num_adapters = C.c_int(-1) - ADL2_Adapter_NumberOfAdapters_Get(self.context, C.byref(num_adapters)) - 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: - 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: - info = ADLMemoryInfo2() - - if ADL2_Adapter_MemoryInfo2_Get(self.context, adapterIndex, C.byref(info)) != ADL_OK: - raise RuntimeError("ADL2: Failed to get MemoryInfo2") - - return info - - def get_dedicated_vram_usage(self, index: int) -> int: - usage = C.c_int(-1) - - if ADL2_Adapter_DedicatedVRAMUsage_Get(self.context, self.devices[index].iAdapterIndex, C.byref(usage)) != ADL_OK: - raise RuntimeError("ADL2: Failed to get DedicatedVRAMUsage") - - return usage.value - - def get_vram_usage(self, index: int) -> int: - usage = C.c_int(-1) - - if ADL2_Adapter_VRAMUsage_Get(self.context, self.devices[index].iAdapterIndex, C.byref(usage)) != ADL_OK: - raise RuntimeError("ADL2: Failed to get VRAMUsage") - - return usage.value diff --git a/modules/dml/optimizer/amd/driver/atiadlxx_apis.py b/modules/dml/optimizer/amd/driver/atiadlxx_apis.py deleted file mode 100644 index 25e6390ef..000000000 --- a/modules/dml/optimizer/amd/driver/atiadlxx_apis.py +++ /dev/null @@ -1,45 +0,0 @@ -import ctypes as C -from platform import platform -from .atiadlxx_structures import * - -if 'Windows' in platform(): - atiadlxx = C.WinDLL("atiadlxx.dll") -else: - atiadlxx = C.CDLL("libatiadlxx.so") # Not tested on Linux system. But will be supported. - -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_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_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_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_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/driver/atiadlxx_defines.py b/modules/dml/optimizer/amd/driver/atiadlxx_defines.py deleted file mode 100644 index c242b9819..000000000 --- a/modules/dml/optimizer/amd/driver/atiadlxx_defines.py +++ /dev/null @@ -1 +0,0 @@ -ADL_OK = 0 \ No newline at end of file diff --git a/modules/dml/optimizer/amd/driver/atiadlxx_structures.py b/modules/dml/optimizer/amd/driver/atiadlxx_structures.py deleted file mode 100644 index a68392ec4..000000000 --- a/modules/dml/optimizer/amd/driver/atiadlxx_structures.py +++ /dev/null @@ -1,87 +0,0 @@ -import ctypes as C - -class _ADLPMActivity(C.Structure): - __slot__ = [ - 'iActivityPercent', - 'iCurrentBusLanes', - 'iCurrentBusSpeed', - 'iCurrentPerformanceLevel', - 'iEngineClock', - 'iMaximumBusLanes', - 'iMemoryClock', - 'iReserved', - 'iSize', - 'iVddc', - ] -_ADLPMActivity._fields_ = [ - ('iActivityPercent', C.c_int), - ('iCurrentBusLanes', C.c_int), - ('iCurrentBusSpeed', C.c_int), - ('iCurrentPerformanceLevel', C.c_int), - ('iEngineClock', C.c_int), - ('iMaximumBusLanes', C.c_int), - ('iMemoryClock', C.c_int), - ('iReserved', C.c_int), - ('iSize', C.c_int), - ('iVddc', C.c_int), -] -ADLPMActivity = _ADLPMActivity - -class _ADLMemoryInfo2(C.Structure): - __slot__ = [ - 'iHyperMemorySize', - 'iInvisibleMemorySize', - 'iMemoryBandwidth', - 'iMemorySize', - 'iVisibleMemorySize', - 'strMemoryType' - ] -_ADLMemoryInfo2._fields_ = [ - ('iHyperMemorySize', C.c_longlong), - ('iInvisibleMemorySize', C.c_longlong), - ('iMemoryBandwidth', C.c_longlong), - ('iMemorySize', C.c_longlong), - ('iVisibleMemorySize', C.c_longlong), - ('strMemoryType', C.c_char * 256) -] -ADLMemoryInfo2 = _ADLMemoryInfo2 - -class _AdapterInfo(C.Structure): - __slot__ = [ - 'iSize', - 'iAdapterIndex', - 'strUDID', - 'iBusNumber', - 'iDeviceNumber', - 'iFunctionNumber', - 'iVendorID', - 'strAdapterName', - 'strDisplayName', - 'iPresent', - 'iExist', - 'strDriverPath', - 'strDriverPathExt', - 'strPNPString', - 'iOSDisplayIndex', - ] -_AdapterInfo._fields_ = [ - ('iSize', C.c_int), - ('iAdapterIndex', C.c_int), - ('strUDID', C.c_char * 256), - ('iBusNumber', C.c_int), - ('iDeviceNumber', C.c_int), - ('iFunctionNumber', C.c_int), - ('iVendorID', C.c_int), - ('strAdapterName', C.c_char * 256), - ('strDisplayName', C.c_char * 256), - ('iPresent', C.c_int), - ('iExist', C.c_int), - ('strDriverPath', C.c_char * 256), - ('strDriverPathExt', C.c_char * 256), - ('strPNPString', C.c_char * 256), - ('iOSDisplayIndex', C.c_int) -] -AdapterInfo = _AdapterInfo -LPAdapterInfo = C.POINTER(_AdapterInfo) - -ADL_CONTEXT_HANDLE = C.c_void_p \ No newline at end of file diff --git a/modules/dml/optimizer/intel/__init__.py b/modules/dml/optimizer/intel/__init__.py deleted file mode 100644 index aaa9e67b1..000000000 --- a/modules/dml/optimizer/intel/__init__.py +++ /dev/null @@ -1,6 +0,0 @@ -from modules.dml.optimizer.optimizer import Optimizer - -class IntelOptimizer(Optimizer): - def memory_stats(index: int): - # DML TODO: Implement or find a general (and also lightweight) way. - return (1073741824, 0) diff --git a/modules/dml/optimizer/nvidia/__init__.py b/modules/dml/optimizer/nvidia/__init__.py deleted file mode 100644 index 5990c11a9..000000000 --- a/modules/dml/optimizer/nvidia/__init__.py +++ /dev/null @@ -1,6 +0,0 @@ -from modules.dml.optimizer.optimizer import Optimizer - -class nVidiaOptimizer(Optimizer): - def memory_stats(index: int): - # DML TODO: Implement or find a general (and also lightweight) way. - return (1073741824, 0) diff --git a/modules/dml/optimizer/optimizer.py b/modules/dml/optimizer/optimizer.py deleted file mode 100644 index db88c6579..000000000 --- a/modules/dml/optimizer/optimizer.py +++ /dev/null @@ -1,8 +0,0 @@ -from abc import * -from typing import * - -class Optimizer(metaclass=ABCMeta): - driver: Any = None - @abstractmethod - 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 deleted file mode 100644 index 79ed11ce5..000000000 --- a/modules/dml/optimizer/unknown/__init__.py +++ /dev/null @@ -1,5 +0,0 @@ -from modules.dml.optimizer.optimizer import Optimizer - -class UnknownOptimizer(Optimizer): - def memory_stats(index: int): - return (1073741824, 0) diff --git a/modules/shared.py b/modules/shared.py index 6e51d6ebb..453b4bf18 100644 --- a/modules/shared.py +++ b/modules/shared.py @@ -11,6 +11,7 @@ import tqdm import requests from modules import errors, ui_components, shared_items, cmd_args from modules.paths_internal import models_path, script_path, data_path, sd_configs_path, sd_default_config, sd_model_file, default_sd_model_file, extensions_dir, extensions_builtin_dir # pylint: disable=W0611 +from modules.dml import directml_init import modules.interrogate import modules.memmon import modules.styles @@ -766,8 +767,8 @@ batch_cond_uncond = opts.always_batch_cond_uncond or not (cmd_opts.lowvram or cm parallel_processing_allowed = not cmd_opts.lowvram and not cmd_opts.medvram mem_mon = modules.memmon.MemUsageMonitor("MemMon", device, opts) mem_mon.start() -if device.type == 'privateuseone': - import modules.dml # pylint: disable=ungrouped-imports +if devices.backend == "directml": + directml_init() def reload_gradio_theme(theme_name=None):