(attempt at) pickle Buffers

Author: matthiasdiener

pyopencl/__init__.py

@@ -2404,4 +2404,80 @@ def fsvm_empty_like(ctx, ary, alignment=None):
     _KERNEL_ARG_CLASSES = (*_KERNEL_ARG_CLASSES, SVM)
 
 
+# {{{ pickling support
+
+import threading
+from contextlib import contextmanager
+
+
+_QUEUE_FOR_PICKLING_TLS = threading.local()
+
+
+@contextmanager
+def queue_for_pickling(queue):
+    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*."""
+    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
+    try:
+        yield None
+    finally:
+        _QUEUE_FOR_PICKLING_TLS.queue = None
+
+
+def _getstate_buffer(self):
+    import pyopencl as cl
+    state = {}
+    state["size"] = self.size
+    state["flags"] = self.flags
+
+    try:
+        queue = _QUEUE_FOR_PICKLING_TLS.queue
+    except AttributeError:
+        queue = None
+
+    if queue is None:
+        raise RuntimeError("CL Buffer 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_buffer(self, state):
+    try:
+        queue = _QUEUE_FOR_PICKLING_TLS.queue
+    except AttributeError:
+        queue = None
+
+    if queue is None:
+        raise RuntimeError("CL Buffer instances can only be unpickled while "
+                            "queue_for_pickling is active.")
+
+    size = state["size"]
+    flags = state["flags"]
+
+    import pyopencl as cl
+
+    a = state["_pickle_data"]
+    Buffer.__init__(self, queue.context, flags | cl.mem_flags.COPY_HOST_PTR, size, a)
+
+
+Buffer.__getstate__ = _getstate_buffer
+Buffer.__setstate__ = _setstate_buffer
+
+# }}}
+
 # vim: foldmethod=marker

test/test_array.py

@@ -2453,6 +2453,26 @@ def test_array_pickling(ctx_factory):
     # }}}
 
 
+def test_buffer_pickling(ctx_factory):
+    context = ctx_factory()
+    queue = cl.CommandQueue(context)
+
+    a = np.array([1, 2, 3, 4, 5]).astype(np.float32)
+    a_gpu = cl.Buffer(context, cl.mem_flags.READ_WRITE, a.nbytes)
+    cl.enqueue_copy(queue, a_gpu, a)
+
+    import pickle
+
+    with pytest.raises(cl.RuntimeError):
+        pickle.dumps(a_gpu)
+
+    with cl.queue_for_pickling(queue):
+        a_gpu_pickled = pickle.loads(pickle.dumps(a_gpu))
+
+    a_new = np.empty_like(a)
+    cl.enqueue_copy(queue, a_new, a_gpu_pickled)
+    assert np.all(a_new == a)
+
 # }}}