GEMM From Naive CUDA to Tensor Cores

A matrix multiplication can return every right number and still waste most of a GPU. One thread mapping scatters memory requests. A better mapping makes adjacent lanes touch adjacent values. Tiling then reuses those values on chip, while Tensor Cores change the arithmetic instruction itself.

The previous GPU compiler lesson separated source language, generated code, and target architecture. In the kernel engineering lab, memory access, occupancy, and synchronization became tuning decisions. Finally, the profiling lesson required correctness before speed. Keep all three contracts active here.

We'll build one operation, $C = AB$, through five kernels. General matrix-matrix multiplication (GEMM) names this family of products. Matrix $A$ has shape $[M,K]$, $B$ has shape $[K,N]$, and $C$ has shape $[M,N]$. All lab buffers use row-major storage, so neighboring columns occupy neighboring addresses.

The complete CUDA source is available as gemm_lab.cu. Its kernels are intentionally educational. A production library has more schedules, layout transforms, epilogues, and hardware-specific paths than one chapter should hide inside a code listing.

Compute one output before launching a warp

Start with matrices small enough to multiply on paper:

$$A = \begin{bmatrix}1 & 2 & 0 \\ -1 & 3 & 2\end{bmatrix}, \qquad B = \begin{bmatrix}2 & -1 \\ 1 & 4 \\ 3 & 0\end{bmatrix}$$

$A$ is $2 \times 3$ and $B$ is $3 \times 2$, so $C$ must be $2 \times 2$. Entry $C_{1,0}$ takes row 1 of $A$ and column 0 of $B$:

$$C_{1,0} = (-1)(2) + (3)(1) + (2)(3) = -2 + 3 + 6 = 7$$

Applying the same dot product to each output gives:

$$C = \begin{bmatrix}4 & 7 \\ 7 & 13\end{bmatrix}$$

GEMM performs $MNK$ multiplications and about the same number of additions. Performance tools conventionally count $2MNK$ floating-point operations (FLOPs), including one multiply and one add for each inner-loop step. Our tiny product therefore counts $2 \cdot 2 \cdot 2 \cdot 3 = 24$ FLOPs.

The script below reproduces the hand calculation and compares two traffic models. Its first model assumes every output reloads two input values for every $k$. An optimistic second model counts each input and output element once, giving a lower bound for a cold, standalone operation rather than a claim about a real cache.

Those intensities are small because the matrices are small. For a square FP32 GEMM with $M=N=K=n$, the unique-byte model becomes:

$$I_{\text{ideal}} = \frac{2n^3}{4(n^2+n^2+n^2)} = \frac{n}{6}\ \text{FLOP/B}$$

Large square GEMMs can have enough reuse to become compute-bound, but only if the kernel captures that reuse. The roofline model bounds attainable compute rate by the lower of peak arithmetic throughput and memory bandwidth multiplied by arithmetic intensity.^{[1]Reference 1Roofline: An Insightful Visual Performance Model for Multicore Architectureshttps://doi.org/10.1145/1498765.1498785} It can't tell us which tile or layout wins. It tells us whether reducing bytes can still move the ceiling.

Why is 0.375 FLOP/B an optimistic bound for the tiny product rather than a measured intensity?

Answer

It counts every A, B, and C element exactly once. A real kernel may fetch cache lines it only partly uses, reload evicted data, or move padding and metadata. A warm cache can also reduce DRAM traffic for one launch, so measure device traffic and record cache conditions.

A correct thread can still request the wrong addresses

The first CUDA kernel assigns one thread to one output. Each thread runs the same $K$-step dot product:

CUDA linearizes threadIdx.x first. Adjacent lanes therefore change row while holding col fixed in this mapping. At one inner-loop value $k$, those lanes read A[row * K + k] with a stride of $K$ and store C[row * N + col] with a stride of $N$. The values are mathematically correct, but a warp's addresses are scattered across memory.

Swap ownership so threadIdx.x selects col:

Now adjacent lanes read adjacent B[k * N + col] values and write adjacent outputs. Threads sharing a row request the same $A$ address instead of 32 unrelated addresses. NVIDIA's current guidance describes coalescing as combining a warp's accesses into the necessary memory transactions and treats adjacent access as a high-priority pattern.^{[2]Reference 2CUDA C++ Best Practices Guidehttps://docs.nvidia.com/cuda/cuda-c-best-practices-guide/}

Coalescing repairs transaction shape. It doesn't preserve an input tile for reuse across multiple inner products. Shared memory handles that next.

Turn global loads into shared tiles

Pick a $16 \times 16$ output tile. One block loads a $16 \times 16$ tile of $A$ and a $16 \times 16$ tile of $B$ into shared memory, computes 256 partial dot products, then advances 16 positions along $K$. Every loaded $A$ value can feed 16 output columns, and every loaded $B$ value can feed 16 output rows.

Two block-wide barriers form correctness boundaries. A first barrier prevents a thread from reading a shared value before another thread writes it. A second prevents a fast thread from overwriting a tile that a slower thread still needs. NVIDIA's GEMM example uses the same shared-memory idea to remove redundant global loads and requires block synchronization when warps consume values written by other warps.^{[2]Reference 2CUDA C++ Best Practices Guidehttps://docs.nvidia.com/cuda/cuda-c-best-practices-guide/}

Notice how edge handling preserves barrier participation. Each cooperative load checks its own A or B coordinate and writes zero only when that source coordinate is out of range. A thread without a valid output may still load an input needed by its neighbors, and it mustn't return before __syncthreads(). Only the final global store is predicated.

For a complete $16 \times 16 \times 16$ step, the block performs $2 \cdot 16^3 = 8192$ FLOPs. It loads two 256-element FP32 tiles, or 2048 bytes, before the output store. Ignoring the store until the full $K$ loop ends, that step exposes $8192/2048=4$ FLOP/B at the global-to-shared boundary. Larger threadblock tiles can raise reuse further, but they also consume more shared memory and registers.

An edge thread has row >= M. Why must it still reach both barriers inside the K loop?

Answer

Other threads in the block may need its B load and its barrier participation. The edge thread still performs both predicated loads, zero-fills only invalid source coordinates, reaches each block-wide barrier, and skips its final out-of-range store.

Let each thread keep more than one answer

The shared kernel assigns one accumulator register to each thread. It spends instruction and address overhead for every output while leaving each thread with little independent arithmetic to schedule. Register blocking gives a thread several nearby outputs.

Our next kernel keeps four accumulators per thread. A $16 \times 16$ block computes a $16 \times 64$ output tile. Each thread owns columns tx, tx + 16, tx + 32, and tx + 48 within the tile:

One shared $A$ value now feeds four thread-local accumulators before the thread moves to the next $k$. The $B$ tile is wider, so the block cooperatively loads it in several passes. The downloadable source includes those loads and guards its right edge.

More registers aren't free. A wider per-thread tile may increase instruction-level parallelism and reuse, but it can reduce occupancy or spill registers into local memory. CUTLASS's efficient GEMM description notes that accumulators can consume at least half of a thread's register budget in blocked kernels.^{[3]Reference 3Efficient GEMM in CUDAhttps://docs.nvidia.com/cutlass/latest/media/docs/cpp/efficient_gemm.html} Check registers per thread, local-memory traffic, active warps, and runtime together. Occupancy alone isn't the objective.

Register blocking changes ownership, not the mathematical product. Tensor Core MMA changes the instruction that performs each tile product.

Hand a tile to Tensor Cores

Tensor Cores are specialized matrix multiply-accumulate units. NVIDIA exposed Volta Tensor Cores through the Warp Matrix Multiply Accumulate (WMMA) API in CUDA 9. A full warp cooperates on a fragment whose element-to-lane mapping is intentionally opaque. The original FP16 path multiplies two $16 \times 16$ fragments and accumulates into a $16 \times 16$ fragment, commonly in FP32.^{[4]Reference 4Programming Tensor Cores in CUDA 9https://developer.nvidia.com/blog/programming-tensor-cores-cuda-9/}

Three rules prevent subtle bugs:

1. Every active lane in the warp must execute the WMMA operations coherently.
2. Fragment layout and leading dimensions must match the backing memory.
3. Edge fragments need padding or a separate predicated path because load_matrix_sync isn't a per-element guarded load.

For FP16 inputs, load_matrix_sync requires a 32-byte-aligned pointer and a leading dimension divisible by 8 elements. The FP32 output store requires the same pointer alignment and a leading dimension divisible by 4 elements. Padding every dimension to a multiple of 16 satisfies those stride rules, while cudaMalloc and the tile-aligned offsets preserve pointer alignment.^{[5]Reference 5CUDA Programming Guide.https://docs.nvidia.com/cuda/cuda-programming-guide/}

The lab uses row-major FP16 inputs, FP32 accumulation, and one warp per $16 \times 16$ output tile:

M isn't used inside the aligned kernel because the host pads all three dimensions to multiples of 16. The runtime launches only complete padded tiles.

Padding is explicit work. For the lab's default $M=257$, $N=263$, and $K=251$, WMMA sees $272 \times 272 \times 256$. Logical work remains $2MNK$, while padding raises executed work to $2M_pN_pK_p$. The lab prints both counts because a padded path can appear inefficient when its extra work stays hidden.

Layout is part of the type

A[row * K + k] describes row-major $A$. B[k * N + col] describes row-major $B$. Changing a fragment declaration to column-major without transforming data doesn't request the same matrix through a faster path. It requests different logical elements.

Libraries often name GEMMs with layout pairs such as NN, NT, or TN. Those letters describe whether operands are consumed in their stored or transposed orientation. Always record:

• logical shapes $[M,K]$, $[K,N]$, and $[M,N]$;
• physical row or column order;
• leading dimensions and batch strides;
• alignment and any padded dimensions;
• input, accumulator, and output dtypes.

An output that looks plausibly random can still be a correct multiplication of the wrong layout.

Correctness needs two references

The FP32 kernels multiply original FP32 inputs. Their CPU reference accumulates those same values in FP64. The WMMA kernel first rounds each input to FP16, so its reference converts the rounded FP16 values back to FP32 and accumulates them in FP64. Comparing WMMA against the unrounded input product would mix input quantization error with accumulation error.

The lab accepts an element when:

$$|C_{\text{gpu}}-C_{\text{ref}}| \leq \text{atol} + \text{rtol}\,|C_{\text{ref}}|$$

Absolute tolerance (atol) protects results near zero. Relative tolerance (rtol) scales with result magnitude. Neither pair is universal. Increase $K$, widen input range, introduce cancellation, change accumulator dtype, or add a fused epilogue, and the error distribution changes.

Use three correctness shapes before timing large squares:

Seed input generation, keep values bounded, include zeros and signed values, and fail on NaN or Inf. Then run NVIDIA Compute Sanitizer to catch memory and synchronization defects that numerical comparison may miss.^{[6]Reference 6Compute Sanitizerhttps://docs.nvidia.com/compute-sanitizer/ComputeSanitizer/}

Run the lab and keep the receipt

Check the installed toolkit and device first. Compile for the GPU you will execute on, not for a favorite architecture copied from another machine:

The default build target is sm_80 only as an explicit starting choice. Use sm_70 for a Volta WMMA lab, sm_80 for Ampere, or the matching target for a newer device. A binary built only for sm_80 isn't a Hopper-specialized TMA kernel.

Exact latency and throughput are device-specific measurements, so the expected output fixes schema and correctness status without inventing performance values:

The program performs five warmups, records CUDA events around repeated launches in one stream, synchronizes the stop event, and reports mean kernel time. For a serious receipt, collect at least median and tail values over multiple process runs. Rotate buffers or state whether warm-cache behavior is intentional.

Use a structured row so someone else can reproduce the comparison:

Don't write 4.2x faster unless the receipt names both kernels, same shape, same device state, same timing method, same precision contract, and equivalent correctness. Compare against a tuned library such as cuBLAS or a CUTLASS profiler kernel before calling a handwritten kernel competitive.

Nsight Compute can then explain a timing change. Start with a small metric set instead of collecting every counter:

Record requested and actual kernel names, memory sectors, achieved occupancy, register count, shared-memory use, tensor-pipe activity, and replay warnings. NVIDIA's profiler documentation also provides roofline analysis for connecting measured arithmetic intensity to compute and bandwidth ceilings.^{[7]Reference 7Nsight Compute Profiling Guidehttps://docs.nvidia.com/nsight-compute/ProfilingGuide/}

Where asynchronous copies begin

The shared and WMMA kernels issue a load, wait, compute, and repeat. A pipelined main loop instead prepares tile $t+1$ while computing tile $t$:

This snippet expresses ownership and overlap. It isn't a portable implementation: copy instructions, barriers, address spaces, alignment, stage count, and MMA schedule all depend on target architecture.

NVIDIA documents two distinct hardware boundaries:

LDGSTS support starts at compute capability 8.0. TMA starts at 9.0 and adds bulk multidimensional copies.^{[5]Reference 5CUDA Programming Guide.https://docs.nvidia.com/cuda/cuda-programming-guide/} Current CUDA documentation also warns that some high-level cuda::memcpy_async calls fall back to synchronous copies when alignment or size requirements aren't met, while lower-level TMA APIs make violating those requirements undefined behavior.^{[8]Reference 8Asynchronous Data Copieshttps://docs.nvidia.com/cuda/cuda-programming-guide/04-special-topics/async-copies.html}

An async-copy kernel returns correct values but matches the synchronous kernel's timing. Why isn't the API name proof that transfer overlapped arithmetic?

Answer

A high-level copy can fall back to a synchronous path when its alignment or size contract isn't met. Even with an asynchronous instruction, too little independent computation, too few pipeline stages, or waits in the wrong place can remove overlap. Inspect the emitted instruction path and stall metrics before changing the stage count.

Hopper's warp-group MMA path adds another boundary. WGMMA is collective across four contiguous warps and targets sm_90a; it has its own fence, commit, and wait protocol.^{[9]Reference 9CUTLASS and CuTe Documentationhttps://docs.nvidia.com/cutlass/latest/overview.html} Replacing wmma::mma_sync with an instruction name doesn't produce a correct Hopper pipeline.

CUTLASS shows why the chapter stops at this line. Its GEMM main loops double-buffer shared-memory tiles and warp fragments so memory movement overlaps arithmetic, then select threadblock schedules and epilogues around the target shape.^{[3]Reference 3Efficient GEMM in CUDAhttps://docs.nvidia.com/cutlass/latest/media/docs/cpp/efficient_gemm.html} Use the CUTLASS profiler to obtain a verified library receipt rather than treating demonstration examples as benchmarks:

CUTLASS kernel availability and command arguments vary by built version and target architecture. Record the repository tag or commit, build configuration, kernel name, and profiler verification result. NVIDIA's current CUTLASS examples page explicitly says examples aren't performance benchmarks and directs measurement work to the profiler.^{[10]Reference 10CUTLASS Programming Exampleshttps://github.com/NVIDIA/cutlass/blob/main/examples/README.md}

Diagnose the first bad signal

Every row starts from a symptom, not a favorite optimization. If correctness fails, stop timing. If timing moves without the expected counter change, revisit the hypothesis.

Keep the artifact honest

A complete lab submission contains:

• source and exact compile command;
• device, driver, toolkit, and target architecture;
• odd, tiled, and long-$K$ correctness shapes;
• reference dtype plus atol and rtol;
• warmup, timing, synchronization, and cache-state method;
• logical and padded operation counts;
• one profiler capture tied to a named kernel;
• a cuBLAS or CUTLASS baseline with the same input and output contract;
• conclusions limited to measured shapes.

The progression now has a stable meaning. Coalescing repairs warp transactions. Shared tiling captures reuse. Register blocking increases per-thread reuse and independent accumulation. WMMA moves a cooperative tile onto Tensor Cores. Async copies and TMA overlap movement with compute, but only after architecture-specific synchronization and layout contracts are correct.