Skip to content

feat: migrate remaining ListArray CUDA kernels to cuda.compute - #4481

Open
ianna wants to merge 13 commits into
scikit-hep:mainfrom
ianna:ianna/cuda-compute-listarray
Open

ianna wants to merge 13 commits into
scikit-hep:mainfrom
ianna:ianna/cuda-compute-listarray

Conversation

@ianna

@ianna ianna commented Oct 25, 2026 •

Copy link
Copy Markdown
Member

ListArray CUDA validity checks currently defer errors even though callers need an immediate answer, and ten slicing/padding/validity kernels still use RawKernel dispatch. Migrate those kernels to cuda.compute while preserving the canonical error, partial-write and output-buffer contracts.

Implementation

  • Implement jagged apply/expand/shrink, next-array and advanced indexing, scalar/range slicing, range spreading, padding and validity with transforms, scans and reductions. Remove their legacy CUDA sources/registrations and update the spec-driven generation/test adapters.
  • Return immediate KernelError results for error-capable kernels, preserving earliest-error priority, original attempted indices, successful prefixes and untouched output suffixes. Forward the current stream, including the existing range carry-length helper. Kernels without error results do not read device scalars to determine launch sizes.
  • Widen offset operands before subtraction in six CPU kernels to match the canonical reference for narrow offsets. Canonical signatures remain unchanged; spec/test data clarify edge cases.
  • Reuse prepared operations/adapters by device/context, dtype, dimensions and slice parameters. Cached proxies hold metadata rather than caller buffers, are rebound under a lock, and are cleared after submission. Each invocation owns its scratch storage.
  • Validate bounded row chunks, reduce first-error positions and write independent output elements in parallel. Small bounded workloads use fused paths. Preserve empty-slice short-circuiting, narrow-offset behavior and sequential last-writer semantics for overlapping spread intervals. Range slicing removes parameter uploads and scans directly into int64 output offsets where possible.
  • Capture CUDA error-context metadata without formatting GPU values or retaining input buffers on successful calls. Both CUDA wrappers retain operation names, array type/length metadata and bounded selectors for deferred diagnostics. Deferred notes use summaries instead of array values; ordinary CPU exception formatting is unchanged.
  • Combine jagged expansion status/first-error/chunk storage, reducing large-path scratch allocations from four to two and removing one buffer binding. Cache reduction workspace sizing.

The awkward-cpp version and root dependency pin remain unchanged.

Della results

On an A100, run 16325085 compared 587eb8e8a against historical baseline 7c403344b. Timings are warmed, unprofiled completed public operations: input/reference preparation and JIT warmup are excluded; stream completion and deferred-error draining are included.

Public slicing Completed latency reduction across tested shapes
Advanced 55–60%
Jagged 37–43%
Spread 35–57%

For jagged 4096:8, completed latency was 2.917 ms → 1.780 ms. All candidate public cases passed. Overall comparison parity is false because the historical baseline misses three expected jagged bounds errors; the candidate raises them. These numbers measure the complete candidate against that baseline and do not isolate individual changes or changes inherited from main.

A separate direct comparison in run 16331686 tested 587eb8e8a against current head 0937802ba three times on one GPU allocation, alternating execution order. Public parity passed in all three comparisons, including bounds errors with unchanged error stages. Jagged completed-time changes ranged from −1.5% to +0.6% (32:8), −1.8% to −0.1% (4096:8), and −0.9% to +0.8% (32:4096). Other workloads shifted similarly: the final scratch/workspace tuning is approximately latency-neutral, with no established jagged-specific public speedup.

Separate instrumented profiles identified eager error-context formatting as a major host cost. After metadata capture, the RegularArray range wrapper's diagnostic duration fell from approximately 1.8 ms to 0.27 ms. Trace/cProfile measurements perturb execution and are not benchmark latencies or pure device execution times. Individual compute kernels, particularly jagged expansion, can remain slower than their historical RawKernel counterparts despite the completed public-operation gains.

Validation and remaining limits

  • Latest focused local suite: 1,363 tests passed; Ruff, Pylint and whitespace checks passed. Earlier native CPU/reference runs covered canonical kernel cases, including rebuilt CPU offset-arithmetic changes.
  • Tests cover dtype combinations, chunk boundaries, first-error priority, partial writes, empty/decreasing/overlapping intervals, spare capacity, prepared setup reuse, caller-buffer lifetimes and concurrent host calls. Added CUDA regressions cover nondefault streams and successful context capture without value formatting or input retention.
  • Della public benchmarks validate the measured slicing workloads and expected errors. They do not establish exhaustive GPU coverage of every migrated kernel, shape or dtype. Full GPU parity, memory-lifetime/error-recovery and stream validation on the latest head remain to be completed/confirmed through the full CUDA suites.
  • Host execution and host Numba compilation of operator bodies do not establish CUDA JIT coverage. Immediate error-capable kernels still require compact status transfers/synchronization; this PR does not claim to eliminate them.
  • Larger capacities, heterogeneous shapes and other GPUs need broader performance measurements. The last allocation tuning has a concrete resource benefit but no demonstrated public latency benefit.

Related work: #4189 and #4461; jagged index-equals-length behavior overlaps #4449; CPU offset widening overlaps #4471, #4474 and #4301. #4483 and #4484 share the returned-KernelError generated-test adapter; combine their kernel-name sets when integrating. The standalone benchmark follow-up is ianna/awkward-matrix#2.

AI assistance was used to implement, review, test and document this contribution, including the benchmark summaries.

@github-actions github-actions Bot added the type/feat PR title type: feat (set automatically) label Oct 25, 2026
@ianna

ianna commented Oct 25, 2026

Copy link
Copy Markdown
Member Author

partially fixes #3793

@codecov

codecov Bot commented Oct 25, 2026 •

Copy link
Copy Markdown

Codecov Report

✅ All modified and coverable lines are covered by tests.
✅ Project coverage is 85.22%. Comparing base (24e8900) to head (f5f362a).

Additional details and impacted files
Files with missing lines Coverage Δ
src/awkward/_backends/cupy.py 97.72% <ø> (ø)
src/awkward/_connect/cuda/__init__.py 91.89% <ø> (ø)
src/awkward/_connect/cuda/_compute.py 54.56% <100.00%> (+8.73%) ⬆️
src/awkward/_connect/cuda/_listarray.py 100.00% <100.00%> (ø)
src/awkward/_connect/cuda/_listarray_prepared.py 100.00% <100.00%> (ø)
src/awkward/_connect/cuda/_listarray_slicing.py 100.00% <100.00%> (ø)
src/awkward/_errors.py 85.65% <100.00%> (+2.06%) ⬆️
src/awkward/_kernels.py 96.50% <100.00%> (+0.69%) ⬆️

... and 11 files with indirect coverage changes

@github-actions

Copy link
Copy Markdown
Contributor

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

@ianna

ianna commented Oct 27, 2026

Copy link
Copy Markdown
Member Author

CUDA performance update (AI-assisted summary):

Della A100 run 16325085, candidate 587eb8e8ac975546ba6c2e66258564035921cbb8 versus historical baseline 7c403344b574980bbb020f98dc2a7b1bd5f80324, measured warmed, unprofiled completed public operations. Inputs/CPU references were prepared outside timing; completion includes stream synchronization and deferred-error draining.

Public slicing Completed latency reduction across tested shapes
Advanced 55–60%
Jagged 37–43%
Spread 35–57%

For jagged 4096:8, completed latency was 2.917 ms → 1.780 ms. All candidate public cases passed. The report's overall parity flag is false because the historical baseline fails to raise the three expected jagged bounds errors; the candidate correctly raises them. Advanced bounds errors are reported during the candidate call rather than at the baseline's later drain.

The error-context change avoids eager GPU value formatting and buffer retention on successful calls, capturing host-only array metadata and bounded selectors for deferred diagnostics. Instrumented tracing showed the RegularArray range wrapper spike falling from approximately 1.8 ms to 0.27 ms; these diagnostic durations are perturbed and are not benchmark latencies. The public numbers above compare the whole candidate against the historical baseline, not an isolated context-change A/B test.

Latest head 0937802ba8a994c0f4b7298688c7f379bda5cade additionally combines jagged expansion scratch allocations (four → two on the large path), removes one buffer binding, and caches reduction workspace sizing. No latency improvement is established for this follow-up: diagnostic runs remain variable and jagged expansion itself remains slower than its historical RawKernel counterpart. A direct unprofiled comparison of 587eb8e8a versus 0937802ba, alternating order on the same GPU, has been submitted on Della and is still queued.

Local validation of the latest tuning: 1,363 focused tests passed, with Ruff, Pylint and whitespace checks passing. This is separate from GPU performance validation. Remaining work is to collect the direct comparison before deciding whether the latest tuning improves latency.

@ianna

ianna commented Oct 27, 2026

Copy link
Copy Markdown
Member Author

Direct follow-up comparison from Della run 16331686 (AI-assisted summary):

Compared 587eb8e8ac975546ba6c2e66258564035921cbb8 against 0937802ba8a994c0f4b7298688c7f379bda5cade in three sequential comparisons on the same GPU allocation, alternating execution order. Both revisions include the CUDA error-context capture fix, so this isolates the subsequent scratch/workspace tuning. Timings are unprofiled completed public operations; JIT warmup and input/reference preparation are excluded, and completion includes stream synchronization and deferred-error draining.

Public parity passed in all three comparisons, including expected bounds errors, with unchanged error stages between these revisions.

Jagged workload Completed-time change across the three comparisons
32:8 −1.5% to +0.6%
4096:8 −1.8% to −0.1%
32:4096 −0.9% to +0.8%

Other workloads shifted similarly. These results do not establish a jagged-specific public latency improvement; the follow-up is approximately latency-neutral at this measurement resolution. Its concrete changes remain fewer large-path scratch allocations (four → two), one fewer buffer binding, and reuse of reduction workspace sizing. Recommendation: retain it for reduced allocation/setup overhead, without claiming a measured public speedup.

This result is separate from the substantial whole-candidate improvements against the historical baseline reported above; it measures only the latest follow-up.

@ianna

ianna commented Oct 27, 2026

Copy link
Copy Markdown
Member Author

Coverage/readiness follow-up (AI-assisted investigation), pushed as db0608c39:

The 45.76% patch report was substantially an attribution problem, not evidence that all those callback paths were untested. The host contract tests AST-load actual implementation definitions to execute CUDA operator bodies with host substitutes. Their namespaces were named after the test module (or lacked an Awkward module name), so --cov=awkward excluded those executions. Their compiled filenames also pointed into the source checkout, while CI covers an installed package.

The tests now read the installed implementation files, compile with those actual filenames, and identify the corresponding Awkward module namespaces. No coverage threshold was lowered and no production lines were excluded.

I also added genuine missing coverage: large-path advanced-index structural errors beyond a chunk boundary, canonical case comparisons through the actual production wrapper dispatch into prepared helpers, descriptor rebinding/clearing, invalid tuning arguments, and low-level/custom diagnostic-context fallbacks.

Validation: 1,701 focused tests passed under coverage run --source=awkward, including a separate package-copy layout outside the checkout to reproduce installed-package attribution. The local helper reports are:

  • _listarray.py: 98%, missing only top-level imports/initialization (lines 10–26).
  • _listarray_prepared.py: 97%, missing only top-level imports/initialization (lines 5–22).
  • _listarray_slicing.py: 96%, missing only top-level imports/initialization (lines 5–15).

All executable function/class-body lines in those three helpers were covered by host execution. GPU CI imports/initializes these modules normally. Host substitutes remain distinct from actual CUDA JIT/execution validation; the previous head 0937802ba passed both CUDA 12/13 jobs, but fresh CI and the combined Codecov report for db0608c39 are pending. Ruff, formatting and whitespace checks passed. Production code and measured performance are unchanged by this follow-up.

Merge readiness: wait for the refreshed checks/coverage result and independent maintainer review. The existing Della public parity/performance results remain applicable to the unchanged implementation.

@ianna
ianna requested a review from TaiSakuma October 27, 2026 21:02

This branch has not been deployed

No deployments
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

type/feat PR title type: feat (set automatically)

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant