pickle SVM, centralize queue_for_pickling, use Buffer/SVM pickling for Arrays

Author: matthiasdiener

pyopencl/__init__.py

@@ -2414,9 +2414,10 @@ def fsvm_empty_like(ctx, ary, alignment=None):
 
 
 @contextmanager
-def queue_for_pickling(queue):
+def queue_for_pickling(queue, alloc=None):
     r"""A context manager that, for the current thread, sets the command queue
-    to be used for pickling and unpickling :class:`Buffer`\ s to *queue*."""
+    to be used for pickling and unpickling :class:`Array`\ s and :class:`Buffer`\ s
+    to *queue*."""
     try:
         existing_pickle_queue = _QUEUE_FOR_PICKLING_TLS.queue
     except AttributeError:
@@ -2427,10 +2428,12 @@ def queue_for_pickling(queue):
                 "inside the context of its own invocation.")
 
     _QUEUE_FOR_PICKLING_TLS.queue = queue
+    _QUEUE_FOR_PICKLING_TLS.alloc = alloc
     try:
         yield None
     finally:
         _QUEUE_FOR_PICKLING_TLS.queue = None
+        _QUEUE_FOR_PICKLING_TLS.alloc = None
 
 
 def _getstate_buffer(self):
@@ -2478,6 +2481,51 @@ def _setstate_buffer(self, state):
 Buffer.__getstate__ = _getstate_buffer
 Buffer.__setstate__ = _setstate_buffer
 
+if get_cl_header_version() >= (2, 0):
+    def _getstate_svm(self):
+        import pyopencl as cl
+
+        state = {}
+        state["size"] = self.size
+
+        try:
+            queue = _QUEUE_FOR_PICKLING_TLS.queue
+        except AttributeError:
+            queue = None
+
+        if queue is None:
+            raise RuntimeError(f"{self.__class__.__name__} instances can only be "
+                               "pickled while queue_for_pickling is active.")
+
+        a = bytearray(self.size)
+        cl.enqueue_copy(queue, a, self)
+
+        state["_pickle_data"] = a
+
+        return state
+
+    def _setstate_svm(self, state):
+        import pyopencl as cl
+
+        try:
+            queue = _QUEUE_FOR_PICKLING_TLS.queue
+        except AttributeError:
+            queue = None
+
+        if queue is None:
+            raise RuntimeError(f"{self.__class__.__name__} instances can only be "
+                               "unpickled while queue_for_pickling is active.")
+
+        size = state["size"]
+
+        a = state["_pickle_data"]
+        SVMAllocation.__init__(self, queue.context, size, alignment=0, flags=0,
+                               queue=queue)
+        cl.enqueue_copy(queue, self, a)
+
+    SVMAllocation.__getstate__ = _getstate_svm
+    SVMAllocation.__setstate__ = _setstate_svm
+
 # }}}
 
 # vim: foldmethod=marker

pyopencl/array.py

@@ -342,39 +342,6 @@ class _copy_queue:  # noqa: N801
 _NOT_PRESENT = object()
 
 
-# {{{ pickling support
-
-import threading
-from contextlib import contextmanager
-
-
-_QUEUE_FOR_PICKLING_TLS = threading.local()
-
-
-@contextmanager
-def queue_for_pickling(queue, alloc=None):
-    r"""A context manager that, for the current thread, sets the command queue
-    to be used for pickling and unpickling :class:`Array`\ s to *queue*."""
-    try:
-        existing_pickle_queue = _QUEUE_FOR_PICKLING_TLS.queue
-    except AttributeError:
-        existing_pickle_queue = None
-
-    if existing_pickle_queue is not None:
-        raise RuntimeError("queue_for_pickling should not be called "
-                "inside the context of its own invocation.")
-
-    _QUEUE_FOR_PICKLING_TLS.queue = queue
-    _QUEUE_FOR_PICKLING_TLS.alloc = alloc
-    try:
-        yield None
-    finally:
-        _QUEUE_FOR_PICKLING_TLS.queue = None
-        _QUEUE_FOR_PICKLING_TLS.alloc = None
-
-# }}}
-
-
 class Array:
     """A :class:`numpy.ndarray` work-alike that stores its data and performs
     its computations on the compute device. :attr:`shape` and :attr:`dtype` work
@@ -742,36 +709,33 @@ def __init__(
 
     def __getstate__(self):
         try:
-            queue = _QUEUE_FOR_PICKLING_TLS.queue
+            queue = cl._QUEUE_FOR_PICKLING_TLS.queue
         except AttributeError:
             queue = None
 
         if queue is None:
             raise RuntimeError("CL Array instances can only be pickled while "
-                               "queue_for_pickling is active.")
+                               "cl.queue_for_pickling is active.")
 
         state = self.__dict__.copy()
 
         del state["allocator"]
         del state["context"]
         del state["events"]
         del state["queue"]
-        del state["base_data"]
-        state["data"] = self.get(queue=queue)
-
         return state
 
     def __setstate__(self, state):
         try:
-            queue = _QUEUE_FOR_PICKLING_TLS.queue
-            alloc = _QUEUE_FOR_PICKLING_TLS.alloc
+            queue = cl._QUEUE_FOR_PICKLING_TLS.queue
+            alloc = cl._QUEUE_FOR_PICKLING_TLS.alloc
         except AttributeError:
             queue = None
             alloc = None
 
         if queue is None:
-            raise RuntimeError("CL Array instances can only be pickled while "
-                               "queue_for_pickling is active.")
+            raise RuntimeError("CL Array instances can only be unpickled while "
+                               "cl.queue_for_pickling is active.")
 
         self.__dict__.update(state)
 
@@ -780,20 +744,6 @@ def __setstate__(self, state):
         self.events = []
         self.queue = queue
 
-        if self.allocator is None:
-            self.base_data = cl.Buffer(self.context, cl.mem_flags.READ_WRITE,
-                                       self.nbytes)
-        else:
-            self.base_data = self.allocator(self.nbytes)
-
-        ary = state["data"]
-
-        # Mimics the stride update in _get() below
-        if ary.strides != self.strides:
-            ary = _as_strided(ary, strides=self.strides)
-
-        self.set(ary, queue=queue)
-
     # }}}
 
     @property

test/test_array.py

@@ -2415,7 +2415,7 @@ def test_array_pickling(ctx_factory):
     with pytest.raises(RuntimeError):
         pickle.dumps(a_gpu)
 
-    with cl_array.queue_for_pickling(queue):
+    with cl.queue_for_pickling(queue):
         a_gpu_pickled = pickle.loads(pickle.dumps(a_gpu))
     assert np.all(a_gpu_pickled.get() == a)
 
@@ -2424,7 +2424,7 @@ def test_array_pickling(ctx_factory):
     a_gpu_tagged = TaggableCLArray(queue, a.shape, a.dtype, tags={"foo", "bar"})
     a_gpu_tagged.set(a)
 
-    with cl_array.queue_for_pickling(queue):
+    with cl.queue_for_pickling(queue):
         a_gpu_tagged_pickled = pickle.loads(pickle.dumps(a_gpu_tagged))
 
     assert np.all(a_gpu_tagged_pickled.get() == a)
@@ -2437,14 +2437,15 @@ def test_array_pickling(ctx_factory):
     from pyopencl.characterize import has_coarse_grain_buffer_svm
 
     if has_coarse_grain_buffer_svm(queue.device):
-        from pyopencl.tools import SVMAllocator, SVMPool
+        from pyopencl.tools import SVMAllocator
 
         alloc = SVMAllocator(context, alignment=0, queue=queue)
-        alloc = SVMPool(alloc)
+        # FIXME: SVMPool is not picklable
+        # alloc = SVMPool(alloc)
 
         a_dev = cl_array.to_device(queue, a, allocator=alloc)
 
-        with cl_array.queue_for_pickling(queue, alloc):
+        with cl.queue_for_pickling(queue, alloc):
             a_dev_pickled = pickle.loads(pickle.dumps(a_dev))
 
         assert np.all(a_dev_pickled.get() == a)