NVIDIA / NVIDIA/cccl

[EPIC] Optimize `thrust::transform` for newer architectures

Open
#1,947 8 comments 4 reactions 1 assignee Claimed by @bernhardmgruber View on GitHub
cub thrust
Dominant language
C++
Stars
2.5k
Forks
486
Avg merge
2d 6h
Merged PRs (30d)
295

Description

**Motivation**
It's increasingly harder to reach SOL on newer GPU architectures, starting with A100 and H100, especially for simple kernels, like:
`thrust::transform(..., thrust::plus{})`, which basically load a few values and perform little compute. CUB algorithms already counter this by processing several elements per thread, but internal research hints at the necessity to further increase the amount of data in flight.

**Use case**
`thrust::transform` is an important primitive for many algorithms and also occurs in BabelStream, i highly relevant HPC benchmark often used to produce representative numbers to compare the performance of hardware architectures. We should therefore dedicate some effort to ensure `thrust::transform` performs well.

**Approach**
The main strategy is to have more "bytes in flight" when reading, with the concrete amount depending on the target architecture (tuning parameter). There are multiple ways to generate more loads. Again, internal research points to using either prefetching or the tensor memory accelerator (TMA, e.g. via `memcpy_async`) on newer architectures. Excessive unrolling and loading to registers works as well, but has the drawback of consuming large amount of registers for architectures requiring a large number of bytes in flight.

**Address stability**
For the loading strategy we have to consider the address stability of data items as well. Users sometimes rely on the ability to retrieve the index inside an input array from the reference of a loaded element:
```c++
transform(par, a, a + n, a, [a,b,c](const T& e) {
const auto i = &e – a; // &e expected to point into global memory
return e + b[i] + c[i];
});
```
Such a user-provided function object inhibits any optimization which loads elements from global memory into registers or shared memory before passing them as arguments, thus only allowing prefetching as optimization. Address oblivious function objects can benefit from a larger variety of optimizations (like TMA or pipelined loading to registers.

**Further concerns**
Furthermore, the computational intensity and shared memory/register consumption of the user provided function object influence the loading strategy. Longer computations seem to require more data in flight. Shared memory is contested by TMA and user-side computation. Register pressure limits unrolling.

**Status quo**
`thrust::transform` (CUDA) is currently built on top of `cub::DeviceFor::Bulk`, which eventually dispatches independently of the uses data types or number of input and output streams. Because `cub::DeviceFor::Bulk` is index based, the involved input and output data streams are not visible and no tuning based on this information is possible. The situation is similar with `cub::DeviceFor::ForEach` et al.

**Strategy**
I propose to add a new CUB algorithm `cub::DeviceTransform` governing transformations of N input streams into a single output stream (maybe M output streams if use cases arrise) and rebasing `thrust::transform` on top of it.

### Future tasks after merging `cub::DeviceTransform`
- [x] https://github.com/NVIDIA/cccl/issues/2091
- [x] Split BabelStream benchmarks by number if iterators
- [x] https://github.com/NVIDIA/cccl/issues/2363
- [x] Drop the fallback_for transform algorithm
- [x] https://github.com/NVIDIA/cccl/issues/2263
- [x] https://github.com/NVIDIA/cccl/issues/2361
- [x] https://github.com/NVIDIA/cccl/issues/2714
- [x] https://github.com/NVIDIA/cccl/issues/2403
- [x] Profile memcpy_async kernel on Ampere workstation GPUs
- [x] ~~Refactor prefetching and TMA loading into new CUB block load algorithms for reuse (requires CUB block load redesign)~~ We have `BlockLoadToSharedNow`.
- [x] Port `thrust::transform_if` to `cub::DeviceTransform` as well
- [x] Review Thrust and try to find more uses for `transform` (especially replacing `for_each`)
- [x] https://github.com/NVIDIA/cccl/issues/2717
- [x] ~~Evaluate the need for the unroll_staged kernel before A100 and for Ampere workstation GPUs~~
- [x] ~~Try to mark non-aliased input streams with `__restrict__`~~ `__restrict__` cannot be combined with PDL in the same kernel, since it introduces data races. So let's skip restrict.
- [x] ~~Port `thrust::transform_if` with stencil to `cub::DeviceTransform` as well~~ Not necessary for this issue.
- [ ] #7507

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.