[EPIC] Improve performance of `DeviceSegmentedSort`
- Dominant language
- C++
- Stars
- 2.5k
- Forks
- 486
- Avg merge
- 2d 6h
- Merged PRs (30d)
- 295
Description
### Context
`cub::DeviceSegmentedSort` sorts N independent segments described by `begin`/`end` offset iterators. Its
current design buckets segments into exactly three size classes and assigns each class a fixed
algorithm. That structure is a good default but, for certain workloads, leaves performance on the table.
### Feature Requests, Prioritization, and Workloads
➡️ **Help us collect and prioritize workloads that we want to focus on for optimizing `DeviceSegmentedSort`**
- https://github.com/NVIDIA/cccl/issues/10661
### Shortcoming of the Current Approach
#### 1. Overprovisioning at the bottom of each size class
Algorithms for each class pad every segment to its class's full tile (the maximum supported size of that class). Because the class boundaries correspond to the tile sizes, the padding factor at the bottom of a class equals the ratio
between consecutive tiles. E.g., for `int` keys, keys-only, on SM 8.6+
([`tuning_segmented_sort.cuh:190`](https://github.com/NVIDIA/cccl/blob/main/cub/cub/device/dispatch/tuning/tuning_segmented_sort.cuh#L190)):
| class | segment sizes | worker | tile | worst-case padding |
| ------ | ------------- | ----------------------------- | ------------- | -----------------: |
| small | 1 … 18 | 2-thread sub-warp merge sort | 2 × 9 = 18 | 18x (at size 1) |
| medium | 19 … 112 | 16-thread sub-warp merge sort | 16 × 7 = 112 | 5.9x (at size 19) |
| large | 113 … 5887 | 256-thread block radix sort | 256 × 23 = 5888 | **52x** (at size 113) |
| huge | ≥ 5888 | multi-pass global-memory radix | streaming | 1x |
The `medium` to `large` cliff is ~50x case: a 113-element segment gets a whole 256-thread CTA running a
`BlockRadixSort` over a 5888-slot tile with a ~23 KiB shared-memory key-exchange buffer, and every padded
slot pays every radix pass.
Related Issues:
- https://github.com/NVIDIA/cccl/issues/6792
- https://github.com/NVIDIA/cccl/issues/895#issue-1984270788
#### 2. Host round-trip in the middle of the algorithm
The partitioned path computes the number of segments in each class on the device, then copies 8 bytes back
to the host and **synchronizes the stream** so the host can size the grids of the two sort kernels
([`dispatch_segmented_sort.cuh:1109-1123`](https://github.com/NVIDIA/cccl/blob/maincub/cub/device/dispatch/dispatch_segmented_sort.cuh#L1109)):
```cpp
local_segment_index_t h_group_sizes[num_selected_groups];
launcher_factory.MemcpyAsync(h_group_sizes, group_sizes.get(), ..., cudaMemcpyDeviceToHost, stream);
SyncStream(stream); // cudaStreamSynchronize
device_segmented_sort_continuation(..., h_group_sizes, ...);
```
Consequences: a full `device -> host -> device` latency on every call, and `DeviceSegmentedSort` cannot be
captured into a CUDA graph.
#### 3. No load balancing for huge segments: work is assigned per segment, not per unit of work
Every size class hands one worker a whole segment, even the huge size class. This is particularly punishing for segment sizes exceeding a block tile, where a segment is sorted by exactly one 256-thread CTA however large it is: a 100M-element segment occupies a single SM while the rest of the GPU may idle. There is no intra-segment parallelism beyond a block. The same imbalance exists in miniature within each class: the CTAs of one KernelLarge grid handle segments anywhere from 113 items to unbounded, and each does its own segment's worth of work.
#### 4. No latency regime
The current tuning policies optimize throughput on a saturated GPU. Users whose whole problem does not
fill the device (few segments, or few total items) pay a fixed pipeline of 4 kernels plus a host round-trip. Tile sizes are too coarse-grained, algorithms are chosen for optimal throughput, while, for latency, another choice may be preferable.
Contributor guide
Research direction
Start with the workloads collected in issue #10661, then read cub/cub/device/dispatch/tuning/tuning_segmented_sort.cuh and cub/cub/device/dispatch/dispatch_segmented_sort.cuh, especially the cited tuning and host-round-trip sections. Define and validate an approach for reducing size-class padding, avoiding the synchronization, improving large-segment balance, and supporting latency-sensitive workloads using relevant performance tests.
Written by the indexing model from the issue text.
Assessment
- Tech stack
- cpp
- Domain
- hpc, performance
- Issue type
- Feature
- Difficulty
- 5/5
- Estimated time
- Over a week
- Activity status
- Quiet
- Clarity
- Needs clarification
- Newbie friendliness
- 25/100