Select few large segments in segmented_top_k with cub::DeviceTopK ins… - #23602
Select few large segments in segmented_top_k with cub::DeviceTopK ins…#23602a-hirota wants to merge 1 commit into
Conversation
f0baa97 to
d814d8b
Compare
Add a cub::DeviceTopK fast path for eligible inputs with at most 64 segments, an average covered segment size of at least 16K rows, and k at most 1/8 of that average. The fast path selects each segment independently, reuses temporary storage, and sorts only the selected indices to preserve the existing per-segment ordering. Empty segments and segments with at most k rows are handled without CUB. Inputs with nulls, floating-point values, unsupported types, more than 64 segments, an average covered segment size below 16K rows, or k greater than one eighth of that average continue to use the existing sort-based implementation.
d814d8b to
0d73d76
Compare
📝 WalkthroughSummary by CodeRabbit
WalkthroughChangesThe segmented top-k implementation adds a CUB-based fast path for eligible large, fixed-width, non-floating-point inputs without nulls. Existing sort-based behavior remains the fallback. Benchmarks and tests cover large-segment configurations and boundary cases. Segmented top-k optimization
Estimated code review effort: 4 (Complex) | ~45 minutes Possibly related PRs
Suggested labels: Suggested reviewers: 🚥 Pre-merge checks | ✅ 5✅ Passed checks (5 passed)
✨ Finishing Touches 💡 1🛠️ Fix failing CI checks 💡
🧪 Generate unit tests (beta)
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: 2
🧹 Nitpick comments (3)
cpp/src/sort/segmented_top_k.cu (3)
222-228: 📐 Maintainability & Code Quality | 🔵 Trivial | 💤 Low valueRun clang-format on this call.
The continuation lines of
stable_segmented_sort_by_keyare not aligned consistently. The{topk_order}and{nulls}arguments use a different indentation than the other arguments.As per coding guidelines: "Format C++ and CUDA code with clang-format."
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@cpp/src/sort/segmented_top_k.cu` around lines 222 - 228, Run clang-format on the stable_segmented_sort_by_key call in the segmented top-k implementation, correcting the continuation-line indentation while preserving the call’s arguments and behavior.Source: Coding guidelines
170-172: 🎯 Functional Correctness | 🔵 Trivial | ⚡ Quick winDocument the non-deterministic tie behavior in the fast path.
The environment requests
determinism::not_guaranteed. When several rows share the key value at thek-th boundary,cub::DeviceTopKcan select a different subset on each run. The sort-based path selects the same subset on every run for the same input. The public API therefore becomes run-to-run non-deterministic for tied inputs on eligible columns only.Add a note in the public
segmented_top_kandsegmented_top_k_orderdocumentation, or usedeterminism::gpu_to_gpuif the cost is acceptable.🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@cpp/src/sort/segmented_top_k.cu` around lines 170 - 172, Document in the public segmented_top_k and segmented_top_k_order APIs that the fast path may select different rows across runs when ties occur at the k-th boundary, while the sort-based path remains deterministic. Alternatively, replace determinism::not_guaranteed with determinism::gpu_to_gpu if its performance cost is acceptable.
292-306: 🚀 Performance & Scalability | 🔵 Trivial | 💤 Low valueNote the added device-to-host synchronization on the eligibility check.
cudf::detail::make_host_vectorsynchronizes the stream. Every null-free fixed-width column with at most 64 segments now pays this synchronization, including columns that later fall back tosort_based_segmented_top_k_orderbecause of the average-size orkchecks. The host offsets are required for the host-side segment loop, so the copy is unavoidable on the fast path. Consider ordering the cheap checks so that fewer fallback cases pay the synchronization, for example by checkingcol.size() / num_segments >= cub_min_avg_segment_sizebefore the copy.As per coding guidelines: "Avoid unnecessary host-device synchronization".
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@cpp/src/sort/segmented_top_k.cu` around lines 292 - 306, Reorder the fast-path eligibility checks in the segmented top-k dispatch so the cheap average-size test using col.size() and num_segments runs before make_host_vector. Only copy segment_offsets to h_offsets when the column passes the segment-count, average-size, and k-fraction checks; preserve h_offsets for dispatch_segmented_topk_fn and retain sort_based_segmented_top_k_order as the fallback.Source: Coding guidelines
🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
Inline comments:
In `@cpp/src/sort/segmented_top_k.cu`:
- Around line 151-166: Validate the host offsets before entering the fast path
in segmented_top_k_order, ensuring they are nondecreasing and each value is
within the input column’s row count. For invalid offsets, route to
sort_based_segmented_top_k_order or reject via CUDF_EXPECTS; only compute
h_out_offsets and allocate indices after validation so negative or oversized
ranges cannot trigger invalid allocations or reads.
- Around line 136-140: Update is_fast_path to reject DECIMAL128 in addition to
null, floating-point, and non-fixed-width exclusions, ensuring DECIMAL128 uses
sort_based_segmented_top_k_order while DECIMAL32 and DECIMAL64 remain eligible
for DeviceTopK.
---
Nitpick comments:
In `@cpp/src/sort/segmented_top_k.cu`:
- Around line 222-228: Run clang-format on the stable_segmented_sort_by_key call
in the segmented top-k implementation, correcting the continuation-line
indentation while preserving the call’s arguments and behavior.
- Around line 170-172: Document in the public segmented_top_k and
segmented_top_k_order APIs that the fast path may select different rows across
runs when ties occur at the k-th boundary, while the sort-based path remains
deterministic. Alternatively, replace determinism::not_guaranteed with
determinism::gpu_to_gpu if its performance cost is acceptable.
- Around line 292-306: Reorder the fast-path eligibility checks in the segmented
top-k dispatch so the cheap average-size test using col.size() and num_segments
runs before make_host_vector. Only copy segment_offsets to h_offsets when the
column passes the segment-count, average-size, and k-fraction checks; preserve
h_offsets for dispatch_segmented_topk_fn and retain
sort_based_segmented_top_k_order as the fallback.
🪄 Autofix
Fix all unresolved CodeRabbit comments on this PR:
- Push a commit to this branch (recommended)
- Create a new PR with the fixes
ℹ️ Review info
⚙️ Run configuration
Configuration used: Path: .coderabbit.yaml
Review profile: CHILL
Plan: Enterprise
Run ID: 3dca77ea-d205-4eea-ae9e-4efc267233cb
📒 Files selected for processing (3)
cpp/benchmarks/sort/segmented_top_k.cppcpp/src/sort/segmented_top_k.cucpp/tests/sort/top_k_tests.cpp
| bool is_fast_path(column_view const& column) | ||
| { | ||
| return !column.has_nulls() && cudf::is_fixed_width(column.type()) && | ||
| !cudf::is_floating_point(column.type()); // requires NaN-aware ordering | ||
| } |
There was a problem hiding this comment.
🗄️ Data Integrity & Integration | 🟠 Major | ⚡ Quick win
🧩 Analysis chain
🏁 Script executed:
#!/bin/bash
# Check whether the tests cover decimal or chrono types on the fast path.
rg -n 'TopKSegmented' -C5 cpp/tests/sort/top_k_tests.cpp | head -120Repository: rapidsai/cudf
Length of output: 5381
🌐 Web query:
CUB DeviceTopK supported key types 128-bit integer __int128 radix
💡 Result:
As of August 10, 2026, cub::DeviceTopK supports both fundamental numeric types and custom user-defined types [1][2]. While the core implementation of DeviceTopK is radix-based [3], it does not have native, specialized support for 128-bit integers (e.g., int128) as a built-in fundamental type in the same manner as standard primitives like int or double [1][2]. However, you can process 128-bit integers by utilizing the support for custom types via a decomposer object [1][2]. To use 128-bit integers with DeviceTopK: 1. You must provide a decomposer object that allows the algorithm to extract the necessary bits or components for the radix-based partitioning [1][2]. 2. The interface requires that for custom types, you use the overload that accepts a decomposer, as the non-decomposer API is strictly limited to fundamental types [2]. The official documentation specifies that DeviceTopK can process built-in C++ numeric primitives, CUDA's 16-bit floating-point types (half and nv_bfloat16), and any user-defined types provided with a corresponding decomposer [1][2].
Citations:
- 1: https://nvidia.github.io/cccl/unstable/cub/api/structcub_1_1DeviceTopK.html
- 2: fbusato/cccl@a0d85a9
- 3: Add support for non-fundamental types to
DeviceTopKcccl#6290
🏁 Script executed:
#!/bin/bash
set -e
printf '%s\n' '--- target implementation ---'
sed -n '1,240p' cpp/src/sort/segmented_top_k.cu
printf '%s\n' '--- CUB DeviceTopK call sites ---'
rg -n -C8 'DeviceTopK|cub_segmented_top_k_order|dispatch_storage_type|is_fast_path' cpp/src cpp/include cpp/tests | head -260
printf '%s\n' '--- decimal128 and test type coverage ---'
rg -n -C4 'DECIMAL128|decimal128|fixed_point|TopKTypes|using.*TypeParam|Types<' cpp/tests/sort cpp/src/sort cpp/include | head -300
printf '%s\n' '--- dependency/version evidence ---'
rg -n -C3 'CCCL|CUB|cub' cmake CMakeLists.txt cpp/include cpp/src | head -240Repository: rapidsai/cudf
Length of output: 50369
🌐 Web query:
Official CCCL cub::DeviceTopK documentation: determine whether the non-decomposer overload accepts __int128_t or decimal128 keys, and identify the exact supported key-type rules and version.
💡 Result:
The CCCL cub::DeviceTopK non-decomposer overload does not officially support int128_t keys [1][2]. Supported Key-Type Rules: cub::DeviceTopK is implemented using radix-based logic, which relies on cub::Traits for bit-level interpretation [3]. The non-decomposer API is strictly limited to built-in C++ numeric primitive types (such as unsigned char, int, double, etc.) and specific CUDA floating-point types, namely half and nv_bfloat16 [1][2]. Types such as int128_t are not considered fundamental numeric primitives within this context and therefore do not function with the non-decomposer overload [3]. For types not supported by the non-decomposer API, developers must use the decomposer-based overload [1][4]. This version allows for custom types by utilizing a decomposer object to interpret the keys, though it explicitly prohibits the use of fundamental types that are already covered by the standard overload [4]. Version Context: Support for non-fundamental types via a decomposer was added to DeviceTopK in recent updates (e.g., via PR #8040, merged around March 2026) [3][4][5]. Despite this expansion, the non-decomposer path remains restricted to the aforementioned built-in types [1][4].
Citations:
- 1: https://nvidia.github.io/cccl/unstable/cub/api/structcub_1_1DeviceTopK.html
- 2: https://gevtushenko.github.io/cccl/cub/api/structcub_1_1DeviceTopK.html
- 3: Add support for non-fundamental types to
DeviceTopKcccl#6290 - 4: fbusato/cccl@a0d85a9
- 5: [EPIC] Add family of top-k algorithms to CUB cccl#5673
🏁 Script executed:
#!/bin/bash
set -e
printf '%s\n' '--- dispatch and path selection ---'
sed -n '220,380p' cpp/src/sort/segmented_top_k.cu
printf '%s\n' '--- segmented top-k test type definitions and decimal coverage ---'
sed -n '1,150p' cpp/tests/sort/top_k_tests.cpp
rg -n -C5 'DECIMAL128|decimal128|fixed_point|TestTypes|ComparableTypes' cpp/tests/sort/top_k_tests.cpp cpp/tests/utilities/type_list.hpp cpp/tests | head -220
printf '%s\n' '--- exact type-dispatch mappings ---'
sed -n '160,230p' cpp/include/cudf/utilities/type_dispatcher.hpp
rg -n -C4 'device_storage_type_t|DECIMAL128' cpp/include/cudf | head -120Repository: rapidsai/cudf
Length of output: 42800
🌐 Web query:
CCCL DeviceTopK MinPairs MaxPairs non-decomposer overload __int128_t compile error decimal128 exact documentation
💡 Result:
The issues and compilation errors you are observing are linked to architectural and structural limitations within the CUDA Core Compute Libraries (CCCL) regarding 128-bit types and the design of cub::DeviceTopK. Regarding int128_t compile errors: The int128_t and uint128_t types are often not supported in device code, leading to compilation errors such as "int128_t contains a 128-bit integer, which is not supported in device code" [1]. If you encounter this, ensure your environment supports it or explicitly disable it using the _LIBCUDACXX_HAS_NO_INT128 macro [1]. CCCL also provides a configuration macro, CCCL_DISABLE_INT128_SUPPORT, to disable this support if needed [2]. Note that hardware and compiler support for 128-bit types (including float128) varies significantly between target architectures (e.g., x86 vs. ARM) and compiler versions [3]. Regarding cub::DeviceTopK and non-decomposer overloads: The cub::DeviceTopK interface is designed primarily for fundamental types [4]. If you are attempting to use a custom decomposer with a fundamental type, you will trigger a static assertion error, as the library explicitly restricts custom decomposers to non-fundamental types [5]. The error message "Custom decomposers are not supported for fundamental types; use the non-decomposer API overload instead" appears because the library expects you to use the standard MinPairs or MaxPairs overloads for basic data types [5]. When working with custom, non-fundamental types, you must provide a decomposer to interpret the type, as the library's internal radix-based partitioning logic otherwise cannot handle them [5][4]. Regarding decimal128: There is no "decimal128 exact" support provided natively within the core CCCL / CUB library. While other systems (like ClickHouse or MongoDB) define their own Decimal128 types for specific purposes [6][7], CCCL does not provide an exact decimal128 arithmetic type. If you are seeing references to "decimal128" in conjunction with CCCL, it is likely related to external library interactions or confusion with the aforementioned 128-bit integer support issues, which are distinct from IEEE 754-2008 decimal floating-point arithmetic.
Citations:
- 1: [BUG] 128bit integers are not supported on device code cccl#1227
- 2: https://nvidia.github.io/cccl/unstable/cccl/development/macro.html
- 3: [BUG]: cuda/std/__cccl/builtin.h", line 430: error: identifier "__float128" is undefined cccl#4114
- 4: Add support for non-fundamental types to
DeviceTopKcccl#6290 - 5: fbusato/cccl@a0d85a9
- 6: https://clickhouse.com/docs/reference/data-types/decimal
- 7: https://github.com/mongodb/specifications/blob/master/source/bson-decimal128/decimal128.md
Exclude DECIMAL128 from the DeviceTopK fast path.
dispatch_storage_type maps DECIMAL128 to __int128_t, but the non-decomposer cub::DeviceTopK::MinPairs and MaxPairs overloads do not support 128-bit keys. Add a DECIMAL128 guard to is_fast_path so these inputs use sort_based_segmented_top_k_order. DECIMAL32 and DECIMAL64 use supported 32-bit and 64-bit storage types.
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
In `@cpp/src/sort/segmented_top_k.cu` around lines 136 - 140, Update is_fast_path
to reject DECIMAL128 in addition to null, floating-point, and non-fixed-width
exclusions, ensuring DECIMAL128 uses sort_based_segmented_top_k_order while
DECIMAL32 and DECIMAL64 remain eligible for DeviceTopK.
| auto const num_segments = static_cast<size_type>(h_offsets.size()) - 1; | ||
|
|
||
| auto h_out_offsets = std::vector<size_type>(num_segments + 1); | ||
| h_out_offsets[0] = 0; | ||
| for (size_type i = 0; i < num_segments; ++i) { | ||
| auto const size = h_offsets[i + 1] - h_offsets[i]; | ||
| h_out_offsets[i + 1] = h_out_offsets[i] + cuda::std::min(size, k); | ||
| } | ||
|
|
||
| // Synchronous copy before any CUB work is queued: h_out_offsets is stack-local and an | ||
| // async copy would defer the read. | ||
| auto offsets = std::make_unique<column>( | ||
| cudf::detail::make_device_uvector(h_out_offsets, stream, mr), rmm::device_buffer{}, 0); | ||
|
|
||
| auto const temp_mr = cudf::get_current_device_resource_ref(); | ||
| auto indices = rmm::device_uvector<size_type>(h_out_offsets[num_segments], stream, temp_mr); |
There was a problem hiding this comment.
🩺 Stability & Availability | 🟠 Major | ⚡ Quick win
Guard the fast path against non-monotonic or out-of-range offsets.
h_out_offsets accumulates min(size, k) for each segment. If h_offsets is not monotonically increasing, size becomes negative and the running total shrinks. The rmm::device_uvector size then becomes negative, and a negative value converted to std::size_t produces an enormous allocation request. The per-segment reads at Line 177 (in + begin) also read out of bounds when an offset exceeds col.size().
The sort-based path does not have the same failure mode, because resolve_segment_indices resolves segments on device against the actual row count. The fast path therefore changes the failure behavior for malformed offsets.
The host offsets are already available at Line 294. Add a cheap validation there and route invalid offsets to sort_based_segmented_top_k_order, or reject them with CUDF_EXPECTS.
🛡️ Proposed eligibility guard in `segmented_top_k_order`
auto const avg_segment_size = (h_offsets.back() - h_offsets.front()) / num_segments;
- if (avg_segment_size >= cub_min_avg_segment_size &&
- k <= avg_segment_size / cub_max_k_fraction) {
+ auto const valid_offsets =
+ h_offsets.front() >= 0 && h_offsets.back() <= col.size() &&
+ std::is_sorted(h_offsets.begin(), h_offsets.end());
+ if (valid_offsets && avg_segment_size >= cub_min_avg_segment_size &&
+ k <= avg_segment_size / cub_max_k_fraction) {As per coding guidelines: "Prevent invalid memory access, including out-of-bounds access" and "Validate inputs such as negative dimensions and null pointers".
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
In `@cpp/src/sort/segmented_top_k.cu` around lines 151 - 166, Validate the host
offsets before entering the fast path in segmented_top_k_order, ensuring they
are nondecreasing and each value is within the input column’s row count. For
invalid offsets, route to sort_based_segmented_top_k_order or reject via
CUDF_EXPECTS; only compute h_out_offsets and allocate indices after validation
so negative or oversized ranges cannot trigger invalid allocations or reads.
Source: Coding guidelines
1 participant
Description
Add a cub::DeviceTopK fast path for eligible inputs with at most 64 segments, an average
covered segment size of at least 16K rows, and k at most 1/8 of that average.
The fast path selects each segment independently, reuses temporary storage, and sorts only
the selected indices to preserve the existing per-segment ordering. Empty segments and
segments with at most k rows are handled without CUB.
Inputs with nulls, floating-point values, unsupported types, more than 64 segments,
an average covered segment size below 16K rows, or k greater than one eighth of that
average continue to use the existing sort-based implementation. Selection among equal
values at the k-th boundary remains unspecified.
This follows the pattern of the non-segmented
top_k(#21582), which routes the sameeligible inputs to
cub::DeviceTopKand falls back to the sort-based path for everythingelse.
Performance
GB10, 16.8M rows, k=100:
Peak temporary memory decreased from 192 MiB to under 1 MiB. FLOAT32 remains on the
fallback path and showed no change.
The guard bounds come from crossover measurements (INT32, 16 segments unless noted; the
sort-based path is k-independent):
452 µs); at 16K the fast path wins 4x (465 µs vs 1.84 ms, also 3.4x at 64 segments).
The segmented sort's small-segment kernels cut off between 8K and 16K.
rows per segment, 16 and 64 segments); at k = avg/2 the margin shrinks to 1.2-1.7x,
since the post-selection sort approaches the full sort this path avoids.
wins at both bounds: 1.8x at 16K rows per segment and 1.5x at k = avg/8 (32K rows per
segment). Below 16K it falls back, which errs toward the unchanged sort-based path.
A
segmented_top_k_few_largebenchmark is added (num_segments = {4, 16, 64}, segment ={16K, 128K, 1M}, k = {100, 2048}) so the fast path's representative shapes and the guard
bounds stay measurable for future re-evaluation.
Checklist