4. 矩阵乘法性能优化

矩阵乘法 C = A × B 是 GPU 上最典型、也最容易被“朴素实现”拖慢的算子。假设 A 为 M×K、B 为 K×N、C 为 M×N，朴素 CUDA 实现让每个线程计算 C 的一个元素，需要读取 A 的一整行和 B 的一整列，共 K 次乘加。问题在于：同一行 A 元素会被 N 个线程重复读取，同一列 B 元素会被 M 个线程重复读取，全局内存访问量被放大了 N 倍和 M 倍。本节围绕“如何让数据被复用、让访存与计算重叠”展开，依次讲清分块、寄存器分块、共享内存分块、双缓冲四种手段，并给出可对照的性能分析框架。

一、分块技术：把大矩阵切成能装进片上存储的小块

分块（Tiling）的核心思想是：把 C 划分为若干 BM×BN 的子块，每个线程块负责一个子块；沿 K 方向再切成 BK 大小的“滑动窗口”，每次只把 A 的 BM×BK 片段和 B 的 BK×BN 片段搬入片上存储，在片上完成这一段的乘加，再滑动到下一段。这样全局内存的读取次数从每个元素 O(N) 或 O(M) 次，降为每个元素约 O(K/BK) 次，复用率大幅提升。

分块尺寸不是越大越好。BM、BN、BK 共同决定片上存储占用与线程块内线程数。以常见的 128×128 输出块、BK=16 为例：A 片段 128×16、B 片段 16×128，若都用共享内存存放，各需 2048 个 float（8 KB），合计 16 KB；线程块可配置为 256 个线程，每个线程负责 8×8=64 个输出元素。这个“每线程多元素”的设计正是寄存器分块的入口。

二、寄存器分块：让每个线程算一小片 C

寄存器分块（Register Tiling）指每个线程不只算 1 个 C 元素，而是算 TM×TN 个小块，例如 8×8。线程把 A 的 8 个值和 B 的 8 个值分别读入寄存器，做 8×8 次乘加，只需 16 次读取就产生 64 次乘加，计算访存比从 1:2 提升到 4:1。实现上通常用 float acc[8][8] 累加器数组，编译器会将其展开到寄存器；外层循环沿 K 逐点推进，内层用两个长度为 8 的寄存器数组 aReg[8]、bReg[8] 完成外积更新：

// 每个线程负责 C 的 8×8 子块（示例参数）
float acc[8][8] = {0.f};
for (int k = 0; k < BK; ++k) {
    float aReg[8], bReg[8];
    #pragma unroll
    for (int i = 0; i < 8; ++i) aReg[i] = sA[ty*8 + i][k];
    #pragma unroll
    for (int j = 0; j < 8; ++j) bReg[j] = sB[k][tx*8 + j];
    #pragma unroll
    for (int i = 0; i < 8; ++i)
        #pragma unroll
        for (int j = 0; j < 8; ++j)
            acc[i][j] += aReg[i] * bReg[j];
}

注意寄存器压力：8×8 累加器加两个 8 元素数组已占约 80 个寄存器，若再增大 TM、TN，可能出现寄存器溢出（spill）到本地内存，反而变慢。工程上常以 8×8、8×4、4×4 作为候选，通过编译选项 -Xptxas -v 查看寄存器用量与 spill 情况来取舍。

三、共享内存分块：把重复读取挡在片上

共享内存分块解决的是“块内线程之间的复用”。仍以 128×128 输出块、BK=16 为例，线程块先把 A 的 128×16 片段和 B 的 16×128 片段协作加载到 __shared__ 数组，__syncthreads() 后所有线程从共享内存取值参与寄存器分块计算，计算完再同步、加载下一个 K 片段。这样每个 K 片段对全局内存只读一次，块内 256 个线程共享这份数据。

共享内存分块有两个必须处理的细节。其一是bank 冲突：若 sA 按 sA[128][16] 声明，同一 warp 的 32 个线程按列访问时可能落到同一 bank。常用做法是给行宽加填充，如 sA[128][16+1]，把冲突打散。其二是边界处理：当 M、N、K 不是 BM、BN、BK 的整数倍时，加载阶段要做越界判断，越界位置填 0，保证乘加结果不受影响。

四、双缓冲技术：让搬数与计算重叠

即便用了共享内存，流程仍是“加载—同步—计算—同步—加载”，加载期间计算单元空转。双缓冲（Double Buffering / 软件流水）把共享内存分成两份缓冲：计算第 t 个 K 片段时，提前把第 t+1 个片段加载到另一份缓冲，从而让全局内存搬运与乘加计算在时间上重叠。实现上可用两个共享数组轮换，或在支持异步拷贝的架构上用 cp.async 把加载指令与计算解耦，配合 cp.async.commit_group 与 cp.async.wait_group 控制流水深度。

双缓冲的收益取决于“搬数时间”与“计算时间”的比例。当 BK 较小、计算量不足时，流水掩盖有限；当寄存器分块使每线程计算量足够大时，双缓冲往往能带来明显提升。它也会增加共享内存占用（两份缓冲），需要在占用与流水收益之间权衡。

五、性能对比分析：看什么指标、怎么判断

对比四种实现时，不能只看“跑得快不快”，而要建立可解释的指标链。第一是有效带宽与访存效率：用 profiler 观察全局内存读写字节数与理论最小值的比值，比值越接近 1，说明分块复用越充分。第二是计算访存比：每线程乘加次数除以加载次数，寄存器分块把它从 1:2 提到 4:1 甚至更高。第三是占用率与寄存器用量：寄存器分块提高复用，但寄存器过多会压低 occupancy，需要看 -Xptxas -v 的寄存器数与 spill 字节。第四是流水掩盖效果：对比开/关双缓冲的耗时差，差值越大说明搬数原本越拖后腿。

一个合理的调优顺序是：先做共享内存分块消除块内重复读取，再加寄存器分块提高单线程计算密度，最后用双缓冲掩盖剩余搬数开销；每一步都用同一组矩阵规模（如 1024×1024×1024、2048×2048×2048）和同一精度（FP32）做对照，避免跨规模比较得出错误结论。若某一步提升不明显，往往说明瓶颈已转移到别处，例如 bank 冲突、边界分支或 occupancy 不足，应回到对应细节排查，而不是继续堆叠优化手段。