Improve pickling support

CHANGELOG.md

@@ -8,8 +8,12 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 
 ### Added
 
+* `dpctl.SyclDevice`, `dpctl.SyclContext`, and `dpctl.SyclQueue` can be pickled. Root devices are restored by filter string, default contexts by one of their devices, and queues by context, device, and properties, with default-property queues restored as the cached queue of their device. `copy.copy` and `copy.deepcopy` return the same object
+
 ### Changed
 
+* Pickled `dpctl.memory.MemoryUSMShared`, `MemoryUSMDevice`, and `MemoryUSMHost` objects keep their queue instead of being restored on the default device. Pickling memory whose queue can not be pickled raises `pickle.PicklingError`
+
 ### Deprecated
 
 ### Removed

dpctl/_sycl_context.pyx

@@ -23,6 +23,7 @@
 """
 
 import logging
+import pickle
 
 from cpython cimport pycapsule
 from cpython.mem cimport PyMem_Free, PyMem_Malloc
@@ -129,6 +130,13 @@ cdef class _SyclContext:
         DPCTLContext_Delete(self._ctxt_ref)
 
 
+def _rebuild_context(SyclDevice dev):
+    """Unpickles a :class:`dpctl.SyclContext` as the default context of
+    ``dev``.
+    """
+    return SyclContext(dev)
+
+
 cdef class SyclContext(_SyclContext):
     """
     SyclContext(arg=None)
@@ -408,6 +416,35 @@ cdef class SyclContext(_SyclContext):
         """
         return DPCTLContext_Hash(self._ctxt_ref)
 
+    def __reduce__(self):
+        """Pickles the default context of a root device by that device.
+
+        Raises:
+            pickle.PicklingError: if the context is not the default context
+                of its devices.
+        """
+        cdef DPCTLSyclContextRef CRef = NULL
+        cdef bint is_default_ctx = False
+        cdef SyclDevice dev = self.get_devices()[0]
+        # sub-devices have no cached context
+        if dev.parent_device is None:
+            CRef = DPCTLDeviceMgr_GetCachedContext(dev.get_device_ref())
+            if CRef is not NULL:
+                is_default_ctx = DPCTLContext_AreEq(CRef, self._ctxt_ref)
+                DPCTLContext_Delete(CRef)
+        if not is_default_ctx:
+            raise pickle.PicklingError(
+                f"Cannot pickle {self!r}: it is not the default context of "
+                "its devices"
+            )
+        return _rebuild_context, (dev,)
+
+    def __copy__(self):
+        return self
+
+    def __deepcopy__(self, memo):
+        return self
+
     cdef DPCTLSyclContextRef get_context_ref(self):
         return self._ctxt_ref
 

dpctl/_sycl_device.pyx

@@ -155,6 +155,7 @@ from ._sycl_platform cimport SyclPlatform
 
 import collections
 import functools
+import pickle
 import warnings
 
 __all__ = [
@@ -371,6 +372,26 @@ def _cached_filter_string(d : SyclDevice):
     return ":".join((br_str, dt_str, str(relId)))
 
 
+def _rebuild_device(str filter_string, str name):
+    """Unpickles a :class:`dpctl.SyclDevice`.
+
+    Raises:
+        pickle.UnpicklingError: if no device matches ``filter_string``, or
+            the device it selects is not named ``name``.
+    """
+    try:
+        dev = SyclDevice(filter_string)
+    except SyclDeviceCreationError as e:
+        raise pickle.UnpicklingError(
+            f"Device '{filter_string}' ('{name}') is not available"
+        ) from e
+    if dev.name != name:
+        raise pickle.UnpicklingError(
+            f"Device '{filter_string}' is '{dev.name}', expected '{name}'"
+        )
+    return dev
+
+
 cdef class SyclDevice(_SyclDevice):
     """ SyclDevice(arg=None)
     A Python wrapper for the ``sycl::device`` C++ class.
@@ -1652,6 +1673,29 @@ cdef class SyclDevice(_SyclDevice):
         """
         return DPCTLDevice_Hash(self._device_ref)
 
+    def __reduce__(self):
+        """Pickles a root device by its filter string and name.
+
+        The filter string's device id is relative to the devices visible
+        under ``ONEAPI_DEVICE_SELECTOR``. If the loading process sees a
+        different set of devices, the pickle may load onto another device
+        of the same name.
+
+        Raises:
+            pickle.PicklingError: if the device is not a root device.
+        """
+        try:
+            filter_string = self.filter_string
+        except ValueError as e:
+            raise pickle.PicklingError(f"Cannot pickle {self!r}: {e}") from e
+        return _rebuild_device, (filter_string, self.name)
+
+    def __copy__(self):
+        return self
+
+    def __deepcopy__(self, memo):
+        return self
+
     cdef tuple create_sub_devices_equally(self, size_t count):
         """ Returns a tuple of sub-devices partitioned from this SYCL device
         based on the ``count`` parameter.

dpctl/_sycl_queue.pyx

@@ -23,6 +23,7 @@
 """
 
 from ._backend cimport (  # noqa: E211
+    DPCTLContext_AreEq,
     DPCTLContext_Create,
     DPCTLContext_Delete,
     DPCTLDefaultSelector_Create,
@@ -103,6 +104,7 @@ from libc.stdlib cimport free, malloc
 
 import collections.abc
 import logging
+import pickle
 import struct
 import sys
 
@@ -775,6 +777,17 @@ cdef DPCTLSyclEventRef _memset_impl(
     return ERef
 
 
+def _rebuild_queue(SyclContext ctx, SyclDevice dev, tuple props):
+    """Unpickles a :class:`dpctl.SyclQueue`. Queues with default properties
+    are restored as the cached queue of ``(ctx, dev)``.
+    """
+    if not props:
+        # local import, _sycl_queue_manager cimports this module
+        from ._sycl_queue_manager import get_device_cached_queue
+        return get_device_cached_queue((ctx, dev))
+    return SyclQueue(ctx, dev, property=props)
+
+
 cdef class _SyclQueue:
     """ Barebone data owner class used by SyclQueue.
     """
@@ -855,6 +868,22 @@ cdef class SyclQueue(_SyclQueue):
       to ensure one-time use. A new named capsule can be constructed by
       using :func:`dpctl.SyclQueue._get_capsule` method.
 
+    A queue can be pickled if its device is a root device and its context is
+    the default context of that device. Unpickling creates or looks up a
+    queue rather than restoring the original one:
+
+    - a queue with default properties is restored as the cached queue of its
+      device, as returned by
+      :func:`dpctl._sycl_queue_manager.get_device_cached_queue`. Distinct
+      queues on the same device, and USM memory allocated on them, are all
+      restored onto that one queue.
+    - a queue with ``"in_order"`` or ``"enable_profiling"`` is restored as a
+      new queue on every unpickling. Only objects pickled in one
+      :func:`pickle.dumps` call that share a :class:`dpctl.SyclQueue` object
+      are restored onto one queue.
+    - the restored queue is equal to the original only if the original is the
+      cached queue.
+
     Args:
         ctx (:class:`dpctl.SyclContext`, optional): Sycl context to create
             :class:`dpctl.SyclQueue` from. If not specified, a single-device
@@ -2190,6 +2219,43 @@ cdef class SyclQueue(_SyclQueue):
         """
         return DPCTLQueue_Hash(self._queue_ref)
 
+    def __reduce__(self):
+        """Pickles the queue's context, device, and properties.
+
+        Raises:
+            pickle.PicklingError: if the queue's device is not a root device,
+                or its context is not the default context of that device.
+        """
+        cdef DPCTLSyclContextRef CRef = NULL
+        cdef bint is_default_ctx = False
+        if self._device.parent_device is not None:
+            raise pickle.PicklingError(
+                f"Cannot pickle {self!r}: its device is not a root device"
+            )
+        CRef = DPCTLDeviceMgr_GetCachedContext(self._device.get_device_ref())
+        if CRef is not NULL:
+            is_default_ctx = DPCTLContext_AreEq(
+                CRef, self._context.get_context_ref()
+            )
+            DPCTLContext_Delete(CRef)
+        if not is_default_ctx:
+            raise pickle.PicklingError(
+                f"Cannot pickle {self!r}: its context is not the default "
+                "context of its device"
+            )
+        props = []
+        if DPCTLQueue_IsInOrder(self._queue_ref):
+            props.append("in_order")
+        if DPCTLQueue_HasEnableProfiling(self._queue_ref):
+            props.append("enable_profiling")
+        return _rebuild_queue, (self._context, self._device, tuple(props))
+
+    def __copy__(self):
+        return self
+
+    def __deepcopy__(self, memo):
+        return self
+
     def _get_capsule(self):
         cdef DPCTLSyclQueueRef QRef = NULL
         QRef = DPCTLQueue_Copy(self._queue_ref)

dpctl/memory/_memory.pyx

@@ -143,18 +143,18 @@ cdef void copy_via_host(void *dest_ptr, SyclQueue dest_queue,
     DPCTLEvent_Delete(E2Ref)
 
 
-def _to_memory(unsigned char[::1] b, str usm_kind):
+def _to_memory(unsigned char[::1] b, str usm_kind, SyclQueue queue=None):
     """
-    Constructs Memory of the same size as the argument
+    Constructs Memory of the same size as the argument on ``queue``
     and copies data into it"""
     cdef _Memory res
 
     if (usm_kind == "shared"):
-        res = MemoryUSMShared(len(b))
+        res = MemoryUSMShared(len(b), queue=queue)
     elif (usm_kind == "device"):
-        res = MemoryUSMDevice(len(b))
+        res = MemoryUSMDevice(len(b), queue=queue)
     elif (usm_kind == "host"):
-        res = MemoryUSMHost(len(b))
+        res = MemoryUSMHost(len(b), queue=queue)
     else:
         raise ValueError(
             f"Unrecognized usm_kind={usm_kind} stored in the pickle"
@@ -373,7 +373,12 @@ cdef class _Memory:
         return self.tobytes()
 
     def __reduce__(self):
-        return _to_memory, (self.copy_to_host(), self.get_usm_type())
+        """Pickles a host copy of the data and the queue. See
+        :class:`dpctl.SyclQueue` for how the queue is restored.
+        """
+        return (
+            _to_memory, (self.copy_to_host(), self.get_usm_type(), self.queue)
+        )
 
     @property
     def __sycl_usm_array_interface__(self):

dpctl/tests/helper/__init__.py

@@ -19,6 +19,7 @@
 from ._helper import (
     create_invalid_capsule,
     get_queue_or_skip,
+    get_sub_device_or_skip,
     has_cpu,
     has_gpu,
     has_sycl_platforms,
@@ -30,4 +31,5 @@
     "has_gpu",
     "has_sycl_platforms",
     "get_queue_or_skip",
+    "get_sub_device_or_skip",
 ]

dpctl/tests/helper/_helper.py

@@ -49,3 +49,13 @@ def get_queue_or_skip(args=()):
     except dpctl.SyclQueueCreationError:
         pytest.skip(f"Queue could not be created from {args}")
     return q
+
+
+def get_sub_device_or_skip():
+    for dev in dpctl.get_devices():
+        for partition in ("next_partitionable", [1, 1]):
+            try:
+                return dev.create_sub_devices(partition=partition)[0]
+            except dpctl.SyclSubDeviceCreationError:
+                pass
+    pytest.skip("No device can be partitioned")

dpctl/tests/test_sycl_context.py

@@ -16,11 +16,14 @@
 
 """Defines unit test cases for the :class:`dpctl.SyclContext` class."""
 
+import copy
+import pickle
+
 import pytest
 
 import dpctl
 
-from .helper import create_invalid_capsule
+from .helper import create_invalid_capsule, get_sub_device_or_skip
 
 list_of_valid_filter_selectors = [
     "opencl",
@@ -341,3 +344,42 @@ def test_context_atomic_fence_scope_capabilities(valid_filter):
     assert all(isinstance(m, dpctl.memory_scope) for m in caps)
     for d in ctx.get_devices():
         assert set(caps).issubset(set(d.atomic_fence_scope_capabilities))
+
+
+def test_context_pickling():
+    devs = dpctl.get_devices()
+    if not devs:
+        pytest.skip("No devices available")
+    for dev in devs:
+        ctx = dpctl.SyclContext(dev)
+        ctx2 = pickle.loads(pickle.dumps(ctx))
+        assert ctx2 == ctx
+        assert dev in ctx2.get_devices()
+
+
+def test_context_pickling_non_default():
+    try:
+        dev = dpctl.SyclDevice()
+    except dpctl.SyclDeviceCreationError:
+        pytest.skip("No default device available")
+    ctx = dpctl.SyclContext([dev])
+    if ctx == dpctl.SyclContext(dev):
+        pytest.skip("Context created from device list is the default context")
+    with pytest.raises(pickle.PicklingError):
+        pickle.dumps(ctx)
+
+
+def test_context_pickling_sub_device():
+    sub_dev = get_sub_device_or_skip()
+    ctx = dpctl.SyclContext(sub_dev)
+    with pytest.raises(pickle.PicklingError):
+        pickle.dumps(ctx)
+
+
+def test_context_copy():
+    try:
+        ctx = dpctl.SyclContext()
+    except dpctl.SyclContextCreationError:
+        pytest.skip("Default context could not be created")
+    assert copy.copy(ctx) is ctx
+    assert copy.deepcopy(ctx) is ctx