diff --git a/pyopencl/array.py b/pyopencl/array.py index d4c78578..153dfeb5 100644 --- a/pyopencl/array.py +++ b/pyopencl/array.py @@ -370,6 +370,7 @@ def queue_for_pickling(queue, alloc=None): yield None finally: _QUEUE_FOR_PICKLING_TLS.queue = None + _QUEUE_FOR_PICKLING_TLS.alloc = None # }}} diff --git a/test/test_array.py b/test/test_array.py index 44c23bc1..01ec811c 100644 --- a/test/test_array.py +++ b/test/test_array.py @@ -2433,19 +2433,21 @@ def test_array_pickling(ctx_factory): # }}} # {{{ SVM test + from pyopencl.characterize import has_coarse_grain_buffer_svm - from pyopencl.tools import SVMAllocator, SVMPool + if has_coarse_grain_buffer_svm(queue.device): + from pyopencl.tools import SVMAllocator, SVMPool - alloc = SVMAllocator(context, alignment=0, queue=queue) - alloc = SVMPool(alloc) + alloc = SVMAllocator(context, alignment=0, queue=queue) + alloc = SVMPool(alloc) - a_dev = cl_array.to_device(queue, a, allocator=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)) + 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 + assert np.all(a_dev_pickled.get() == a) + assert a_dev_pickled.allocator is alloc # }}}