NVIDIA / NVIDIA/cccl

[FEA]: Add segmented-sort partitioning with in-place small-segment processing

Open
#10,804 1 comment 0 reactions 1 assignee Claimed by @colin-mcd View on GitHub
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

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.