XPU: register pinned host memory via Level Zero ext_memmap

accelerator/xpu_accelerator.py

@@ -3,6 +3,7 @@
 
 # DeepSpeed Team
 
+import ctypes
 import torch
 from deepspeed.accelerator.abstract_accelerator import DeepSpeedAccelerator
 import functools
@@ -15,6 +16,107 @@
 except ImportError as e:
     oneccl_imported_p = False
 
+# Host-memory registration on XPU goes through Level Zero's
+# ZE_extension_external_memmap_sysmem extension: zeMemAllocHost with an
+# ze_external_memmap_sysmem_ext_desc_t chained onto the host-alloc descriptor
+# maps an existing page-locked host buffer into the device page tables (the
+# returned pointer equals the input pointer), and zeMemFree releases the
+# mapping without touching the host memory itself. This is the XPU equivalent
+# of cudaHostRegister/cudaHostUnregister.
+#
+# Values from ze_api.h (stable ABI): structure types, result codes, and the
+# descriptor layouts the loader dispatches on.
+_ZE_STRUCTURE_TYPE_CONTEXT_DESC = 0xD
+_ZE_STRUCTURE_TYPE_HOST_MEM_ALLOC_DESC = 0x16
+_ZE_STRUCTURE_TYPE_EXTERNAL_MEMMAP_SYSMEM_EXT_DESC = 0x00020037
+_ZE_MAX_EXTENSION_NAME = 256
+_ZE_RESULT_SUCCESS = 0x0
+_ZE_INIT_FLAG_GPU_ONLY = 0x1
+
+
+class _ZeContextDesc(ctypes.Structure):
+
+    _fields_ = [("stype", ctypes.c_uint32), ("pNext", ctypes.c_void_p), ("flags", ctypes.c_uint32)]
+
+
+class _ZeExternalMemmapSysmemDesc(ctypes.Structure):
+
+    _fields_ = [("stype", ctypes.c_uint32), ("pNext", ctypes.c_void_p), ("pSystemMemory", ctypes.c_void_p),
+                ("size", ctypes.c_uint64)]
+
+
+class _ZeHostMemAllocDesc(ctypes.Structure):
+
+    _fields_ = [("stype", ctypes.c_uint32), ("pNext", ctypes.c_void_p), ("flags", ctypes.c_uint32)]
+
+
+class _ZeDriverExtensionProperties(ctypes.Structure):
+
+    _fields_ = [("name", ctypes.c_char * _ZE_MAX_EXTENSION_NAME), ("version", ctypes.c_uint32)]
+
+
+@functools.lru_cache(maxsize=None)
+def _l0_host_registration():
+    """(loader, context) handles for host-memory registration, or None when unsupported.
+
+    The mapping is created in our own Level Zero context; device page tables
+    are shared across the driver's contexts, so buffers registered here are
+    device-accessible from torch's runtime as well.
+    """
+    try:
+        ze = ctypes.CDLL("libze_loader.so.1")
+    except OSError:
+        return None
+
+    ze.zeInit.argtypes = [ctypes.c_uint32]
+    ze.zeInit.restype = ctypes.c_uint32
+    if ze.zeInit(_ZE_INIT_FLAG_GPU_ONLY) != _ZE_RESULT_SUCCESS:
+        return None
+
+    ze.zeDriverGet.argtypes = [ctypes.POINTER(ctypes.c_uint32), ctypes.POINTER(ctypes.c_void_p)]
+    ze.zeDriverGet.restype = ctypes.c_uint32
+    count = ctypes.c_uint32(0)
+    if ze.zeDriverGet(ctypes.byref(count), None) != _ZE_RESULT_SUCCESS or count.value == 0:
+        return None
+    drivers = (ctypes.c_void_p * count.value)()
+    if ze.zeDriverGet(ctypes.byref(count), drivers) != _ZE_RESULT_SUCCESS:
+        return None
+    driver = drivers[0]
+
+    # zeMemAllocHost-with-descriptor is an optional extension, so probe before
+    # the first registration instead of failing per buffer.
+    ze.zeDriverGetExtensionProperties.argtypes = [
+        ctypes.c_void_p,
+        ctypes.POINTER(ctypes.c_uint32),
+        ctypes.POINTER(_ZeDriverExtensionProperties)
+    ]
+    ze.zeDriverGetExtensionProperties.restype = ctypes.c_uint32
+    if ze.zeDriverGetExtensionProperties(driver, ctypes.byref(count), None) != _ZE_RESULT_SUCCESS:
+        return None
+    extensions = (_ZeDriverExtensionProperties * count.value)()
+    if ze.zeDriverGetExtensionProperties(driver, ctypes.byref(count), extensions) != _ZE_RESULT_SUCCESS:
+        return None
+    names = [extensions[i].name.decode() for i in range(count.value)]
+    if "ZE_extension_external_memmap_sysmem" not in names:
+        return None
+
+    context = ctypes.c_void_p()
+    desc = _ZeContextDesc(stype=_ZE_STRUCTURE_TYPE_CONTEXT_DESC, pNext=None, flags=0)
+    ze.zeContextCreate.argtypes = [ctypes.c_void_p, ctypes.POINTER(_ZeContextDesc), ctypes.POINTER(ctypes.c_void_p)]
+    ze.zeContextCreate.restype = ctypes.c_uint32
+    if ze.zeContextCreate(driver, ctypes.byref(desc), ctypes.byref(context)) != _ZE_RESULT_SUCCESS:
+        return None
+
+    ze.zeMemAllocHost.argtypes = [
+        ctypes.c_void_p,
+        ctypes.POINTER(_ZeHostMemAllocDesc), ctypes.c_size_t, ctypes.c_size_t,
+        ctypes.POINTER(ctypes.c_void_p)
+    ]
+    ze.zeMemAllocHost.restype = ctypes.c_uint32
+    ze.zeMemFree.argtypes = [ctypes.c_void_p, ctypes.c_void_p]
+    ze.zeMemFree.restype = ctypes.c_uint32
+    return ze, context
+
 
 class XPU_Accelerator(DeepSpeedAccelerator):
 
@@ -239,6 +341,43 @@ def _torch_empty_pinned(self, tensor, shape):
     def _torch_is_pinned(self, tensor):
         return tensor.is_pinned(device=self.current_device_name())
 
+    def register_host_memory(self, address, num_bytes):
+        handles = _l0_host_registration()
+        if handles is None:
+            from deepspeed.utils import logger
+            logger.warning_once(
+                "Level Zero host-memory registration is unavailable (libze_loader.so.1 or the "
+                "ZE_extension_external_memmap_sysmem extension is missing); native pinned memory stays mlock-only.")
+            return False
+        ze, context = handles
+        memmap_desc = _ZeExternalMemmapSysmemDesc(stype=_ZE_STRUCTURE_TYPE_EXTERNAL_MEMMAP_SYSMEM_EXT_DESC,
+                                                  pNext=None,
+                                                  pSystemMemory=address,
+                                                  size=num_bytes)
+        host_desc = _ZeHostMemAllocDesc(stype=_ZE_STRUCTURE_TYPE_HOST_MEM_ALLOC_DESC,
+                                        pNext=ctypes.cast(ctypes.byref(memmap_desc), ctypes.c_void_p),
+                                        flags=0)
+        mapped = ctypes.c_void_p()
+        rc = ze.zeMemAllocHost(context, ctypes.byref(host_desc), num_bytes, 0, ctypes.byref(mapped))
+        if rc == _ZE_RESULT_SUCCESS and mapped.value != address:
+            # The spec guarantees the mapping preserves the virtual address; a
+            # non-identity mapping is unusable (callers hold the original
+            # pointer), so release it and report failure so they fall back to mlock.
+            ze.zeMemFree(context, mapped.value)
+            from deepspeed.utils import logger
+            logger.warning_once(
+                "Level Zero returned a non-identity mapping for host registration "
+                "(%#x != %#x); falling back to mlock-only pinned memory.", mapped.value, address)
+            return False
+        return rc == _ZE_RESULT_SUCCESS
+
+    def unregister_host_memory(self, address):
+        handles = _l0_host_registration()
+        if handles is None:
+            return None
+        ze, context = handles
+        ze.zeMemFree(context, address)
+
     def op_builder_dir(self):
         try:
             # is op_builder from deepspeed or a 3p version? this should only succeed if it's deepspeed

tests/unit/v1/pin_memory/test_pin_memory.py

@@ -2,13 +2,16 @@
 
 # DeepSpeed Team
 
+import ctypes
 import pytest
 import torch
 
 from deepspeed.accelerator import npu_accelerator
 from deepspeed.accelerator.cpu_accelerator import CPU_Accelerator
 from deepspeed.accelerator.cuda_accelerator import CUDA_Accelerator
 from deepspeed.accelerator.npu_accelerator import NPU_Accelerator
+from deepspeed.accelerator import xpu_accelerator
+from deepspeed.accelerator.xpu_accelerator import XPU_Accelerator
 from deepspeed.utils.pin_memory import NativePinnedMemory
 
 
@@ -211,6 +214,63 @@ def cudaHostUnregister(self, address):
     assert errors == [0, 0]
 
 
+def test_xpu_device_registration_calls_level_zero(monkeypatch):
+
+    calls = []
+
+    class _FakeLoader:
+
+        def zeMemAllocHost(self, context, host_desc, size, alignment, out):
+            host = ctypes.cast(host_desc, ctypes.POINTER(xpu_accelerator._ZeHostMemAllocDesc))[0]
+            memmap = ctypes.cast(host.pNext, ctypes.POINTER(xpu_accelerator._ZeExternalMemmapSysmemDesc))[0]
+            calls.append(("alloc", memmap.pSystemMemory, memmap.size, size))
+            # byref() without argtypes arrives as a CArgObject; cast it back.
+            ctypes.cast(out, ctypes.POINTER(ctypes.c_void_p))[0] = memmap.pSystemMemory
+            return xpu_accelerator._ZE_RESULT_SUCCESS
+
+        def zeMemFree(self, context, address):
+            calls.append(("free", address))
+            return xpu_accelerator._ZE_RESULT_SUCCESS
+
+    monkeypatch.setattr(xpu_accelerator, "_l0_host_registration", lambda: (_FakeLoader(), "ctx"))
+    accelerator = XPU_Accelerator.__new__(XPU_Accelerator)
+
+    assert accelerator.register_host_memory(1234, 4096) is True
+    accelerator.unregister_host_memory(1234)
+    assert calls == [("alloc", 1234, 4096, 4096), ("free", 1234)]
+
+
+def test_xpu_registration_rejects_changed_address(monkeypatch):
+    """A mapping that comes back at another address is a failed registration
+    and the unusable mapping is released."""
+
+    freed = []
+
+    class _FakeLoader:
+
+        def zeMemAllocHost(self, context, host_desc, size, alignment, out):
+            ctypes.cast(out, ctypes.POINTER(ctypes.c_void_p))[0] = 0xDEAD0000
+            return xpu_accelerator._ZE_RESULT_SUCCESS
+
+        def zeMemFree(self, context, address):
+            freed.append(address)
+            return xpu_accelerator._ZE_RESULT_SUCCESS
+
+    monkeypatch.setattr(xpu_accelerator, "_l0_host_registration", lambda: (_FakeLoader(), "ctx"))
+    accelerator = XPU_Accelerator.__new__(XPU_Accelerator)
+    assert accelerator.register_host_memory(1234, 4096) is False
+    assert freed == [0xDEAD0000]
+
+
+def test_xpu_device_registration_without_level_zero(monkeypatch):
+    """Missing loader or extension must degrade to mlock-only, not raise."""
+    monkeypatch.setattr(xpu_accelerator, "_l0_host_registration", lambda: None)
+    accelerator = XPU_Accelerator.__new__(XPU_Accelerator)
+
+    assert accelerator.register_host_memory(1234, 4096) is False
+    assert accelerator.unregister_host_memory(1234) is None
+
+
 def test_cpu_native_pin_with_register_env_on(monkeypatch, native_pins):
     """CPU accelerator has no register hook; native pin still works with default-on."""
     monkeypatch.setenv("DS_PIN_MEMORY_REGISTER_DEVICE", "1")