NVIDIA / NVIDIA/cccl

[FEA]: Add a runtime-width group merge-sort primitive for segmented sort

Open
#10,944 3 comments 0 reactions 1 assignee Claimed by @kaatish View on GitHub
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

Open the contributing guide

Assessment

This issue has not been assessed yet.

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.