From 7a4ad7c825b5ec82febef7a2f7f9eb38aa5c5168 Mon Sep 17 00:00:00 2001 From: giadarol Date: Thu, 15 Oct 2026 21:17:50 +0200 Subject: [PATCH] Restore CPU kernel dispatch after unpickling contexts --- tests/test_kernel.py | 11 ++++++++++- xobjects/context_cpu.py | 2 ++ 2 files changed, 12 insertions(+), 1 deletion(-) diff --git a/tests/test_kernel.py b/tests/test_kernel.py index 5f6cebd..5f76dd9 100644 --- a/tests/test_kernel.py +++ b/tests/test_kernel.py @@ -4,6 +4,7 @@ # ########################################### # import cffi +import pickle import pytest import sysconfig from pathlib import Path @@ -14,7 +15,8 @@ from xobjects.test_helpers import for_all_test_contexts, requires_context -def test_kernel_cpu(): +@pytest.mark.parametrize('pickle_context', [False, True]) +def test_kernel_cpu(pickle_context): ctx = xo.ContextCpu() src_code = r""" double my_mul(const int n, const double* x1, @@ -45,6 +47,13 @@ def test_kernel_cpu(): assert y == 285.0 + if pickle_context: + restored = pickle.loads(pickle.dumps(ctx)) + assert not restored.kernels + assert ctx.kernels.my_mul(n=len(a1), x1=a1, x2=a2) == 285.0 + restored.add_kernels(sources=[src_code], kernels=kernel_descriptions) + assert restored.kernels.my_mul(n=len(a1), x1=a1, x2=a2) == 285.0 + @for_all_test_contexts def test_kernel_pointer_none_maps_to_null(test_context): diff --git a/xobjects/context_cpu.py b/xobjects/context_cpu.py index fd67bc4..26184dd 100644 --- a/xobjects/context_cpu.py +++ b/xobjects/context_cpu.py @@ -96,6 +96,7 @@ def kernel_compilation_help_message(): from .context import ( Kernel, + KernelDict, ModuleNotAvailable, SourceType, XBuffer, @@ -821,6 +822,7 @@ def __getstate__(self): def __setstate__(self, state): self.__dict__.update(state) + self._kernels = KernelDict(self._kernels) self._buffers = weakref.WeakSet()