Shared Memory Tiling for Matrix Multiplication

Multiply $C = A \times B$ where $A$ is $M \times K$, $B$ is $K \times N$, $C$ is $M \times N$. Each output element is a dot product:

$C_{ij} = \sum_{k=0}^{K-1} A_{ik} \cdot B_{kj}$

For a square $N \times N$ case: $N^2$ outputs, each costing $N$ multiply-adds. $2N^3$ FLOPs total, but only $3N^2$ bytes of input/output. The arithmetic intensity grows linearly with $N$ — matmul is the textbook compute-dense workload, if you don't blow the memory hierarchy. Naive matmul does blow it. That's what tiling fixes.

The obvious mapping: each thread computes one element of $C$. It reads a full row of $A$ and a full column of $B$ from global memory.

__global__ void matmulNaive(const float* A, const float* B, float* C, int N) {
  int row = blockIdx.y * blockDim.y + threadIdx.y;
  int col = blockIdx.x * blockDim.x + threadIdx.x;
  if (row >= N || col >= N) return;
  float acc = 0.0f;
  for (int k = 0; k < N; ++k) {
    acc += A[row * N + k] * B[k * N + col];
  }
  C[row * N + col] = acc;
}

It's correct. It's also painfully slow. Let's see why.

Each thread does $2N$ global loads ($N$ from A, $N$ from B) to produce one output. Across all $N^2$ threads that's $2N^3$ global loads.

But every element of $A$ is genuinely needed only $N$ times (once per column of $B$). Every element of $B$ is needed $N$ times (once per row of $A$). The minimum number of distinct loads is $2N^2$.

So the naive kernel reads each input value $N$ times from global memory. For $N = 1024$ that's a 1024× overcount. The L2 cache absorbs some of it but not enough. The kernel ends up memory-bound on a chip that could be doing 80 TFLOPS.

Break $C$ into blocks of $\text{TILE} \times \text{TILE}$ outputs (say, 32×32). Each block of threads computes one tile of $C$.

Observation: to compute a TILE×TILE region of $C$, you only need strips of $A$ and $B$ of width TILE. The threads in the block all need the same strips. So:

1. Have the block cooperatively load a TILE×TILE slab of $A$ and a TILE×TILE slab of $B$ into shared memory — once.
2. Now every thread can do its part of the dot product reading from shared memory, not global.
3. Slide along $K$, loading the next pair of slabs, accumulating.

Each element of $A$ and $B$ now gets loaded from global memory TILE times fewer than the naive version.

One block produces a TILE-by-TILE chunk of C by loading paired slabs.

#define TILE 32

__global__ void matmulTiled(const float* A, const float* B, float* C, int N) {
  __shared__ float As[TILE][TILE];
  __shared__ float Bs[TILE][TILE];

  int row = blockIdx.y * TILE + threadIdx.y;
  int col = blockIdx.x * TILE + threadIdx.x;
  float acc = 0.0f;

  for (int t = 0; t < N / TILE; ++t) {
    // Each thread loads one element of each slab
    As[threadIdx.y][threadIdx.x] = A[row * N + (t * TILE + threadIdx.x)];
    Bs[threadIdx.y][threadIdx.x] = B[(t * TILE + threadIdx.y) * N + col];
    __syncthreads();          // wait for all loads

    for (int k = 0; k < TILE; ++k) {
      acc += As[threadIdx.y][k] * Bs[k][threadIdx.x];
    }
    __syncthreads();          // wait before overwriting tiles
  }
  C[row * N + col] = acc;
}

Launch with dim3 block(TILE, TILE); dim3 grid(N/TILE, N/TILE);. Add boundary guards if $N$ isn't a multiple of TILE.

Two __syncthreads() calls bookend the work in each iteration of the outer loop. They're load-bearing:

• First barrier — after all threads load their elements into As and Bs, before any thread reads from them. Without it, some threads would read stale or uninitialised tile entries while others were still writing.
• Second barrier — after every thread finishes its $K$-loop on the current slab, before any thread overwrites As/Bs with the next slab. Without it, fast threads would clobber the data slower threads were still reading.

__syncthreads() only synchronises threads in the same block. There is no built-in way to sync across blocks within a kernel; that's what kernel boundaries are for.

For $N = 1024$: naive does 2048 loads/output; tiled does 64. 32× less global memory traffic. Bigger tiles win bigger — until you hit the shared memory budget per block (~48 KB) or register pressure capping occupancy.

In practice, on an A100 this single optimisation takes matmul from ~5% of peak FLOPS to ~50%. Going further (warp-level tiling, double-buffering, tensor cores) is how cuBLAS hits 95%+. But the fundamental idea is right here: trade one expensive global load for TILE cheap shared-memory reads.

A few traps people hit on their first tiled matmul:

• Boundary handling. If $N$ isn't a multiple of TILE, the loads can run off the edge of $A$ or $B$. Either pad the matrices or guard each load with if and write zeros.
• Bank conflicts. Reading Bs[k][threadIdx.x] is column access — the textbook fix is __shared__ float Bs[TILE][TILE+1]; (the +1 padding shifts banks).
• Choosing TILE. 32×32 = 1024 threads per block is the max; 16×16 = 256 threads runs more blocks per SM (higher occupancy) and is often faster on smaller GPUs. Profile both.
• Forgetting the second __syncthreads(). Causes race conditions that only show up at certain matrix sizes or compiler optimisation levels. Painful to debug.