[cudax] Harden cooperative Reduce reuse - #11209
Draft
tpn wants to merge 2 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-exchange-shuffle
branch
from
September 4, 2026 18:10
7587282 to
03793c5
Compare
tpn
force-pushed
the
codex/cuda-coop-cudax-reduce-reuse-prereq
branch
from
September 4, 2026 18:11
ee61c0a to
64d2fcd
Compare
Mapped warp groups synchronize independently, but reduce used one shared scratch object. Concurrent siblings could corrupt CUB storage, partials, and broadcast results. Passing owning groups by value could also destroy copied synchronizer state. Give each physical block warp its own temporary and value slot. Use the exact static block extent when available and the architectural 32-warp bound otherwise. Reject noncontiguous mappings until a safe physical membership traversal is available. Add the collective barriers required before block, cluster, grid, and mapped-group scratch is reused by a later call. Add repeated regressions for direct, nested, and viewed mappings with static and dynamic extents, plus block, cluster, and grid scratch reuse. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com> Signed-off-by: Trent Nelson <trent@trent.me>
Embedded compiler providers cannot always include cooperative_groups.h. Add an opt-out that suppresses the include and conversion overloads while preserving the default interoperability behavior. Signed-off-by: Trent Nelson <trent@trent.me>
tpn
force-pushed
the
codex/cuda-coop-cudax-reduce-reuse-prereq
branch
from
September 4, 2026 18:20
64d2fcd to
be03724
Compare
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 next
cuda.coopNumba-CUDA-MLIR layer lowers hierarchy-aware reductionsthrough
cuda::experimental::coop::reduce. Consecutive reductions cannotsafely reuse CUB scratch until every member has finished, and independently
synchronized mapped Warp groups must not alias one shared scratch object.
Embedded compiler providers also need to consume CUDAX group headers without
pulling in
cooperative_groups.h.This is a separate prerequisite so the CUDAX correctness change can be
reviewed independently from the Python Reduce API.
Reviewer-visible behavior
mapped-group scratch is reused.
mapped groups, and static and dynamic block extents.
_CUDAX_DISABLE_COOPERATIVE_GROUPS_INTEROPas an opt-out for embeddedcompilation while preserving the existing default interop behavior.
Validation
Focused CUDAX targets passed on an RTX PRO 6000 with CUDA 13.3 under both
C++17 and C++20:
Changed-file pre-commit, including clang-format, gersemi, codespell, and
secret scanning, and
git diff --checkpassed.Stack and scope
This draft is stacked on #11208. It does not add a Python API or compiler
lowering. Grid Reduce still requires a per-launch workspace policy before a
Numba backend can expose it safely.