[FEA]: Add segmented-sort partitioning with in-place small-segment processing
- Dominant language
- C++
- Stars
- 2.5k
- Forks
- 487
- Avg merge
- 2d 7h
- Merged PRs (30d)
- 296
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.
`DeviceSegmentedSort` inspects segment sizes and route segments to suitable sorting specializations.
For small segments, writing the segment to a partitioning queue and launching another kernel may cost more than sorting the segment itself. The partitioning kernel has already loaded the segment offsets and identified its size, so it can process sufficiently small segments in place and enqueue only the remaining segments.
CUB does not currently have a generalized partitioning building block that supports this fused processing.
This is part of the roadmap for improving `DeviceSegmentedSort`: https://github.com/NVIDIA-dev/cccl_private/issues/829
It depends on the single-thread unstable odd-even merge sort building block: https://github.com/NVIDIA/cccl/issues/10792
### Describe the solution you'd like
This issue can be closed with a PR that does in-place/small-segment processing and queuing of remaining segments in a fused partitioning kernel.
The initial contribution may use a fixed number of size classes. The representation should not prevent the generic tuning design from later selecting the number of classes, thresholds, partitioning algorithm, and in-place threshold.
Because the proposed in-place building block is unstable, the first integration should target the unstable segmented-sort path. Stable entry points should either disable in-place processing or use a stable specialization.
The implementation should include:
- Correctness tests for empty, small, boundary-size, and large segments.
- Key-only and key-value tests.
- Tests for independent begin/end offset iterators and segments with gaps.
- Validation that every non-in-place segment is emitted exactly once.
- Temporary-storage and overflow validation
cc @colin-mcd
### Describe alternatives you've considered
_No response_
### Additional context
_No response_
Contributor guide
Assessment
This issue has not been assessed yet.