Skip to content

cuda.coop: add Numba Scan prefix callbacks - #11217

Draft
tpn wants to merge 9 commits into
codex/cuda-coop-numba-base-scanfrom
codex/cuda-coop-numba-scan-prefix-callbacks
Draft

cuda.coop: add Numba Scan prefix callbacks#11217
tpn wants to merge 9 commits into
codex/cuda-coop-numba-base-scanfrom
codex/cuda-coop-numba-scan-prefix-callbacks

Conversation

@tpn

@tpn tpn commented Sep 4, 2026 •
edited
Loading

Copy link
Copy Markdown
Contributor

Why this is needed

The base Scan PR deliberately leaves prefix callbacks out so its operator,
result, and storage contracts can be reviewed independently. This follow-up
adds qualified Block Scan prefix callbacks through the same registered family
and generic operator interfaces without changing activation, root dispatch,
caching, or rewrite orchestration.

Prefix callbacks are a compiler-specific extension. The portable API and its
import surface remain unchanged.

Reviewer-visible behavior

  • Adds qualified Block Scan prefix callbacks to Scan, InclusiveScan,
    ExclusiveScan, InclusiveSum, and ExclusiveSum.
  • Uses the canonical prefix_op keyword. The removed
    block_prefix_callback_op compatibility alias is rejected.
  • Supports stateless callbacks and immutable StatefulFunction descriptors
    with explicit numeric, one-item-per-thread state.
  • Allows state and scan dtypes to differ while requiring exact descriptor
    matching.
  • Uses CUB's first-Warp callback ABI and makes thread 0's state authoritative
    after the call. Callers initialize every participating state cell equally.
  • Covers implicit, caller-owned, and dynamic temporary storage, including
    manual reuse synchronization.
  • Rejects prefix callbacks combined with initial values, aggregate results,
    Warp groups, nonnumeric state, or stateful scan operators before provider
    compilation.
  • Preserves the portable Scan signature ordering in the qualified surface;
    callback parameters are additive qualified extensions.

Validation

At exact head 9b299e894feeaeeb98867009f72cb244b3822b5f:

  • 1,440 source-tree tests passed with the GPU hidden.
  • An exact-head wheel installed outside the source tree passed 607 contract,
    packaging, public-surface, and activation tests.
  • The installed wheel passed 50 real NVRTC LTO/nvJitLink compile tests.
  • All 366 installed-wheel runtime tests passed on an RTX PRO 6000.
  • Changed-file pre-commit and git diff --check passed.

Stack and scope

This draft is stacked on #11212. Portable callbacks, Warp callbacks, stateful
binary scan operators, structured state, and grid workspace policy remain
deferred.

The inherited private compatibility shim remains limited to the supported
Numba-CUDA-MLIR 0.5.x contract. CUDA 12/13 L4 and Windows CI remain external
readiness gates.

@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 force-pushed the codex/cuda-coop-numba-base-scan branch 2 times, most recently from 1eef1a3 to 8a594e5 Compare September 4, 2026 18:47
tpn added 9 commits September 4, 2026 11:52
@tpn
Model Scan prefix callbacks in core contracts
a01dac5 
Represent stateless and stateful prefix callbacks in normalized Scan
semantics and place their parameters in CUB overload order.

Reject invalid prefix combinations and unsupported scopes before a
backend materializes a provider.

Signed-off-by: Trent Nelson <trent@trent.me>
@tpn
Compile explicit-state Numba operators
ceecfe6 
Expose an immutable qualified descriptor for device callables and
state dtypes. Lower the generic core contract to a target-aware
pointer ABI with callable-based semantic identity and stable cache
keys.

Signed-off-by: Trent Nelson <trent@trent.me>
@tpn
Add Numba Scan prefix callbacks
9e43334 
Expose Block-only prefix callbacks through the qualified Scan API.
Canonicalize the compatibility spelling, validate explicit per-thread
state before inference, and preserve runtime operands through provider
lowering. Cover every spelling, algorithm, storage mode, and the
two-phase compiler path.

Signed-off-by: Trent Nelson <trent@trent.me>
@tpn
Type Numba Scan callbacks precisely
ca41e5f 
Correlate callback arguments and results with scalar and per-thread
array payloads. Carry state and value types through StatefulFunction
while retaining functor-class support through its bound call shape.

Allow explicit None forwarding for every qualified runtime parameter.
Reject nonnumeric, wrong-arity, and mismatched callback forms during
strict consumer checks.

Signed-off-by: Trent Nelson <trent@trent.me>
@tpn
Verify Scan callback wheel contents
b60fef0 
Require the qualified StatefulFunction implementation and declaration in
the cuda-coop wheel while preserving the deferred-module exclusions.

Exercise lazy import, callable identity, runtime signatures, and package
origins from an isolated installed-wheel process.

Signed-off-by: Trent Nelson <trent@trent.me>
@tpn
Document Numba Scan prefix callbacks
1c07837 
Describe qualified Block Scan callback spellings, state lifetime, and
thread-zero authority. Explain callable exclusions and the interaction
between persistent prefix state, temporary storage, and reuse barriers.

Signed-off-by: Trent Nelson <trent@trent.me>
@tpn
Use string selectors in Scan callback examples
099f52b 
Keep callback coverage aligned with portable Scan by spelling shared
algorithm controls as lowercase strings. Remove stale enum claims from
API documentation so qualified extensions do not imply alternate
selector types.

Signed-off-by: Trent Nelson <trent@trent.me>
@tpn
Remove the Scan prefix callback compatibility alias
91782c2 
Make prefix_op the only qualified Scan keyword for stateless and
stateful prefix callbacks. Remove duplicate overloads and legacy planner
canonicalization while preserving the callback provider paths.

Keep signature, wheel, typing, compiler, and runtime coverage focused on
the canonical spelling, with a negative regression for the removed
keyword.

Signed-off-by: Trent Nelson <trent@trent.me>
@tpn
Preserve qualified Scan surface ordering
9b299e8 
Keep the qualified export list in portable order. Append backend-only
names as one coherent extension block.

Assert that every shared Scan parameter retains its portable contract.
Account explicitly for the positional state and keyword callback
extensions.

Signed-off-by: Trent Nelson <trent@trent.me>
@tpn
tpn force-pushed the codex/cuda-coop-numba-scan-prefix-callbacks branch from d28b8cc to 9b299e8 Compare September 4, 2026 19:00
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