[ianna]

Migrate IndexedArray simplification from RawKernel to cuda.compute for all nine outer/inner index dtype combinations, preserving the canonical immediate-error and partial-write contract.

Inputs of at most 64 entries use one prepared operation that validates and writes sequentially, stopping before the first invalid read. The large path now fuses validation and first-error selection into a prepared reduction over CountingIterator row numbers, followed by the existing guarded parallel gather and attempted-index packing. It needs no separate validation transform, chunk-results buffer or TransformIterator.

The custom binary operator maps valid rows to the length sentinel and invalid rows to their own position, then selects the minimum. The mapping is idempotent, making the operation associative and commutative on row numbers plus the sentinel. Negative outer indices map to -1; negative inner indices are copied unchanged. The failing row and output suffix stay untouched. One 16-byte status read preserves immediate errors and the original attempted index.

Prepared operations/adapters and workspace sizing are cached by device/context, dtypes and declared lengths. Per-call status and reduction workspace remain independent. Cached proxies hold metadata rather than array owners; rebinding/launch is locked and proxies are cleared afterward. Remove the legacy registration/source and update the spec-driven generation/test adapters. Canonical signatures and reference semantics remain unchanged. No C++ implementation, dependency, awkward-cpp version or root pin changes.

Performance evidence and limits

The original 249aba447 passed canonical GPU parity but regressed substantially. Prepared setup and cached workspace sizing reduced that overhead. Della run 16371839, candidate 712d13f0b versus historical baseline 7c403344b, passed all 24 completed public field-selection cases, but valid large cases remained 13–14% slower and bounds cases 8–9% slower. Both timings include stream completion and error draining; the candidate reports errors during the call while baseline reports them during draining.

A benchmark-only fused prototype on Della run 16372118 reduced completed prepared-call latency from approximately 137 µs to 122–123 µs, a 10.4–10.8% reduction, consistently across valid inputs and first/middle/last errors. The prototype checks full output buffers, first error and attempted index. These timings include scratch, rebinding, status transfer and completion, but exclude public/backend dispatch and error-context capture.

Revision e0eab8138375d3b78bbd30968ba2d99223e978c8 also captures CUDA error context using bounded host metadata, avoiding GPU value reads and eager array formatting on successful calls. This is the same context-capture implementation used in #4481; the reduction algorithm is unchanged by this final optimization.

Della run 16374446 compared this revision with baseline 7c403344b574980bbb020f98dc2a7b1bd5f80324 on an NVIDIA A100. All 24 public-operation cases passed CPU-reference parity. Inputs and CPU reference results were prepared outside timing, JIT warmup was excluded, and completed timings include the public call, result construction, stream completion and deferred-error draining. Error classification and output comparisons were outside timing.

Representative completed public-operation medians (baseline → candidate):

Projection	Shape (rows:width)	Completed µs	Candidate / baseline
indexed→indexed	32:8	731.589 → 209.680	0.287
indexed→option	32:8	639.295 → 208.553	0.326
indexed→indexed	4096:8	664.447 → 252.962	0.381
indexed→option	4096:8	643.138 → 254.675	0.396
option→indexed	4096:8	673.238 → 257.936	0.383
option→option	4096:8	602.069 → 255.221	0.424
indexed→indexed	32:4096	626.206 → 210.406	0.336
indexed→option	32:4096	603.182 → 211.573	0.351

Across all 12 valid cases, completed ratios were 0.287–0.424 (2.4–3.5× faster). The large indexed→option and indexed→indexed candidate timings are effectively equal. The paired diagnostics in run 16372993 showed one simplify dispatch and one context capture for each projection, with essentially identical status-read costs; they did not identify additional option-specific kernel dispatch.

The final revision 131b98a18d38f7a0410462dbf5a3e618099cb0af also reuses the CUDA metadata snapshot when decorating immediate exceptions. Previously the original public error context formatted array values again during unwinding. CPU exceptions without a CUDA snapshot retain value formatting; error type/message and immediate reporting are preserved.

Final public-operation results

Della run 16376376, final revision 131b98a18 versus the same historical baseline 7c403344b, passed all 24 CPU-reference parity cases. The methodology is unchanged: unprofiled medians, JIT warmup excluded, inputs/reference preparation and comparisons outside timing, and stream completion/error draining included.

Projection	Shape	Valid completed µs, baseline → candidate	Bounds completed µs, baseline → candidate
indexed→indexed	32:8	717.707 → 210.010	722.756 → 197.697
indexed→option	32:8	640.556 → 211.278	646.578 → 200.633
indexed→indexed	4096:8	666.646 → 255.551	669.476 → 242.873
indexed→option	4096:8	641.613 → 258.362	649.418 → 245.423
option→indexed	4096:8	669.125 → 258.487	673.919 → 245.387
option→option	4096:8	602.745 → 258.732	607.610 → 247.270
indexed→indexed	32:4096	624.371 → 211.809	627.883 → 199.806
indexed→option	32:4096	605.821 → 212.721	608.065 → 201.435

All 12 valid cases improved, with completed ratios 0.293–0.429 (2.3–3.4× faster). All 12 bounds-error cases also improved, with ratios 0.274–0.407 (2.5–3.7× faster). The previous bounds-path slowdown is resolved in this measurement. Candidate bounds errors occur during the call; baseline errors occur during final draining (250/250 samples for each case).

Remaining limitations: these public-operation results do not replace the separate kernel comparison or establish that its review gate passes. GPU CI still needs independent confirmation. Instrumented diagnostic timings are not benchmark latencies or pure GPU times, and a prepared reduction call may launch multiple device kernels. Full run artifacts remain in results/16376376/public-indexed/ on Della; the benchmark harness is in ianna/awkward-matrix. Results here were transcribed from the supplied report; raw JSON/Nsight artifacts have not been uploaded to this PR.

Validation

• 1,296 local tests passed, covering native CPU parity, actual host-executed callbacks, production wrapper dispatch, all dtype combinations, extreme/unsigned indices, chunk/size boundaries, first errors, partial writes, untouched capacity, setup/workspace-query reuse, buffer lifetimes and concurrent host calls.
• Algebra tests exercise the actual fused binary operator's associativity and commutativity. The standalone benchmark tests also check arbitrary reduction trees against native first-error results.
• Expanded GPU tests cover all nine index dtype combinations at lengths 64, 65, 129, 4096 and 4097, on default and nonblocking streams, with repeated rebinding and valid/first/middle/last error cases. GPU execution was skipped locally because CuPy/hardware is unavailable.
• Host coverage attribution uses actual installed source paths/module names. Host simulation does not establish CUDA JIT coverage. Ruff, Pylint and whitespace checks pass.

This branch is independent of #4481/#4483. They share the returned-KernelError generated-test adapter: combine the kernel-name sets when integrating. Benchmark tools live in ianna/awkward-matrix on ianna/cuda-migration-validation.

AI assistance was used to implement, review, test and document this contribution.

[ianna]

partially fixes #3793

[codecov]

Codecov Report

❌ Patch coverage is 99.46809% with 1 line in your changes missing coverage. Please review.
✅ Project coverage is 84.52%. Comparing base (24e8900) to head (da37d62).
⚠️ Report is 3 commits behind head on main.

Files with missing lines	Patch %	Lines
src/awkward/_connect/cuda/_compute.py	97.14%	1 Missing ⚠️

Additional details and impacted files

Files with missing lines	Coverage Δ
src/awkward/_backends/cupy.py	97.72% <ø> (ø)
src/awkward/_connect/cuda/_indexedarray.py	100.00% <100.00%> (ø)
src/awkward/_errors.py	85.71% <100.00%> (+2.12%)	⬆️
src/awkward/_kernels.py	95.80% <100.00%> (ø)
src/awkward/_connect/cuda/_compute.py	50.53% <97.14%> (+4.70%)	⬆️

... and 9 files with indirect coverage changes

[github-actions]

The documentation preview is ready to be viewed at https://preview.awkward-array.org/PR4484/