Matmul I: tiling and shared memory

1. The kernel that decides everything

The problem, fixed for this chapter and the next: C = A · B with A of shape [M, K], B of shape [K, N], C of shape [M, N], all row-major, all FP32. The BLAS name is SGEMM (single-precision general matrix multiply). Each of the M·N outputs is a dot product of length K, so 2MNK FLOPs in total. The benchmark size is 4096 × 4096 × 4096: 137 GFLOPs, 4.4 ms at your card’s 30.9 TFLOPS FP32 peak if every cycle of every FMA unit were used.

Nobody reaches the peak, so the course measures against the best real implementation: cuBLAS, NVIDIA’s BLAS, with dozens of hand-tuned kernels per architecture. Lab 6 calls it in plain FP32 mode (no TF32 tensor cores, which would be a different contest; chapter 8) and reports every kernel as a percentage of cuBLAS. Expect cuBLAS itself at somewhere around 60–85% of the nominal peak: the peak assumes the boost clock, and under a long, power-hungry matmul the card clocks lower.

The road through these two chapters is well mapped. Simon Boehm’s worklog, the standard modern account, took SGEMM on an RTX A6000 through ten kernels; his measurements, as a percentage of cuBLAS, are the landmarks we will predict against:

Look at the shape of that column before reading on. The ideas of this chapter, the ones every CUDA tutorial teaches, get you to about an eighth of cuBLAS. The big steps come after, and the reason they come after is the most important thing this chapter has to explain.

2. The naive kernel

One thread per output element, each running the dot product in a loop. You met it in lab 5 as matmul_naive; here is the general version with a 2D launch:

__global__ void sgemm_naive(const float* A, const float* B, float* C, int M, int N, int K)
{
    int row = /* which row of C this thread owns */;
    int col = /* which column */;
    if (row >= M || col >= N) return;
    float acc = 0.0f;
    for (int k = 0; k < K; ++k)
        acc += A[row * K + k] * B[k * N + col];     // one FMA, two global loads
    C[row * N + col] = acc;
}

Every iteration of the loop does one fused multiply-add (FFMA, 2 FLOPs) and two loads from global memory. If each load went to DRAM, that would be 2 FLOPs per 8 bytes, an intensity of 0.25, and a ceiling of 168 GFLOPS: half a percent of peak (chapter 4). Before measuring anything, the question is how the two blank lines at the top are filled in, because that decides what a warp’s 32 loads look like.

3. Which way the warp lies decides coalescing

A warp is 32 threads with consecutive threadIdx.x (chapter 2). So the mapping from threadIdx.x to the output decides which 32 outputs a warp computes, and therefore which addresses its loads hit at each step k. There are two natural choices:

// v1: x walks the ROWS of C              // v2: x walks the COLUMNS of C
int row = blockIdx.x * 32 + threadIdx.x;  int col = blockIdx.x * 32 + threadIdx.x;
int col = blockIdx.y * 8  + threadIdx.y;  int row = blockIdx.y * 8  + threadIdx.y;

v2’s A load deserves a second look. All 32 lanes read the same address, and the hardware sends one request and broadcasts the value to every lane. That is as cheap as a load can be, and it will matter again when the same pattern appears in shared memory.

4. Where the naive kernel’s loads really go

v2 does 2 FLOPs per 8 bytes loaded. If DRAM served every byte, it could not exceed 168 GFLOPS. Yet it will run at several TFLOP/s on your card, ten or more times faster. So most of its loads never reach DRAM. Look at who else needs the same data at the same moment:

• With 32×8 blocks, the 8 warps of a block compute 8 rows of the same 32 columns. At step k they all load the same 128 bytes of row k of B. One of them brings it into L1; the other seven hit.
• Each warp’s A load is one float of its row, and consecutive k are consecutive floats: one 32-byte sector serves 8 iterations, so 7 of every 8 A loads hit L1.
• Blocks in the same column of C, possibly on other SMs, read the same B columns; blocks in the same row read the same A rows. L2 (48 MB) catches much of that reuse between blocks.

So the caches recover a good part of the reuse the math allows, by luck of timing. But notice what they do not change: every FMA still issues two load instructions, each of which has to travel through the L1/texture pipeline, be checked against the cache tags and come back hundreds of cycles later if it missed. The kernel spends its time waiting for loads and issuing loads. The fix is to stop relying on luck: bring the data the block will reuse into on-chip storage once, on purpose, and read it from there.

5. Tiling: load once, use T times

Give each block a T×T tile of C, one thread per element as before. To compute it the block needs a strip of T full rows of A and a strip of T full columns of B, far too much to hold at once. So it walks along K in phases of width T. In phase p:

1. every thread loads one element of the T×T A tile (the block’s rows, columns pT…pT+T−1) and one element of the B tile (rows pT…, the block’s columns) into shared memory;
2. the block waits at a barrier until both tiles are complete;
3. every thread does T multiply-adds, reading a row of the A tile and a column of the B tile from shared memory;
4. the block waits again, so that nobody overwrites a tile someone is still reading.

Count the global loads. The naive kernel issues 2K loads per output. The tiled kernel issues 2 per thread per phase, and there are K/T phases: 2K/T per output, T times fewer. If every one of them went to DRAM, the intensity would be T/4 FLOP/B: 4 for 16×16 tiles, 8 for 32×32 (chapter 4’s numbers). Walk through it:

6. The tiled kernel, line by line

template <int T>
__global__ void sgemm_tiled(const float* A, const float* B, float* C, int M, int N, int K)
{
    __shared__ float As[T][T];
    __shared__ float Bs[T][T];
    const int tx = threadIdx.x, ty = threadIdx.y;
    const int row = blockIdx.y * T + ty;          // this thread's element of C
    const int col = blockIdx.x * T + tx;

    float acc = 0.0f;
    for (int k0 = 0; k0 < K; k0 += T) {                                              // one phase
        As[ty][tx] = (row < M && k0 + tx < K) ? A[row * K + k0 + tx]   : 0.0f;       // (1)
        Bs[ty][tx] = (k0 + ty < K && col < N) ? B[(k0 + ty) * N + col] : 0.0f;       // (2)
        __syncthreads();                                                             // (3)
        #pragma unroll
        for (int k = 0; k < T; ++k)
            acc += As[ty][k] * Bs[k][tx];                                            // (4)
        __syncthreads();                                                             // (5)
    }
    if (row < M && col < N) C[row * N + col] = acc;                                  // (6)
}
// launch: dim3 block(T, T), grid(ceil(N / T), ceil(M / T))

1. (1) and (2), the loads, are both coalesced. In (1), consecutive tx read consecutive columns of a row of A; in (2), consecutive tx read consecutive columns of a row of B. That is the point of choosing which thread loads which element: the thread that loads an element is not the one that uses it, so loads can be arranged for coalescing independently of the math.
2. Out-of-range elements load as 0, which adds nothing to a dot product. That is how one kernel handles any M, N, K. Threads outside C do not return early: they still load and still reach both barriers (chapter 3’s rule), and only the final store (6) is guarded.
3. (3) and (5) are chapter 3’s two barriers: read-after-write (the tile must be complete before anyone reads it) and write-after-read (nobody overwrites it for the next phase until everyone is done). On a GPU a missing barrier shows up only sometimes, depending on timing, which is what makes it dangerous. (The CPU emulator these labs were checked with, dev/cudasim in the repository, runs each thread as far as it can before switching to the next, which turns a missing barrier into a wrong answer every time.)
4. (4), the inner loop, is conflict-free. With T = 32, a warp is one row of the block (one ty, tx = 0…31). As[ty][k] is the same word for all 32 lanes: a broadcast. Bs[k][tx] is 32 consecutive words, one per bank. With T = 16 a warp spans two rows: two distinct As words 16 banks apart, and the two half-warps read the same 16 Bs words, still one pass each.

7. Why tiling stops near a tenth of peak

Tiling divides global traffic by T. Look at the inner loop again, though, and count per FMA: one load from As, one load from Bs. The loads now come from shared memory instead of L1/L2, which is faster and never misses, but there are still two of them per multiply-add. Now compare the two pipes they compete for, per SM, per clock cycle:

• FMA: 128 FP32 lanes, so up to 4 warp-wide FFMA instructions per cycle (one per quadrant scheduler).
• Shared memory: 32 banks of 4 bytes, so one warp-wide load per cycle (one pass through the banks, whether it is 32 different words or one broadcast).

So the tiled kernel’s ceiling is one eighth of peak, about 3.9 TFLOP/s on your card, and it does not depend on T. One caveat, in the spirit of chapter 5’s “check the SASS”: in the unrolled inner loop a thread reads As[ty][k], As[ty][k+1], … four consecutive, aligned floats, and the compiler may merge them into a single 16-byte LDS.128. That makes one A load per four FMAs, 1.25 shared loads per FMA instead of 2, and raises the ceiling to about a fifth of peak. make sass in lab 6 counts the instructions so you can see which one your compiler produced. Either way the ceiling is 12–20% of peak, and real kernels land below their ceiling (barriers, latency, the loads themselves take issue slots).

Here is the whole picture for any tile size: occupancy, the global-traffic ceiling, and the shared-memory ceiling:

The ratio that remains

To go further, each value loaded from shared memory must be used by more than one FMA. There is only one place faster than shared memory to keep it: the thread’s own registers. If a thread computes not one output but a small block of them, say TM rows by TN columns, then at each step k it can load TM values of A and TN values of B into registers once and combine every pair: TM·TN FMAs for TM + TN loads.

That is register blocking, and it is where chapter 8 starts. It is also why the table in section 1 jumps from 13% to 69% of cuBLAS in two steps.

8. Predictions, and what the profiler should show

At 4096³, before you run anything. These ranges come from the mechanisms above and from Boehm’s ratios, not from measurements on your card:

The two stall names for the tiled kernels are new. mio_throttle means a warp could not issue because the memory-input/output queue, which feeds shared memory, was full; short_scoreboard means a warp is waiting for the result of a shared-memory load (a short-latency dependency, as opposed to long_scoreboard’s global loads). Together they are the profiler’s way of saying “shared memory is the bottleneck”: exactly figure 3.

9. Lab 6: from naive to tiled

cd ~/unibo-lessons/mlsys/labs/lab06-matmul-tiling
make test        # every kernel on every shape: PASS
make bench       # the table; v1 at 4096 takes a few seconds
make sass        # FFMA and LDS counts per kernel: is it 2 loads per FMA, or 1.25?
make ncu-table   # sectors per request, shared loads, FFMAs, stall reasons at 2048 cubed

Check your understanding

Sources and further reading

• S. Boehm, How to Optimize a CUDA Matmul Kernel for cuBLAS-like Performance: a Worklog (2022): the kernel sequence and the A6000 percentages quoted in section 1.
• Hwu, Kirk, El Hajj, Programming Massively Parallel Processors, 4th ed. (2023), ch. 3 (the naive matmul), ch. 5 (tiling with shared memory, boundary conditions), ch. 6 (coalescing and thread coarsening).
• NVIDIA, CUDA C++ Best Practices Guide, “Shared Memory in Matrix Multiplication (C=AB)”: the coalescing and reuse argument for the tiled kernel.
• NVIDIA, CUDA C++ Programming Guide: shared memory, bank behaviour, __syncthreads() semantics; Nsight Compute Profiling Guide: the mio_throttle and short_scoreboard stall reasons.
• V. Volkov, J. Demmel, Benchmarking GPUs to Tune Dense Linear Algebra (SC 2008): the original argument that register blocking, not more threads, is what feeds the math units.
• NVIDIA, cuBLAS documentation: cublasSgemm, column-major conventions, math modes.