From 06be34799dc58564fcf8a69be806017ebb30ba6b Mon Sep 17 00:00:00 2001 From: Ashwin Srinath Date: Fri, 4 Sep 2026 09:43:02 -0400 Subject: [PATCH] Handle state appropriately in TransformIterator ops --- python/cuda_cccl/cuda/compute/_jit.py | 5 + .../cuda_cccl/cuda/compute/iterators/_base.py | 45 ++++-- .../cuda/compute/iterators/_transform.py | 50 +++++- python/cuda_cccl/cuda/compute/op.py | 10 ++ python/cuda_cccl/tests/compute/test_raw_op.py | 142 ++++++++++++++++++ 5 files changed, 231 insertions(+), 21 deletions(-) diff --git a/python/cuda_cccl/cuda/compute/_jit.py b/python/cuda_cccl/cuda/compute/_jit.py index 2dd14b8078d0..4b4b0f5afa1a 100644 --- a/python/cuda_cccl/cuda/compute/_jit.py +++ b/python/cuda_cccl/cuda/compute/_jit.py @@ -1015,6 +1015,11 @@ def __init__(self, func, state): def get_state(self): return self._state.to_bytes() + @property + def state_alignment(self) -> int: + # State is packed device-array pointers; see _compile_stateful_op. + return np.dtype(np.intp).alignment + def compile(self, input_types, output_type=None) -> Op: transformed_func = _transform_function_ast(self._func, self._state.names) return _compile_stateful_op( diff --git a/python/cuda_cccl/cuda/compute/iterators/_base.py b/python/cuda_cccl/cuda/compute/iterators/_base.py index aafdd18ae019..e9f8b50636b7 100644 --- a/python/cuda_cccl/cuda/compute/iterators/_base.py +++ b/python/cuda_cccl/cuda/compute/iterators/_base.py @@ -246,35 +246,29 @@ def _deterministic_suffix(kind: Hashable) -> str: return hashlib.sha256(kind_str.encode()).hexdigest()[:16] -def compose_iterator_states( - iterators: list[IteratorBase], +def compose_state_blobs( + blobs: list[tuple[bytes, int]], ) -> tuple[bytes, int, list[int]]: """ - Concatenate multiple iterator states with proper alignment. - - This is used by composite iterators (like ZipIterator and PermutationIterator) - that need to store multiple child iterator states in their own state. + Concatenate raw (state_bytes, state_alignment) blobs with proper padding. Args: - iterators: List of child iterators whose states should be composed + blobs: List of (state_bytes, state_alignment) pairs to compose Returns: Tuple of: - combined_state_bytes: Concatenated state bytes with padding - combined_alignment: Maximum alignment requirement - - offsets: List of byte offsets for each iterator's state + - offsets: List of byte offsets for each blob """ - if not iterators: + if not blobs: return (b"", 1, []) - states = [bytes(memoryview(it.state)) for it in iterators] - alignments = [it.state_alignment for it in iterators] - offsets = [] current_offset = 0 combined = b"" - for state, align in zip(states, alignments): + for state, align in blobs: # Add padding to meet alignment requirement padding = (align - (current_offset % align)) % align combined += b"\x00" * padding @@ -284,10 +278,33 @@ def compose_iterator_states( combined += state current_offset += len(state) - max_alignment = max(alignments) + max_alignment = max(align for _, align in blobs) return (combined, max_alignment, offsets) +def compose_iterator_states( + iterators: list[IteratorBase], +) -> tuple[bytes, int, list[int]]: + """ + Concatenate multiple iterator states with proper alignment. + + This is used by composite iterators (like ZipIterator and PermutationIterator) + that need to store multiple child iterator states in their own state. + + Args: + iterators: List of child iterators whose states should be composed + + Returns: + Tuple of: + - combined_state_bytes: Concatenated state bytes with padding + - combined_alignment: Maximum alignment requirement + - offsets: List of byte offsets for each iterator's state + """ + return compose_state_blobs( + [(bytes(memoryview(it.state)), it.state_alignment) for it in iterators] + ) + + cache_with_registered_key_functions.register(IteratorBase, lambda it: it.kind) __all__ = ["IteratorBase"] diff --git a/python/cuda_cccl/cuda/compute/iterators/_transform.py b/python/cuda_cccl/cuda/compute/iterators/_transform.py index e43ed8cd5cd6..60444a5f0a8b 100644 --- a/python/cuda_cccl/cuda/compute/iterators/_transform.py +++ b/python/cuda_cccl/cuda/compute/iterators/_transform.py @@ -12,7 +12,7 @@ from .._cpp_compile import compile_cpp_op_code, make_variable_declaration from ..op import make_op_adapter from ..types import TypeDescriptor, signature_from_annotations -from ._base import IteratorBase +from ._base import IteratorBase, compose_state_blobs from ._common import CUDA_PREAMBLE, ensure_iterator @@ -46,6 +46,7 @@ class TransformIterator(IteratorBase): "_value_type", "_is_input", "_compiled_op", + "_op_state_offset", ] def __init__( @@ -98,9 +99,30 @@ def __init__( assert value_type is not None self._value_type = value_type + # A stateful transform_op's state is folded into this + # iterator's own device state, alongside the underlying iterator's + # state, so that the generated deref glue (see _make_input_deref_op / + # _make_output_deref_op) has a pointer to hand the op as its `state` + # argument. + op_state = self._transform_op.get_state() + self._op_state_offset: int | None + if op_state: + underlying_state = bytes(memoryview(self._underlying.state)) + state_bytes, state_alignment, offsets = compose_state_blobs( + [ + (underlying_state, self._underlying.state_alignment), + (op_state, self._transform_op.state_alignment), + ] + ) + self._op_state_offset = offsets[1] + else: + state_bytes = bytes(self._underlying.state) + state_alignment = self._underlying.state_alignment + self._op_state_offset = None + super().__init__( - state_bytes=bytes(self._underlying.state), - state_alignment=self._underlying.state_alignment, + state_bytes=state_bytes, + state_alignment=state_alignment, value_type=value_type, ) @@ -160,16 +182,23 @@ def _make_input_deref_op(self) -> Op | None: symbol = self._make_input_deref_symbol() temp_decl = make_variable_declaration(self._underlying.value_type, "temp") + if compiled_op.operator_type == OpKind.STATEFUL: + op_decl = f'extern "C" __device__ void {compiled_op.name}(void* state, void* input, void* output);' + op_call = f"{compiled_op.name}(static_cast(state) + {self._op_state_offset}, &temp, result);" + else: + op_decl = f'extern "C" __device__ void {compiled_op.name}(void* input, void* output);' + op_call = f"{compiled_op.name}(&temp, result);" + source = dedent(f""" {CUDA_PREAMBLE} extern "C" __device__ void {child_op.name}(void* state, void* result); - extern "C" __device__ void {compiled_op.name}(void* input, void* output); + {op_decl} extern "C" __device__ void {symbol}(void* state, void* result) {{ {temp_decl} {child_op.name}(state, &temp); - {compiled_op.name}(&temp, result); + {op_call} }} """).strip() @@ -200,15 +229,22 @@ def _make_output_deref_op(self) -> Op | None: symbol = self._make_output_deref_symbol() temp_decl = make_variable_declaration(self._underlying.value_type, "temp") + if compiled_op.operator_type == OpKind.STATEFUL: + op_decl = f'extern "C" __device__ void {compiled_op.name}(void* state, void* input, void* output);' + op_call = f"{compiled_op.name}(static_cast(state) + {self._op_state_offset}, value, &temp);" + else: + op_decl = f'extern "C" __device__ void {compiled_op.name}(void* input, void* output);' + op_call = f"{compiled_op.name}(value, &temp);" + source = dedent(f""" {CUDA_PREAMBLE} extern "C" __device__ void {child_op.name}(void* state, void* value); - extern "C" __device__ void {compiled_op.name}(void* input, void* output); + {op_decl} extern "C" __device__ void {symbol}(void* state, void* value) {{ {temp_decl} - {compiled_op.name}(value, &temp); + {op_call} {child_op.name}(state, &temp); }} """).strip() diff --git a/python/cuda_cccl/cuda/compute/op.py b/python/cuda_cccl/cuda/compute/op.py index 009d89fbe68d..c427c77c15e2 100644 --- a/python/cuda_cccl/cuda/compute/op.py +++ b/python/cuda_cccl/cuda/compute/op.py @@ -51,6 +51,11 @@ def get_state(self) -> bytes: """ return b"" + @property + def state_alignment(self) -> int: + """Return the alignment requirement of the op's state bytes.""" + return 1 + def get_return_type(self, input_types): """Get the return type for this op given input types.""" raise NotImplementedError( @@ -167,6 +172,11 @@ def get_state(self) -> bytes: """Return the op's state bytes.""" return self._state + @property + def state_alignment(self) -> int: + """Return the alignment requirement of the op's state bytes.""" + return self._state_alignment + @property def _identity(self): return ( diff --git a/python/cuda_cccl/tests/compute/test_raw_op.py b/python/cuda_cccl/tests/compute/test_raw_op.py index d2f7ca3615b6..56509e545c9c 100644 --- a/python/cuda_cccl/tests/compute/test_raw_op.py +++ b/python/cuda_cccl/tests/compute/test_raw_op.py @@ -523,3 +523,145 @@ def test_cpp_stateful_op_select_with_counter(): assert np.array_equal(selected_values, expected_selected), ( "Selected values don't match" ) + + +def test_cpp_stateful_op_with_transform_output_iterator(): + """Regression test: a stateful RawOp used as a TransformOutputIterator's + transform op must receive its state. + + This mirrors computing a mean as sum(x) * (1/n): the scale factor lives + in the RawOp's state and is applied to the reduction result by a + TransformOutputIterator. + """ + from cuda.compute import OpKind, TransformOutputIterator + + scale_factor = np.array([3], dtype=np.int32) + state_data = scale_factor.tobytes() + state_alignment = np.dtype(np.int32).alignment + + cpp_source = """ + extern "C" __device__ void scale_by_state(void* state, void* input, void* output) { + int factor = *static_cast(state); + *static_cast(output) = *static_cast(input) * factor; + } + """ + + op = make_cpp_stateful_op(cpp_source, state_data, "scale_by_state", state_alignment) + + num_items = 10 + h_input = np.arange(num_items, dtype=np.int32) + d_input = DeviceArray.from_numpy(h_input) + d_output = DeviceArray.empty(1, np.int32) + + output_iterator = TransformOutputIterator( + d_output, op, output_value_type=types.int32 + ) + + h_init = np.array(0, dtype=np.int32) + cuda.compute.reduce_into( + d_in=d_input, + d_out=output_iterator, + num_items=num_items, + op=OpKind.PLUS, + h_init=h_init, + ) + + result = d_output.copy_to_host()[0] + expected = int(np.sum(h_input)) * 3 + assert result == expected, f"Expected {expected}, got {result}" + + +def test_cpp_stateful_op_with_transform_input_iterator(): + """Regression test: a stateful RawOp used as a TransformIterator's + (input-side) transform op must receive its state.""" + from cuda.compute import OpKind, TransformIterator + + scale_factor = np.array([3], dtype=np.int32) + state_data = scale_factor.tobytes() + state_alignment = np.dtype(np.int32).alignment + + cpp_source = """ + extern "C" __device__ void scale_by_state(void* state, void* input, void* output) { + int factor = *static_cast(state); + *static_cast(output) = *static_cast(input) * factor; + } + """ + + op = make_cpp_stateful_op(cpp_source, state_data, "scale_by_state", state_alignment) + + num_items = 10 + h_input = np.arange(num_items, dtype=np.int32) + d_input = DeviceArray.from_numpy(h_input) + d_output = DeviceArray.empty(1, np.int32) + + transform_iter = TransformIterator(d_input, op, value_type=types.int32) + + h_init = np.array(0, dtype=np.int32) + cuda.compute.reduce_into( + d_in=transform_iter, + d_out=d_output, + num_items=num_items, + op=OpKind.PLUS, + h_init=h_init, + ) + + result = d_output.copy_to_host()[0] + expected = int(np.sum(h_input)) * 3 + assert result == expected, f"Expected {expected}, got {result}" + + +def test_cpp_stateful_op_transform_nested_in_zip(): + """Regression test: a stateful RawOp's state must still be correctly + composed when its TransformIterator is nested inside another compound + iterator (here, ZipIterator), which treats its children's `.state` / + `.state_alignment` opaquely.""" + from cuda.compute import OpKind, TransformIterator, ZipIterator + + # Stateful op: scales its input by a state-provided factor. + scale_factor = np.array([3], dtype=np.int32) + state_data = scale_factor.tobytes() + state_alignment = np.dtype(np.int32).alignment + + scale_source = """ + extern "C" __device__ void scale_by_state(void* state, void* input, void* output) { + int factor = *static_cast(state); + *static_cast(output) = *static_cast(input) * factor; + } + """ + scale_op = make_cpp_stateful_op( + scale_source, state_data, "scale_by_state", state_alignment + ) + + # Stateless op: sums the two fields of the zipped pair. + add_source = """ + struct Pair { int field_0; int field_1; }; + extern "C" __device__ void add_pair(void* input, void* output) { + Pair* p = static_cast(input); + *static_cast(output) = p->field_0 + p->field_1; + } + """ + add_op = make_cpp_op(add_source, "add_pair") + + num_items = 10 + h_x = np.arange(num_items, dtype=np.int32) + h_y = np.arange(num_items, dtype=np.int32) * 100 + d_x = DeviceArray.from_numpy(h_x) + d_y = DeviceArray.from_numpy(h_y) + + scaled_x = TransformIterator(d_x, scale_op, value_type=types.int32) + zipped = ZipIterator(scaled_x, d_y) + combined = TransformIterator(zipped, add_op, value_type=types.int32) + + d_output = DeviceArray.empty(1, np.int32) + h_init = np.array(0, dtype=np.int32) + cuda.compute.reduce_into( + d_in=combined, + d_out=d_output, + num_items=num_items, + op=OpKind.PLUS, + h_init=h_init, + ) + + result = d_output.copy_to_host()[0] + expected = int((h_x * 3 + h_y).sum()) + assert result == expected, f"Expected {expected}, got {result}"