[BUG]: stateful op in a TransformIterator (algorithm input) fails (get_return_type NotImplementedError / cudaErrorLaunchFailure)
- Dominant language
- C++
- Stars
- 2.5k
- Forks
- 486
- Avg merge
- 2d 6h
- Merged PRs (30d)
- 295
Description
## Bug
A **stateful op** (one that closes over a device array) works fine as a *direct* algorithm op — e.g. `select`'s `cond`, `unary_transform`'s `op` (these auto-capture device arrays as state arrays). But the same kind of op **fails when wrapped in a `TransformIterator` that is used as an algorithm's input** (`reduce_into` / `segmented_reduce` / `histogram_even`). Two distinct symptoms:
1. **Type inference** (no return annotation):
```
NotImplementedError: get_return_type not implemented for _StatefulOp
```
2. **Execution** (with a return annotation to get past inference):
```
cudaErrorLaunchFailure: unspecified launch failure
```
## Reproducer
```python
import numpy as np, cupy as cp
from cuda.compute import reduce_into, OpKind
from cuda.compute.iterators import CountingIterator, TransformIterator
N = 1000; src = cp.arange(N, dtype=cp.float32); out = cp.empty(1, cp.float32); h = np.zeros(1, np.float32)
def stateless(i) -> np.float32: return np.float32(i) * np.float32(2.0)
def stateful(i) -> np.float32: return src[i] * np.float32(2.0) # closes over a device array
def run(tag, op):
try:
reduce_into(d_in=TransformIterator(CountingIterator(np.int64(0)), op),
d_out=out, op=OpKind.PLUS, num_items=N, h_init=h)
cp.cuda.Device(0).synchronize(); print(tag, "OK", float(out[0]))
except Exception as e:
print(tag, "FAIL", type(e).__name__ + ":", str(e)[:120])
run("stateless:", stateless) # stateless: OK 999000.0
run("stateful: ", stateful) # stateful: FAIL CUDARuntimeError: cudaErrorLaunchFailure
```
Output:
```
stateless: OK 999000.0
stateful: FAIL CUDARuntimeError: cudaErrorLaunchFailure: unspecified launch failure
```
(Without the `-> np.float32` annotation, the `stateful` case instead raises `NotImplementedError: get_return_type not implemented for _StatefulOp` during construction.)
## Impact
This blocks **fusing an index-gather computation into a reduction or histogram**. The per-element transform can only enter these algorithms through the input iterator (`segmented_reduce`'s `op` is the binary reducer; `histogram_even` has no op at all). So patterns like:
- gather selected values into `histogram_even` (`d_samples = TransformIterator(idx, lambda i: src[i])`), or
- compute a per-pair metric by indexing other arrays and feed it into a `segmented_reduce`
cannot be expressed without first materializing the intermediate array — defeating the point of the iterator. Stateless transform-iterator ops fuse correctly, so this is specifically about stateful (device-array-capturing) ops in iterators.
## Possibly related
#2479 ([FEA] Support fancy iterators in cuda.parallel) is broader; this report is a specific correctness/crash bug for `TransformIterator`, which already works for stateless ops. (Note: `PermutationIterator` covers pure gather; this is about general stateful transform ops.)
## Environment
- `cuda-cccl` 3.5.0.dev296+g8c905c4ca
- CUDA 13.3, NVIDIA RTX 6000 Ada, driver 580.167.08
- Python 3.12, cupy-cuda13x 14.1.1, numba-cuda 0.30.2
Contributor guide
Research direction
Start by running the supplied TransformIterator and reduce_into reproducer, then compare the stateless and stateful cases. Trace TransformIterator through reduce_into, segmented_reduce, and histogram_even, focusing on state capture, return-type inference, and kernel launch behavior; done means stateful transforms work without materialization, inference errors, or launch failures.
Written by the indexing model from the issue text.
Assessment
- Tech stack
- cpp, python
- Domain
- backend-api-design, hpc
- Issue type
- Bug
- Difficulty
- 4/5
- Estimated time
- 3-5 days
- Activity status
- Quiet
- Clarity
- Mostly clear
- Newbie friendliness
- 48/100