Add cuda.coop Merge Sort with Numba-CUDA-MLIR by tpn · Pull Request #11510 · NVIDIA/cccl

tpn · 2026-09-18T23:12:12Z

Why this is needed

Kernels need to order the values already held by a cooperative group and
keep associated data, such as original indices, attached to each key.
This adds common and Numba-CUDA-MLIR Merge Sort APIs:

ordered_keys, ordered_indices = coop.merge_sort_pairs(block, keys, indices)

Behavior

Add merge_sort_keys and merge_sort_pairs for Block, physical Warp,
and logical Warp groups. Inputs are preserved and results use blocked
layout. Equal-key order is unspecified. The common API accepts read-only
ThreadData inputs.
Support ascending/descending order, partial tiles with explicit sentinels,
and independent scratch for each physical or logical Warp group.
Support explicit Block TempStorage; the qualified namespace also accepts
fixed local arrays and custom comparison callbacks.
Integrate through the existing family registry, planning, and rewrite
hooks, with public stubs and host/compiler/GPU coverage.

Stacked on #11217. User guides, expanded API documentation, and executable
documentation examples are added in the downstream docs PR #11320.

Validation

Fresh installed wheels passed 685 host/packaging checks on each of
Python 3.10 and 3.14,
56 focused unit/compiler checks with GPUs hidden, and 57 GPU tests
on CUDA 13.3 with Numba-CUDA-MLIR 0.5.2. Source, wheel, installed module,
and bundled-header bytes were compared. Wheel validation and changed-file
pre-commit passed.

Remote CI is separate from these local results.

copy-pr-bot · 2026-09-18T23:12:16Z

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 · 2026-09-18T23:35:14Z

@coderabbitai full review

Please review this PR against its stacked base, focusing on correctness,
compiler integration, synchronization, input/output contracts, and meaningful
regressions. Merge Sort is based on #11217. Radix Rank/Sort (#11511) and TopK (#11512) follow it.

Expanded API docstrings, user guides, and runnable documentation examples are
in the downstream docs PR #11320; launch-bounds inference follows in #11392.
Consider those layers when assessing apparent omissions, and distinguish an
actual defect in this PR from an intentional downstream addition. Keep optional
refactoring and unrelated features separate from actionable findings.

coderabbitai · 2026-09-18T23:35:26Z

important: I will review the PR against its stacked base. I will focus on correctness, compiler integration, synchronization, input/output contracts, and regressions. I will treat the documentation and launch-bounds work in #11320 and #11392 as downstream scope.

✅ Action performed

Full review finished.

Copilot

Copilot review overview

🟡 Changes recommended

Unresolved API read-only handling and packaging/export validation gaps remain.

Get a fresh assessment by requesting another Copilot review.

Review effort: Lite
Findings: 1

Open (1)

Allow read-only payloads in merge sort key validation · New

What changed in this PR

Adds portable and Numba-CUDA-MLIR Merge Sort APIs for block and warp groups, including key/pair sorting, partial tiles, scratch storage, and custom comparisons.

Changes:

Adds backend-neutral semantics, planning, and CUB providers.
Integrates lowering, exports, typing, and packaging.
Adds contract, compiler, runtime, GPU, and packaging tests.

Review findings:

Moderate (2 votes): Allow read-only thread data in the Merge Sort API checks.
Moderate (1 vote): Add Merge Sort files and headers to the independent wheel validator.
Nit (1 vote): Require both new exports in the installed-wheel coop.__all__ probe.

File	Reviewed change
`python/cuda_coop/tests/packaging/typing/portable_consumer.py`	Portable typing coverage
`python/cuda_coop/tests/packaging/typing/numba_consumer.py`	Qualified typing coverage
`python/cuda_coop/tests/packaging/typing/invalid_consumer.py`	Invalid API typing cases
`python/cuda_coop/tests/packaging/test_wheel.py`	Wheel payload coverage
`python/cuda_coop/tests/packaging/test_metadata.py`	Package metadata coverage
`python/cuda_coop/tests/contracts/core/test_core_merge_sort.py`	Core planning contracts
`python/cuda_coop/tests/backends/numba_mlir/unit/test_public_foundation.py`	Backend surface checks
`python/cuda_coop/tests/backends/numba_mlir/unit/test_merge_sort_planning.py`	Planning diagnostics
`python/cuda_coop/tests/backends/numba_mlir/runtime/test_merge_sort.py`	GPU runtime tests
`python/cuda_coop/tests/backends/numba_mlir/compile/test_merge_sort_compile.py`	Compilation tests
`python/cuda_coop/cuda/coop/numba_mlir/_lowering/_merge_sort.py`	Provider generation
`python/cuda_coop/cuda/coop/numba_mlir/_group_merge_sort.pyi`	Qualified type stubs
`python/cuda_coop/cuda/coop/numba_mlir/_group_merge_sort.py`	Qualified entry points
`python/cuda_coop/cuda/coop/numba_mlir/_compiler/_rewrite_merge_sort.py`	Payload inference
`python/cuda_coop/cuda/coop/numba_mlir/_compiler/_group_merge_sort.py`	Backend planning and lowering
`python/cuda_coop/cuda/coop/numba_mlir/__init__.pyi`	Qualified exports
`python/cuda_coop/cuda/coop/numba_mlir/__init__.py`	Lazy qualified exports
`python/cuda_coop/cuda/coop/_core/warp/merge_sort.py`	Warp semantics
`python/cuda_coop/cuda/coop/_core/warp/__init__.py`	Warp exports
`python/cuda_coop/cuda/coop/_core/group/merge_sort.py`	Group planning
`python/cuda_coop/cuda/coop/_core/group/__init__.py`	Group exports
`python/cuda_coop/cuda/coop/_core/block/merge_sort.py`	Block semantics
`python/cuda_coop/cuda/coop/_core/block/__init__.py`	Block exports
`python/cuda_coop/cuda/coop/_core/api/merge_sort.pyi`	Portable type stubs
`python/cuda_coop/cuda/coop/_core/api/merge_sort.py`	Portable entry points
`python/cuda_coop/cuda/coop/_core/api/__init__.pyi`	API stub exports
`python/cuda_coop/cuda/coop/_core/api/__init__.py`	API runtime exports
`python/cuda_coop/cuda/coop/_core/__init__.py`	Core exports
`python/cuda_coop/cuda/coop/__init__.pyi`	Top-level type exports

💡 Configure MCP servers for context-aware, tailored reviews. Learn more in the docs.

coderabbitai · 2026-09-18T23:42:35Z

No actionable comments were generated in the recent review. 🎉

ℹ️ Recent review info

⚙️ Run configuration

Configuration used: Repository: NVIDIA/cccl/.coderabbit.yaml

Review profile: CHILL

Plan: Enterprise

Run ID: aba2e5c4-598a-4e0a-a7cb-4e3c5863ee6d

📥 Commits

Reviewing files that changed from the base of the PR and between ab4123b and e574f2e.

📒 Files selected for processing (3)

python/cuda_coop/cuda/coop/_core/api/merge_sort.py
python/cuda_coop/tests/contracts/core/test_core_merge_sort.py
python/cuda_coop/tests/packaging/typing/portable_consumer.py

Included review availability: Your plan provides up to 12 included reviews per hour; 7 remain after this review.

📝 Summary

Summary by CodeRabbit

New Features
- Added cooperative CUDA merge sort for keys and key/value pairs.
- Supports block and warp groups, ascending or descending order, partial inputs, custom comparators, and optional temporary storage.
- Preserves original inputs while returning sorted results.
- Added public APIs and typing support across portable and Numba-MLIR interfaces.
Tests
- Added comprehensive coverage for sorting behavior, validation, compilation, runtime execution, typing, and packaging.

Walkthrough

Added cooperative CUDA merge-sort APIs for keys and key/value pairs. The implementation supports block and warp groups, full and partial tiles, custom comparators, typed payloads, CUB lowering, and Numba-MLIR compilation.

Changes

Merge Sort

Layer / File(s)	Summary
Merge sort contracts and group planning `python/cuda_coop/cuda/coop/_core/{block,group,warp}/...`, `python/cuda_coop/cuda/coop/_core/api/...`, `python/cuda_coop/cuda/coop/{__init__.pyi,_core/__init__.py}`	Added block, warp, and group merge-sort semantics. Added validation for payloads, tile sizes, dimensions, validity counts, defaults, comparison operators, and supported group widths. Added typed public exports and overloads.
Compiler registration and CUDA lowering `python/cuda_coop/cuda/coop/numba_mlir/...`	Added keys and pairs entry points, payload inference, operation registration, comparator handling, input-preserving result construction, and CUB providers for block and warp full or partial sorts.
Compile, runtime, and contract validation `python/cuda_coop/tests/{contracts,backends,packaging}/...`	Added tests for planning diagnostics, CUB specialization, runtime sorting, partial-count guards, custom comparators, input preservation, typing surfaces, public exports, and wheel contents.

Priority: ➖ Normal

Change: Feature

Merge Risk: ⚪ Minimal · up to e574f

No merge-blocking issue is established for the added merge-sort API surface.

_{Comment @coderabbitai help to get the list of available commands.}

Add input-preserving key and key/value Merge Sort operations with blocked results, partial tiles, and independent storage for logical warps. Qualified Numba calls also accept local arrays and stateless comparators. Validate runtime tile counts before narrowing them to the CUB ABI and retain inferred payload dtypes across chained operations. Signed-off-by: Trent Nelson <trent@trent.me>

tpn · 2026-09-19T00:00:19Z

@coderabbitai review

The read-only input validation finding is fixed in the owning Merge Sort commit, with 12 host regressions and readable-payload typing coverage.

Please review the current head incrementally. Expanded documentation remains in #11320, followed by launch-bounds inference in #11392. Focus on credible correctness or contract findings within this PR.

coderabbitai · 2026-09-19T00:00:32Z

@tpn important: I will review the current head incrementally. I will focus on credible correctness, input/output contract, compiler integration, synchronization, and regression findings in this PR. I will exclude the downstream documentation and launch-bounds scope.

✅ Action performed

Review finished.

Note: CodeRabbit is an incremental review system and does not re-review already reviewed commits. This command is applicable only when automatic reviews are paused.

github-actions · 2026-09-19T00:29:48Z

🥳 CI Workflow Results

🟩 Finished in 28m 21s: Pass: 100%/8 | Total: 47m 12s | Max: 11m 46s

See results here.

github-project-automation Bot added this to CCCL Sep 18, 2026

github-project-automation Bot moved this to Todo in CCCL Sep 18, 2026

cccl-authenticator-app Bot moved this from Todo to In Progress in CCCL Sep 18, 2026

tpn mentioned this pull request Sep 18, 2026

Add cuda.coop Radix Rank and Sort with Numba-CUDA-MLIR #11511

Open

tpn marked this pull request as ready for review September 18, 2026 23:34

tpn requested a review from a team as a code owner September 18, 2026 23:34

tpn requested review from NaderAlAwar and a lite review from Copilot and removed request for a team September 18, 2026 23:34

cccl-authenticator-app Bot moved this from In Progress to In Review in CCCL Sep 18, 2026

Copilot started reviewing on behalf of tpn September 18, 2026 23:34 View session

tpn mentioned this pull request Sep 18, 2026

Add cuda.coop TopK with Numba-CUDA-MLIR #11512

Open

Copilot AI reviewed Sep 18, 2026

View reviewed changes

Comment thread python/cuda_coop/cuda/coop/_core/api/merge_sort.py

tpn force-pushed the codex/cuda-coop-numba-merge-sort branch from ab4123b to e574f2e Compare September 18, 2026 23:58

Navigation Menu

Add cuda.coop Merge Sort with Numba-CUDA-MLIR - #11510

Add cuda.coop Merge Sort with Numba-CUDA-MLIR#11510
tpn wants to merge 1 commit into
codex/cuda-coop-numba-scan-prefix-callbacksfrom
codex/cuda-coop-numba-merge-sort

tpn commented Sep 18, 2026 •

edited

Loading

Uh oh!

copy-pr-bot Bot commented Sep 18, 2026

Uh oh!

tpn commented Sep 18, 2026

Uh oh!

coderabbitai Bot commented Sep 18, 2026 •

edited

Loading

Uh oh!

Copilot AI left a comment

Uh oh!

Uh oh!

coderabbitai Bot commented Sep 18, 2026 •

edited

Loading

Summary by CodeRabbit

Walkthrough

Changes

Uh oh!

tpn commented Sep 19, 2026

Uh oh!

coderabbitai Bot commented Sep 19, 2026 •

edited

Loading

Uh oh!

github-actions Bot commented Sep 19, 2026

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

2 participants

Conversation

tpn commented Sep 18, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Why this is needed

Behavior

Validation

Uh oh!

copy-pr-bot Bot commented Sep 18, 2026

Uh oh!

tpn commented Sep 18, 2026

Uh oh!

coderabbitai Bot commented Sep 18, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

Copilot AI left a comment

Choose a reason for hiding this comment

Copilot review overview

🟡 Changes recommended

Uh oh!

Uh oh!

coderabbitai Bot commented Sep 18, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Summary by CodeRabbit

Walkthrough

Changes

Uh oh!

tpn commented Sep 19, 2026

Uh oh!

coderabbitai Bot commented Sep 19, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

github-actions Bot commented Sep 19, 2026

🥳 CI Workflow Results

🟩 Finished in 28m 21s: Pass: 100%/8 | Total: 47m 12s | Max: 11m 46s

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

2 participants

tpn commented Sep 18, 2026 •

edited

Loading

coderabbitai Bot commented Sep 18, 2026 •

edited

Loading

coderabbitai Bot commented Sep 18, 2026 •

edited

Loading

coderabbitai Bot commented Sep 19, 2026 •

edited

Loading