Warp-Level Matrix Multiplication with `nvcuda::wmma`

The nvcuda::wmma C++ API

The Warp Matrix Multiply-Accumulate (WMMA) API is a C++ template library, accessible by including the mma.h header file. It lives within the nvcuda::wmma namespace. Its purpose is to provide a high-level, programmatic interface to the Tensor Core hardware, abstracting away the tedious and error-prone details of writing raw PTX instructions and manually managing register data layouts.

The core idea remains the same as what we saw with PTX: operations are performed at the warp level. When you call a wmma function, you are directing all 32 threads in the warp to act in concert.

The central data structure in this API is the nvcuda::wmma::fragment.

The fragment: An Opaque Container

A fragment is a C++ object that represents a piece of a matrix distributed across the private registers of an entire warp. The key word here is opaque. As a programmer, you don't know—and shouldn't care—which thread holds which specific element of the matrix. You only need to know that the fragment object as a whole contains the matrix data for the warp.

A fragment is a templated type, and its parameters define the matrix operation:
nvcuda::wmma::fragment<Use, M, N, K, T, Layout>

• Use: Specifies the fragment's role in the MMA operation. It can be nvcuda::wmma::matrix_a, nvcuda::wmma::matrix_b, or nvcuda::wmma::accumulator.
• M, N, K: These integers define the shape of the warp-level MMA operation, $D_{M\times N}=A_{M\times K}\times B_{K\times N}+C_{M\times N}$. A common shape for many architectures is 16x16x16.
• T: The data type of the elements, such as __half (for FP16) or float.
• Layout: Can be nvcuda::wmma::col_major or nvcuda::wmma::row_major. This parameter tells the API about the logical layout of the matrix in memory (e.g., global memory), from which data will be loaded. The API then handles the complex task of transforming this memory layout into the specific, opaque register layout required by the hardware.

The Core WMMA Workflow

Using the API involves a clear, three-step workflow for each MMA computation: load, compute, and store.

Let's watch a brief video that walks through this exact workflow in a simplified code example.

Now, let's dissect each of those steps in more detail by examining a more complete code implementation.

1. Loading Data: load_matrix_sync

Before you can compute, you must load matrix tiles from memory (global or shared) into your fragment objects.

Function Signature:
nvcuda::wmma::load_matrix_sync(fragment, memory_pointer, leading_dimension);

• The first argument is the destination fragment object (e.g., a_frag, b_frag).
• The second is a pointer to the top-left element of the matrix tile in memory.
• The third, leading_dimension (or stride), is crucial. It tells the function how many elements are in a full row (for row-major) or column (for column-major) of the matrix in memory. This allows it to correctly calculate the address of subsequent rows/columns when fetching the tile.

In the reference code, you can see these calls inside the for loop:

// Load the mma matrix inputs.
nvcuda::wmma::load_matrix_sync(a_frag, matrix_mma_a_mptr, lda);
nvcuda::wmma::load_matrix_sync(b_frag, matrix_mma_b_mptr, ldb);

The _sync suffix signifies this is a synchronous, collective operation. All threads in the warp must execute this call, and the warp will not proceed until the load is complete for all threads.

2. Performing MMA: mma_sync

This is the function that actually invokes the mma.sync PTX instruction on the Tensor Cores.

Function Signature:
nvcuda::wmma::mma_sync(accumulator_D, fragment_A, fragment_B, accumulator_C);

This performs the operation $D=A\times B+C$. The result is written into the accumulator_D fragment. A common pattern for accumulation is to use the same fragment for accumulator_D and accumulator_C, effectively performing C = A * B + C.

In the reference code, this is the core of the computation:

// Perform the matrix multiplication
nvcuda::wmma::mma_sync(acc_frag, a_frag, b_frag, acc_frag);

Before this loop begins, acc_frag is initialized to zero using nvcuda::wmma::fill_fragment(acc_frag, 0.0f);. This ensures the first MMA operation is an accumulation onto zero.

3. Storing Results: store_matrix_sync

After the computation is complete (typically after the loop over the K-dimension), the final result held in the accumulator fragment is written back to memory.

Function Signature:
nvcuda::wmma::store_matrix_sync(memory_pointer, fragment, leading_dimension);

• The first argument is a pointer to the top-left element of the destination tile in memory.
• The second is the source accumulator fragment containing the results.
• The third is the leading dimension of the output matrix in memory.

In the example code, this happens after the for loop:

// Store the output
nvcuda::wmma::store_matrix_sync(matrix_mma_c_mptr, c_frag, ldc, nvcuda::wmma::mem_col_major);

The code also demonstrates the full GEMM capability by first loading the original C matrix into a fragment, performing element-wise scaling (alpha * AB + beta * C), and then storing the final result.

Tiling Strategy and Data Flow

The wmma API operates on fixed-size tiles (e.g., 16x16). To multiply larger matrices, we use a tiled approach. The provided code demonstrates a common strategy:

1. The entire output matrix C is partitioned into tiles.
2. The grid of CUDA blocks is laid out such that each warp is responsible for computing one 16x16 tile of the output matrix.
3. To compute its output tile, each warp must iterate along the K-dimension, loading corresponding 16x16 tiles from matrix A and matrix B, and accumulating the results using mma_sync.

The image below visualizes this hierarchical breakdown. A large problem is decomposed into thread block tiles, which are further decomposed into warp tiles. The warp then loads fragments of matrices A and B to compute its warp tile result.

Conclusion

You have now moved from the raw PTX assembly of the last lesson to a practical C++ API for Tensor Core programming. You've seen that the nvcuda::wmma library provides a powerful abstraction that hides the complexities of register management while still giving you direct control over the warp-level MMA operations.

Key Takeaways:

• The nvcuda::wmma API in the mma.h header is the standard C++ interface for programming Tensor Cores.
• The nvcuda::wmma::fragment is an opaque object representing a matrix tile distributed across a warp's registers.
• The core workflow consists of three main synchronous, warp-wide functions:
  1. load_matrix_sync: Loads data from memory into a fragment.
  2. mma_sync: Executes the $D=A\times B+C$ operation on the Tensor Cores.
  3. store_matrix_sync: Writes data from a fragment back to memory.
• Large matrix multiplications are built by tiling the problem, where each warp computes a small, fixed-size output tile by looping over the K-dimension and accumulating partial results.

We've focused on moving data directly from global memory to registers for the wmma operations. However, this is not the most performant approach, as it leads to many redundant global memory reads. In our next lesson, we will combine the wmma API with the shared memory tiling techniques you learned in Module 9. You will learn how to construct a complete, high-performance GEMM kernel that stages data from global memory into shared memory, and then from shared memory into fragments, creating an efficient data pipeline to truly feed the beast.