diff --git a/CHANGELOG.md b/CHANGELOG.md index 0aa6be7eab..2f1f261fcb 100644 --- a/CHANGELOG.md +++ b/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 diff --git a/dpctl/_sycl_context.pyx b/dpctl/_sycl_context.pyx index f48b6d34c3..2160202150 100644 --- a/dpctl/_sycl_context.pyx +++ b/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 diff --git a/dpctl/_sycl_device.pyx b/dpctl/_sycl_device.pyx index 624575a5d8..6cd14f7f02 100644 --- a/dpctl/_sycl_device.pyx +++ b/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. diff --git a/dpctl/_sycl_queue.pyx b/dpctl/_sycl_queue.pyx index 3fbf9c7bf6..b60b943984 100644 --- a/dpctl/_sycl_queue.pyx +++ b/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) diff --git a/dpctl/memory/_memory.pyx b/dpctl/memory/_memory.pyx index 6b2050bb9e..b53586752b 100644 --- a/dpctl/memory/_memory.pyx +++ b/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): diff --git a/dpctl/tests/helper/__init__.py b/dpctl/tests/helper/__init__.py index be5f8d2e33..bb80f1d7bb 100644 --- a/dpctl/tests/helper/__init__.py +++ b/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", ] diff --git a/dpctl/tests/helper/_helper.py b/dpctl/tests/helper/_helper.py index 5857b23195..23630f0d56 100644 --- a/dpctl/tests/helper/_helper.py +++ b/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") diff --git a/dpctl/tests/test_sycl_context.py b/dpctl/tests/test_sycl_context.py index 058bb7ad7f..ddf272ebec 100644 --- a/dpctl/tests/test_sycl_context.py +++ b/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 diff --git a/dpctl/tests/test_sycl_device.py b/dpctl/tests/test_sycl_device.py index dd1a579272..c3441966f7 100644 --- a/dpctl/tests/test_sycl_device.py +++ b/dpctl/tests/test_sycl_device.py @@ -16,12 +16,15 @@ """Defines unit test cases for the SyclDevice class.""" +import copy +import pickle + import pytest import dpctl from dpctl import SyclDeviceCreationError -from .helper import get_queue_or_skip +from .helper import get_queue_or_skip, get_sub_device_or_skip def test_standard_selectors(device_selector, check): @@ -455,3 +458,45 @@ def test_peer_access_value_keyword_validation(): bad_value = "wrong" with pytest.raises(ValueError): dev0.can_access_peer(dev1, value=bad_value) + + +def test_device_pickling(): + devs = dpctl.get_devices() + if not devs: + pytest.skip("No devices available") + for dev in devs: + dev2 = pickle.loads(pickle.dumps(dev)) + assert dev2 == dev + assert dev2.filter_string == dev.filter_string + + +def test_device_pickling_sub_device(): + sub_dev = get_sub_device_or_skip() + with pytest.raises(pickle.PicklingError): + pickle.dumps(sub_dev) + + +def test_device_unpickling_name_mismatch(): + from dpctl._sycl_device import _rebuild_device + + try: + dev = dpctl.SyclDevice() + except SyclDeviceCreationError: + pytest.skip("No default device available") + with pytest.raises(pickle.UnpicklingError): + _rebuild_device(dev.filter_string, dev.name + " (other)") + + +def test_device_unpickling_missing_device(): + from dpctl._sycl_device import _rebuild_device + + with pytest.raises(pickle.UnpicklingError, match="is not available"): + _rebuild_device("opencl:gpu:999", "Missing device") + + +def test_device_copy(): + sub_dev = get_sub_device_or_skip() + root_dev = sub_dev.parent_device + for dev in (root_dev, sub_dev): + assert copy.copy(dev) is dev + assert copy.deepcopy(dev) is dev diff --git a/dpctl/tests/test_sycl_queue.py b/dpctl/tests/test_sycl_queue.py index af0fe71dcd..4d9e2fcabe 100644 --- a/dpctl/tests/test_sycl_queue.py +++ b/dpctl/tests/test_sycl_queue.py @@ -16,14 +16,22 @@ """Defines unit test cases for the SyclQueue class.""" +import copy import ctypes +import pickle +import subprocess import sys import pytest import dpctl +from dpctl._sycl_queue_manager import get_device_cached_queue -from .helper import create_invalid_capsule +from .helper import ( + create_invalid_capsule, + get_queue_or_skip, + get_sub_device_or_skip, +) def test_standard_selectors(device_selector, check): @@ -402,3 +410,79 @@ def test_cython_api(dpctl_cython_extension): except dpctl.SyclDeviceCreationError: pytest.skip("Default-construction of SyclDevice failed") assert q.sycl_device == d + + +@pytest.mark.parametrize( + "props", + [ + (), + ("in_order",), + ("enable_profiling",), + ("in_order", "enable_profiling"), + ], +) +def test_queue_pickling(props): + try: + q = dpctl.SyclQueue(property=props or "default") + except dpctl.SyclQueueCreationError: + pytest.skip("Default-constructed queue could not be created") + q2 = pickle.loads(pickle.dumps(q)) + assert isinstance(q2, dpctl.SyclQueue) + assert q2.sycl_device == q.sycl_device + assert q2.sycl_context == q.sycl_context + assert q2.is_in_order == ("in_order" in props) + assert q2.has_enable_profiling == ("enable_profiling" in props) + assert (q2 == get_device_cached_queue(q.sycl_device)) == (not props) + + +def test_queue_unpickling_cached_queue(): + dev = get_queue_or_skip().sycl_device + cached = get_device_cached_queue(dev) + q0, q1 = dpctl.SyclQueue(dev), dpctl.SyclQueue(dev) + assert q0 != q1 + assert pickle.loads(pickle.dumps(q0)) == cached + assert pickle.loads(pickle.dumps(q1)) == cached + q2 = dpctl.SyclQueue(dev, property="in_order") + assert pickle.loads(pickle.dumps(q2)) != pickle.loads(pickle.dumps(q2)) + + +def test_queue_pickling_non_default_context(): + q = get_queue_or_skip() + dev = q.sycl_device + ctx = dpctl.SyclContext([dev]) + if ctx == q.sycl_context: + pytest.skip("Context created from device list is the default context") + with pytest.raises(pickle.PicklingError, match="SyclQueue.*context"): + pickle.dumps(dpctl.SyclQueue(ctx, dev)) + + +def test_queue_pickling_sub_device(): + sub_dev = get_sub_device_or_skip() + q = dpctl.SyclQueue(sub_dev) + with pytest.raises(pickle.PicklingError, match="not a root device"): + pickle.dumps(q) + + +def test_queue_unpickling_in_subprocess(): + q = get_queue_or_skip() + q = dpctl.SyclQueue(q.sycl_device, property="in_order") + script = ( + "import pickle, sys\n" + "q = pickle.loads(sys.stdin.buffer.read())\n" + "print(q.sycl_device.filter_string, q.is_in_order)\n" + ) + res = subprocess.run( + [sys.executable, "-c", script], + input=pickle.dumps(q), + capture_output=True, + check=True, + ) + assert res.stdout.decode().split() == [q.sycl_device.filter_string, "True"] + + +def test_queue_copy(): + q = get_queue_or_skip() + assert copy.copy(q) is q + assert copy.deepcopy(q) is q + d = copy.deepcopy({"queue": q}) + assert d["queue"] is q diff --git a/dpctl/tests/test_sycl_usm.py b/dpctl/tests/test_sycl_usm.py index 398a07b216..59417e7baf 100644 --- a/dpctl/tests/test_sycl_usm.py +++ b/dpctl/tests/test_sycl_usm.py @@ -179,19 +179,75 @@ def test_pickling(memory_ctor): ), "Pickling/unpickling should be changing pointer" -def test_pickling_reconstructor_invalid_type(memory_ctor): +def test_pickling_reconstructor_invalid_type(): + from dpctl.memory._memory import _to_memory + + with pytest.raises(ValueError): + _to_memory(bytearray(8), "unknown") + + +def test_pickling_preserves_queue(memory_ctor): + import pickle + + aspect = f"usm_{expected_usm_type_str(memory_ctor)}_allocations" + devs = [ + d for d in dpctl.get_devices() if getattr(d, f"has_aspect_{aspect}") + ] + if not devs: + pytest.skip(f"No SYCL devices with {aspect} available") + for dev in devs: + q = dpctl.SyclQueue(dev, property="in_order") + mobj = memory_ctor(64, queue=q) + mobj2 = pickle.loads(pickle.dumps(mobj)) + q2 = mobj2.sycl_queue + assert q2.sycl_device == dev + assert q2.sycl_context == q.sycl_context + assert q2.is_in_order + + +def test_pickling_restores_onto_cached_queue(memory_ctor): import pickle + from dpctl._sycl_queue_manager import get_device_cached_queue + try: - mobj = memory_ctor(1024, alignment=64) + dev = dpctl.SyclDevice() except dpctl.SyclDeviceCreationError: pytest.skip("No SYCL devices available") - good_pickle_bytes = pickle.dumps(mobj) - usm_types = expected_usm_type_str(memory_ctor).encode("utf-8") - i = good_pickle_bytes.rfind(usm_types) - bad_pickle_bytes = good_pickle_bytes[:i] + b"u" + good_pickle_bytes[i + 1 :] - with pytest.raises(ValueError): - pickle.loads(bad_pickle_bytes) + m0 = memory_ctor(64, queue=dpctl.SyclQueue(dev)) + m1 = memory_ctor(64, queue=dpctl.SyclQueue(dev)) + m2 = pickle.loads(pickle.dumps(m0)) + m3 = pickle.loads(pickle.dumps(m1)) + assert m2.sycl_queue == m3.sycl_queue == get_device_cached_queue(dev) + + +def test_pickling_non_default_context(memory_ctor): + import pickle + + try: + dev = dpctl.SyclDevice() + except dpctl.SyclDeviceCreationError: + pytest.skip("No SYCL devices available") + ctx = dpctl.SyclContext([dev]) + if ctx == dpctl.SyclContext(dev): + pytest.skip("Context created from device list is the default context") + mobj = memory_ctor(64, queue=dpctl.SyclQueue(ctx, dev)) + with pytest.raises(pickle.PicklingError, match="SyclQueue"): + pickle.dumps(mobj) + + +def test_pickling_reconstructor_without_queue(memory_ctor): + from dpctl.memory._memory import _to_memory + + try: + dev = dpctl.SyclDevice() + except dpctl.SyclDeviceCreationError: + pytest.skip("No SYCL devices available") + usm_type = expected_usm_type_str(memory_ctor) + mobj = _to_memory(bytearray(b"abcdefgh"), usm_type) + assert type(mobj) is memory_ctor + assert mobj.sycl_device == dev + assert mobj.tobytes() == b"abcdefgh" @pytest.fixture(params=[MemoryUSMShared, MemoryUSMDevice, MemoryUSMHost])