NVIDIA / NVIDIA/cccl

[FEA]: Add a fixed size segmented sort overload

Open
#10,518 4 comments 0 reactions 0 assignees 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.

Today, `cub::DeviceSegmentedSort` requires the user to specify the begin and end offsets for each segment, which allows us to support segmented arrays with different size segments. Internally, it relies on `ThreeWayPartition` to group the segments into small, medium, and large depending on their size, and then uses a different cooperative algorithm to do the sorting based on the size.

https://github.com/NVIDIA/cccl/pull/3969 added an overload to `cub::DeviceSegmentedReduce` which accepts `segment_size` as an argument, so we would need a similar API here

### Describe the solution you'd like

The solution should be relatively simple. The fixed size overload would not need the `ThreeWayPartition` kernel since the sizes of the segments are already known (and uniform). Instead, the dispatch logic will simply call the segmented sort kernel, passing in the segment size and using the same cooperative algorithm for all segments

The initial PR should reuse the existing tunings, which means using the same threshold to decide whether segments are small, medium, or large, and using the same cooperative algorithms as the non fixed size overload. Later work could include adding specific tunings for the fixed size overload

### Describe alternatives you've considered

_No response_

### Additional context

_No response_

Contributor guide

Open the contributing guide

Research direction

Start by comparing cub::DeviceSegmentedSort with the fixed-size overload added to cub::DeviceSegmentedReduce in PR 3969. Trace the existing ThreeWayPartition and segmented sort kernel dispatch, including the small, medium, and large thresholds. Done means a fixed-size API uses the existing tunings and cooperative algorithms without the ThreeWayPartition kernel.

Written by the indexing model from the issue text.

Assessment

Tech stack
cpp
Domain
data
Issue type
Feature
Difficulty
4/5
Estimated time
3-5 days
Activity status
Quiet
Clarity
Mostly clear
Newbie friendliness
52/100

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.