Skip to content

Commit

Permalink
misc fixes
Browse files Browse the repository at this point in the history
  • Loading branch information
matthiasdiener committed Sep 18, 2024
1 parent 25ebc7b commit 3713741
Show file tree
Hide file tree
Showing 2 changed files with 11 additions and 8 deletions.
1 change: 1 addition & 0 deletions pyopencl/array.py
Original file line number Diff line number Diff line change
Expand Up @@ -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

# }}}

Expand Down
18 changes: 10 additions & 8 deletions test/test_array.py
Original file line number Diff line number Diff line change
Expand Up @@ -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

# }}}

Expand Down

0 comments on commit 3713741

Please sign in to comment.