Skip to content

cuda.coop: add Numba hierarchy Reduce - #11210

Draft
tpn wants to merge 17 commits into
codex/cuda-coop-cudax-reduce-reuse-prereqfrom
codex/cuda-coop-numba-hierarchy-reduce
Draft

cuda.coop: add Numba hierarchy Reduce#11210
tpn wants to merge 17 commits into
codex/cuda-coop-cudax-reduce-reuse-prereqfrom
codex/cuda-coop-numba-hierarchy-reduce

Conversation

@tpn

@tpn tpn commented Sep 4, 2026 •
edited
Loading

Copy link
Copy Markdown
Contributor

Why this is needed

The preceding cuda.coop stack establishes generic family, provider, result,
storage, and synchronization seams. This PR adds Reduce as the first
hierarchy-wide collective on those seams, covering both storage-free CUDAX
providers and storage-bearing CUB providers without adding primitive
conditionals to generic activation, dispatch, caching, or rewrite phases.

Reviewer-visible behavior

  • Adds portable and qualified reduce and sum for thread, physical-Warp,
    logical-Warp, block, mapped-Warp, and cluster groups.
  • Portable and qualified calls share payload protocols, lowercase operator and
    algorithm strings, broadcast, and dtype inference for untyped
    ThreadData.
  • Qualified execution additionally accepts recognized operator and NumPy
    aliases and stateless Python device callbacks.
  • Uses cuda::experimental::coop::reduce for full-group built-in operations,
    including broadcast and root-only projections without temporary storage.
  • Uses deterministic CUB BlockReduce and WarpReduce providers for valid
    prefixes, explicit algorithms, and qualified callback forms.
  • Compiles stateless callbacks to target-specific LTO-IR and rejects target
    mismatches before provider generation.
  • Preserves array and ThreadData inputs, isolates logical-group storage, and
    emits block or exact-mask Warp reuse synchronization from provider metadata.
  • Guards nonmembers of partial mapped groups and rejects unsupported grid,
    stateful-callback, nondeterministic-algorithm, and invalid-prefix forms with
    structured diagnostics.
total = coop.sum(coop.this_block(), value)
root_max = coop.reduce(
    coop.this_block(),
    values,
    binary_op="max",
    broadcast=False,
    algorithm="warp_reductions",
)

Validation

The exact final-stack head 9b299e894f, which includes this layer, passed:

  • 1,440 source-tree tests with the GPU hidden.
  • 607 installed-wheel contract, packaging, public-surface, and activation
    tests outside the source tree.
  • 50 installed-wheel real NVRTC LTO/nvJitLink compile tests.
  • 366 installed-wheel runtime tests on an RTX PRO 6000.

Stack and scope

This draft is stacked on #11209. Stateful callbacks, Scan, prefix callbacks,
and a safe per-launch grid workspace policy remain deferred. The guarded
Numba-CUDA-MLIR compatibility shim remains inherited technical debt.

@github-project-automation github-project-automation Bot moved this to Todo in CCCL Sep 4, 2026
@copy-pr-bot

copy-pr-bot Bot commented Sep 4, 2026

Copy link
Copy Markdown
Contributor

Auto-sync is disabled for draft pull requests in this repository. Workflows must be run manually.

Contributors can view more details about this message here.

@cccl-authenticator-app cccl-authenticator-app Bot moved this from Todo to In Progress in CCCL Sep 4, 2026
@tpn tpn mentioned this pull request Sep 4, 2026
Draft
@tpn
tpn force-pushed the codex/cuda-coop-cudax-reduce-reuse-prereq branch 2 times, most recently from 64d2fcd to be03724 Compare September 4, 2026 18:20
tpn added 16 commits September 4, 2026 11:20
@tpn
Support storage-free raw group providers
8a61205 
Add a compiler-local C ABI invocable for providers that own storage.

Admit GROUP execution only for storage-free lowering plans.

Require provider metadata to declare no storage or synchronization too.

Signed-off-by: Trent Nelson <trent@trent.me>
@tpn
Plan hierarchy-aware cooperative reductions
ddf2cbe 
Add portable Reduce and Sum contracts across supported group scopes.

Route full built-in reductions to storage-free CUDAX plans.

Keep prefixes, algorithms, and custom operators on CUB plans.

Correct mapped topology at physical warp boundaries.

Signed-off-by: Trent Nelson <trent@trent.me>
@tpn
Compile stateless Numba reduction operators
da1b2ca 
Compile stateless Python operators to device LTO-IR.

Adapt scalar and aggregate values to the generated C ABI.

Keep stateful callbacks outside this slice.

Signed-off-by: Trent Nelson <trent@trent.me>
@tpn
Lower Numba hierarchy reductions
65a873c 
Register lazy Reduce and Sum lowering for supported thread groups.
Route full built-in collectives through storage-free CUDAX providers and
retain deterministic CUB block and warp paths for prefixes, arrays, and
stateless callbacks.

Signed-off-by: Trent Nelson <trent@trent.me>
@tpn
Compile Numba hierarchy reduction providers
61f41c0 
Exercise storage-free CUDAX hierarchy sources and deterministic CUB
block and warp providers with the real CUDA toolchain. Verify stateless
callback LTO artifacts link with the generated wrapper.

Signed-off-by: Trent Nelson <trent@trent.me>
@tpn
Qualify Numba reduction runtime behavior
5fd699c 
Exercise hierarchy and direct-CUB reductions on representative payloads,
groups, algorithms, and result-visibility modes. Isolate invalid-prefix
trap probes so the parent CUDA context remains usable.

Signed-off-by: Trent Nelson <trent@trent.me>
@tpn
Verify Reduce wheel and typing contracts
86990bc 
Ship the complete Reduce implementation and its required headers in the
universal wheel. Exercise lazy registration and narrow portable and
qualified type signatures to the supported numeric closure.

Signed-off-by: Trent Nelson <trent@trent.me>
@tpn
Canonicalize complete mapped group topology
7c5ec83 
Use canonical ranks whenever a mapping covers its physical parent, even
when it was declared nonexhaustive. This lets storage-bearing logical
WarpReduce fallbacks allocate and select their per-group slices.

Signed-off-by: Trent Nelson <trent@trent.me>
@tpn
Document hierarchy reduction behavior
61160b3 
Describe CUDAX and CUB selection, result visibility, valid prefixes,
storage, synchronization, and deferred functionality. Add a runnable
portable block-sum example with a per-thread payload.

Signed-off-by: Trent Nelson <trent@trent.me>
@tpn
Qualify Python callback LTO by target
e01d516 
Compile stateless device callbacks for the exact compute capability and
include that target in cache and coalescing identities. Reject callback
artifacts before linking them with a provider for another target.

Signed-off-by: Trent Nelson <trent@trent.me>
@tpn
Load Numba reduction providers lazily
85b76c1 
Keep provider registration dormant during backend activation and marker
access. Import and register the Reduce providers only when the first
Reduce or Sum call is planned.

Signed-off-by: Trent Nelson <trent@trent.me>
@tpn
Type qualified reduction callable aliases
f17a31e 
Accept recognized NumPy ufunc and operator callable shapes for
full-group qualified reductions. Keep dtype-specific custom callbacks
on the direct-CUB, root-only overloads.

Signed-off-by: Trent Nelson <trent@trent.me>
@tpn
Constrain callback reduction algorithms
e5e0e5a 
Recognized built-in aliases may use every deterministic block algorithm.

Keep arbitrary callbacks off the commutativity-only algorithm.

Signed-off-by: Trent Nelson <trent@trent.me>
@tpn
Expand Numba reduction path coverage
a71556e 
Exercise callback paths for block arrays and logical-warp prefixes.

Cover public operator aliases and the remaining bitwise built-in path.

Signed-off-by: Trent Nelson <trent@trent.me>
@tpn
Validate the complete cuda-coop wheel surface
27475b4 
Advance the standalone wheel validator from the Load and Store
foundation to the cumulative Exchange, Shuffle, and Reduce surface.

Require every shipped module and header while continuing to reject
deferred Scan, stateful callback, dataclass, and CUTLASS code.

Signed-off-by: Trent Nelson <trent@trent.me>
@tpn
Align Reduce selectors across public APIs
6b15a85 
Use canonical string validation for portable and qualified Reduce
algorithms and built-in operators before provider planning. Keep
qualified callable extensions explicit and internal CUB enums behind the
lowering boundary.

Accept read-only portable payload protocols in qualified reduction stubs
so input contracts match the portable API without weakening mutable
result types.

Signed-off-by: Trent Nelson <trent@trent.me>
@tpn
tpn force-pushed the codex/cuda-coop-numba-hierarchy-reduce branch from ee48a4f to 6b15a85 Compare September 4, 2026 18:21
@tpn
Route Reduce IR access through the compatibility shim
e7ad96e 
Import the validated IR facade from rewrite support. This keeps private
Numba-CUDA-MLIR compiler access inside the guarded compatibility
boundary.

Signed-off-by: Trent Nelson <trent@trent.me>
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Reviewers

Copilot code review Copilot
Copilot will automatically review once the pull request is marked ready for review

Assignees

No one assigned

Labels

None yet

Projects

CCCL
Status: In Progress

Milestone

No milestone

Development

Successfully merging this pull request may close these issues.

1 participant

@tpn