前幾天一路從 GPU 的 Thread、Warp、Memory Hierarchy 講到 Roofline Model。

今天終於可以真的來寫一個 CUDA kernel。

這次直接拿 LeetGPU 的 General Matrix Multiplication（GEMM）Challenge 來練習。

LeetGPU 給的 template 長這樣：

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

// 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
) {}

我們要做的事情其實就是 GEMM：

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

其中：

A: M × K
B: K × N
C: M × N

而 half 就是 FP16。

一個 Thread 算一個 C

先不要想任何 optimization。

矩陣乘法裡，C 的每一個元素都可以獨立計算。

例如：

C[0][0]  C[0][1]  C[0][2]
C[1][0]  C[1][1]  C[1][2]
C[2][0]  C[2][1]  C[2][2]

最直覺的方法就是：

一個 CUDA Thread 負責一個 C 的元素。

例如某個 Thread 負責 C[row][col]，它只需要拿：

A 的第 row 列
×
B 的第 col 欄

做一次 dot product。

所以第一版 kernel 可以直接寫成：

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

__global__ void gemm_naive(
    const half* A,
    const half* B,
    half* C,
    int M,
    int N,
    int K,
    float alpha,
    float beta
) {
    int row = blockIdx.y * blockDim.y + threadIdx.y;
    int col = blockIdx.x * blockDim.x + threadIdx.x;

    if (row >= M || col >= N) {
        return;
    }

    float sum = 0.0f;

    for (int k = 0; k < K; ++k) {
        float a = __half2float(A[row * K + k]);
        float b = __half2float(B[k * N + col]);

        sum += a * b;
    }

    float old_c = __half2float(C[row * N + col]);

    C[row * N + col] =
        __float2half(alpha * sum + 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(16, 16);

    dim3 blocks(
        (N + threads.x - 1) / threads.x,
        (M + threads.y - 1) / threads.y
    );

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

這裡剛好可以把 Day 2 的東西重新串回來。

我們設定：

dim3 threads(16, 16);

代表一個 Block 有：

16 × 16 = 256 Threads

每個 Thread 再透過：

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

找到自己要負責的 C element。

所以：

Thread (0, 0) → C[0][0]
Thread (0, 1) → C[0][1]
Thread (0, 2) → C[0][2]
...

這就是前幾天講的：

Grid
 ↓
Blocks
 ↓
Threads

第一次真的用到了。

為什麼累加用 float？

雖然 A、B、C 都是：

half

但我們中間的 accumulator 是：

float sum = 0.0f;

也就是用 FP32 累加。

因為每個 C element 都要做 K 次：

A[row][k] × B[k][col]

再全部加起來。

如果一路都用 FP16 累加，數值誤差會比較容易累積。

所以這裡會：

FP16 input
   ↓
轉成 FP32
   ↓
Multiply + Accumulate
   ↓
最後轉回 FP16

這種 mixed precision 的做法之後看 Tensor Core 時還會再遇到。

但這個版本其實很浪費

這個 kernel 看起來已經把大量工作平行化了。

但看看相鄰兩個 Threads。

Thread 0 算：

C[0][0]

需要：

A[0][0]
A[0][1]
A[0][2]
...

Thread 1 算：

C[0][1]

結果它也需要：

A[0][0]
A[0][1]
A[0][2]
...

也就是：

Thread 0 ──→ A[0][k]
Thread 1 ──→ A[0][k]
Thread 2 ──→ A[0][k]
Thread 3 ──→ A[0][k]

很多 Threads 都在使用相同的資料。

但在這個 naive kernel 裡，我們沒有主動把這些資料搬到可以重複使用的地方。

這就是 Day 3 講的 Data Movement 又回來了。

用了很多 Threads，不代表就有效利用了 GPU。

如果 Threads 一直在等 Global Memory，那 GPU 有再多計算能力也不一定用得上。

下一步：能不能把資料搬近一點？

既然同一個 Block 裡的 Threads 會重複使用 A 和 B 的部分資料，那很自然會想到：

能不能先把這些資料從 Global Memory 搬進 Shared Memory？

例如：

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

這就是 Tiling 的基本想法。

我們不是每次需要一個數字就重新去 Global Memory 拿，而是把 Matrix 切成一塊一塊，先搬一塊進來，再把它盡可能多用幾次。

這件事會：

Data Reuse ↑
Data Movement ↓
Arithmetic Intensity ↑

剛好又和昨天的 Roofline Model 接起來。

所以今天先完成最直覺的版本：

One Thread, One Output Element。