NVIDIA / NVIDIA/cccl

[BUG]: stateful op in a TransformIterator (algorithm input) fails (get_return_type NotImplementedError / cudaErrorLaunchFailure)

Open
#9,627 1 comment 0 reactions 0 assignees View on GitHub
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

Open the contributing 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

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.