cuda.coop: add Numba Block and Warp Scan - #11212
Draft
tpn wants to merge 11 commits into
Draft
Conversation
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. |
tpn
force-pushed
the
codex/cuda-coop-numba-hierarchy-reduce
branch
from
September 4, 2026 18:21
ee48a4f to
6b15a85
Compare
tpn
force-pushed
the
codex/cuda-coop-numba-base-scan
branch
from
September 4, 2026 18:30
cf1672c to
1eef1a3
Compare
Define backend-neutral Block and Warp Scan semantics, including out-of-place results, aggregate side outputs, prefix bounds, and storage contracts. Expose five portable group-first spellings while keeping backend-only controls out of the root API. Signed-off-by: Trent Nelson <trent@trent.me>
Materialize scalar and array BlockScan plus scalar physical and logical WarpScan providers through the existing core adapter. Preserve callback LTO, CUB storage, and scope-derived reuse barriers. Add an explicit dependent type-expression placeholder so partial exclusive sums inject a typed zero. Signed-off-by: Trent Nelson <trent@trent.me>
Recognize all five Scan spellings through the family registry and lower them from exact launch topology. Carry inferred payload, initial, aggregate, prefix-width, and storage facts into providers without changing the shared compiler phases. Signed-off-by: Trent Nelson <trent@trent.me>
Exercise every numeric dtype through real provider compilation and cover Block, physical-Warp, and logical-Warp execution. Cover aggregate, initial, prefix, storage, callback, and isolated invalid-prefix paths so every concrete ABI runs end to end. Signed-off-by: Trent Nelson <trent@trent.me>
Require complete portable and qualified Scan APIs in installed wheels while preserving lazy compiler-family loading. Exercise valid and invalid static call shapes, including Python initial literals and the explicit N6 exclusions. Signed-off-by: Trent Nelson <trent@trent.me>
Describe Scan modes, operators, result ownership, partial-Warp rules, storage synchronization, and the prefix-callback deferral. Add a runnable portable Block Load, Scan, and Store example and expose the family from the Python documentation index. Signed-off-by: Trent Nelson <trent@trent.me>
Extend the standalone wheel validator with every portable and Numba Scan module plus the bundled CUB BlockScan and WarpScan headers. Keep stateful prefix callbacks, dataclasses, CUTLASS, and AOT payloads out of the supported wheel surface. Signed-off-by: Trent Nelson <trent@trent.me>
Require string selectors for Scan across the portable and qualified surfaces, and normalize shared aliases through one path. Keep callable operators as a qualified-only extension. Use common read-only payload, storage, and contextual initial-value typing contracts, and remove the public BlockScanAlgorithm enum. Signed-off-by: Trent Nelson <trent@trent.me>
Reject string-backed enums for Scan modes and algorithms. Apply the same rule to built-in operators. Enforce it at portable, planner, and provider boundaries. Match shared runtime signatures while retaining qualified extensions. Prevent the removed public enum API from returning through coercions. Signed-off-by: Trent Nelson <trent@trent.me>
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>
tpn
force-pushed
the
codex/cuda-coop-numba-base-scan
branch
from
September 4, 2026 18:47
1eef1a3 to
8a594e5
Compare
Retain separate variables for the portable descriptor and the qualified descriptor viewed through their shared protocol. This keeps the strict consumer valid as Scan adds its cross-surface storage check. Signed-off-by: Trent Nelson <trent@trent.me>
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
1 participant
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Why this is needed
The preceding
cuda.coopstack establishes declarative group topology,storage, result projection, and stateless operator compilation. This PR adds
Scan as another independently registered family without changing activation,
root dispatch, caching, or rewrite phase orchestration.
Reviewer-visible behavior
InclusiveSum, and ExclusiveSum entry points.
lowercase modes, operators and algorithms, storage descriptors, and dtype
inference for untyped
ThreadData.valid_items, aggregate output, and compiler-resolved callback or operatoraliases.
ThreadDatapayloads and scalar physical orlogical Warp payloads.
implicit/caller/dynamic Block storage, and compiler-owned per-Warp storage.
only outputs within the valid prefix are defined.
Validation
The exact final-stack head
9b299e894f, which includes this layer, passed:tests outside the source tree.
Stack and scope
This draft is stacked on #11210. Stateful and prefix Scan callbacks and any
grid workspace policy remain deferred. CUDA 12/13 L4 and Windows CI remain
external readiness gates.