`cub::detail::WarpScanSmem` could be optimized/cleaned up
- Dominant language
- C++
- Stars
- 2.5k
- Forks
- 487
- Avg merge
- 2d 7h
- Merged PRs (30d)
- 296
Description
This is a rather low-priority issue given that this specialization of `cub::WarpScan` is only used for `LOGICAL_WARP_THREADS` is not a power of 2 which I assume is quite exotic. But while adding new member functions in #5379 I observed quite a few potential inefficiencies. Most of them come from the fact that when CUB was sprinkled with `__syncwarp()` for Volta/Independent Thread Scheduling in 83ab80618c866079d5ed5d3a9ed9414a1572c432, there was little regard for how else the code could be modified to fit the new paradigm.
1. All loads and store from/to shared memory are done with `ThreadLoad`/`ThreadStore` since long before the Volta changes. Doing volatile loads/stores should not be necessary in the presence of the explicit warp synchronization.
2. Afaik there is at least one occurrence of redundant synchronization in line 417 here: https://github.com/NVIDIA/cccl/blob/3d67fa8da087651514402e9c0b41e06a32acff29/cub/cub/warp/specializations/warp_scan_smem.cuh#L411-L423
The second store is explicitly done with an offset such that the aggregate is not overwritten.
3. There is an explicit optimization for `cuda::std::plus` on primitive types based on knowing the identity value. That optimization could easily be done for other operators as well (`cuda::std::multiplies`, `cuda::minimum`, `cuda::maximum`, `cuda::std::bit_or`, `cuda::std::bit_and`, `cuda::std::bit_xor`) in https://github.com/NVIDIA/cccl/blob/3d67fa8da087651514402e9c0b41e06a32acff29/cub/cub/warp/specializations/warp_scan_smem.cuh#L134-L139 `HAS_IDENTITY` could be true for all of these operators. See also `identity_v` from #5317
4. Another optimization currently only done for `cuda::std::plus` on integers is calculating the exclusive scan from the inclusive scan as `exclusive = inclusive - input` instead of communicating through shared memory. The same could be done for `cuda::std::bit_xor` and in theory `cuda::std::multiplies` although integer division might not actually be faster than going through shared memory.
5. Remove `CellT` and all the C-style casts associated with it. https://github.com/NVIDIA/cccl/blob/3d67fa8da087651514402e9c0b41e06a32acff29/cub/cub/warp/specializations/warp_scan_smem.cuh#L86-L87
Contributor guide
Assessment
This issue has not been assessed yet.