[FEA]: CUB Sort Deviceless Scratch Size
- 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.
We (at OpenXLA) would like to be able to obtain the amount of temporary storage needed for the CUB sorting operations, without needing a device to be attached.
In particular we'd need a deviceless API that provides us the scratch space required for the following operations:
* `cub::DeviceRadixSort::SortPairs` / `SortPairsDescending`
* `cub::DeviceSegmentedRadixSort::SortPairs` / `SortPairsDescending`
* `cub::DeviceRadixSort::SortKeys` / `SortKeysDescending`
* `cub::DeviceSegmentedRadixSort::SortKeys` / `SortKeysDescending`
It is enough if this solution only works for Hopper+ devices.
### Describe the solution you'd like
We originally created the [following doc](https://docs.google.com/document/d/1OWekzUxP7Lc3GbMrmTPnuy8AqfqaBzjw2B5YpPkPdGY/edit?usp=sharing) proposing a way to pass the properties of the device we are targeting, so that CUB doesn't need to query the actual device to obtain the information it needs to calculate the temporary storage size (very similar to the cuDNN AOT mode).
@danielfrg mentioned that CUB may already support deviceless friendly algorithms for these cases (e.g. OneSweep for `cub::DeviceRadixSort`), that we could select using the Tuning API.
### Describe alternatives you've considered
_No response_
### Additional context
XLA needs to calculate the scratch sizes during compile time, so this would allow us to compile models using CUB sort without needing a GPU to be available.
Contributor guide
Assessment
This issue has not been assessed yet.