Fixes various issues in partial tile merge sort - #10733
Conversation
📝 WalkthroughSummary by CodeRabbit
WalkthroughChangesBlock merge sort now supports partial tiles without sentinel values. It adds padded shared-memory storage, clamped merge rounds, valid-item-only Partial tile merge sort
Assessment against linked issues
Out-of-scope changes
Suggested reviewers: Merge Risk: 🔵 Low · up to The sorting fix is localized, but test helpers still need to follow required API conventions, and the custom-type size check is weaker than intended. The PR is mergeable with explicit owner awareness and follow-up on these bounded issues. Comment |
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
There was a problem hiding this comment.
Actionable comments posted: 1
🧹 Nitpick comments (5)
cub/cub/block/block_merge_sort.cuh (5)
225-233: 🚀 Performance & Scalability | 🔵 Trivial | 💤 Low valuesuggestion: the padding looks larger than required. Runs are tile-aligned, so
keys1_end <= start + size <= ITEMS_PER_TILE - sizeandsize >= ItemsPerThread.serial_mergeadvances at mostItemsPerThreadpositions pastkeys1_beg_loc, so the largest key index read isITEMS_PER_TILE, andkeys2reads stop atkeys2_end <= ITEMS_PER_TILE.ITEMS_PER_TILE + 1then covers every read, and the extraItemsPerThread - 1slots increase shared memory for allBlockMergeSortusers, including full-tile sorts. If you keep the larger padding, state in the comment which access actually needs more than one slot.Source: Path instructions
468-502: 📐 Maintainability & Code Quality | 🔵 Trivial | 💤 Low valuesuggestion: add the
@rst .. versionadded::block used by every neighboring overload, so the new public entry points carry a version marker in the generated docs.
844-890: 📐 Maintainability & Code Quality | 🔵 Trivial | 💤 Low valuesuggestion: honor
_Unrolland mark the indexconst. The merge loops use_CCCL_PRAGMA_UNROLL(_Unroll ? ItemsPerThread : 1), but these loops force full unrolling and ignore thepolicy.unrollknob.idxis never modified at lines 849 and 871.Proposed change
- _CCCL_PRAGMA_UNROLL_FULL() + _CCCL_PRAGMA_UNROLL(_Unroll ? ItemsPerThread : 1) for (int item = 0; item < ItemsPerThread; ++item) { - int idx = ItemsPerThread * linear_tid + item; + const int idx = ItemsPerThread * linear_tid + item; temp_storage.keys_shared[idx] = keys[item]; }As per coding guidelines: "All variables that are not modified must be declared
const".Source: Coding guidelines
895-898: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚖️ Poor tradeoffsuggestion:
MergeRoundsandMergeRoundsClampedduplicate roughly 50 lines that differ only in the run and diagonal clamping. Consider one implementation parameterized by abool Clampednon-type template parameter, with the clamped boundary computation behindif constexpr. That keeps the two paths from drifting apart in later fixes.
954-986: 🚀 Performance & Scalability | 🔵 Trivial | ⚡ Quick winsuggestion: exit the round loop early for small
valid_items. OnceItemsPerThread * merged_threads_number >= valid_items, run 2 of the only non-empty group is empty, so the round is an identity copy plus twoSync()calls, and all later rounds are identical.valid_itemsandsizeare uniform across the block, so the exit is block-uniform and does not create a barrier divergence hazard. For a nearly empty final tile this removes most rounds.Proposed change
const int merged_threads_number = target_merged_threads_number / 2; const int mask = target_merged_threads_number - 1; + + // All remaining rounds merge a single non-empty run with an empty one: nothing left to do. + if (ItemsPerThread * merged_threads_number >= valid_items) + { + break; + }
ℹ️ Review info
⚙️ Run configuration
Configuration used: Path: .coderabbit.yaml
Review profile: CHILL
Plan: Enterprise
Run ID: 108c0304-ec7e-4a04-8f80-d43b449aba44
📒 Files selected for processing (2)
cub/cub/agent/agent_merge_sort.cuhcub/cub/block/block_merge_sort.cuh
This comment has been minimized.
This comment has been minimized.
This comment has been minimized.
This comment has been minimized.
|
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. |
- proto_ballot_merge.cu + BALLOT_MERGE_RESULTS.md: three-way stable-merge comparison on B200 (umbriel-b200-022). Rank-augmented bitonic merger wins both axes at all sizes (stable 64-elem merge in 341 cyc keys / 518 pairs, 4-5x the smem co-rank baseline); the ballot-routed merge-path design works and beats smem on latency but is collective-issue-bound and strictly dominated - documented as a principled negative. Recommendation: rank augmentation is the route to WarpStableMerge / stable WarpBitonicSort. - WMS_STATIC_SWITCH_RESULTS.md + proto_wms_static.cu + wms_static/: the MERGE_SORT_SEARCH_STATIC switch productization study (warp -14..-24% latency, block -10..-14% for tiles <= 512, occupancy cost quantified; default stays dynamic). - MERGE_SORT_IMPLEMENTATION_REPORT.md + MERGE_SORT_PARTIAL_TILE_UNINIT_READS.md: walkthrough and analysis behind PR NVIDIA#10733 / issue NVIDIA#5327. - proto_merge_fix.cu + merge_fix/: dual-build validation harness and patched headers for the partial-tile fix branch. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
b623edc to
05a0af2
Compare
… and B Real-header three-way on B200 (umbriel-b200-037), today's main. Full and near-full tiles: all three within +-3-7% (no regression; FIX-A even leads full-tile throughput slightly). Overload A pays for its now-delivered suffix contract on partials (full-tile work regardless of valid_items: ~1.5x latency, ~0.5x throughput at valid=N/8 vs old impl). Overload B beats the old implementation outright on partials (-9..-30% latency, up to +60% valid-normalized throughput at N/8; early exit) and ties at full. Correctness gate green on both builds; FIX-A delivers suffix==oob everywhere, BASE does not (stock gap reconfirmed vs today's main). Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
|
Note GitHub couldn't provide a complete incremental comparison for this pull request, so CodeRabbit is performing a full review instead. This review may take a little longer. |
There was a problem hiding this comment.
Actionable comments posted: 1
🧹 Nitpick comments (3)
c2h/include/c2h/custom_type.h (1)
58-71: 🎯 Functional Correctness | 🔵 Trivial | ⚡ Quick winsuggestion: Add compile-time layout checks for
huge_datapolicy packs.Current instantiations use one
huge_datapolicy and one-byte remaining policies. A secondhuge_datapolicy or a larger policy makesobject_bytesdiffer fromsizeof(custom_type_t<...>). Add compile-time size assertions or enforce this policy-pack restriction.Source: Coding guidelines
cub/test/catch2_test_block_merge_sort.cu (2)
195-200: 🎯 Functional Correctness | 🔵 Trivial | 🏗️ Heavy liftsuggestion: Keep the no-sentinel parameterizations independent of sentinel construction. The new actions ignore
oob_default, but bothblock_merge_sorthelpers still evaluatecuda::std::numeric_limits<KeyT>::max()at lines 168 and 182. A key type without a sentinel value cannot instantiate these tests. Add a no-sentinel helper and kernel path that does not construct or acceptoob_default.
As per path instructions, this CUB test should verify the real API contract and relevant correctness coverage. The supplied production contract states that the no-sentinel overload does not require a sentinel value.Also applies to: 216-221, 266-273
Source: Path instructions
195-200: 🩺 Stability & Availability | 🔵 Trivialimportant: Run the targeted key and pair partial-tile tests with
compute-sanitizer --tool initcheck. Confirm that thecompute_init_lid0CI job passes. Also run the linkedDeviceReduceByKeyandDeviceScanByKeytests before merge. As per coding guidelines, build and run relevant tests, prefer targeted builds, and report the results.Also applies to: 216-221, 232-232, 266-273, 303-303
Source: Coding guidelines
ℹ️ Review info
⚙️ Run configuration
Configuration used: Path: .coderabbit.yaml
Review profile: CHILL
Plan: Enterprise
Run ID: 6ceb6a24-531b-4197-8bce-1020e384d448
📒 Files selected for processing (4)
c2h/include/c2h/custom_type.hcub/cub/agent/agent_merge_sort.cuhcub/cub/block/block_merge_sort.cuhcub/test/catch2_test_block_merge_sort.cu
🚧 Files skipped from review as they are similar to previous changes (2)
- cub/cub/agent/agent_merge_sort.cuh
- cub/cub/block/block_merge_sort.cuh
Included review availability: Your plan provides up to 12 included reviews per hour; 11 remain after this review.
⏱️ CCCL compile-time benchmark comparison: Public headers compile-time benchResult: 0 regression row(s), 3 improvement row(s) above threshold.
Artifacts: reports and traces Direct file processing
🟢 Direct file processing — Improvements
|
| // ordering relative to the valid keys therefore cannot affect the sorted prefix, so callers | ||
| // that violate the documented ordering requirement still get their valid items sorted. Keys | ||
| // beyond the boundary are left unspecified. | ||
| SortPartialTile<true>(keys, items, compare_op, valid_items, oob_default); |
There was a problem hiding this comment.
To be sure, its intentional that this now calls MergeRounds<true> in the end?
There was a problem hiding this comment.
Yeah, it's intentional - for now - (even though I dislike it).
After discussing with Georgii, we agreed to separate the two concerns:
a. Fix compute-sanitizer of DeviceMergeSort by introducing the new BlockMergeSort interface that does not take an oob_default - as that may not be generally available, e.g., for custom comparators (as you had pointed out.
b. For the existing overload that takes an oob_default, i.e., Sort(..., valid_items, oob_default). We would like to apply our typical pad-oob-with-oob_default then run the algorithm on the full tile.
We agreed to defer (b) for now, even though I would really like to do it. But it's not as as pressing as (a) - plus, we would change the behavior for existing users. Though, worth noting, we would only change behavior for existing users that did not comply with our existing contract(!). So that may give us the freedom to do (b) - but as follow-up.
There was a problem hiding this comment.
🧹 Nitpick comments (1)
c2h/include/c2h/custom_type.h (1)
70-75: 🎯 Functional Correctness | 🔵 Trivial | ⚡ Quick winsuggestion: Enforce the size-contributing-policy invariant instead of checking only the estimated filler bound. A pack such as
custom_type_t<huge_data<512>::type, huge_data<1024>::type>can pass this assertion because eachhuge_datacounts the other policy as one byte, whileobject_bytesno longer matches the composed type size. Reject additional size-contributing policies or validate their actual layout with a compile-time check.
ℹ️ Review info
⚙️ Run configuration
Configuration used: Path: .coderabbit.yaml
Review profile: CHILL
Plan: Enterprise
Run ID: 1111fa38-ac11-45cf-b0d7-ec1b8221a530
📒 Files selected for processing (2)
c2h/include/c2h/custom_type.hcub/test/catch2_test_block_merge_sort.cu
Included review availability: Your plan provides up to 12 included reviews per hour; 11 remain after this review.
🥳 CI Workflow Results🟩 Finished in 3h 14m: Pass: 100%/272 | Total: 13d 05h | Max: 3h 13m | Hits: 9%/1503668See results here. |
640 configs (keys+pairs, 7-8 types, 2^16..2^28, 2 entropies) on B200: mean -0.06%/-0.17%, median ~0 - no systematic regression. Wins up to -5.7% on expensive-comparator/small types at 2^24-2^28 (clamp-free full-tile path); +1-3% at 2^16 from overload B on nearly-full last tiles (bounded by tile-count dilution; the partial-tile wins need non-pow2 sizes, covered by the collective eval). ncu: block-sort kernel unchanged (40 regs, 75% occ, +0.49% cycles) - attribution closes against the +0.12% end-to-end at F32 2^28. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
|
2 participants
Description
Closes #5327
The core of the issue was that we were previously conflating two code paths in
BlockMergeSort, a version that: (a) pre-pads oob items withMAX_KEY(oob_default, i.e., something that compares GEQ than any provided item) and subsequently simply do a full sort and (b) a code path for an interface that only takesvalid_itemsand introduces the oob-guards where needed - for that code path, nooob_defaultis needed here. This is basically the design thatWarpBitonicSortfollows, I like, and that I introduce the distinction for in this PR.A more detailed assessment of the underlying issues, I've shared in #5327 (comment).