NVIDIA / NVIDIA/cccl

[FEA]: Prepare productization of BlockLoadToShared

Open
#6,315 0 comments 0 reactions 1 assignee Claimed by @elstehle View on GitHub
Dominant language
C++
Stars
2.5k
Forks
486
Avg merge
2d 6h
Merged PRs (30d)
295

Description

## Background

This issue is part of a broader initiative that aims to achieve two main goals:
1. Providing users with an abstraction that makes it easy to leverage modern hardware features (e.g., TMA), to efficiently load data from global memory to shared memory.
2. Introduce utilities that facilitate the use of modern hardware features for data movement in CUB algorithms

## Next steps

### 1. Determining the right interface for `BlockLoadToShared`
- [x] Consider dropping the default constructor of `BlockLoadToShared`, as it relies on shared memory that is shared by all instances of the same `BlockLoadToShared` template specialization, which is problematic. Alternatively, we could consider identifying whether the `mbarrier` is used by multiple instances and add an assertion. However, I would start out with a safe and defensive approach of removing the default constructor for now.
- [x] We want to consider encouraging the right sequence of member function invocations, particularly `ctor()->CopyAsync()*->Commit()->Wait()`. Similar to `barrier_arrive_tx`, this could be encouraged by having `Commit` return a token that needs to be passed to `Wait()`.

### 2. Integrating `BlockLoadToShared` into a first set of CUB algorithms

- [[FEA]: Use BlockLoadToShared in cub::DeviceFind #6412](https://github.com/NVIDIA/cccl/issues/6412)
- [[FEA] Optimize cub::DeviceMerge by using cub::detail::BlockLoadToShared #6005](https://github.com/NVIDIA/cccl/issues/6005)

### 3. Determining an optimal implementation for the TMA code path
With the device algorithms and stand-alone benchmarking of TMA we can explore and evaluate whether there's further room for improvement of the current implementation or whether we can use the implementation as is and start tuning the device algorithms that `BlockLoadToShared` was integrated into.

### 4. Explore a simple `BlockLoad` specialization that loads via shared memory to registers
We want to explore whether this is a expedient low-effort approach to achieve performance improvements across the board in CUB on recent architectures. The idea is to have a `BlockLoad` specialization that loads via shared memory to registers (using `BlockLoadToShared` under the hood) that can serve as drop-in replacement for existing `BlockLoad` usage in CUB's device-scope algorithms.

### 5. Design an abstraction that can support optimal performance
Long-term we want to explore a new processing scheme that allows leveraging pipelining processing and TMA data loading together with work stealing.

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.