[FEA]: Smart API for DeviceReduce/DeviceSegmentedReduce
- 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.
The implementation of `DeviceReduce `and `DeviceSegmentedReduce` causes a performance gap for small segment numbers. This means that for small number of segments, it is more efficient to loop independent submissions of `DeviceReduce` instead of using `DeviceSegmentedReduce`. It can be hard to zero in on the exact point that this performance trade off occurs, and is a nuanced performance issue that most users may not recognize or know how to optimize for.
The Following runtimes shows the divergence in performance for a 2M element 1D tensor, using both `DeviceReduce` and `DeviceSegmentedReduce` on an A100 with CUDA 12.6.
DeviceSegmentedReduce : 1.2ms
DeviceReduce: 21us
### Describe the solution you'd like
We can address this usability hurdle by providing a "smart" API that will automatically choose to loop individual `DeviceReduce` calls or call `DeviceSegmentedReduce` for various numbers of segments, depending on the performance of the underlying CUB kernels. I believe this API could also simplify the denotation of DeviceSegmented versus DeviceReduce by having a unified API for all dimensionality, with the trade off of a less configurable API (specifically losing the ability to define arbitrary length segments).
It is a much more common use pattern to have fixed size segments per dimension (standard dense matrix use pattern), which this API would most-closely support. this would ensure maximum performance with the existing kernels without the user needing detailed knowledge and profiling for their given use sizes. It also significantly reduces the amount of code required by the user to reach an optimal implementation with CUB.
### Describe alternatives you've considered
The current solution path to this problem is to develop control mechanisms in user code that switches between `DeviceSegmentedReduce` and `DeviceReduce` based on the size of data and dimensionality of the problem. This is not a good solution as it requires a high level of understanding on the user to maximize performance, causes repeated work and code bloat in user code, and leaves significant latitude for varying "maximum performance" based on the user's own implementation and robustness of the code.
The current pattern leaves CUB susceptible to many non-optimal implementations by users without significant research/profiling from the user.
### Additional context
_No response_
Contributor guide
Assessment
This issue has not been assessed yet.