diff --git a/modules/dml/__init__.py b/modules/dml/__init__.py index 3c8d8f0e1..a36af5da3 100644 --- a/modules/dml/__init__.py +++ b/modules/dml/__init__.py @@ -1,11 +1,40 @@ -import os +from platform import system import torch from typing import NamedTuple, Callable, Optional from modules.sd_hijack_utils import CondFunc +memory_providers = ["None", "atiadlxx (AMD only)"] +default_memory_provider = "None" +if system() == "Windows": + memory_providers.append("Performance Counter") + default_memory_provider = "Performance Counter" do_nothing = lambda: None +def _set_memory_provider(): + from modules.shared import opts, cmd_opts, log + + if opts.directml_memory_provider == "Performance Counter": + from .backend import pdh_mem_get_info + from .memory import MemoryProvider + torch.dml.mem_get_info = pdh_mem_get_info + if torch.dml.memory_provider is not None: + del torch.dml.memory_provider + torch.dml.memory_provider = MemoryProvider() + elif opts.directml_memory_provider == "atiadlxx (AMD only)": + device_name = torch.dml.get_device_name(cmd_opts.device_id) + if "AMD" not in device_name and "Radeon" not in device_name: + log.warning(f"Memory stats provider is changed to None because the current device is not AMDGPU. Current Device: {device_name}") + opts.directml_memory_provider = "None" + _set_memory_provider() + return + from .backend import amd_mem_get_info + torch.dml.mem_get_info = amd_mem_get_info + else: + from .backend import mem_get_info + torch.dml.mem_get_info = mem_get_info + torch.cuda.mem_get_info = torch.dml.mem_get_info + def directml_init(): from modules.dml.backend import DirectML # pylint: disable=ungrouped-imports # Alternative of torch.cuda for DirectML. @@ -29,10 +58,6 @@ def directml_init(): torch.Tensor.directml = lambda self: self.to(torch.dml.current_device()) - mem_bound = os.environ.get("DML_GPU_MEMORY_BOUND", None) - if mem_bound is not None: - torch.dml.set_gpu_memory_bound(int(mem_bound)) - def directml_do_hijack(): import modules.dml.hijack from modules.devices import device @@ -42,6 +67,8 @@ def directml_do_hijack(): lambda orig_func, *args, **kwargs: orig_func(args[0].astype('float32')), lambda *args, **kwargs: args[1].dtype == float) + _set_memory_provider() + class OverrideItem(NamedTuple): value: str condition: Optional[Callable] @@ -61,9 +88,9 @@ def directml_override_opts(): count = 0 for key in opts_override_table: - count += 1 item = opts_override_table[key] if getattr(shared.opts, key) != item.value and (item.condition is None or item.condition(shared.opts)): + count += 1 setattr(shared.opts, key, item.value) if item.message is not None: shared.log.warning(item.message) @@ -71,3 +98,5 @@ def directml_override_opts(): if count > 0: shared.log.info(f'{count} options are automatically overriden. If you want to keep them from overriding, run with --experimental argument.') + + _set_memory_provider() diff --git a/modules/dml/backend.py b/modules/dml/backend.py index feb0a07bd..ef46f288a 100644 --- a/modules/dml/backend.py +++ b/modules/dml/backend.py @@ -1,13 +1,24 @@ # pylint: disable=no-member,no-self-argument,no-method-argument -from typing import Optional +from typing import Optional, Callable import torch import torch_directml # pylint: disable=import-error import modules.dml.amp as amp -from .memctl.unknown import UnknownMemoryControl from .utils import rDevice, get_device from .device import device from .device_properties import DeviceProperties +from .memory_amd import AMDMemoryProvider +from .memory import MemoryProvider + +def amd_mem_get_info(device: Optional[rDevice]=None) -> tuple[int, int]: + return AMDMemoryProvider.mem_get_info(get_device(device).index) + +def pdh_mem_get_info(device: Optional[rDevice]=None) -> tuple[int, int]: + mem_info = DirectML.memory_provider.get_memory(get_device(device).index) + return (mem_info["total_committed"] - mem_info["dedicated_usage"], mem_info["total_committed"]) + +def mem_get_info(device: Optional[rDevice]=None) -> tuple[int, int]: + return (8589934592, 8589934592) class DirectML: amp = amp @@ -15,29 +26,10 @@ class DirectML: context_device: Optional[torch.device] = None - __gpu_memory_bound: Optional[int] = None - is_autocast_enabled = False autocast_gpu_dtype = torch.float16 - def __get_memory_control(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 .memctl.nvidia import nVidiaMemoryControl as memory_control - elif 'AMD' in device_name or 'Radeon' in device_name: - from .memctl.amd import AMDMemoryControl as memory_control - elif 'Intel' in device_name: - from .memctl.intel import IntelMemoryControl as memory_control - else: - return UnknownMemoryControl - return memory_control - except Exception: - return UnknownMemoryControl - - def set_gpu_memory_bound(bound: Optional[int]): - DirectML.__gpu_memory_bound = bound + memory_provider: Optional[MemoryProvider] = None def is_available() -> bool: return torch_directml.is_available() @@ -73,15 +65,7 @@ class DirectML: "num_alloc_retries": mem_stat_fill, } - def mem_get_info(device: Optional[rDevice]=None) -> tuple[int, int]: - device = get_device(device) - memory_control = DirectML.__get_memory_control(device) - mem_info = memory_control.mem_get_info(device.index) - if DirectML.__gpu_memory_bound is None: - return mem_info - used = mem_info[1] - mem_info[0] - available = DirectML.__gpu_memory_bound - used - return (0 if available < 0 else available, DirectML.__gpu_memory_bound) + mem_get_info: Callable = mem_get_info def memory_allocated(device: Optional[rDevice]=None) -> int: return sum(torch_directml.gpu_memory(get_device(device).index)) * (1 << 20) diff --git a/modules/dml/memctl/amd/__init__.py b/modules/dml/memctl/amd/__init__.py deleted file mode 100644 index b2de1c1bf..000000000 --- a/modules/dml/memctl/amd/__init__.py +++ /dev/null @@ -1,8 +0,0 @@ -from modules.dml.memctl.memctl import MemoryControl -from .driver.atiadlxx import ATIADLxx - -class AMDMemoryControl(MemoryControl): - driver: ATIADLxx = ATIADLxx() - def mem_get_info(index): - usage = AMDMemoryControl.driver.get_dedicated_vram_usage(index) * (1 << 20) - return (AMDMemoryControl.driver.iHyperMemorySize - usage, AMDMemoryControl.driver.iHyperMemorySize) diff --git a/modules/dml/memctl/intel/__init__.py b/modules/dml/memctl/intel/__init__.py deleted file mode 100644 index 24909c5bf..000000000 --- a/modules/dml/memctl/intel/__init__.py +++ /dev/null @@ -1,6 +0,0 @@ -from modules.dml.memctl.memctl import MemoryControl - -class IntelMemoryControl(MemoryControl): - def mem_get_info(index: int): - # DML TODO: Implement or find a general (and also lightweight) way. - return (1073741824, 1073741824) diff --git a/modules/dml/memctl/memctl.py b/modules/dml/memctl/memctl.py deleted file mode 100644 index bda31dd20..000000000 --- a/modules/dml/memctl/memctl.py +++ /dev/null @@ -1,8 +0,0 @@ -from abc import * -from typing import * - -class MemoryControl(metaclass=ABCMeta): - driver: Any = None - @abstractmethod - def mem_get_info(index: int) -> Tuple[int, int]: - pass diff --git a/modules/dml/memctl/nvidia/__init__.py b/modules/dml/memctl/nvidia/__init__.py deleted file mode 100644 index 3334b3518..000000000 --- a/modules/dml/memctl/nvidia/__init__.py +++ /dev/null @@ -1,6 +0,0 @@ -from modules.dml.memctl.memctl import MemoryControl - -class nVidiaMemoryControl(MemoryControl): - def mem_get_info(index: int): - # DML TODO: Implement or find a general (and also lightweight) way. - return (1073741824, 1073741824) diff --git a/modules/dml/memctl/unknown/__init__.py b/modules/dml/memctl/unknown/__init__.py deleted file mode 100644 index 31fda2bcd..000000000 --- a/modules/dml/memctl/unknown/__init__.py +++ /dev/null @@ -1,5 +0,0 @@ -from modules.dml.memctl.memctl import MemoryControl - -class UnknownMemoryControl(MemoryControl): - def mem_get_info(index: int): - return (1073741824, 1073741824) diff --git a/modules/dml/memory.py b/modules/dml/memory.py new file mode 100644 index 000000000..af2d8060f --- /dev/null +++ b/modules/dml/memory.py @@ -0,0 +1,31 @@ +from os import getpid +from collections import defaultdict + +from modules.dml.pdh import HQuery, HCounter, expand_wildcard_path + +class MemoryProvider: + hQuery: HQuery + hCounters: defaultdict[str, list[HCounter]] + + def __init__(self): + self.hQuery = HQuery() + self.hCounters = defaultdict(list) + + def get_memory(self, device_id: int) -> dict[str, int]: + if len(self.hCounters) == 0: + pid = getpid() + paths_dedicated = expand_wildcard_path(f"\\GPU Process Memory(pid_{pid}_*_phys_{device_id})\\Dedicated Usage") + paths_committed = expand_wildcard_path(f"\\GPU Process Memory(pid_{pid}_*_phys_{device_id})\\Total Committed") + for path in paths_dedicated: + self.hCounters["dedicated_usage"].append(self.hQuery.add_counter(path)) + for path in paths_committed: + self.hCounters["total_committed"].append(self.hQuery.add_counter(path)) + self.hQuery.collect_data() + result = defaultdict(int) + for key in self.hCounters: + for hCounter in self.hCounters[key]: + result[key] += hCounter.get_formatted_value(int) + return dict(result) + + def __del__(self): + self.hQuery.close() diff --git a/modules/dml/memory_amd/__init__.py b/modules/dml/memory_amd/__init__.py new file mode 100644 index 000000000..9928d5bc5 --- /dev/null +++ b/modules/dml/memory_amd/__init__.py @@ -0,0 +1,7 @@ +from .driver.atiadlxx import ATIADLxx + +class AMDMemoryProvider: + driver: ATIADLxx = ATIADLxx() + def mem_get_info(index): + usage = AMDMemoryProvider.driver.get_dedicated_vram_usage(index) * (1 << 20) + return (AMDMemoryProvider.driver.iHyperMemorySize - usage, AMDMemoryProvider.driver.iHyperMemorySize) diff --git a/modules/dml/memctl/amd/driver/atiadlxx.py b/modules/dml/memory_amd/driver/atiadlxx.py similarity index 100% rename from modules/dml/memctl/amd/driver/atiadlxx.py rename to modules/dml/memory_amd/driver/atiadlxx.py diff --git a/modules/dml/memctl/amd/driver/atiadlxx_apis.py b/modules/dml/memory_amd/driver/atiadlxx_apis.py similarity index 100% rename from modules/dml/memctl/amd/driver/atiadlxx_apis.py rename to modules/dml/memory_amd/driver/atiadlxx_apis.py diff --git a/modules/dml/memctl/amd/driver/atiadlxx_defines.py b/modules/dml/memory_amd/driver/atiadlxx_defines.py similarity index 100% rename from modules/dml/memctl/amd/driver/atiadlxx_defines.py rename to modules/dml/memory_amd/driver/atiadlxx_defines.py diff --git a/modules/dml/memctl/amd/driver/atiadlxx_structures.py b/modules/dml/memory_amd/driver/atiadlxx_structures.py similarity index 100% rename from modules/dml/memctl/amd/driver/atiadlxx_structures.py rename to modules/dml/memory_amd/driver/atiadlxx_structures.py diff --git a/modules/dml/pdh/__init__.py b/modules/dml/pdh/__init__.py new file mode 100644 index 000000000..0dcd466cb --- /dev/null +++ b/modules/dml/pdh/__init__.py @@ -0,0 +1,85 @@ +from ctypes import * +from ctypes.wintypes import * +from typing import NamedTuple, TypeVar + +from .apis import PdhExpandWildCardPathW, PdhOpenQueryW, PdhAddEnglishCounterW, PdhCollectQueryData, PdhGetFormattedCounterValue, PdhGetFormattedCounterArrayW, PdhCloseQuery +from .structures import PDH_HQUERY, PDH_HCOUNTER, PDH_FMT_COUNTERVALUE, PPDH_FMT_COUNTERVALUE_ITEM_W +from .defines import * +from .msvcrt import malloc +from .errors import PDHError + +class __InternalAbstraction(NamedTuple): + flag: int + attr_name: str + +_type_map = { + int: __InternalAbstraction(PDH_FMT_LARGE, "largeValue"), + float: __InternalAbstraction(PDH_FMT_DOUBLE, "doubleValue"), +} + +def expand_wildcard_path(path: str) -> list[str]: + listLength = DWORD(0) + if PdhExpandWildCardPathW(None, LPCWSTR(path), None, byref(listLength), PDH_NOEXPANDCOUNTERS) != PDH_MORE_DATA: + raise PDHError("Something went wrong.") + expanded = (WCHAR * listLength.value)() + if PdhExpandWildCardPathW(None, LPCWSTR(path), expanded, byref(listLength), PDH_NOEXPANDCOUNTERS) != PDH_OK: + raise PDHError(f"Couldn't expand wildcard path '{path}'") + result = list() + cur = str() + for chr in expanded: + if chr == '\0': + result.append(cur) + cur = str() + else: + cur += chr + result.pop() + return result + +T = TypeVar("T", *_type_map.keys()) + +class HCounter(PDH_HCOUNTER): + def get_formatted_value(self, type: T) -> T: + if type not in _type_map: + raise PDHError(f"Invalid value type: {type}") + flag, attr_name = _type_map[type] + value = PDH_FMT_COUNTERVALUE() + if PdhGetFormattedCounterValue(self, DWORD(flag | PDH_FMT_NOSCALE), None, byref(value)) != PDH_OK: + raise PDHError("Couldn't get formatted counter value.") + return getattr(value.u, attr_name) + + def get_formatted_dict(self, type: T) -> dict[str, T]: + if type not in _type_map: + raise PDHError(f"Invalid value type: {type}") + flag, attr_name = _type_map[type] + bufferSize = DWORD(0) + itemCount = DWORD(0) + if PdhGetFormattedCounterArrayW(self, DWORD(flag | PDH_FMT_NOSCALE), byref(bufferSize), byref(itemCount), None) != PDH_MORE_DATA: + raise PDHError("Something went wrong.") + itemBuffer = cast(malloc(c_size_t(bufferSize.value)), PPDH_FMT_COUNTERVALUE_ITEM_W) + if PdhGetFormattedCounterArrayW(self, DWORD(flag | PDH_FMT_NOSCALE), byref(bufferSize), byref(itemCount), itemBuffer) != PDH_OK: + raise PDHError("Couldn't get formatted counter array.") + result: dict[str, T] = dict() + for i in range(0, itemCount.value): + item = itemBuffer[i] + result[item.szName] = getattr(item.FmtValue.u, attr_name) + return result + +class HQuery(PDH_HQUERY): + def __init__(self): + super(HQuery, self).__init__() + if PdhOpenQueryW(None, None, byref(self)) != PDH_OK: + raise PDHError("Couldn't open PDH query.") + + def add_counter(self, path: str) -> HCounter: + hCounter = HCounter() + if PdhAddEnglishCounterW(self, LPCWSTR(path), None, byref(hCounter)) != PDH_OK: + raise PDHError("Couldn't add counter query.") + return hCounter + + def collect_data(self): + if PdhCollectQueryData(self) != PDH_OK: + raise PDHError("Couldn't collect query data.") + + def close(self): + if PdhCloseQuery(self) != PDH_OK: + raise PDHError("Couldn't close PDH query.") diff --git a/modules/dml/pdh/apis.py b/modules/dml/pdh/apis.py new file mode 100644 index 000000000..87c1d1204 --- /dev/null +++ b/modules/dml/pdh/apis.py @@ -0,0 +1,36 @@ +from ctypes import * +from ctypes.wintypes import * +from typing import Callable + +from .structures import * +from .defines import * + +pdh = CDLL("pdh.dll") + +PdhExpandWildCardPathW: Callable = pdh.PdhExpandWildCardPathW +PdhExpandWildCardPathW.restype = PDH_FUNCTION +PdhExpandWildCardPathW.argtypes = [LPCWSTR, LPCWSTR, PZZWSTR, LPDWORD, DWORD] + +PdhOpenQueryW: Callable = pdh.PdhOpenQueryW +PdhOpenQueryW.restype = PDH_FUNCTION +PdhOpenQueryW.argtypes = [LPCWSTR, DWORD_PTR, POINTER(PDH_HQUERY)] + +PdhAddEnglishCounterW: Callable = pdh.PdhAddEnglishCounterW +PdhAddEnglishCounterW.restype = PDH_FUNCTION +PdhAddEnglishCounterW.argtypes = [PDH_HQUERY, LPCWSTR, DWORD_PTR, POINTER(PDH_HCOUNTER)] + +PdhCollectQueryData: Callable = pdh.PdhCollectQueryData +PdhCollectQueryData.restype = PDH_FUNCTION +PdhCollectQueryData.argtypes = [PDH_HQUERY] + +PdhGetFormattedCounterValue: Callable = pdh.PdhGetFormattedCounterValue +PdhGetFormattedCounterValue.restype = PDH_FUNCTION +PdhGetFormattedCounterValue.argtypes = [PDH_HCOUNTER, DWORD, LPDWORD, PPDH_FMT_COUNTERVALUE] + +PdhGetFormattedCounterArrayW: Callable = pdh.PdhGetFormattedCounterArrayW +PdhGetFormattedCounterArrayW.restype = PDH_FUNCTION +PdhGetFormattedCounterArrayW.argtypes = [PDH_HCOUNTER, DWORD, LPDWORD, LPDWORD, PPDH_FMT_COUNTERVALUE_ITEM_W] + +PdhCloseQuery: Callable = pdh.PdhCloseQuery +PdhCloseQuery.restype = PDH_FUNCTION +PdhCloseQuery.argtypes = [PDH_HQUERY] diff --git a/modules/dml/pdh/defines.py b/modules/dml/pdh/defines.py new file mode 100644 index 000000000..a5ea1d479 --- /dev/null +++ b/modules/dml/pdh/defines.py @@ -0,0 +1,22 @@ +from ctypes import * +from ctypes.wintypes import * + +PDH_FUNCTION = c_int +PDH_OK = 0x00000000 +PDH_MORE_DATA = -2147481646#0x800007D2 + +DWORD_PTR = POINTER(DWORD) +PWSTR = POINTER(WCHAR) +PZZWSTR = POINTER(WCHAR) + +PDH_NOEXPANDCOUNTERS = 1 +PDH_NOEXPANDINSTANCES = 2 +PDH_REFRESHCOUNTERS = 4 + +PDH_FMT_LONG = 0x00000100 +PDH_FMT_DOUBLE = 0x00000200 +PDH_FMT_LARGE = 0x00000400 + +PDH_FMT_NOSCALE = 0x00001000 +PDH_FMT_1000 = 0x00002000 +PDH_FMT_NOCAP100 = 0x00008000 diff --git a/modules/dml/pdh/errors.py b/modules/dml/pdh/errors.py new file mode 100644 index 000000000..60d9ab8f7 --- /dev/null +++ b/modules/dml/pdh/errors.py @@ -0,0 +1,3 @@ +class PDHError(Exception): + def __init__(self, message: str): + super(PDHError, self).__init__(message) diff --git a/modules/dml/pdh/msvcrt.py b/modules/dml/pdh/msvcrt.py new file mode 100644 index 000000000..bc5d93031 --- /dev/null +++ b/modules/dml/pdh/msvcrt.py @@ -0,0 +1,11 @@ +from ctypes import * + +msvcrt = CDLL("msvcrt") + +malloc = msvcrt.malloc +malloc.restype = c_void_p +malloc.argtypes = [c_size_t] + +free = msvcrt.free +free.restype = None +free.argtypes = [c_void_p] diff --git a/modules/dml/pdh/structures.py b/modules/dml/pdh/structures.py new file mode 100644 index 000000000..8fb09e6cb --- /dev/null +++ b/modules/dml/pdh/structures.py @@ -0,0 +1,41 @@ +from ctypes import * +from ctypes.wintypes import * + +PDH_HQUERY = HANDLE +PDH_HCOUNTER = HANDLE + +class PDH_FMT_COUNTERVALUE_U(Union): + _fields_ = [ + ("longValue", LONG), + ("doubleValue", c_double), + ("largeValue", c_longlong), + ("AnsiStringValue", LPCSTR), + ("WideStringValue", LPCWSTR), + ] + + longValue: int + doubleValue: float + largeValue: int + AnsiStringValue: LPCSTR + WideStringValue: LPCWSTR + +class PDH_FMT_COUNTERVALUE(Structure): + _anonymous_ = ("u",) + _fields_ = [ + ("CStatus", DWORD), + ("u", PDH_FMT_COUNTERVALUE_U), + ] + + CStatus: DWORD + u: PDH_FMT_COUNTERVALUE_U +PPDH_FMT_COUNTERVALUE = POINTER(PDH_FMT_COUNTERVALUE) + +class PDH_FMT_COUNTERVALUE_ITEM_W(Structure): + _fields_ = [ + ("szName", LPWSTR), + ("FmtValue", PDH_FMT_COUNTERVALUE), + ] + + szName: str + FmtValue: PDH_FMT_COUNTERVALUE +PPDH_FMT_COUNTERVALUE_ITEM_W = POINTER(PDH_FMT_COUNTERVALUE_ITEM_W) diff --git a/modules/shared.py b/modules/shared.py index 711a77802..e7e585223 100644 --- a/modules/shared.py +++ b/modules/shared.py @@ -12,7 +12,7 @@ import requests import fasteners 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_do_hijack +from modules.dml import memory_providers, default_memory_provider, directml_do_hijack import modules.interrogate import modules.memmon import modules.styles @@ -390,6 +390,7 @@ options_templates.update(options_section(('cuda', "Compute Settings"), { "cuda_compile_verbose": OptionInfo(False, "Model compile verbose mode"), "cuda_compile_errors": OptionInfo(True, "Model compile suppress errors"), "disable_gc": OptionInfo(True, "Disable Torch memory garbage collection"), + "directml_memory_provider": OptionInfo(default_memory_provider, '[DirectML] Memory stats provider', gr.Dropdown, lambda: {"choices": memory_providers}), })) options_templates.update(options_section(('diffusers', "Diffusers Settings"), {