NVIDIA / NVIDIA/cccl

[FEA]: Smart API for DeviceReduce/DeviceSegmentedReduce

Open
#2,738 0 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.

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

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.