咖啡图片
正在将巧克力泡入咖啡
ntainer" style="display: none">
文章

GEMM优化

整理 GPU GEMM 优化的一般分块思路：global memory 到 shared memory、shared memory 到 register、register 分块和 prefetch

发表于 更新于
Preview Image
作者 Hatrix
26 分钟阅读
GEMM优化

针对NVIDIA GPU的矩阵乘法（GEMM）优化已有非常成熟的套路，此处进行梳理。主要介绍GEMM中的数据分块和如何在多级存储进行数据搬运。这也是HPC优化的核心思想，怎么样让数据放在更近的存储上来掩盖计算的延时，从而减少存储墙的影响。

主要参考：

GEMM 的问题形式

GEMM（General Matrix Multiplication）一般写作：

\[C = A \times B\]

其中：

  • $A \in \mathbb{R}^{m \times k}$
  • $B \in \mathbb{R}^{k \times n}$
  • $C \in \mathbb{R}^{m \times n}$

对 $C$ 中任意一个元素：

\[C_{i,j} = \sum_{t=0}^{k-1} A_{i,t}B_{t,j}\]

naive CUDA 写法通常让一个线程负责 $C$ 的一个元素。这样一共开启 $m \times n$ 个线程，每个线程读取 $A$ 的一行和 $B$ 的一列。只看 global memory 访问，读操作约为 $2mnk$ 次，写操作为 $mn$ 次。问题不在计算公式，而在数据复用太差：同一行 $A$ 和同一列 $B$ 会被许多线程反复从 global memory 读取。

理想 GEMM 时间估算

从 roofline 的角度看，理想 GEMM 时间可以粗略写成：

\[T_{\text{ideal}} = \max(T_{\text{io}}, T_{\text{compute}})\]

其中 IO 时间按一次读入 $A$、一次读入 $B$、一次写回 $C$ 估算：

\[T_{\text{io}} = \frac{mk + nk + mn}{B_{\text{mem}}}\]

这里的 $mk + nk + mn$ 先按元素个数记；如果要换成字节数，还要乘以单个元素的字节数。$B_{\text{mem}}$ 是硬件内存带宽。

计算时间要看硬件算力的标注口径。如果把一次乘加按一个 MAC 计算：

\[T_{\text{compute}} = \frac{mnk}{P_{\text{mac}}}\]

如果硬件峰值按 FLOP/s 标，GEMM 中一次乘加通常记作 2 FLOPs，此时写成：

\[T_{\text{compute}} = \frac{2mnk}{P_{\text{flops}}}\]

因此，$T_{\text{io}} > T_{\text{compute}}$ 时主要卡在访存，$T_{\text{compute}} > T_{\text{io}}$ 时主要卡在计算。分块、预取和双缓冲做的事，就是尽量降低有效 IO 开销，让计算单元少等数据。

GEMM 优化的主线就是逐层提高数据复用：

1

global memory -> shared memory -> register -> FFMA

越靠近计算单元，容量越小，但延迟更低、带宽更高。优化的重点是让数据在高层存储中被多次使用，而不是每次 FFMA 都去等 global memory。

从 global memory 到 shared memory

先把 $C$ 分成很多 block tile。设一个 thread block 负责计算 $C$ 的一个 $b_m \times b_n$ 小块；在 $K$ 维度上每次只处理 $b_k$ 个元素。那么每轮迭代需要加载：

  • $A$ 的一个 $b_m \times b_k$ tile
  • $B$ 的一个 $b_k \times b_n$ tile

这两个 tile 先从 global memory 搬到 shared memory，然后 block 内所有线程复用它们计算同一个 $C$ tile。

如果 $m,n,k$ 都能整除对应 tile size，则 tile 网格可以写成：

  • $M = m / b_m$
  • $N = n / b_n$
  • $K = k / b_k$

GPU 启动 $M \times N$ 个 thread block，每个 block 负责一个 $C$ tile。每个 block 在 $K$ 方向迭代 $K$ 次，每次读取 $b_m b_k + b_k b_n$ 个 float。总 global memory 读取量约为：

\[MNK(b_m b_k + b_k b_n) = mnk\left(\frac{1}{b_n} + \frac{1}{b_m}\right)\]

相比 naive 的 $2mnk$ 次读取，比例为：

\[\frac{1}{2}\left(\frac{1}{b_m} + \frac{1}{b_n}\right)\]

若取 $b_m=b_n=b$，global memory 读取量大约降为 naive 的 $1/b$。这就是第一层分块的收益：把 global memory 上的重复读取转成 shared memory 中的块内复用。

从 shared memory 到 register

只做 block 级 tile 还不够。进入一个 block 后，如果仍然让一个线程只计算 $C$ 的一个元素，那么 shared memory 也会被大量重复读取。下一步是在 thread 级别继续分块：让每个线程负责一个 $r_m \times r_n$ 的小矩阵，而不是一个标量。

一个 thread block 负责 $b_m \times b_n$ 个 $C$ 元素。若每个线程负责 $r_m \times r_n$ 个元素，则 block 内线程数为：

\[\frac{b_m}{r_m} \times \frac{b_n}{r_n}\]

对一次 $b_k$ 方向的 tile 迭代来说，每个线程每个内层 $k$ 步只需要从 shared memory 读取：

  • $r_m$ 个 $A$ 元素
  • $r_n$ 个 $B$ 元素

然后在寄存器中做 $r_m \times r_n$ 次 FFMA。整个 block 的 shared memory 读取量约为：

\[\frac{b_m}{r_m}\frac{b_n}{r_n} b_k (r_m+r_n) = b_m b_n b_k \left(\frac{1}{r_n}+\frac{1}{r_m}\right)\]

如果不做 thread 级分块，每个 $C$ 元素都要读 $b_k$ 个 $A$ 和 $b_k$ 个 $B$，shared memory 读取量约为 $2b_m b_n b_k$。所以 thread 级分块后的 shared memory 读取比例为：

\[\frac{1}{2}\left(\frac{1}{r_m}+\frac{1}{r_n}\right)\]

若 $r_m=r_n=r$，读取量大约降为原来的 $1/r$。代价是每个线程要占用更多寄存器，包括：

  • $r_m$ 个来自 $A$ 的片段
  • $r_n$ 个来自 $B$ 的片段
  • $r_m r_n$ 个累加器

这里开始出现一个取舍：$r_m,r_n$ 取大，单线程数据复用更好，但寄存器占用上升，可能降低 occupancy。

shared memory 和 register 的 bank conflict

bank conflict 的基本原因很简单：硬件里的 shared memory 或 register file 不是一个整体，而是拆成多个 bank。一次并行访问中，如果多个请求落到不同 bank，可以同时服务；如果很多请求挤到同一个 bank，就会被拆成多次访问，吞吐下降。

shared memory bank conflict 发生在：

1

shared memory -> register

NVIDIA GPU 的 shared memory 常见是 32 个 bank。对 32-bit float 来说，可以粗略理解为：

\[bank\_id = (address / 4) \bmod 32\]

一个 warp 中 32 个线程同时访问 shared memory。若线程访问 shared[threadIdx.x] 这类连续地址，通常能分散到不同 bank；若访问 shared[threadIdx.x * 32]，就很容易让所有线程打到同一个 bank。前者接近并行访问，后者会被串行化。

在 GEMM 里，这个问题常出现在按列读 tile、做转置、或者 stride 刚好是 32 的倍数时。常见处理方式是 padding：

1

__shared__ float tile[32][33];

多出来的一列不是为了存有效数据，而是为了让下一行的起始地址错开 bank 映射。

register bank conflict 发生在：

1

register -> FFMA

一条 FFMA 指令需要读取多个源操作数，例如：

1

acc = acc + a * b;

如果 acc、a、b 对应的物理寄存器落到同一个 register bank，而硬件同一周期不能从这个 bank 读出足够多的操作数，就会拖慢指令发射。和 shared memory 不同，普通 CUDA C++ 通常不能直接控制变量最终分配到哪个物理寄存器。只有写 PTX/SASS 或做很底层的 kernel 排布时，才会专门考虑寄存器编号、累加器布局和 FFMA 顺序。

这两个 conflict 对应的层级不同：

  • shared memory bank conflict：数据已经从 global memory 搬进 shared memory，但线程从 shared memory 取 tile 片段时不顺。
  • register bank conflict：数据已经进入寄存器，但 FFMA 读源操作数时不顺。

普通 GEMM 优化先看 global/shared/register 分块是否合理；继续逼近峰值算力时，bank conflict、指令顺序和汇编排布才会变成主要问题。

register 分块和 bank conflict

在寄存器层，每个线程真正执行的是一组 FFMA 指令。以一个线程负责 $r_m \times r_n$ 个输出为例，某个内层 $k$ 步需要做：

\[r_m r_n\]

条 FFMA。伪代码可以写成：

1
2
3
4
5

for (int i = 0; i < r_m; ++i) {
  for (int j = 0; j < r_n; ++j) {
    acc[i][j] += frag_a[i] * frag_b[j];
  }
}

这里的 frag_a、frag_b 和 acc 都在寄存器中。问题是 NVIDIA GPU 的 register file 也存在 bank。若一条 FFMA 指令的多个源操作数落到同一个 register bank，就可能产生 bank conflict，指令需要重发射，吞吐下降。

原文用 Maxwell 架构举例：register bank 数为 4，可以粗略按 register_id % 4 判断寄存器属于哪个 bank。若累加器按非常朴素的顺序排列，连续 FFMA 很容易反复读到同一个 bank。解决思路是重新排列寄存器编号和计算顺序，让同一条 FFMA 的源寄存器尽量分布在不同 bank 上。

这部分有两个要点：

  • register 分块不只是「让一个线程算多个元素」，还包括寄存器布局和指令顺序的安排。
  • 越接近算力上限，register bank conflict 这类低层细节越重要；普通 CUDA C++ 写法未必能完全控制它，所以后续文章才会讨论汇编层面的优化。

prefetch 和双缓冲

前面的分块减少了访存次数，但没有完全消除访存等待。一个 block 在计算当前 $C$ tile 时，需要不断做这样的流程：

  1. 从 global memory 读取下一组 $A/B$ tile 到 shared memory。
  2. 从 shared memory 读取当前小片段到 register。
  3. 用 register 中的数据执行 FFMA。

如果这三件事串行发生，计算单元会频繁等待数据。prefetch 的目标是把「搬下一批数据」和「算当前数据」重叠起来。

做法是多开一份 buffer，也就是常说的 double buffering：

  • shared memory 中准备两份 tile buffer，一份读、一份写。
  • register 中也准备两份 fragment buffer，一份参与计算，一份装下一轮数据。
  • 当前轮计算使用 read buffer，同时把下一轮需要的数据预取到 write buffer。
  • 下一轮开始时交换 read/write 角色。

以 global memory 到 shared memory 为例：当 block 正在用 shared memory 中的 $A_t/B_t$ 计算时，提前把 $A_{t+1}/B_{t+1}$ 从 global memory 搬进另一块 shared memory。等当前计算完成，下一轮可以直接使用已经搬好的数据。

shared memory 到 register 也类似：在用当前 register fragment 做 FFMA 前，提前发起下一组 shared memory load，用计算指令尽量掩盖这段 load latency。

prefetch 的代价也很明确：

  • shared memory 需要多一份 buffer。
  • register 也需要多一份 fragment。
  • 资源占用变高，active block 数可能下降。

所以 double buffering 不是无脑打开就一定更快。它需要和 $b_m,b_n,b_k,r_m,r_n$ 一起调参。好的参数组合能提高数据复用并掩盖 latency；差的参数组合可能因为寄存器和 shared memory 占用过高，反而降低并发度。

参数含义速记

参数含义影响
$b_m$一个 block 负责的 $C$ tile 高度影响 global memory 复用、shared memory 占用
$b_n$一个 block 负责的 $C$ tile 宽度影响 global memory 复用、shared memory 占用
$b_k$每次沿 $K$ 维度加载的 tile 深度影响每轮搬运量、同步频率和 prefetch 粒度
$r_m$一个线程负责的输出 tile 高度影响 shared memory 复用和寄存器占用
$r_n$一个线程负责的输出 tile 宽度影响 shared memory 复用和寄存器占用

常见教学参数会取类似：

\[b_m=128,\quad b_n=128,\quad b_k=8,\quad r_m=8,\quad r_n=8\]

若矩阵大小为 $2048 \times 2048$，则：

  • block 数为 $(2048/128) \times (2048/128)=256$
  • 每个 block 负责 $128 \times 128=16384$ 个 $C$ 元素
  • 每个 block 线程数为 $(128/8)\times(128/8)=256$
  • 每个线程负责 $8\times8=64$ 个输出元素
  • 每个 block 在 $K$ 方向有 $2048/8=256$ 次大迭代

这个例子很适合理解多级分块：block 级 tile 控制 global memory 到 shared memory 的复用，thread 级 tile 控制 shared memory 到 register 的复用，register 级布局控制 FFMA 的实际发射效率。

CUDA kernel 的实现落点

第二篇文章主要把前面的分块思路落到 CUDA C++ kernel 中。它不使用手写汇编，也不使用 Tensor Core，而是用模板参数、向量化搬运、寄存器分块和双缓冲，把普通 FP32 GEMM 尽量写到接近 cuBLAS 的方向。

常见实现会把 tile 参数写成编译期常量：

1
2
3
4
5

BLOCK_SIZE_M = 128;
BLOCK_SIZE_N = 128;
BLOCK_SIZE_K = 8;
THREAD_SIZE_Y = 8;
THREAD_SIZE_X = 8;

由此可以推出：

  • 一个 block 计算 $128 \times 128$ 的 $C$ tile。
  • 一个线程计算 $8 \times 8$ 的输出小块。
  • 一个 block 中线程数为 $(128/8)\times(128/8)=256$。
  • thread block 形状通常是 dim3(16, 16)，blockIdx.y 选择 $M$ 方向的 tile，blockIdx.x 选择 $N$ 方向的 tile。

对线程 (tx, ty) 来说，它负责的 $C$ 子块大致是：

\[row = 128 \cdot blockIdx.y + 8 \cdot ty + [0, 8)\] \[col = 128 \cdot blockIdx.x + 8 \cdot tx + [0, 8)\]

也就是说，前文的 $b_m,b_n,b_k,r_m,r_n$ 在代码里分别对应 BLOCK_SIZE_M、BLOCK_SIZE_N、BLOCK_SIZE_K、THREAD_SIZE_Y、THREAD_SIZE_X。

global memory 到 shared memory 的协作搬运

每一轮 $K$ tile 迭代中，block 需要搬运：

  • $A$ 的 $128 \times 8$ tile
  • $B$ 的 $8 \times 128$ tile

两者都是 1024 个 float。若一个 block 有 256 个线程，每个线程用 float4 一次搬 4 个 float，那么每个线程刚好负责 $A$ 的一组连续 4 个元素和 $B$ 的一组连续 4 个元素。

代码里会先把二维线程编号压成一维：

1

tid = ty * THREAD_X_PER_BLOCK + tx;

然后用几个派生参数决定每个线程搬哪一段：

  • A_TILE_THREAD_PER_ROW = BLOCK_SIZE_K / 4：搬运 $A$ tile 的一行需要几个线程。$BLOCK_SIZE_K=8$ 时，一行需要 2 个线程。
  • B_TILE_THREAD_PER_ROW = BLOCK_SIZE_N / 4：搬运 $B$ tile 的一行需要几个线程。$BLOCK_SIZE_N=128$ 时，一行需要 32 个线程。
  • A_TILE_ROW_START / B_TILE_ROW_START：当前线程负责搬运的起始行。
  • A_TILE_COL / B_TILE_COL：当前线程负责搬运的起始列，按 4 个 float 对齐。
  • A_TILE_ROW_STRIDE / B_TILE_ROW_STRIDE：当一个线程需要搬多行时，每次跨多少行。

这里用 float4 的意义不是改变算法，而是减少 load/store 指令数，并让 global memory 访问更容易合并成连续事务。对应的限制也很直接：地址最好 16 字节对齐，矩阵维度和 tile 偏移要能按 4 个 float 对齐。示例代码为了聚焦优化主线，通常假设 $M,N,K$ 能被对应 tile size 整除；真正写通用 kernel 时，还要补边界判断和残块处理。

还有一个容易忽略的细节：$A$ 读入 shared memory 时通常会转置存放。概念上是把 global memory 中的 $A[M][K]$ tile 存成 shared memory 中的 As[K][M]。这样后续每个线程从 shared memory 取 frag_a 时，可以沿着连续的 $M$ 方向读一组元素，和 float4 读取更匹配。$B$ tile 则保持 Bs[K][N] 的布局，方便沿 $N$ 方向连续读取。

shared memory 到 register 的计算小块

进入计算阶段后，每个线程维护三类寄存器数据：

  • accum[THREAD_SIZE_Y][THREAD_SIZE_X]：该线程负责的 $8 \times 8$ 个输出累加器。
  • frag_a[...][THREAD_SIZE_Y]：当前 $k$ 步需要的 8 个 $A$ 片段。
  • frag_b[...][THREAD_SIZE_X]：当前 $k$ 步需要的 8 个 $B$ 片段。

对 $BLOCK_SIZE_K=8$ 的一轮大 tile 来说，内部有 8 次小迭代。每次小迭代从 shared memory 取一组 frag_a 和 frag_b，然后执行：

1
2
3
4
5

for (int y = 0; y < THREAD_SIZE_Y; ++y) {
  for (int x = 0; x < THREAD_SIZE_X; ++x) {
    accum[y][x] += frag_a[y] * frag_b[x];
  }
}

这就是 thread 级 register 分块的代码形态：每个线程每个 $k$ 步只取 $8+8$ 个操作数，却做 $8\times8=64$ 次 FFMA。计算结束后，accum 再按 float4 分段写回 $C$。

两级 double buffer

第二篇文章里的 double buffer 不只发生在 shared memory，也发生在 register fragment 上。

shared memory 层通常写成：

1
2

__shared__ float As[2][BLOCK_SIZE_K][BLOCK_SIZE_M];
__shared__ float Bs[2][BLOCK_SIZE_K][BLOCK_SIZE_N];

write_stage_idx 指向下一轮要写入的 shared memory buffer，load_stage_idx = write_stage_idx ^ 1 指向当前正在参与计算的 buffer。主循环大致做四件事：

  1. 先把下一轮 $A/B$ tile 从 global memory 预取到临时寄存器 ldg_a_reg / ldg_b_reg。
  2. 用当前 shared memory buffer 中的数据继续做 FFMA。
  3. 当前大 tile 快算完时，把临时寄存器里的下一轮 tile 写入另一个 shared memory buffer。
  4. __syncthreads() 后翻转 write_stage_idx，下一轮直接读刚写好的 buffer。

这层 double buffer 的目的，是让 global memory load 和当前 tile 的计算尽量形成重叠。代码顺序上仍然是一串指令，但因为预取的数据和当前计算用的数据互不依赖，GPU 调度器有机会把访存等待藏到后续计算指令后面。

register 层也用两份 fragment：

1
2

float frag_a[2][THREAD_SIZE_Y];
float frag_b[2][THREAD_SIZE_X];

在小迭代 kk 中，代码先预取 kk+1 要用的 frag_a / frag_b，同时用 kk 已经在寄存器里的 fragment 做 FFMA。这里依赖关系被错开了：load 写的是 (kk + 1) % 2，compute 读的是 kk % 2。这能掩盖 shared memory 到 register 的一部分延迟。

#pragma unroll 和资源占用

这些循环的边界都由模板参数决定，所以常配合 #pragma unroll 展开。展开后可以减少循环控制开销，也给编译器更多指令调度空间。但展开、寄存器分块和 double buffer 会一起抬高资源占用：

  • accum[8][8] 已经需要 64 个累加寄存器。
  • frag_a[2][8] 和 frag_b[2][8] 又增加了双缓冲 fragment。
  • global memory 预取还需要 ldg_a_reg / ldg_b_reg。

所以参数不能只看访存复用。$r_m,r_n$ 变大后，单线程算得更多，但寄存器压力也更高；shared memory 双缓冲会提升流水线机会，但也会让每个 block 占更多片上资源。最后的性能由访存复用、指令级并行、occupancy 和 bank conflict 一起决定。

C++ 版本的优化边界

这类 kernel 已经覆盖了普通 CUDA C++ 能做的主要优化：

  • block 级 tiling 减少 global memory 读取。
  • thread 级 tiling 提高 shared memory 和 register 复用。
  • float4 做向量化搬运。
  • shared memory 中调整 $A$ tile 布局，避免后续读取变成低效的跨列访问。
  • shared memory 和 register fragment 双缓冲，用预取掩盖延迟。
  • 循环展开提高调度空间。

再往后继续逼近 cuBLAS，问题会变得更底层：shared memory bank conflict、register bank conflict、warp/lane 到输出子块的映射、FFMA 指令顺序，甚至 SASS 级寄存器编号都会影响吞吐。也就是说，第二篇的价值在于把「分块思想」变成一个完整 CUDA C++ kernel；它能解释大部分性能提升，但还不是手写汇编级 GEMM 的终点。

总结

GEMM 优化可以按存储层次拆开看：

  • global memory 分块：让 $A/B$ tile 进入 shared memory，在 block 内复用，减少 global memory 访问。
  • shared memory 分块：让每个线程负责一个小输出矩阵，把 shared memory 数据复用到多个 FFMA，同时避开 shared memory bank conflict。
  • register 分块：把 frag_a、frag_b 和 accumulator 放在寄存器里，同时注意 register bank conflict。
  • prefetch / double buffering：用读写分离把下一轮数据搬运和当前轮计算重叠起来。
  • CUDA C++ 实现：用模板参数固化 tile shape，用 float4 协作搬运，用双层 buffer 把 global/shared/register 三段数据流串成流水线。

这套思路的核心不是某个固定参数，而是沿着 global -> shared -> register -> compute 这条路径减少重复搬运、提高局部性，并用更多片上资源去换取更少的等待时间。

本文由作者按照 CC BY 4.0 进行授权
/body>