昨天在 LeetGPU 上寫了最簡單的 GEMM：

C = alpha * (A × B) + beta * C

做法也很直覺：

一個 Thread 算一個 C element。

但問題是，相鄰的 Threads 其實會一直使用到相同的 A、B 資料。Naive 版本卻讓大家各自從 Global Memory 重新讀取。

所以今天先做 GEMM 最經典的第一個優化：Shared Memory Tiling。

把資料搬近一點

想法很簡單。

不要每做一次乘法就跑去 Global Memory 拿資料，而是把 A、B 切成 Tile：

Global Memory
      ↓
 A Tile + B Tile
      ↓
 Shared Memory
      ↓
 Threads 重複使用

例如 TILE = 16，一個 16 × 16 的 Block 負責計算一個 16 × 16 的 C Tile。

A、B 的 Tile 先由 Threads 一起搬到 Shared Memory，之後再重複使用。

核心就是前幾天一直提到的：

Data Reuse ↑
Global Memory Access ↓

Tiled GEMM

繼續使用 LeetGPU 的 template：

#include <cuda_fp16.h>
#include <cuda_runtime.h>

#define TILE 16

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

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

    float sum = 0.0f;

    for (int t = 0; t < (K + TILE - 1) / TILE; ++t) {
        int a_col = t * TILE + threadIdx.x;
        int b_row = t * TILE + threadIdx.y;

        As[threadIdx.y][threadIdx.x] =
            (row < M && a_col < K)
                ? A[row * K + a_col]
                : __float2half(0.0f);

        Bs[threadIdx.y][threadIdx.x] =
            (b_row < K && col < N)
                ? B[b_row * N + col]
                : __float2half(0.0f);

        __syncthreads();

        for (int k = 0; k < TILE; ++k) {
            sum += __half2float(As[threadIdx.y][k]) *
                   __half2float(Bs[k][threadIdx.x]);
        }

        __syncthreads();
    }

    if (row < M && col < N) {
        float old_c = __half2float(C[row * N + col]);
        C[row * N + col] =
            __float2half(alpha * sum + beta * old_c);
    }
}

extern "C" void solve(
    const half* A,
    const half* B,
    half* C,
    int M, int N, int K,
    float alpha, float beta
) {
    dim3 threads(TILE, TILE);
    dim3 blocks(
        (N + TILE - 1) / TILE,
        (M + TILE - 1) / TILE
    );

    gemm_tiled<<<blocks, threads>>>(
        A, B, C, M, N, K, alpha, beta
    );
}

__syncthreads() 的用途也很直覺：要等同一個 Block 裡的 Threads 都把 Tile 搬完，大家才能開始使用它。

H100 實測

我在 LeetGPU 的 H100 上測了幾個 Tile Size：

Tile = 8 幾乎沒差，但到了 16 和 32，改善就明顯很多。

目前最快的是 Tile = 32：

0.53 ms → 0.34 ms
≈ 1.56× speedup

至少可以確認一件事：同樣的 GEMM，只是改變資料搬移和 reuse 的方式，就能有明顯差距。

那為什麼不直接 TILE = 64？

很自然會想到：

8 → 16 → 32 都變快了
那 64 呢？

但我們現在的寫法是：

dim3 threads(TILE, TILE);

所以：

TILE = 8  →   64 Threads
TILE = 16 →  256 Threads
TILE = 32 → 1024 Threads
TILE = 64 → 4096 Threads

H100 一個 Thread Block 最多只能有 1024 Threads。

因此 TILE = 32 已經剛好碰到上限，TILE = 64 會直接變成非法的 4096 Threads / Block。

不過要注意：

不能的是 64 × 64 Threads，不是 64 × 64 Output Tile。

這兩件事不一樣。

一個 Thread 其實可以算不只一個 C

我們到現在都假設：

1 Thread → 1 C element

但其實完全可以：

1 Thread → 2 × 2 C elements

甚至：

1 Thread → 4 × 4 C elements

例如一個 64 × 64 C Tile 總共有 4096 個 outputs。

如果每個 Thread 算 4 × 4 = 16 個 outputs：

4096 / 16 = 256 Threads

那就不需要 4096 Threads 了。

而這些 C 的 partial sums 可以放在 Registers 裡一路累加：

Global Memory
      ↓
Shared Memory
      ↓
Registers
      ↓
一個 Thread 算多個 Outputs

這就是 Thread Tiling / Register Tiling 的基本想法。

資料從 Global Memory 搬進 Shared Memory 後，還可以再進一步在 Register 裡 reuse。

當然，一個 Thread 也不能無限算更多元素，因為 Register 用量會跟著增加，最後又會影響能同時留在 SM 上的 Threads / Warps。

GEMM optimization 基本上就是一直在這些 trade-off 中找平衡。

附錄：Thread Tiling kernel

0.25ms

#include <cuda_fp16.h>
#include <cuda_runtime.h>

#define BM 64   // output tile rows
#define BN 64   // output tile cols
#define BK 16   // K dimension tile

#define TM 4    // outputs per thread in M dimension
#define TN 4    // outputs per thread in N dimension

__global__ void gemm_register_tiled(
    const half* A,
    const half* B,
    half* C,
    int M,
    int N,
    int K,
    float alpha,
    float beta
) {
    __shared__ half As[BM][BK];
    __shared__ half Bs[BK][BN];

    // 16 x 16 = 256 threads
    int tx = threadIdx.x;
    int ty = threadIdx.y;

    int tid = ty * blockDim.x + tx;
    int num_threads = blockDim.x * blockDim.y;

    // Each thread owns a 4 x 4 output tile
    int row_base = blockIdx.y * BM + ty * TM;
    int col_base = blockIdx.x * BN + tx * TN;

    // 4 x 4 partial sums kept in registers
    float acc[TM][TN] = {0.0f};

    for (int kb = 0; kb < K; kb += BK) {

        // --------------------------------
        // Cooperatively load A tile
        // A tile: 64 x 16 = 1024 elements
        // --------------------------------
        for (int idx = tid; idx < BM * BK; idx += num_threads) {
            int r = idx / BK;
            int k = idx % BK;

            int global_r = blockIdx.y * BM + r;
            int global_k = kb + k;

            if (global_r < M && global_k < K)
                As[r][k] = A[global_r * K + global_k];
            else
                As[r][k] = __float2half(0.0f);
        }

        // --------------------------------
        // Cooperatively load B tile
        // B tile: 16 x 64 = 1024 elements
        // --------------------------------
        for (int idx = tid; idx < BK * BN; idx += num_threads) {
            int k = idx / BN;
            int c = idx % BN;

            int global_k = kb + k;
            int global_c = blockIdx.x * BN + c;

            if (global_k < K && global_c < N)
                Bs[k][c] = B[global_k * N + global_c];
            else
                Bs[k][c] = __float2half(0.0f);
        }

        __syncthreads();

        // --------------------------------
        // Compute
        //
        // One thread computes:
        //
        // C00 C01 C02 C03
        // C10 C11 C12 C13
        // C20 C21 C22 C23
        // C30 C31 C32 C33
        //
        // --------------------------------
        #pragma unroll
        for (int k = 0; k < BK; ++k) {

            float a_frag[TM];
            float b_frag[TN];

            // Load A values into registers
            #pragma unroll
            for (int i = 0; i < TM; ++i) {
                a_frag[i] =
                    __half2float(As[ty * TM + i][k]);
            }

            // Load B values into registers
            #pragma unroll
            for (int j = 0; j < TN; ++j) {
                b_frag[j] =
                    __half2float(Bs[k][tx * TN + j]);
            }

            // Outer product:
            // 4 A values x 4 B values = 16 FMAs
            #pragma unroll
            for (int i = 0; i < TM; ++i) {
                #pragma unroll
                for (int j = 0; j < TN; ++j) {
                    acc[i][j] += a_frag[i] * b_frag[j];
                }
            }
        }

        __syncthreads();
    }

    // --------------------------------
    // Write the 4 x 4 output tile
    // --------------------------------
    #pragma unroll
    for (int i = 0; i < TM; ++i) {
        #pragma unroll
        for (int j = 0; j < TN; ++j) {

            int row = row_base + i;
            int col = col_base + j;

            if (row < M && col < N) {
                float old_c = __half2float(C[row * N + col]);

                C[row * N + col] =
                    __float2half(
                        alpha * acc[i][j]
                        + beta * old_c
                    );
            }
        }
    }
}


// A, B, and C are device pointers

extern "C" void solve(
    const half* A,
    const half* B,
    half* C,
    int M,
    int N,
    int K,
    float alpha,
    float beta
) {
    dim3 threads(
        BN / TN,   // 64 / 4 = 16
        BM / TM    // 64 / 4 = 16
    );

    dim3 blocks(
        (N + BN - 1) / BN,
        (M + BM - 1) / BM
    );

    gemm_register_tiled<<<blocks, threads>>>(
        A, B, C,
        M, N, K,
        alpha, beta
    );
}