[FEA]: Add a runtime-width group merge-sort primitive 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.
Fine-grained segmented sort may assign a configurable group of cooperating threads to each segment. The same group merge-sort building block applies whether that group contains fewer than one warp, exactly one warp, or multiple warps; the cases differ only in the number of threads cooperating on a segment.
Treating multi-warp sorting as a separate primitive would duplicate the underlying building block and would not match the runtime-width selection already expressed by `call_group_merge_runtime`.
The primitive should therefore take `threads_per_segment` at runtime and use the same group merge-sort implementation for all supported widths.
### Describe the solution you'd like
Add a reusable runtime-width group merge-sort primitive based on the building block used by `call_group_merge_runtime` that:
- Accepts the supported cooperating-group width at runtime.
- Uses the same interface and implementation structure for sub-warp (`threads_per_segment < 32`), warp (`threads_per_segment == 32`), and multi-warp (`threads_per_segment > 32`) groups.
- Sorts within independent groups of consecutive threads.
- Supports a runtime valid-item count for partially filled groups.
- Prevents synchronization and shared-memory state from crossing group boundaries.
- Supports keys in the initial implementation and provides a path to key-value pairs.
- Supports ascending and descending order.
- Provides correctness coverage across supported sub-warp, warp, and multi-warp widths, including partial groups, duplicate keys, and multiple independent groups in one CTA.
This issue can be closed once the runtime-width primitive is integrated into CUB with tests and an interface suitable for use by specialized segmented-sort kernels. A separate compile-time-width variant is not required to close this issue and is out of scope.
### Describe alternatives you've considered
A compile-time `threads_per_segment` variant was considered. For the current contribution, the runtime-width primitive covers the required sub-warp, warp, and multi-warp cases, so a compile-time variant will not be implemented as part of this issue.
### Additional context
Runtime group-width selection is part of the primitive delivered by this issue rather than a later optimization.
Contributor guide
Assessment
This issue has not been assessed yet.