Abdi Moalim

Choosing shmem or warp shuffles

Warp shuffles work with registers without sync besides implicit lockstep execution which makes them ideal for intra-warp loads/stores. On the other hand, shared memory is required once loads/stores cross warp boundaries. In many cases, the required information is known at compile time and we would be interested in using the least unit of scheduling.

A simple approach is to encode the cooperation topology in a type.

template<int X, int Y = 1, int Z = 1>
struct tile
{
  static constexpr int size = X * Y * Z;
};

Obviously, the unused branch is eliminated because the tile size is known.

template<typename Tile> __device__ float reduce(float value) {
  if constexpr (Tile::size <= warpSize) {
    /* intra-warp/shuffle */
  } else {
    /* inter-warp/shmem */
  }
}

The same interface can be used regardless of the cooperation size.

auto a = reduce(tile<32>{}, value);
auto b = reduce(tile<128>{}, value);

The first instantiation selects warp shuffles whereas the second selects shmem without virtual dispatch or SFINAE.

Any collective operation whose implementation depends on said layout can use the same compile-time selection.

reduce<tile<64>>(value);
scan<tile<64>>(value);
broadcast<tile<64>>(value);
exchange<tile<64>>(value);

I only like this because it's very declarative which is ironic because kernels write to memory and never return. Refer to NVIDIA's Cooperative Groups library for related ideas.