[FEA]: Investigate Fine Grained Bucket Sizes for Segmented Sort
- Dominant language
- C++
- Stars
- 2.5k
- Forks
- 486
- Avg merge
- 2d 6h
- Merged PRs (30d)
- 295
Description
### Is this a duplicate?
- [x] I confirmed there appear to be no [duplicate issues](https://github.com/NVIDIA/cccl/issues) for this request and that I agree to the [Code of Conduct](CODE_OF_CONDUCT.md)
### Area
CUB
### Is your feature request related to a problem? Please describe.
Currently, CUB::SegmentedSort uses a coarse-grained small/medium/large strategy for binning segment lengths. Due to the coarseness of the binning, a significant amount wasted work is performed when an input's distribution of segment lengths is far from the maximum size of a bin, significantly decreasing performance. This problem cannot be solved by library-level tuning, as changing the size of a bin increases performance for some segment lengths at the cost of others.
### Describe the solution you'd like
There are two avenues of granularity which should be explored:
Kernel Level Granularity: Investigate the performance impact of increasing granularity to the [logarithmic radix binning](https://ieeexplore.ieee.org/document/8547581) scheme used by [Hou et al.](https://dl.acm.org/doi/10.1145/3079079.3079105), [Kobus et al.](https://dl.acm.org/doi/10.1007/978-3-031-39698-4_45), and [GPUSorting](https://github.com/b0nes164/GPUSorting/tree/cleanup/GPUSortingCUDA/SegSort/SplitSort). While increasing kernel granularity improves performance by reducing wasted work, it also risks higher kernel launch overheads and occupancy tail costs. Although the current CUB use model precludes multiple cudaStreams, PDL could potentially enable concurrent kernel execution; this is currently being investigated in #6790.
Dynamic/Runtime Granularity: Investigate the impact of dynamically choosing size-specialized algorithm variants within a kernel. For algorithms that are generally threadblock uniform, this should always yield a performance increase. For algorithms that aren't, dynamically branching on segment length increases efficiency at the risk of thrashing the instruction cache, so the advantage is less clear.
### Describe alternatives you've considered
_No response_
### Additional context
An issue with existing implementations of "fine-grain" segmented sort is that histogramming during the binning pass is done by atomics. This creates an unstable ordering of segments within the bin, which can cause cache locality/memory access inefficiencies in the downstream sort.
The solution to this is to replace pure atomic histogramming with a multisplit/warp_match_any based approach. This preserves a stable ordering of segments, and because the number of bins is typically less than 16, only 4 bits are required to match on, ensuring low ALU overhead during the multisplit/warp_match_any.
Functionally, this kernel would operate almost identically to a single round of OneSweep sort.
Contributor guide
Assessment
This issue has not been assessed yet.