Handle state appropriately in TransformIterator ops - #11213
Conversation
📝 SummarySummary by CodeRabbit
WalkthroughChangesRawOp state propagation
Assessment against linked issues
Suggested reviewers: Merge Risk: 🔵 Low · up to Stateful transform operations now receive their runtime state correctly, but an invalid zero alignment on a stateful RawOp can fail iterator construction on the host. Validate alignment before merge. Comment |
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
There was a problem hiding this comment.
Actionable comments posted: 1
ℹ️ Review info
⚙️ Run configuration
Configuration used: Path: .coderabbit.yaml
Review profile: CHILL
Plan: Enterprise
Run ID: 74764057-0aa3-4e75-a828-61024513c362
📒 Files selected for processing (5)
python/cuda_cccl/cuda/compute/_jit.pypython/cuda_cccl/cuda/compute/iterators/_base.pypython/cuda_cccl/cuda/compute/iterators/_transform.pypython/cuda_cccl/cuda/compute/op.pypython/cuda_cccl/tests/compute/test_raw_op.py
Included review availability: Your plan provides up to 12 included reviews per hour; 11 remain after this review.
| @property | ||
| def state_alignment(self) -> int: | ||
| """Return the alignment requirement of the op's state bytes.""" | ||
| return self._state_alignment |
There was a problem hiding this comment.
🩺 Stability & Availability | 🟡 Minor | ⚡ Quick win
important: Validate state_alignment before returning it. A RawOp with nonempty state and state_alignment=0 reaches compose_state_blobs() through TransformIterator; its padding calculation then raises ZeroDivisionError. Raise ValueError for non-positive or non-integer alignments in RawOp.__init__.
| ) | ||
|
|
||
|
|
||
| def test_cpp_stateful_op_with_transform_output_iterator(): |
There was a problem hiding this comment.
Important: please add similar tests to the numba path as well (since these changes support that)
| # 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() |
There was a problem hiding this comment.
Critical: different values of the op state still trigger a full recompile. TransformIterator.kind includes self._transform_op, and RawOp._identity includes self._state, so two iterators that differ only in the state bytes hash differently and we build the whole algorithm again. With the scale_by_state op from the new tests, same state reuses the cached reducer in 0.00s and a different state rebuilds in about 1.1s. For the mean use case in #11142 every distinct n recompiles, which defeats the point of passing it as state.
Since the state is now rebound on every call, only its length and alignment affect the generated code (they fix the offset baked into the deref glue). Could we replace self._state in RawOp._identity with len(self._state)? The Numba path has the same problem inside a TransformIterator for a different reason (_StatefulOp.__eq__ compares _JitOpState by identity)
| 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);" | ||
|
|
There was a problem hiding this comment.
nit: this block is a copy of the one in _make_input_deref_op (line 185), with only the argument order swapped. A small helper along the lines of _op_decl_and_call(compiled_op, in_expr, out_expr) returning the declaration and call strings would keep the two prototypes in one place, so a future change to the stateful signature only needs to happen once.
| ) | ||
| self._op_state_offset = offsets[1] | ||
| else: | ||
| state_bytes = bytes(self._underlying.state) |
There was a problem hiding this comment.
nit: line 110 above goes through bytes(memoryview(...)) while this one calls bytes(...) directly on the same IteratorState. Both work, but it reads like they are doing different things. Could we pick one form and use it in both branches?
😬 CI Workflow Results🟥 Finished in 1h 25m: Pass: 98%/102 | Total: 1d 10h | Max: 57m 09sSee results here. AI failure analysis1. cuda.compute double-buffer radix sort worker crash on Windows CTK 12.0 · 1 jobExplanation: The sysctk job lost an xdist worker while the int64/1024 DoubleBuffer radix-sort case was running, with no Python exception or native exit code in the log. The PR does not modify radix-sort code and its added stateful TransformIterator tests pass in this job, so a crash dump or isolated reproduction is needed to distinguish a platform/toolkit defect from a latent binding bug. Evidence: Copy this prompt into a coding agentJobs: 2. cuda.compute large radix sort exits with Windows status 0xC0000409 · 1 jobExplanation: The pinned-package job terminates the serial large-test process during the uint8/1,048,576-key radix sort; -1073740791 corresponds to Windows status 0xC0000409, and no traceback or crash dump was emitted. The PR does not modify radix-sort code and all added stateful TransformIterator tests had already passed, so the evidence is insufficient to attribute the native termination to the PR. Evidence: Copy this prompt into a coding agentJobs: |
2 participants
Description
TransformIteratorcurrently ignores itsopstate, instead using only the underlying iterator's state to compose its state bytes. This PR makes it so that if the op of aTransformIteratoris stateful, we combine the op state with the underlying iterator's state. The dereference operator is modified to decompose the states and handle them appropriately.Closes #11142
Checklist