Handle state appropriately in TransformIterator ops

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(

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"]

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<char*>(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<char*>(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()

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 (

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<int*>(state);
+        *static_cast<int*>(output) = *static_cast<int*>(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<int*>(state);
+        *static_cast<int*>(output) = *static_cast<int*>(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<int*>(state);
+        *static_cast<int*>(output) = *static_cast<int*>(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<Pair*>(input);
+        *static_cast<int*>(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}"