[FEA]: k-way merge API
- 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
Thrust
### Is your feature request related to a problem? Please describe.
Problem description provided by @mfoerster:
In the cuPyNumeric implementation of multi-GPU sort, after the local data sorting we have to re-distribute the data according to split positions. Each rank gets data from all other ranks that are pre-sorted and wants to merge them into a single new local chunk. Currently this is a sequence of `ranks-1` 2-way `thrust::merge`s for standard sort, or `thrust::merge_by_key`s for argsort. Even when doing this in a hierarchical pattern we end up accessing all data `log(ranks)` times.
In more detail:
* The initial local sort is done via cub whenever possible (fallback to thrust for complex types) where we use `DeviceRadixSort` or `DeviceSegmentedRadixSort`.
* The distribution is done via sampling, so all ranks agree on a set of split positions to all2all the data.
* Each rank ends up with `#ranks` chunks that are individually sorted, and need to be merged.
* During merge, we have either a primitive `VAL` type or a tuple of primitives (for segmented sort) as sort criteria (comparison operator `thrust::less>()`).
* The merge cost grows linearly with `log(ranks)` as we have to do multiple 2-way passes.
* We currently launch the merges of each level sequentially.
Having a k-way merge API in thrust would slot in very naturally for this step, and can have mild to significant performance implications, depending on implementation (see below).
### Describe the solution you'd like
Most of this information provided by @pauleonix.
The simplest implementation of a k-way merge API might involve reusing the "segmented/batched device-merge" part from thrust's mergesort, which IIUC is a sequence of `log2(#CTAs)` batched 2-way merge kernels. With this, cuPyNumeric would immediately get the benefits of fewer kernel launches, and maybe even better work distribution.
If a "true" k-way merge algorithm becomes available in the future, the API could transparently switch to using a lower number of higher-k batched merges. In the best-case scenario, this would speed up the merge phase by a factor of `log2(ranks)`. Such a "true" k-way merge could also be used internally in thrust's mergesort.
### Describe alternatives you've considered
We can consider using different algorithms altogether, e.g. a partition-based sort (instead of local sort, shuffle and merge), such as an MSD radix sort.
### Additional context
There appears to be an issue already asking for distributed sort primitives https://github.com/NVIDIA/cccl/issues/5683, but it doesn't give many details on the required primitives. This ask could certainly be considered as one of those primitives.
Contributor guide
Assessment
This issue has not been assessed yet.