From 4febc90988f57d5dfd7e12d8c8156d22d52ddd7e Mon Sep 17 00:00:00 2001 From: brandon-b-miller Date: Tue, 27 May 2025 07:04:42 -0700 Subject: [PATCH 1/6] implement __reduce__ --- cuda_core/cuda/core/experimental/_module.py | 14 ++++++++++++++ 1 file changed, 14 insertions(+) diff --git a/cuda_core/cuda/core/experimental/_module.py b/cuda_core/cuda/core/experimental/_module.py index 021165c5c..9731e84c8 100644 --- a/cuda_core/cuda/core/experimental/_module.py +++ b/cuda_core/cuda/core/experimental/_module.py @@ -302,6 +302,9 @@ def _init(cls, module, code_type, *, symbol_mapping: Optional[dict] = None): return self + def __reduce__(self): + return _rebuild_objectcode_instance(self._module, self._code_type, self._sym_map) + @staticmethod def from_cubin(module: Union[bytes, str], *, symbol_mapping: Optional[dict] = None) -> "ObjectCode": """Create an :class:`ObjectCode` instance from an existing cubin. @@ -397,3 +400,14 @@ def handle(self): handle, call ``int(ObjectCode.handle)``. """ return self._handle + + +def _rebuild_objectcode_instance( + module: Union[bytes, str], + code_type: str, + symbol_mapping: Optional[dict] = None, +) -> ObjectCode: + """Rebuild an ObjectCode instance from its serialized form.""" + obj = ObjectCode._init(module, code_type, symbol_mapping=symbol_mapping) + obj._handle = None # handle will be lazily loaded + return obj From f074acecd059165c9253b6f249ff942f1bb6e481 Mon Sep 17 00:00:00 2001 From: brandon-b-miller Date: Tue, 27 May 2025 07:43:33 -0700 Subject: [PATCH 2/6] updates --- cuda_core/cuda/core/experimental/_module.py | 4 +--- cuda_core/tests/test_module.py | 12 ++++++++++++ 2 files changed, 13 insertions(+), 3 deletions(-) diff --git a/cuda_core/cuda/core/experimental/_module.py b/cuda_core/cuda/core/experimental/_module.py index 9731e84c8..8c7e57a27 100644 --- a/cuda_core/cuda/core/experimental/_module.py +++ b/cuda_core/cuda/core/experimental/_module.py @@ -408,6 +408,4 @@ def _rebuild_objectcode_instance( symbol_mapping: Optional[dict] = None, ) -> ObjectCode: """Rebuild an ObjectCode instance from its serialized form.""" - obj = ObjectCode._init(module, code_type, symbol_mapping=symbol_mapping) - obj._handle = None # handle will be lazily loaded - return obj + return ObjectCode._init(module, code_type, symbol_mapping=symbol_mapping) diff --git a/cuda_core/tests/test_module.py b/cuda_core/tests/test_module.py index cc8620906..4e022f234 100644 --- a/cuda_core/tests/test_module.py +++ b/cuda_core/tests/test_module.py @@ -2,6 +2,7 @@ # SPDX-License-Identifier: Apache-2.0 import ctypes +import pickle # nosec B403, B301 import warnings import pytest @@ -245,3 +246,14 @@ def test_num_args_error_handling(deinit_all_contexts_function, cuda12_prerequisi with pytest.raises(CUDAError): # assignment resolves linter error "B018: useless expression" _ = krn.num_arguments + + +def test_module_serialization_roundtrip(get_saxpy_kernel): + _, objcode = get_saxpy_kernel + result = pickle.loads(pickle.dumps(objcode)) # nosec B403, B301 + + assert isinstance(result, ObjectCode) + assert result._handle is None + assert objcode.code == result.code + assert objcode._sym_map == result._sym_map + assert objcode._code_type == result._code_type From a15c43494eb3ded51e52a04de08a5c26ef11ca28 Mon Sep 17 00:00:00 2001 From: brandon-b-miller Date: Wed, 28 May 2025 06:56:50 -0700 Subject: [PATCH 3/6] inline function --- cuda_core/cuda/core/experimental/_module.py | 11 +---------- 1 file changed, 1 insertion(+), 10 deletions(-) diff --git a/cuda_core/cuda/core/experimental/_module.py b/cuda_core/cuda/core/experimental/_module.py index 8c7e57a27..0138899e3 100644 --- a/cuda_core/cuda/core/experimental/_module.py +++ b/cuda_core/cuda/core/experimental/_module.py @@ -303,7 +303,7 @@ def _init(cls, module, code_type, *, symbol_mapping: Optional[dict] = None): return self def __reduce__(self): - return _rebuild_objectcode_instance(self._module, self._code_type, self._sym_map) + return ObjectCode._init(self._module, self._code_type, symbol_mapping=self._sym_map) @staticmethod def from_cubin(module: Union[bytes, str], *, symbol_mapping: Optional[dict] = None) -> "ObjectCode": @@ -400,12 +400,3 @@ def handle(self): handle, call ``int(ObjectCode.handle)``. """ return self._handle - - -def _rebuild_objectcode_instance( - module: Union[bytes, str], - code_type: str, - symbol_mapping: Optional[dict] = None, -) -> ObjectCode: - """Rebuild an ObjectCode instance from its serialized form.""" - return ObjectCode._init(module, code_type, symbol_mapping=symbol_mapping) From d35346f1703169f7f653befc3a8270129f2d6c4d Mon Sep 17 00:00:00 2001 From: brandon-b-miller Date: Wed, 28 May 2025 07:44:26 -0700 Subject: [PATCH 4/6] fix reduce --- cuda_core/cuda/core/experimental/_module.py | 6 +++++- 1 file changed, 5 insertions(+), 1 deletion(-) diff --git a/cuda_core/cuda/core/experimental/_module.py b/cuda_core/cuda/core/experimental/_module.py index 0138899e3..aad8bc43e 100644 --- a/cuda_core/cuda/core/experimental/_module.py +++ b/cuda_core/cuda/core/experimental/_module.py @@ -303,7 +303,11 @@ def _init(cls, module, code_type, *, symbol_mapping: Optional[dict] = None): return self def __reduce__(self): - return ObjectCode._init(self._module, self._code_type, symbol_mapping=self._sym_map) + def _reconstruct(moduel, code_type, symbol_mapping): + # just for forwarding kwargs + return ObjectCode._init(moduel, code_type, symbol_mapping=symbol_mapping) + + return _reconstruct, (self._module, self._code_type, self._sym_map) @staticmethod def from_cubin(module: Union[bytes, str], *, symbol_mapping: Optional[dict] = None) -> "ObjectCode": From f0e86322281cafb18e30220240bd3c37c2cd3bda Mon Sep 17 00:00:00 2001 From: brandon-b-miller Date: Thu, 29 May 2025 06:55:36 -0700 Subject: [PATCH 5/6] apply patch --- cuda_core/cuda/core/experimental/_module.py | 11 ++++++----- 1 file changed, 6 insertions(+), 5 deletions(-) diff --git a/cuda_core/cuda/core/experimental/_module.py b/cuda_core/cuda/core/experimental/_module.py index aad8bc43e..afb229538 100644 --- a/cuda_core/cuda/core/experimental/_module.py +++ b/cuda_core/cuda/core/experimental/_module.py @@ -302,12 +302,13 @@ def _init(cls, module, code_type, *, symbol_mapping: Optional[dict] = None): return self - def __reduce__(self): - def _reconstruct(moduel, code_type, symbol_mapping): - # just for forwarding kwargs - return ObjectCode._init(moduel, code_type, symbol_mapping=symbol_mapping) + @classmethod + def _reduce_helper(self, module, code_type, symbol_mapping): + # just for forwarding kwargs + return ObjectCode._init(module, code_type, symbol_mapping=symbol_mapping) - return _reconstruct, (self._module, self._code_type, self._sym_map) + def __reduce__(self): + return ObjectCode._reduce_helper, (self._module, self._code_type, self._sym_map) @staticmethod def from_cubin(module: Union[bytes, str], *, symbol_mapping: Optional[dict] = None) -> "ObjectCode": From 77e51deb3c27dbf4a2674b217f069d5af1a18c93 Mon Sep 17 00:00:00 2001 From: brandon-b-miller Date: Thu, 29 May 2025 06:56:10 -0700 Subject: [PATCH 6/6] address reviews --- cuda_core/tests/test_module.py | 1 - 1 file changed, 1 deletion(-) diff --git a/cuda_core/tests/test_module.py b/cuda_core/tests/test_module.py index 4e022f234..d85a4745e 100644 --- a/cuda_core/tests/test_module.py +++ b/cuda_core/tests/test_module.py @@ -253,7 +253,6 @@ def test_module_serialization_roundtrip(get_saxpy_kernel): result = pickle.loads(pickle.dumps(objcode)) # nosec B403, B301 assert isinstance(result, ObjectCode) - assert result._handle is None assert objcode.code == result.code assert objcode._sym_map == result._sym_map assert objcode._code_type == result._code_type