Improve pickling support

CHANGELOG.md

@@ -12,11 +12,15 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 * Added `SyclQueue_DeferRelease` C-API function, which defers a release callback until work submitted to a queue completes, and through which `dpctl::utils::keep_args_alive` keeps its arguments alive [gh-2359](https://github.com/IntelPython/dpctl/pull/2359)
 * Added `dpctl::utils::defer_release` to `dpctl4pybind11.hpp`, which calls a release callback once the given events have completed [gh-2359](https://github.com/IntelPython/dpctl/pull/2359)
 
+* `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
 * Object lifetime during offload is now managed by a list of releases deferred until the events that gate them have completed, rather than by a `host_task`. `dpctl` keeps one thread for looking over the list and running the releases that have come due, started by the first release deferred and again in the child of a `fork`, which takes the GIL to do so and lets a release call into the interpreter and free USM allocations [gh-2359](https://github.com/IntelPython/dpctl/pull/2359)
 * `import_dpctl()` now returns 0 on success and -1 with a Python exception set on failure, and `dpctl4pybind11.hpp` raises that exception rather than failing on the first C-API call [gh-2359](https://github.com/IntelPython/dpctl/pull/2359)
 * The event returned by `dpctl.SyclQueue.keep_args_alive`, `dpctl::utils::keep_args_alive` and `dpctl.SyclQueue._submit_keep_args_alive` is now for an empty kernel that gates the deferred release rather than for a `host_task` that performs it, so its completion means that the objects are no longer in use rather than that they have been released [gh-2359](https://github.com/IntelPython/dpctl/pull/2359)
 
+* 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
 * Deprecated `dpctl.SyclQueue._submit_keep_args_alive` in favor of `dpctl.SyclQueue.keep_args_alive`, which it now calls and which behaves identically [gh-2359](https://github.com/IntelPython/dpctl/pull/2359)
 * Deprecated the order manager's `add_event_pair`, `host_task_events` and `num_host_task_events`, as `host_task` is no longer used for managing object lifetimes, in favor of `add_event`, `add_cleanup_event`, `cleanup_events` and `num_cleanup_events`; they do not yet emit a `DeprecationWarning` [gh-2359](https://github.com/IntelPython/dpctl/pull/2359)

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,
@@ -105,6 +106,7 @@ import atexit
 import collections.abc
 import logging
 import os
+import pickle
 import struct
 import sys
 import warnings
@@ -790,6 +792,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.
     """
@@ -870,6 +883,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
@@ -2273,6 +2302,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"
@@ -394,7 +394,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