[FEA]: Experiment with decoupled look-back optimization for primitive types
- 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.
Decoupled look-back implementation looks as follows:
```cpp
T process_window(tile_id) {
return WarpReduce(storage).TailSegmentedReduce(wait_for_valid(tile_id));
}
while (__all_sync(predecessor_status != inclusive)) {
exclusive_prefix = scan_op(process_window(predecessor_idx), exclusive_prefix);
predecessor_idx -= 32;
}
```
In the case of commutative operations, tail reduction is not needed on reading every window.
### Describe the solution you'd like
When the operation is commutative, we can accumulate partial sum in thread and reduce the previous scheme to:
```cpp
T process_window(tile_id) {
return wait_for_valid(tile_id);
}
while (__all_sync(predecessor_status != inclusive)) {
exclusive_prefix = scan_op(process_window(predecessor_idx), exclusive_prefix);
predecessor_idx -= 32;
}
WarpReduce(storage).TailSegmentedReduce(exclusive_prefix);
```
This modification removes N collective operations from the loop and adds one collective operation after it. Preliminary results show ~10% speedup on 4090 when computing prefix sum of int128_t. The change do not lead to significant improvements for different integer types on 4090.
More research is needed. We should benchmark the change on H100 and also try re-tuning scan with the proposed scheme.
### Describe alternatives you've considered
_No response_
### Additional context
_No response_
Contributor guide
Assessment
This issue has not been assessed yet.