Naive GEMM vs Tiled GEMM

The same multiplication, two performance worlds

GEMM (C = A·B) is where 'just throw threads at it' stops working. The naive version assigns one thread to one output element and loops over the K dimension. The tiled version cooperates within a block to load chunks of A and B into shared memory, then reuses each chunk many times before the next DRAM trip. Same answer, dramatically different speed.

Naive GEMM — feel the pain

Each thread computes one C[m][n]:

for k in 0..K: C[m][n] += A[m][k] * B[k][n]

Every FMA requires two global-memory loads (one from A, one from B). No reuse, no shared memory. The kernel is bandwidth-bound and runs at ~5% of peak FLOPs.

Tiled GEMM — keep hot data in shared memory

Per output tile, threads in a block cooperatively load A-tile and B-tile into shared memory. Each thread then computes its partial output by reading from shared memory — many cheap reads per one global load. Inner loop:

Load next A-tile and B-tile (global → shared).
__syncthreads() to ensure all threads see the loaded data.
Multiply-accumulate into a register: Csub += As[ty][k] * Bs[k][tx].
Move to next K-slice; repeat.

Numbers from RTX 4090 (FP32, non-Tensor-Core path):

Kernel	Time (ms)	TFLOP/s	Comment
gemm-naive (512³)	0.65	0.41	One DRAM read per FMA
gemm-tiled (2048³)	4.50	3.8	Shared-mem reuse, register blocking
cuBLAS sgemm	~3.9	~7.0	Multi-stage tiling, Tensor Cores via cublasGemmEx

The 9× speedup from naive to tiled is purely about data reuse. cuBLAS adds another ~2× by using Tensor Cores and multi-stage software pipelining (Track 8).

Code

Naive GEMM (CUDA) — one thread per output element·cuda

// One thread per output element. Two global loads per FMA. Bandwidth-bound.
__global__ void gemm_naive(const float* A, const float* B, float* C,
                           int M, int N, int K) {
    int row = blockIdx.y * blockDim.y + threadIdx.y;
    int col = blockIdx.x * blockDim.x + threadIdx.x;
    if (row >= M || col >= N) return;

    float c = 0.f;
    for (int k = 0; k < K; ++k) {
        c += A[row * K + k] * B[k * N + col];   // 2 loads, 1 FMA
    }
    C[row * N + col] = c;
}

Tiled GEMM (CUDA) — block-cooperative tile loads·cuda

// 32×32 tiles. Each block computes one 32×32 chunk of C.
// Inside the block, threads cooperatively load A and B tiles into shared
// memory, sync, then multiply-accumulate from shared.
#define TILE 32

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

    int tx = threadIdx.x, ty = threadIdx.y;
    int row = blockIdx.y * TILE + ty;
    int col = blockIdx.x * TILE + tx;

    float c = 0.f;
    for (int kt = 0; kt < (K + TILE - 1) / TILE; ++kt) {
        int a_col = kt * TILE + tx;
        int b_row = kt * TILE + ty;
        As[ty][tx] = (row < M && a_col < K) ? A[row * K + a_col] : 0.f;
        Bs[ty][tx] = (b_row < K && col < N) ? B[b_row * N + col] : 0.f;
        __syncthreads();

        #pragma unroll
        for (int k = 0; k < TILE; ++k) {
            c += As[ty][k] * Bs[k][tx];          // both reads from smem
        }
        __syncthreads();
    }
    if (row < M && col < N) C[row * N + col] = c;
}

// Launch:
//   dim3 block(TILE, TILE);
//   dim3 grid((N+TILE-1)/TILE, (M+TILE-1)/TILE);
//   gemm_tiled<<<grid, block>>>(dA, dB, dC, M, N, K);

Exercise

Implement both naive and tiled GEMM (32×32 tile) for a 1024×1024 FP32 multiply on whatever GPU you have. Verify both against numpy.dot. Time them — you should see roughly 5–10× speedup from naive to tiled. Then increase to 2048×2048 — the gap usually widens because the larger problem makes the naive version even more bandwidth-starved.