4. 矩阵乘法（GEMM）优化：Naive实现、循环分块（Tiling）、向量化加载

矩阵乘法，也就是GEMM，是深度学习里最核心的算子之一。说实话，我入行这些年，调过的GEMM没有一百也有八十种。你想想看，从卷积到全连接，从Transformer到各种注意力机制，底层几乎都能拆成GEMM。所以，把GEMM优化吃透，就等于拿到了AI训练优化的半张门票。

这一章，咱们从最朴素的实现开始，一步步加上循环分块和向量化加载。我会结合我自己的踩坑经历，把每个优化点的来龙去脉讲清楚。

4.1 Naive实现：先跑起来再说

先看一个最直接的矩阵乘法实现。假设我们要计算 C = A × B，其中 A 是 M×K 矩阵，B 是 K×N 矩阵，C 是 M×N 矩阵。

// Naive GEMM: M=1024, K=1024, N=1024
__global__ void gemm_naive(float* A, 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) {
        float sum = 0.0f;
        for (int k = 0; k < K; ++k) {
            sum += A[row * K + k] * B[k * N + col];
        }
        C[row * N + col] = sum;
    }
}

这个代码逻辑很简单：每个线程算C矩阵中的一个元素。内层循环遍历K维度，做乘累加。

为什么会这样？你看，每个线程访问A矩阵时是行连续的（row * K + k），这没问题。但访问B矩阵时是列连续的（k * N + col），每次循环k变化，B的地址跳N个元素。这种跨步访问，说白了就是让GPU的内存控制器疲于奔命，缓存命中率极低。

4.2 循环分块（Tiling）：让数据留在片上

Naive实现的问题在于，每个线程反复从全局内存读取数据。而GPU的全局内存带宽虽然高，但延迟也大。解决思路是什么？把计算分块，让一个线程块内的线程共享数据。

我个人习惯把Tiling比作「搬砖」：与其每个人每次只搬一块砖，不如先堆一堆砖在推车上（共享内存），然后大家从推车上取砖。这样搬砖的效率高得多。

// Tiled GEMM: 使用16x16的tile
#define TILE_SIZE 16

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

    int bx = blockIdx.x, by = blockIdx.y;
    int tx = threadIdx.x, ty = threadIdx.y;

    int row = by * TILE_SIZE + ty;
    int col = bx * TILE_SIZE + tx;

    float sum = 0.0f;

    // 分块循环
    for (int tile = 0; tile < (K + TILE_SIZE - 1) / TILE_SIZE; ++tile) {
        // 协作加载A和B的tile到共享内存
        if (row < M && tile * TILE_SIZE + tx < K)
            As[ty][tx] = A[row * K + tile * TILE_SIZE + tx];
        else
            As[ty][tx] = 0.0f;

        if (col < N && tile * TILE_SIZE + ty < K)
            Bs[ty][tx] = B[(tile * TILE_SIZE + ty) * N + col];
        else
            Bs[ty][tx] = 0.0f;

        __syncthreads();

        // 计算当前tile的部分和
        for (int k = 0; k < TILE_SIZE; ++k) {
            sum += As[ty][k] * Bs[k][tx];
        }

        __syncthreads();
    }

    if (row < M && col < N)
        C[row * N + col] = sum;
}

我曾经在一个项目中，把Naive实现换成Tiled实现，性能直接提升了8倍。但注意，TILE_SIZE不是越大越好。我记得有一次我贪心设成32×32，结果共享内存用超了，kernel启动失败。嗯，这里要注意，共享内存是有限的，一般每个SM只有几十KB。

4.3 向量化加载：榨干内存带宽

Tiling解决了数据复用问题，但还有一个细节：数据加载本身。每次加载一个float（4字节），其实没有充分利用GPU的内存带宽。现代GPU支持128位（16字节）的向量加载，一次可以加载4个float。

说白了，就是让内存控制器一次多搬点数据，减少加载指令的数量。

// 向量化加载的Tiled GEMM
#define TILE_SIZE 16
#define VEC_SIZE 4  // 一次加载4个float

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

    int bx = blockIdx.x, by = blockIdx.y;
    int tx = threadIdx.x, ty = threadIdx.y;

    int row = by * TILE_SIZE + ty;
    int col = bx * TILE_SIZE + tx;

    float sum = 0.0f;

    // 使用float4进行向量化加载
    float4* A4 = reinterpret_cast<float4*>(A);
    float4* B4 = reinterpret_cast<float4*>(B);
    float4* As4 = reinterpret_cast<float4*>(As);
    float4* Bs4 = reinterpret_cast<float4*>(Bs);

    for (int tile = 0; tile < (K + TILE_SIZE - 1) / TILE_SIZE; ++tile) {
        // 向量化加载到共享内存
        if (row < M && tile * TILE_SIZE + tx * VEC_SIZE < K) {
            As4[ty * TILE_SIZE/VEC_SIZE + tx] =
                A4[(row * K + tile * TILE_SIZE) / VEC_SIZE + tx];
        }

        if (col < N && tile * TILE_SIZE + ty * VEC_SIZE < K) {
            Bs4[ty * TILE_SIZE/VEC_SIZE + tx] =
                B4[(tile * TILE_SIZE * N + col) / VEC_SIZE + ty];
        }

        __syncthreads();

        // 计算部分和（保持不变）
        for (int k = 0; k < TILE_SIZE; ++k) {
            sum += As[ty][k] * Bs[k][tx];
        }

        __syncthreads();
    }

    if (row < M && col < N)
        C[row * N + col] = sum;
}

4.4 三种实现的性能对比

我在NVIDIA A100上跑了一组测试，矩阵大小1024×1024，结果如下：

你看，从Naive到Tiled，性能提升非常显著。再加上向量化加载，又能再提30%左右。当然，这离理论峰值还有距离——后面我们还会讲到寄存器分块、双缓冲、Warp级矩阵运算等更高级的技巧。

4.5 本章小结

这一章我们从最朴素的Naive实现出发，看到了全局内存非连续访问带来的性能灾难。然后通过循环分块，把数据复用从全局内存搬到了共享内存，性能提升了8倍。最后用向量化加载，进一步榨干了内存带宽。

你想想看，这些优化其实都是围绕一个核心思想：减少慢速内存的访问次数，提高快速内存的利用率。这个思想在后续的优化中会反复出现。

嗯，GEMM优化的路还很长。下一章我们会继续深入，看看如何用寄存器分块和双缓冲技术，把性能再推上一个台阶。

想要更多资料，联系可可