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.
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.
|
This comment has been minimized.
This comment has been minimized.
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>
a987aae to
f793bd9
Compare
14e0f01 to
031509d
Compare
🥳 CI Workflow Results🟩 Finished in 25m 14s: Pass: 100%/8 | Total: 1h 19m | Max: 20m 05sSee 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
Fresh installed-wheel checks at
031509df2a3c:Wheel validation and source, installed-package, and bundled-header comparisons passed.
Remote CI passed at
031509df2a3c: GitHub Actions run. All exact-head checks and statuses, including pre-commit, are successful or conditionally skipped.