Matrix multiplication

Look again at the two barriers in the shared-memory kernel above. The second one exists only because there is one buffer: no thread may begin staging the next tile until every thread has finished reading the current one. That makes the mainloop strictly alternate between a load phase in which the math units have nothing to do and a compute phase in which the memory pipe has nothing to do, and the escape from an alternation like that is to hold two of whatever is contended. CUTLASS calls the technique software pipelining and implements it by double buffering at two scopes: threadblock-scoped shared-memory tiles, where two tiles are allocated and one loads data for the current matrix operation while the other buffers data from global memory for the next mainloop iteration, and warp-scoped matrix fragments, where two fragments are allocated in registers and one is passed to the CUDA and tensor cores while the other receives shared-memory fetch returns for the next warp-level operation.

The reason a GEMM kernel has to arrange this by hand connects back to the occupancy argument above. CUTLASS’s own explanation is that the blocked structure demands a large storage allocation in each thread’s registers, with the accumulator elements alone typically occupying at least half of a thread’s register budget, so occupancy is relatively low compared to other classes of GPU workload, and that limits the GPU’s ability to hide memory latency and other stalls by switching to another resident warp. There are fewer warps to switch to, so the latency has to be hidden inside the thread instead.

The other half of the technique is the route the bytes take. Prior to cuda::memcpy_async, copying from global to shared memory was a two-step process: the thread block copied data from global memory into registers and then from registers into shared memory, which NVIDIA describes as a long journey through the memory hierarchy. Ampere added a direct path. With an asynchronous copy the block no longer stages data through registers, which frees it from the task of moving data and frees those registers to be used by computations; because the copy uses no registers, register pressure falls and occupancy improves, and the copy is emitted without depending on the compiler’s dynamic loop unrolling.

The instruction underneath is LDGSTS, available from compute capability 8.0. It copies 4, 8, or 16 bytes: at 4 or 8 bytes the data is also cached in L1, while the 16-byte form enables an L1 bypass mode that leaves the cache unpolluted. Global to shared is the only direction it supports, and best performance wants both the shared and the global address aligned to 128 bytes. Completion is signalled through a shared-memory barrier or a pipeline, and the default is narrower than it looks: each thread waits only for its own copies, so a block whose threads prefetch data that other threads will read still needs a __syncthreads() after the completion wait. Substituting the asynchronous copy into a staging loop is otherwise mechanical:

#include <cooperative_groups.h>
#include <cooperative_groups/memcpy_async.h>

template <typename T>
__global__ void example_kernel(T * global1, T * global2, size_t subset_count)
{
    extern __shared__ T shared[];
    auto group = cooperative_groups::this_thread_block();

    for (size_t subset = 0; subset < subset_count; ++subset) {
        cooperative_groups::memcpy_async(group, shared,
                                         &global1[subset * group.size()], sizeof(T) * group.size());
        cooperative_groups::memcpy_async(group, shared + group.size(),
                                         &global2[subset * group.size()], sizeof(T) * group.size());

        cooperative_groups::wait(group); // Wait for all copies to complete

        compute(shared);

        group.sync();
    }
}

That version is a drop-in replacement: cooperative_groups::memcpy_async stands in for the assignment and cooperative_groups::wait stands in for the barrier that used to guard it, and the loop keeps the shape the matmul mainloop has, copy the subset, wait, compute, synchronize. What it does not do yet is overlap anything. It shortens the route the bytes take and hands the registers back, but the block still waits for the copy it just issued before it computes on the result.

Overlap needs stages. Because asynchronous data movement lets multiple batches be in flight at the same time, a block can pipeline its way through a large structure by submitting N stages of asynchronous copies, waiting for the oldest stage to complete, computing with that batch of shared memory, and submitting a new stage before waiting on the next oldest. At N of 2 that is exactly double buffering: stage k+1 is being fetched while stage k is being computed. cuda::pipeline names the four points where the handoff has to be arranged:

template<int block_dim, int num_stages>
__global__ void in_between_pipe_thread(int* dest, int const* src, size_t size) {
    // Read blockDim.x integers per pipeline stage
    __shared__ int smem[num_stages][block_dim];

    // Grid stride loop:
    int offset = blockIdx.x * blockDim.x;
    size_t stride = gridDim.x * blockDim.x;

    // No pipeline::shared_state needed
    cuda::pipeline<cuda::thread_scope_thread> pipe = cuda::make_pipeline();

    // Load all pipeline stages.
    for (int stage = 0; stage < num_stages; ++stage) {
        pipe.producer_acquire();
        size_t idx = offset + stage * stride + threadIdx.x;
        if (idx < size) {
            cuda::memcpy_async(&smem[stage][threadIdx.x], &src[idx], sizeof(int), pipe);
        }
        pipe.producer_commit();
    }

    // At this point, there are `num_stages` commited into the pipeline. This is a loop.
    // invariant that is upheld throughout the loop.
    int stage = 0;
    for (size_t block_idx = offset; block_idx < size; block_idx += stride) {
        // Wait for the first stage to have completed loading, or equivalently: wait until
        // at most `num_stages - 1` stages are still loading.
        cuda::pipeline_consumer_wait_prior<num_stages - 1>(pipe);

        // __syncthreads is necessary if other threads want to read this thread's loaded data.
        __syncthreads();

        // ... compute on smem[stage][..] ...

        // __syncthreads is necessary if other threads are reading data that this thread
        // is about to overwrite below.
        __syncthreads();

        // Release the consumed stage.
        pipe.consumer_release();

        // Pre-load data for `num_stages` into the future.
        pipe.producer_acquire();
        // To ensure that the number of commited stages into the pipeline remains constant,
        // producer_acquire and producer_commit are called even if the load is out-of-bounds.
        size_t idx = block_idx + num_stages * stride + threadIdx.x;
        if (idx < size) {
            cuda::memcpy_async(&smem[stage][threadIdx.x], &src[idx], sizeof(int), pipe);
        }
        pipe.producer_commit();

        stage = (stage + 1) % num_stages;
    }
}

producer_acquire claims a stage before the copy is issued and producer_commit publishes it; pipeline_consumer_wait_prior with a template argument of num_stages - 1 waits until at most that many stages are still loading, which is the same statement as waiting for the oldest; and consumer_release hands the stage back. The stage index advancing modulo num_stages is what turns the shared-memory allocation into a circular buffer. The loop invariant the comments name is the part worth copying into a kernel of your own: exactly num_stages stages stay committed at the top of every iteration, which is why the acquire and commit pair still runs when the load it guards is out of bounds and copies nothing.

CUTLASS generalizes those four calls into its Pipeline classes, where producer_acquire blocks until a consumer releases the stage it wants, producer_commit notifies waiting consumers without blocking, consumer_wait blocks until the producers have committed a stage, and consumer_release frees the stage again. Its stated reason for having the abstraction is worth weighing before hand-rolling one: a persistent GEMM manages dozens of different kinds of asynchronously executing operations synchronizing through multiple barriers organized as a circular list, complexity CUTLASS calls too much for human programmers to manage by hand. It is also the vocabulary the next section needs. Hopper’s warp specialization is this pipeline with the two roles handed to different warps rather than to different phases of one loop: a producer warp group waits for a shared-memory buffer to be signalled empty and issues a TMA load into it, the TMA updates that stage’s barrier on completion, and a consumer warp group waits on the barrier, runs the tensor-core matmuls, and releases the buffer for the next load.