support pickling allocator, test SVM

Author: matthiasdiener

pyopencl/array.py

@@ -352,7 +352,7 @@ class _copy_queue:  # noqa: N801
 
 
 @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:`Array`\ s to *queue*."""
     try:
@@ -365,6 +365,7 @@ 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:
@@ -749,6 +750,8 @@ def __getstate__(self):
                                "queue_for_pickling is active.")
 
         state = self.__dict__.copy()
+
+        del state["allocator"]
         del state["context"]
         del state["events"]
         del state["queue"]
@@ -760,14 +763,18 @@ def __getstate__(self):
     def __setstate__(self, state):
         try:
             queue = _QUEUE_FOR_PICKLING_TLS.queue
+            alloc = _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.")
 
         self.__dict__.update(state)
+
+        self.allocator = alloc
         self.context = queue.context
         self.events = []
         self.queue = queue

test/test_array.py

@@ -2419,6 +2419,8 @@ def test_array_pickling(ctx_factory):
         a_gpu_pickled = pickle.loads(pickle.dumps(a_gpu))
     assert np.all(a_gpu_pickled.get() == a)
 
+    # {{{ subclass test
+
     a_gpu_tagged = TaggableCLArray(queue, a.shape, a.dtype, tags={"foo", "bar"})
     a_gpu_tagged.set(a)
 
@@ -2428,6 +2430,25 @@ def test_array_pickling(ctx_factory):
     assert np.all(a_gpu_tagged_pickled.get() == a)
     assert a_gpu_tagged_pickled.tags == a_gpu_tagged.tags
 
+    # }}}
+
+    # {{{ SVM test
+
+    from pyopencl.tools import SVMAllocator, SVMPool
+
+    alloc = SVMAllocator(context, alignment=0, queue=queue)
+    alloc = SVMPool(alloc)
+
+    a_dev = cl_array.to_device(queue, a, allocator=alloc)
+
+    with cl_array.queue_for_pickling(queue, alloc):
+        a_dev_pickled = pickle.loads(pickle.dumps(a_dev))
+
+    assert np.all(a_dev_pickled.get() == a)
+    assert a_dev_pickled.allocator is alloc
+
+    # }}}
+
 
 # }}}