Conversation
|
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. |
|
@coderabbitai full review Please review this PR against its stacked base, focusing on correctness, Expanded API docstrings, user guides, and runnable documentation examples are |
|
✅ Action performedFull review finished. |
|
No actionable comments were generated in the recent review. 🎉 ℹ️ Recent review info⚙️ Run configurationConfiguration used: Repository: NVIDIA/cccl/.coderabbit.yaml Review profile: CHILL Plan: Enterprise Run ID: 📒 Files selected for processing (2)
Included review availability: Your plan provides up to 12 included reviews per hour; 5 remain after this review. 📝 SummarySummary by CodeRabbit
WalkthroughChangesThe pull request adds block TopK selection for minimum and maximum keys or key-value pairs. It adds CUB-backed planning and lowering, public exports, typed declarations, runtime controls, temporary storage support, and compilation, runtime, contract, packaging, and typing tests. Block TopK
Priority: ➖ Normal Change: Feature Merge Risk: ⚪ Minimal · up to The added TopK planning contract coverage matches the recorded uniform participation requirements. No actionable merge risk remains in the reviewed changes. Comment |
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
There was a problem hiding this comment.
🧹 Nitpick comments (1)
python/cuda_coop/cuda/coop/numba_mlir/_compiler/_group_topk.py (1)
52-55: 📐 Maintainability & Code Quality | 🔵 Trivial | 💤 Low valuesuggestion: guard the dtype during payload inference.
_infer_payloadruns before_lower_topk. If both dtype lookups remainNone, validation can fail during inference. Add the diagnostic before_validate_common_numeric_dtype:if dtype is None: raise GroupRewriteError(f"topk could not infer {name} dtype")
ℹ️ Review info
⚙️ Run configuration
Configuration used: Repository: NVIDIA/cccl/.coderabbit.yaml
Review profile: CHILL
Plan: Enterprise
Run ID: c07daef7-5afc-4974-a0f2-6b9d84983760
📒 Files selected for processing (25)
python/cuda_coop/cuda/coop/__init__.pyipython/cuda_coop/cuda/coop/_core/__init__.pypython/cuda_coop/cuda/coop/_core/api/__init__.pypython/cuda_coop/cuda/coop/_core/api/__init__.pyipython/cuda_coop/cuda/coop/_core/api/topk.pypython/cuda_coop/cuda/coop/_core/api/topk.pyipython/cuda_coop/cuda/coop/_core/block/__init__.pypython/cuda_coop/cuda/coop/_core/block/topk.pypython/cuda_coop/cuda/coop/_core/group/__init__.pypython/cuda_coop/cuda/coop/_core/group/topk.pypython/cuda_coop/cuda/coop/numba_mlir/__init__.pypython/cuda_coop/cuda/coop/numba_mlir/__init__.pyipython/cuda_coop/cuda/coop/numba_mlir/_compiler/_group_topk.pypython/cuda_coop/cuda/coop/numba_mlir/_group_topk.pypython/cuda_coop/cuda/coop/numba_mlir/_group_topk.pyipython/cuda_coop/cuda/coop/numba_mlir/_lowering/_topk.pypython/cuda_coop/tests/backends/numba_mlir/compile/test_topk_compile.pypython/cuda_coop/tests/backends/numba_mlir/runtime/test_topk.pypython/cuda_coop/tests/backends/numba_mlir/unit/test_public_foundation.pypython/cuda_coop/tests/backends/numba_mlir/unit/test_topk_planning.pypython/cuda_coop/tests/contracts/core/test_core_block_topk.pypython/cuda_coop/tests/packaging/test_wheel.pypython/cuda_coop/tests/packaging/typing/invalid_consumer.pypython/cuda_coop/tests/packaging/typing/numba_consumer.pypython/cuda_coop/tests/packaging/typing/portable_consumer.py
Included review availability: Your plan provides up to 12 included reviews per hour; 9 remain after this review.
8e1bbbc to
ff7a94c
Compare
ec442e5 to
a987aae
Compare
|
@coderabbitai review This PR was restacked over the read-only input validation correction in Merge Sort #11510. Its own primitive implementation is unchanged; preserve the dispositions from the completed full review. 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. |
|
important: I will review the current head incrementally. I will preserve the completed review dispositions and limit new findings to credible correctness or contract regressions in this PR. ✅ Action performedReview finished.
|
There was a problem hiding this comment.
Copilot review overview
🟡 Changes recommended
Unresolved critical contract and moderate packaging and validation issues must be addressed before approval.
Get a fresh assessment by requesting another Copilot review.
Review effort: Lite
Findings: 1
Open (3)
What changed in this PR
Adds block TopK key and key-pair operations with CUB and Numba-CUDA-MLIR integration, validation, typing, and GPU coverage.
Changes:
- Adds min/max TopK APIs and backend lowering.
- Adds planning, storage, and count validation.
- Expands tests, exports, typing, and wheel checks.
Final findings include one critical uniformity-contract issue and moderate issues involving temp_storage validation and independent packaging allowlists.
| File | Summary |
|---|---|
python/cuda_coop/tests/packaging/typing/portable_consumer.py |
Portable TopK typing coverage. |
python/cuda_coop/tests/packaging/typing/numba_consumer.py |
Qualified TopK typing coverage. |
python/cuda_coop/tests/packaging/typing/invalid_consumer.py |
Invalid-call diagnostics. |
python/cuda_coop/tests/packaging/test_wheel.py |
Wheel membership and header checks. |
python/cuda_coop/tests/contracts/core/test_core_block_topk.py |
Core TopK contract tests. |
python/cuda_coop/tests/backends/numba_mlir/unit/test_topk_planning.py |
Planning validation tests. |
python/cuda_coop/tests/backends/numba_mlir/unit/test_public_foundation.py |
Public operation registration checks. |
python/cuda_coop/tests/backends/numba_mlir/runtime/test_topk.py |
GPU runtime behavior tests. |
python/cuda_coop/tests/backends/numba_mlir/compile/test_topk_compile.py |
Provider compilation tests. |
python/cuda_coop/cuda/coop/numba_mlir/_lowering/_topk.py |
TopK provider factories. |
python/cuda_coop/cuda/coop/numba_mlir/_group_topk.pyi |
Qualified TopK type declarations. |
python/cuda_coop/cuda/coop/numba_mlir/_group_topk.py |
Qualified TopK API. |
python/cuda_coop/cuda/coop/numba_mlir/_compiler/_group_topk.py |
TopK planning and lowering. |
python/cuda_coop/cuda/coop/numba_mlir/__init__.pyi |
Qualified export typing. |
python/cuda_coop/cuda/coop/numba_mlir/__init__.py |
Qualified runtime exports. |
python/cuda_coop/cuda/coop/_core/group/topk.py |
Group planning semantics. |
python/cuda_coop/cuda/coop/_core/group/__init__.py |
Group semantic exports. |
python/cuda_coop/cuda/coop/_core/block/topk.py |
CUB TopK specification and validation. |
python/cuda_coop/cuda/coop/_core/block/__init__.py |
Block specification exports. |
python/cuda_coop/cuda/coop/_core/api/topk.pyi |
Portable TopK declarations. |
python/cuda_coop/cuda/coop/_core/api/topk.py |
Portable TopK API. |
python/cuda_coop/cuda/coop/_core/api/__init__.pyi |
API type exports. |
python/cuda_coop/cuda/coop/_core/api/__init__.py |
API runtime exports. |
python/cuda_coop/cuda/coop/_core/__init__.py |
Core symbol exports. |
python/cuda_coop/cuda/coop/__init__.pyi |
Public TopK exports. |
💡 Add a code-review agent skill or configure MCP servers for context-aware, tailored reviews. Learn more in the docs.
Select minimum or maximum keys and pairs with blocked per-thread results. Preserve inputs, validate wide count controls, and support caller-managed temporary storage. Isolate the private CUB TopK dependency in the family provider. Signed-off-by: Trent Nelson <trent@trent.me>
Record dtypes inferred from indexed ThreadData writes on their original constructors so copied results remain typed in subsequent group calls. Exercise core and qualified key/pair chains with irregular block sizes and add positive and negative public typing consumer coverage. Signed-off-by: Trent Nelson <trent@trent.me>
ff7a94c to
14e0f01
Compare
|
@coderabbitai review Please review the current head incrementally. The only change since your completed review is the TopK participation contract recording uniform Focus on credible correctness and contract regressions in this PR relative to #11511. Downstream #11320 remains documentation, and #11392 handles launch bounds. The two packaging comments have been addressed with evidence of the existing real-wheel CI checks. |
|
✅ Action performedReview finished.
|
🥳 CI Workflow Results🟩 Finished in 29m 16s: Pass: 100%/8 | Total: 56m 13s | Max: 18m 12sSee results here. |
Why this is needed
A kernel often needs only the smallest or largest keys in its block's tile.
This adds common and Numba-CUDA-MLIR TopK operations without requiring a
full sort:
Behavior
payloads and static or runtime
kandvalid_itemscontrols.The core participation contract records that
kand suppliedvalid_itemsmust be uniform across the block.min(k, valid_items)blocked output positions are defined. The selectionis unordered, and ties have no ordering or selection guarantee.
TempStorage, with the existingreuse-synchronization contract. The qualified API also accepts local arrays.
compiler, GPU, and isolated runtime-trap coverage.
The provider isolates the bundled private
cub::detail::block_topkimplementation behind one compatibility shim. It uses the current CUB
signature; bit-range controls and Warp TopK are not exposed.
Stacked on #11511. User guides, expanded API documentation, and
executable examples are added in downstream #11320.
Validation
A fresh installed wheel at the final stack tip passed 751 host/packaging
checks on each of Python 3.10 and 3.14, plus 54 focused TopK and backend
unit/compiler checks with GPUs hidden. Six host regressions reproduced
the missing uniform-count preconditions before the correction.
The family previously passed 56 GPU tests on CUDA 13.3 with
Numba-CUDA-MLIR 0.5.2. The uniform-count correction changes declarative
planner metadata; provider and compiler lowering code are unchanged.
Source, wheel, installed module, and bundled-header bytes were compared.
Wheel validation and changed-file pre-commit passed.