From 2a077414bcf93969592b6768e9de7a183e1eaa3b Mon Sep 17 00:00:00 2001 From: Seunghoon Lee Date: Fri, 21 Mar 2025 14:56:58 +0900 Subject: [PATCH] zluda experimental torch.compile --- modules/zluda.py | 4 +-- modules/zluda_hijacks.py | 14 ++++++++-- modules/zluda_installer.py | 54 ++++++++++++++++++++++++++++++++------ 3 files changed, 60 insertions(+), 12 deletions(-) diff --git a/modules/zluda.py b/modules/zluda.py index 0203a6398..431ab2c8c 100644 --- a/modules/zluda.py +++ b/modules/zluda.py @@ -3,7 +3,8 @@ from typing import Union import torch from torch._prims_common import DeviceLikeType import onnxruntime as ort -from modules import shared, devices +from modules import shared, devices, zluda_installer +from modules.zluda_installer import core, default_agent # pylint: disable=unused-import from modules.onnx_impl.execution_providers import available_execution_providers, ExecutionProvider @@ -32,7 +33,6 @@ def initialize_zluda(): from modules.zluda_hijacks import do_hijack do_hijack() - from modules import zluda_installer torch.backends.cudnn.enabled = zluda_installer.MIOpen_available if not zluda_installer.MIOpen_available: torch.backends.cuda.enable_cudnn_sdp(False) diff --git a/modules/zluda_hijacks.py b/modules/zluda_hijacks.py index 9f622be94..1b3f750a9 100644 --- a/modules/zluda_hijacks.py +++ b/modules/zluda_hijacks.py @@ -1,5 +1,6 @@ import torch -from modules import rocm +import torch._dynamo.device_interface +from modules import rocm, zluda _topk = torch.topk @@ -10,7 +11,7 @@ def topk(input: torch.Tensor, *args, **kwargs): # pylint: disable=redefined-buil class DeviceProperties: - PROPERTIES_OVERRIDE = {"regs_per_multiprocessor": 65535} + PROPERTIES_OVERRIDE = {"regs_per_multiprocessor": 65535, "gcnArchName": "UNKNOWN ARCHITECTURE"} internal: torch._C._CudaDeviceProperties def __init__(self, props: torch._C._CudaDeviceProperties): @@ -27,11 +28,20 @@ def torch_cuda__get_device_properties(device): return DeviceProperties(__get_device_properties(device)) +_cuda_getCurrentRawStream = torch._C._cuda_getCurrentRawStream # pylint: disable=protected-access +def torch__C__cuda_getCurrentRawStream(device): + return zluda.core.to_hip_stream(_cuda_getCurrentRawStream(device)) + + def do_hijack(): torch.version.hip = rocm.version torch.topk = topk + if zluda.default_agent is not None: + DeviceProperties.PROPERTIES_OVERRIDE["gcnArchName"] = zluda.default_agent.name torch.cuda._get_device_properties = torch_cuda__get_device_properties # pylint: disable=protected-access + torch._C._cuda_getCurrentRawStream = torch__C__cuda_getCurrentRawStream # pylint: disable=protected-access + torch._dynamo.device_interface.CudaInterface.get_raw_stream = staticmethod(torch__C__cuda_getCurrentRawStream) # pylint: disable=protected-access try: import triton _get_device_properties = triton.runtime.driver.active.utils.get_device_properties diff --git a/modules/zluda_installer.py b/modules/zluda_installer.py index 4a0b9ffa4..35a84cfd2 100644 --- a/modules/zluda_installer.py +++ b/modules/zluda_installer.py @@ -17,7 +17,6 @@ DLL_MAPPING = { 'nvrtc.dll': 'nvrtc64_112_0.dll', } HIPSDK_TARGETS = ['rocblas.dll', 'rocsolver.dll', 'hipfft.dll',] -ZLUDA_TARGETS = ('nvcuda.dll', 'nvml.dll',) hipBLASLt_available = False MIOpen_available = False @@ -27,7 +26,49 @@ default_agent: Union[rocm.Agent, None] = None hipBLASLt_enabled = False nightly = os.environ.get("ZLUDA_NIGHTLY", "0") == "1" -skip_arch_test = os.environ.get("ZLUDA_SKIP_ARCH_TEST", "0") == "1" + + +class ZLUDAResult(ctypes.Structure): + _fields_ = [ + ('return_code', ctypes.c_int), + ('value', ctypes.c_ulonglong), + ] + + +class ZLUDALibrary: + internal: ctypes.WinDLL + + def __init__(self, internal: ctypes.WinDLL): + self.internal = internal + + +class Core(ZLUDALibrary): + internal: ctypes.WinDLL + + def __init__(self, internal: ctypes.WinDLL): + internal.zluda_get_hip_object.restype = ZLUDAResult + internal.zluda_get_hip_object.argtypes = [ctypes.c_void_p, ctypes.c_int] + + internal.zluda_get_nightly_flag.restype = ctypes.c_int + internal.zluda_get_nightly_flag.argtypes = [] + + super().__init__(internal) + + def to_hip_stream(self, zluda_object: ctypes.c_void_p): + return self.internal.zluda_get_hip_object(zluda_object, 1).value + + def get_nightly_flag(self) -> int: + return self.internal.zluda_get_nightly_flag().value + + +core = None +ml = None + + +def load_core_modules(): + global core, ml # pylint: disable=global-statement + core = Core(ctypes.windll.LoadLibrary(os.path.join(path, 'nvcuda.dll'))) + ml = ZLUDALibrary(ctypes.windll.LoadLibrary(os.path.join(path, 'nvml.dll'))) def set_default_agent(agent: rocm.Agent): @@ -36,10 +77,8 @@ def set_default_agent(agent: rocm.Agent): is_nightly = False try: - nvcuda = ctypes.windll.LoadLibrary(os.path.join(path, 'nvcuda.dll')) - nvcuda.zluda_get_nightly_flag.restype = ctypes.c_int - nvcuda.zluda_get_nightly_flag.argtypes = [] - is_nightly = nvcuda.zluda_get_nightly_flag() == 1 + load_core_modules() + is_nightly = core.get_nightly_flag() except Exception: pass @@ -113,10 +152,9 @@ def load() -> None: os.environ["ZLUDA_COMGR_LOG_LEVEL"] = "1" os.environ["ZLUDA_NVRTC_LIB"] = os.path.join([v for v in site.getsitepackages() if v.endswith("site-packages")][0], "torch", "lib", "nvrtc64_112_0.dll") + load_core_modules() for v in HIPSDK_TARGETS: ctypes.windll.LoadLibrary(os.path.join(rocm.path, 'bin', v)) - for v in ZLUDA_TARGETS: - ctypes.windll.LoadLibrary(os.path.join(path, v)) for v in DLL_MAPPING.values(): ctypes.windll.LoadLibrary(os.path.join(path, v))