NVIDIA / NVIDIA/cccl

[FEA]: Experiment with decoupled look-back optimization for primitive types

Open
#512 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.

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

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.