CuPy raw kernels → cuda.compute: status and migration plan

[ianna]

CuPy raw kernels → cuda.compute: status and migration plan

Status: working document · Scope: src/awkward/_connect/cuda/ (the CUDA backend's kernel dispatch), audited 2026-07.

How dispatch actually works

Three layers decide what runs for a CUDA kernel call:

1. _kernel_signatures.py (auto-generated by dev/generate-kernel-signatures.py) maps every (name, dtypes) key either to a host wrapper function (which launches CuPy RawModule kernels compiled from cuda_kernels/*.cu) or to None.
2. CupyBackend.__getitem__ (_backends/cupy.py) tries that table first; on None it falls through to _supports_cuda_compute + _get_cuda_compute_impl, which dispatch to _connect/cuda/_compute.py (cuda.compute / CCCL).
3. The generator's cuda_kernels_impl list in dev/generate-kernel-signatures.py is the source of truth: uncommented names are emitted as CuPy wrappers; commented-out names are emitted as None and therefore compute-dispatched.

So "migrated" ≠ "the .cu file is gone" — a kernel is migrated when its generator entry is commented out, regardless of whether the .cu file still sits in the directory.

Current state (counts, audited against the code)

category	count
CUDA kernel families in the signature table	143
already compute-dispatched (emitted as None)	108
still CuPy-dispatched (live .cu + wrapper)	35
— of which: _compute.py implementation already written, just not enabled	17
— of which: no compute implementation yet	18
.cu files on disk	109
dead .cu files (family compute-dispatched, file never used)	74

The 108 migrated include all value reducers (sum/prod/min/max/argmin/argmax/count/countnonzero + bool/complex variants — no .cu ever existed for these), awkward_sort, and ~74 structural kernels whose .cu remains on disk as dead code. Every dead .cu is still concatenated into the RawModule source blob at initialize_cuda_kernels() and compiled by NVRTC — pure startup cost.

The migration recipe (per kernel)

1. Implement in _compute.py (offsets-based, matching current kernel-specification.yml; keep the legacy "... in compiled CUDA code ()" error-message suffix — generated tests match on it).
2. Add the name to _supports_cuda_compute and _get_cuda_compute_impl in _backends/cupy.py.
3. Comment the name out of cuda_kernels_impl in dev/generate-kernel-signatures.py and regenerate (python dev/generate-kernel-signatures.py).
4. Delete the .cu file; run tests-cuda-kernels (+ tests-cuda) on GPU.

Phase 1 — flip the 17 that are already implemented

These have working-looking defs in _compute.py but are still uncommented in the generator and absent from the support table. Migration = steps 2–4 only, one PR each or batched by family, validated on GPU:

IndexedArray_simplify, ListArray_getitem_jagged_expand, ListArray_getitem_next_array, ListArray_getitem_next_array_advanced, ListArray_getitem_next_at, ListArray_getitem_next_range, ListArray_getitem_next_range_spreadadvanced, ListArray_rpad_axis1, ListArray_validity, ListOffsetArray_drop_none_indexes, ListOffsetArray_flatten_offsets, ListOffsetArray_rpad_axis1, ListOffsetArray_toRegularArray, RegularArray_getitem_next_array_regularize, UnionArray_flatten_combine, UnionArray_flatten_length, UnionArray_nestedfill_tags_index.

(Some of these are the same names as generator entries still uncommented — the implementation predates the flip. Treat each as "review impl, enable, test".)

Phase 1 results (della, 2026-07)

11 of the 17 flipped cleanly (0.6–3.9 ms per operation in bench_flip17.py):
IndexedArray_simplify, ListArray_getitem_jagged_expand, ListArray_getitem_next_array,
ListArray_getitem_next_array_advanced, ListArray_getitem_next_range_spreadadvanced,
ListArray_validity, ListOffsetArray_drop_none_indexes, ListOffsetArray_flatten_offsets,
ListOffsetArray_toRegularArray, RegularArray_getitem_next_array_regularize,
UnionArray_nestedfill_tags_index.

6 were reverted to CuPy dispatch (marked in cuda_kernels_impl):

• ListArray_getitem_next_range, ListArray_rpad_axis1, ListOffsetArray_rpad_axis1
  — ~1.83 s per call (identical across all three ⇒ per-call numba JIT: their
  _compute.py ops are closures over per-call state, defeating cuda.compute's
  op cache). Fix before re-flipping: define ops at module level / precompiled
  numba device functions / OpKind, so the cache keys stay stable across calls.
• ListArray_getitem_next_at, UnionArray_flatten_length, UnionArray_flatten_combine
  — TypeError: len() of unsized object at call time (needs a GPU traceback to
  pinpoint; likely the op adapter probing a captured scalar/0-d object).

Lesson for Phases 1–2: any _compute.py implementation whose op is a closure
over call arguments must be validated for cache-stable compilation (call
twice, assert the second call is fast) before flipping.

Before/after comparison (della, 200k and 2M lists): startup
initialize_cuda_kernels improved 59.9 → 39.6 ms (0.66×) from deleting
.cu files — strong evidence for Phase 3's dead-code deletion. The kept
kernels measure 0.99–1.03× (parity) at both sizes.

Two more reverts after the 2M-list scale test:
ListArray_getitem_next_array (1.41× → 1.45×) and ListArray_validity
(1.42× → 1.52×) did not amortize — these kernels are launch-bound (their
raw-kernel time barely grows over 10× data), so the eager host-side
error-check sync (cp.any() + boolean elementwise kernels) that replaced the
raw kernels' deferred error queue is a permanent ~1.4–1.5× tax at every
size. Re-flip criterion: replace the eager sync with a device-side error
flag checked through the existing synchronize_cuda machinery.

Final Phase 1 tally: 9 kept, 8 reverted. Kept:
IndexedArray_simplify, ListArray_getitem_jagged_expand,
ListArray_getitem_next_array_advanced†,
ListArray_getitem_next_range_spreadadvanced,
ListOffsetArray_drop_none_indexes, ListOffsetArray_flatten_offsets,
ListOffsetArray_toRegularArray, RegularArray_getitem_next_array_regularize,
UnionArray_nestedfill_tags_index.
(† _advanced shares the reverted getitem_next_array's eager-error-check
pattern but was not directly measured; it fires only on rarer advanced-index
paths — revert it too if tests-cuda timing flags it.)

The lesson generalizes: launch-bound kernels with mandatory error checks are
the worst candidates for the current _compute.py style; bandwidth-bound
structural kernels flip cleanly. Phase 2 should implement error checks as
device-side flags from the start.

Phase 2 — implement the remaining 18, grouped by CCCL primitive

2a. Trivial maps (unary_transform, one thread per element/bin):

• NumpyArray_reduce_adjust_starts_64 — toptr[k] -= starts[k] (post-offsets-migration form; the rawkward HIP port is a direct reference)
• IndexedArray_fill, IndexedArray_fill_count — gather/copy with base offset
• NumpyArray_rearrange_shifted — pure map (canonical fromparents form)
• RegularArray_reduce_local_nextparents_64 — nextparents[k] = k / size
• ListOffsetArray_reduce_nonlocal_nextstarts_64 — conditional first-of-run scatter
• ListOffsetArray_reduce_nonlocal_outstartsstops_64 — per-bin map over offsets

2b. Select / scan (the segmented_select pattern already proven in awkward._connect.cuda._segment_algorithms and _compute.py's IndexedArray_reduce_next_64):

• ByteMaskedArray_reduce_next_64 — identical shape to the migrated IndexedArray_reduce_next_64 (its current CuPy wrapper is already a 3-phase mark/cumsum/write; replace with one select + ZipIterator index capture and a per-bin nextoffsets searchsorted)
• IndexedArray_index_of_nulls — stream compaction of null positions (select on index < 0, then subtract starts[bin] via a transform)
• IndexedArray_local_preparenext_64 — per-bin pairing of offsets/nextoffsets (map over bins; rawkward's parallel HIP port is the reference)
• ListArray_getitem_jagged_apply, ListArray_getitem_jagged_shrink — carrylen scan + gather (mirror the already-migrated jagged_carrylen/jagged_descend/jagged_numvalid)
• ListOffsetArray_reduce_local_nextparents_64 — this is exactly offsets_to_segment_ids from _segment_algorithms.py (searchsorted); reuse it
• ListOffsetArray_reduce_local_outoffsets_64 — inverse direction: searchsorted(parents, arange(outlength+1))

2c. Bin-major restructure:

• RegularArray_reduce_nonlocal_preparenext_64 — emit nextcarry/nextoffsets bin-major; the rawkward CPU/HIP implementations of the same kernel are line-for-line references

2d. Combinations:

• RegularArray_combinations_64 — adapt the already-migrated ListArray_combinations compute implementation to fixed size

2e. Run-length family (last):

• sorting_ranges, sorting_ranges_length — boundary detection parents[i-1] != parents[i]. No run_length_encode in cuda.compute yet (same gap noted in _segment_algorithms.py's TODO), but expressible today as select over a ZipIterator of adjacent elements + CountingIterator. If cuda.compute grows unique_by_key/RLE (cf. [FEA]: Load-balanced segmented reduce NVIDIA/cccl#6171 for segmented_reduce), simplify then.

Phase 2 batch 1 (implemented, pending GPU validation)

15 of the 18 implemented in _compute.py as pure-CuPy vectorized
functions (searchsorted / cumsum / gather arithmetic — no numba closures, so
cache-stable by construction) and flipped in the dispatch tables + generator:
adjust_starts_64, IndexedArray_fill, fill_count, index_of_nulls,
IndexedArray_local_preparenext_64, ListOffsetArray_reduce_local_nextparents_64,
reduce_local_outoffsets_64, reduce_nonlocal_nextstarts_64,
reduce_nonlocal_outstartsstops_64, NumpyArray_rearrange_shifted,
RegularArray_reduce_local_nextparents_64,
RegularArray_reduce_nonlocal_preparenext_64, ByteMaskedArray_reduce_next_64,
sorting_ranges, sorting_ranges_length.

All 15 numerically validated against loop references transcribed from the
.cu sources (run with cp = numpy in CI-less environments). The .cu
files are intentionally left on disk this round: dispatch is controlled by
the generator alone, so a revert is a one-line uncomment + regenerate; delete
the .cu files only after tests-cuda-kernels + bench pass on GPU.

Deferred to batch 2 (complex logic / error paths, need device-side error
flags): RegularArray_combinations_64, ListArray_getitem_jagged_apply,
ListArray_getitem_jagged_shrink.

Remaining CuPy-dispatched after batch 1: 11 = 3 batch-2 + 8 Phase-1
reverts.

Pre-existing silent coverage gaps (discovered by the batch-1 bench)

Twelve kernel families map to None in the signature table but were never
added to the support table — they raise "not supported on CUDA backend" at
head, independent of this migration. Two common-path ones are now fixed
with compute implementations (validated against loop references, 30
randomized layouts for preparenext):

• awkward_argsort — ak.argsort was broken on CUDA; now a key/value
  segmented_sort with within-segment positions as values.
• awkward_ListOffsetArray_reduce_nonlocal_preparenext_64 — ak.sum/prod/
  min/max(jagged, axis=0) were broken on CUDA; now a lexsort-based bin-major
  gather (same head semantics as the rawkward CPU/HIP ports).

Measured on della (200k lists): all 15 batch-1 kernels healthy (sum_axis1
1.00 ms, argmin 1.13 ms, masked reduce 1.98 ms, regular nonlocal 1.59 ms,
option concat 1.78 ms); the two gap fixes run at argsort 1.63 ms and
sum_axis0 6.10 ms. sum_axis0 is the slowest row — the lexsort inside
preparenext is O(n log n) with no prior baseline (the op was unsupported);
if it matters in practice, replace the lexsort with a counting/radix
bucketing over the bounded nextbin keys.

Still open (documented, not blocking common paths):
ListOffsetArray_reduce_nonlocal_nextshifts_64 (ak.argmin/argmax(jagged,
axis=0); genuinely sequential per-bin state — needs a segmented-scan
formulation or a raw kernel), IndexedArray_unique_next_index_and_offsets_64,
argsort_strings, sort_asstrings_uint8, unique_strings,
prepare_utf8_to_utf32_padded, utf8_to_utf32_padded, unique_offsets,
unique_ranges, unique_ranges_bool (string/unique long tail).

Phase 3 — cleanup and retirement

1. Delete the 74 dead .cu files now (independent of Phases 1–2): they are never dispatched, but every GPU session pays NVRTC compilation for them. Also drop their populate_kernel_errors entries. This is the cheapest win in the whole plan — measure initialize_cuda_kernels() time before/after.
2. As Phases 1–2 land, delete each newly-dead .cu.
3. When cuda_kernels_impl is empty: remove the RawModule/NVRTC machinery (initialize_cuda_kernels, fetch_specialization, template plumbing, cuda_common.cu), the synchronize_cuda error-queue can shrink to compute-side host checks, and _kernel_signatures.py degenerates to a name→None set that can be replaced by a simple frozenset.

Risks and invariants

• Error semantics: raw kernels report via the invocation_index/err_code device machinery; compute implementations check host-side. The exact legacy message text must be preserved (generated tests-cuda-kernels match on it).
• Dtype coverage: each flipped kernel must handle every specialization the signature table emitted (int32/uint32/int64 index variants especially).
• Stream correctness: _compute.py helpers must keep taking the current CuPy stream, as segmented_sort/reducers do.
• Testing: tests-cuda-kernels (generated from kernel-test-data.json) is the per-kernel gate; tests-cuda the integration gate. Both require GPU (della).
• Small-input overhead: cuda.compute has per-call host overhead (op adapter, temp-storage sizing). The reducer migration already accepted this; for hot tiny kernels (e.g. fill_count, scalar-ish maps) verify with the bench_ak_funcs-style comparison before deleting the .cu.

Relation to rawkward

The offsets-based APIs these implementations need are exactly what the parents → offsets migration standardized across rawkward's cpu/simd/cuda/hip/metal trees — for 2a–2c, the rawkward CPU and HIP ports are direct references. Conversely, once _compute.py reaches full coverage, its per-kernel patterns become the specification for rawkward's staged cuda/ sketches.