Merge pull request #2331 from zxue2/enable_ipc_mem

Add Python bindings for SYCL IPC memory via dpctl

Author: ndgrigorian

CMakeLists.txt

@@ -115,7 +115,7 @@ install(DIRECTORY
 )
 
 # find Python before enabling pybind11
-find_package(Python REQUIRED COMPONENTS Development.Module)
+find_package(Python REQUIRED COMPONENTS Interpreter Development.Module)
 
 # Define CMAKE_INSTALL_xxx: LIBDIR, INCLUDEDIR
 include(GNUInstallDirs)

docs/doc_sources/api_reference/dpctl/memory.rst

@@ -66,6 +66,27 @@ Should the USM allocation fail, the following Python exception will be raised:
 
     USMAllocationError
 
+.. rubric:: IPC (Inter-Process Communication) Memory
+
+Classes and functions for sharing USM device allocations across processes
+using IPC memory handles.
+
+.. autosummary::
+    :toctree: generated
+    :template: autosummary/usmmemory.rst
+    :nosignatures:
+
+    IPCMemoryHandle
+    MemoryIPCDevice
+
+.. autosummary::
+    :toctree: generated
+    :nosignatures:
+
+    SyclIPCGetMemHandle
+    SyclIPCOpenMemHandle
+    SyclIPCCloseMemHandle
+
 .. toctree::
     :hidden:
 

dpctl/_backend.pxd

@@ -101,6 +101,7 @@ cdef extern from "syclinterface/dpctl_sycl_enum_types.h":
         _emulated                           "emulated",
         _is_component                       "is_component",
         _is_composite                       "is_composite",
+        _ext_oneapi_ipc_memory              "ext_oneapi_ipc_memory",
 
     ctypedef enum _partition_affinity_domain_type \
             "DPCTLPartitionAffinityDomainType":
@@ -601,6 +602,7 @@ cdef extern from "syclinterface/dpctl_sycl_extension_interface.h":
         DPCTLSyclWorkGroupMemoryRef Ref)
 
     cdef bint DPCTLWorkGroupMemory_Available()
+    cdef bint DPCTLIPCMem_Available()
 
     cdef struct DPCTLOpaqueRawKernelArg
     ctypedef DPCTLOpaqueRawKernelArg *DPCTLSyclRawKernelArgRef
@@ -612,3 +614,19 @@ cdef extern from "syclinterface/dpctl_sycl_extension_interface.h":
         DPCTLSyclRawKernelArgRef Ref)
 
     cdef bint DPCTLRawKernelArg_Available()
+
+cdef extern from "syclinterface/dpctl_sycl_ipc_memory_interface.h":
+    cdef int DPCTLIPCMem_GetHandle(
+        DPCTLSyclUSMRef Ptr,
+        DPCTLSyclContextRef CRef,
+        char **DataOut,
+        size_t *SizeOut)
+    cdef DPCTLSyclUSMRef DPCTLIPCMem_OpenHandle(
+        const char *HandleData,
+        size_t HandleDataSize,
+        DPCTLSyclContextRef CRef,
+        DPCTLSyclDeviceRef DRef)
+    cdef void DPCTLIPCMem_CloseHandle(
+        DPCTLSyclUSMRef MappedPtr,
+        DPCTLSyclContextRef CRef)
+    cdef void DPCTLIPCMem_FreeHandleData(char *Data)

dpctl/_sycl_device.pyx

@@ -885,6 +885,18 @@ cdef class SyclDevice(_SyclDevice):
         cdef _aspect_type AT = _aspect_type._is_composite
         return DPCTLDevice_HasAspect(self._device_ref, AT)
 
+    @property
+    def has_aspect_ext_oneapi_ipc_memory(self):
+        """ Returns ``True`` if this device supports inter-process
+        communication (IPC) memory handles, ``False`` otherwise.
+
+        Returns:
+            bool:
+                Indicates if device supports IPC memory.
+        """
+        cdef _aspect_type AT = _aspect_type._ext_oneapi_ipc_memory
+        return DPCTLDevice_HasAspect(self._device_ref, AT)
+
     @property
     def image_2d_max_width(self):
         """ Returns the maximum width of a 2D image or 1D image in pixels.

dpctl/memory/__init__.py

@@ -30,17 +30,27 @@
 """
 
 from ._memory import (
+    IPCMemoryHandle,
+    MemoryIPCDevice,
     MemoryUSMDevice,
     MemoryUSMHost,
     MemoryUSMShared,
+    SyclIPCCloseMemHandle,
+    SyclIPCGetMemHandle,
+    SyclIPCOpenMemHandle,
     USMAllocationError,
     as_usm_memory,
 )
 
 __all__ = [
+    "IPCMemoryHandle",
+    "MemoryIPCDevice",
     "MemoryUSMDevice",
     "MemoryUSMHost",
     "MemoryUSMShared",
     "USMAllocationError",
     "as_usm_memory",
+    "SyclIPCGetMemHandle",
+    "SyclIPCOpenMemHandle",
+    "SyclIPCCloseMemHandle",
 ]

dpctl/memory/_memory.pxd

@@ -90,3 +90,12 @@ cdef public api class MemoryUSMDevice(_Memory) [
     object PyMemoryUSMDeviceObject, type PyMemoryUSMDeviceType
 ]:
     pass
+
+
+cdef class MemoryIPCDevice(MemoryUSMDevice):
+    @staticmethod
+    cdef object create_ipc_from_usm_pointer_size_qref(
+        DPCTLSyclUSMRef USMRef,
+        Py_ssize_t nbytes,
+        DPCTLSyclQueueRef QRef,
+    )