[EPIC] Productize WSPRO algorithm as CUB warp primitive
- Dominant language
- C++
- Stars
- 2.5k
- Forks
- 486
- Avg merge
- 2d 6h
- Merged PRs (30d)
- 295
Description
## Overview
Integrate `cub::WarpReduceBatched` as a new warp-level primitive in CUB, implementing the Warp Shuffle Parallel Reduction Optimization (WSPRO) algorithm. This primitive performs batched reductions of N arrays with M elements each across M lanes of a logical warp (power of two size <= 32) in O(N) steps (each containing a shuffle and an application of the reduction operator). The basic idea comes from realizing that the butterfly reduction algorithm has a lot of redundant communication and computation that could be used more productively. Butterfly warp reductions use XOR-warp-shuffles for communication and result in all threads arriving at the same reduction-result (assuming associativity and commutativity of the operator). In contrast to that WSPRO distributes the results of the batched reductions among the (logical) warp lanes while the result of `cub::WarpReduce` is only returned in lane 0 which makes it a bad fit for such batched operation.
For comparison N sequential applications of a butterfly reduction (or `cub::WarpReduce`) would take **N × log₂(M)** steps.
- For N being a divisible by M, WSPRO takes exactly **N - N / M** steps, so
- For N = 2 × M the number of steps is **N - 2**,
- For N = M the number of steps is **N - 1**,
- Generally it is the sum of N **ceiling**-divided by powers of two [2, ..., M]
- For N = 4, M = 8 the number of steps is 4/2 + 4/4 + ⌈4/8⌉ = 2 + 1 + 1 = 4,
- For N = 8, M = 32 the number of steps is 8/2 + 8/4 + 8/8 + ⌈8/16⌉ + ⌈8/32⌉ = 4 + 2 + 1 + 1 + 1 = 9,
- ...
- For the N = 2 the number of steps is **log₂(M)** (meaning we are still doing two reductions for the "price" of one).
- Uneven N always takes the same number of steps as the next bigger even N' = N + 1.
Contributor guide
Assessment
This issue has not been assessed yet.