GEMM优化原理

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

问题定义

NVIDIA 的存储资源分为 global memory、shared memory 和 register 三种，三者访问速度依次变快，容量依次减小。如果反复从 global memory 中读取，数据复用太差，就会导致大量时间浪费在等待数据搬运，计算单元长期空转，也就是所谓的存储墙问题。因此需要对矩阵分块将子矩阵的数据搬运到 shared memory 中，再对子矩阵分块将小矩阵的数据搬运到 register 中，同时利用 double buffer（ping-pong buffer）让访存和计算同步进行，通过这些方式将访存延迟隐藏在计算延迟中，这就是大致的优化思路。关于分块的具体的分析后续展开，此处先定义相关变量。

GEMM（General Matrix Multiplication）一般写作：

$$\begin{aligned} \mathbf{C} &= \mathbf{A} \times \mathbf{B}\\ \mathbf{C}_{i, j} &= \sum_{t = 0}^{k-1} \mathbf{A}_{i, t}\mathbf{B}_{t, j} \end{aligned}$$

其中 $A、B、C$ 分别为 $m \times k、k \times n、m \times n$ 的矩阵。将 $A、B、C$ 三个矩阵分别划分为多个 $b_m \times b_k、b_k \times b_n、bm \times b_n$ 的子矩阵，三个矩阵分别形成 $M \times K、K \times N、M \times N$ 的子矩阵网格。同理再将 $\mathbf{C}$ 的子矩阵划分为多个 $r_m \times r_n$ 的小矩阵，形成 $X \times Y$ 的小矩阵网格。

有如下关系成立（此处讨论可以对齐整除的情况）：

$$\begin{aligned} M &= \frac{m}{b_m}\\ N &= \frac{n}{b_n}\\ K &= \frac{k}{b_k}\\ X &= \frac{b_m}{r_m}\\ Y &= \frac{b_n}{r_n}\\ \end{aligned}$$

此外，优化的边界限定为一次 GEMM 计算，主要的优化对象是 MAC 计算所依赖的读取环节，写入延迟在这里并不是瓶颈，因此不作讨论。

从 global memory 到 shared memory 再到 register

基于 global memory 计算

naive CUDA 的常规写法是，启动 $m \times n$ 个线程，一个线程负责 $\mathbf{C}$ 的一个元素。每个线程读取 $\mathbf{A}$ 的一行和 $\mathbf{B}$ 的一列。这样计算一个元素就需要从 global memory 中读取 $2k$ 个元素，算完整个矩阵就需要对 global memory 进行 $2mnk$ 次读取。这样的问题是数据复用太差，同一行 $\mathbf{A}$ 和同一列 $\mathbf{B}$ 会被许多线程反复从 global memory 读取，而 global memory 的访存速度相对于 MAC 计算速度是极慢的，因此就出现了访存瓶颈。

计算整个矩阵所需的读取操作：

$$2mnk 次\text{global memory}$$

基于 shared memory 计算

将矩阵分块为子矩阵后，启动 $M \times N$ 个 thread block，一个 thread block 内的线程数为 $b_m \times b_n$。一个 thread block 负责计算 $\mathbf{C}$ 的一个子矩阵（$b_m \times b_n$ 个元素），本质上还是一个线程负责 $\mathbf{C}$ 的一个元素，区别在于这一个 block 内的 thread 可以复用 shared memory 中的数据。得到 $\mathbf{C}$ 中的每个子矩阵需要将 $K$ 个 $\mathbf{A}、\mathbf{B}$ 子矩阵相乘的结果累加，即经过 $K$ 次迭代得到，这个迭代在后文统称为大迭代。

对于每一个 $\mathbf{A}、\mathbf{B}$ 子矩阵相乘，由于子矩阵读取到 shared memory 中可以复用，而不是像 global memory 中那样每次都需要重复读取，因此计算一个子矩阵只需要从 global memory 中将 $\mathbf{A}、\mathbf{B}$ 子矩阵中全部元素读取到 shared memory，即需要从 global memory 中读取 $b_m b_k + b_k b_n$ 次，然后再从 shared memory 中读取 $2 b_m b_n b_k$ 次。算完整个矩阵（$K$ 次迭代）总共需要从 global memory 中读取 $MNK (b_m b_k + b_k b_n) = mnk (\frac{1}{b_m} + \frac{1}{b_n})$ 次，再从 shared memory 中读取 $MNK \times 2 b_m b_n b_k=2mnk$ 次。

计算整个矩阵所需的读取操作：

$$mnk (\frac{1}{b_m} + \frac{1}{b_n})次\text{global memory} + 2mnk 次\text{shared memory}$$

基于 register 计算

和上面同理，做完 block 级的 tile 分块之后，shared memory 也存在大量重复读取，于是可以进一步将需要复用的数据搬到 register 上。将子矩阵分块为小矩阵，每个线程负责一个小矩阵（$r_m \times r_n$ 个元素）而不是一个元素的计算，一个 thread block 仍然负责一个子矩阵（$b_m \times b_n$ 个元素），因此一个 block 内线程数为 $\frac{b_m}{r_m} \times \frac{b_n}{r_n}=X \times Y$，总共启动 $M \times N$ 个 thread block。得到一个小矩阵 $r_m \times r_n$，需要完成一个 $r_m \times b_k$ 矩阵和一个 $b_k \times r_n$ 矩阵相乘的计算；具体来说会依次完成 $b_k$ 次 $r_m \times 1$ 列矩阵和 $1 \times r_n$ 行矩阵的相乘，每次会得到一个小矩阵 $r_m \times r_n$ 的中间结果，$b_k$ 次中间结果累加迭代即得到要算的小矩阵，这个迭代在后文统称为小迭代。另外，由于每个线程负责的元素变为了 $r_m \times r_n$ 个，因此会出现一个取舍，$r_m、r_n$ 取大，单线程数据复用更好，但寄存器占用上升，可能降低 occupancy。

对于每一个小矩阵的计算，和上面同理，只需要从 shared memory 中读取 $r_m b_k + b_k r_n$ 次，然后再从 register 中读取 $2 r_m r_n b_k$ 次。算完整个子矩阵的一次迭代总共需要从 shared memory 中读取$X Y \times b_k (r_m+r_n) = b_m b_n b_k \left(\frac{1}{r_n}+\frac{1}{r_m}\right)$，再从register中读取$XY\times 2 r_m r_n b_k = 2 b_m b_n b_k$次。

计算整个矩阵所需的读取操作：

$$mnk (\frac{1}{b_m} + \frac{1}{b_n})次\text{global memory} + mnk (\frac{1}{r_m} + \frac{1}{r_n})次\text{shared memory} + 2mnk 次\text{register}$$

寄存器重排

对一个 thread 计算小矩阵的某一次迭代而言，拥有 $r_m$ 个 $\mathbf{A}$ 矩阵的值、$r_n$ 个 $\mathbf{B}$ 矩阵的值以及 $r_m \times r_n$ 个 $\mathbf{C}$ 矩阵的值，完成本次迭代需要 $r_m \times r_n$ 条 FFMA（Floating-point Fused Multiply-Add，浮点融合乘加指令）指令。在每一条 FFMA 执行过程中，如果直接进行 naive 地执行，会存在大量的 register bank conflict，因此需要对寄存器进行重排。

bank 机制是为了并行访问而把一块存储切分成多个小存储的设计，存储的物理硬件是一整块单口 SRAM，bank 设计的目的是用可接受的成本提供高并行的访问带宽。多个线程可以并行访问不同的 bank，但访问同一个 bank 就可能发生冲突而使得访存串行化。对 shared memory bank 而言是把地址映射到 bank 上，对 register bank 而言是把源寄存器映射到 bank 上。存储和 bank 的映射关系在不同的硬件架构上可能也不同，以 32bit-word、32 个 bank 的 shared memory 为例，一个典型的映射是 bank_id = (address / 4) % 32。在 GEMM 里，这个问题常出现在按列读 tile、做转置、或者 stride 刚好是 32 的倍数时。常见处理方式是 padding：__shared__ float tile[32][33];，多出来的一列不存有效数据，而是为了让下一行的起始地址错开 bank 映射。

回到这里的计算，如果一条指令的源寄存器有两个以上来自同一 bank，就会产生 conflict 导致指令重发射而浪费一个 cycle，这对于计算密集型的算子尤其重要。因此需要对寄存器进行重排。CUDA C++ 通常不能直接控制变量最终分配到哪个物理寄存器。只有写 PTX/SASS 或做很底层的 kernel 排布时，才会专门考虑寄存器编号、累加器布局和 FFMA 顺序。普通 GEMM 优化先看 global/shared/register 分块是否合理；继续逼近峰值算力时，bank conflict、指令顺序和汇编排布才会变成主要问题。

prefetch 和双缓冲

按照上述的优化，GEMM 计算会不断进行下面的流程

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

如果三者串行执行，计算单元显然会频繁等待数据，因此需要做数据预取，并行进行下一批数据的搬运和当前批数据的计算。具体方式是使用双缓冲（double buffer / ping-pong buffer），双缓冲的思想非常朴素，即准备两份存储轮流负责预取搬运和当前计算，这样就可以提前发起下一批数据搬运，用计算指令尽量掩盖数据搬运的延迟。

当然，prefetch 也存在代价，一方面是存储资源的翻倍，另一方面是资源占用变高之后 active block 数可能下降，因此 prefetch 需要和分块参数一起调参。好的参数组合能提高数据复用并掩盖 latency；差的参数组合可能因为 register 和 shared memory 占用过高，反而降低并发度。理想情况下，如果访存延迟可以完全隐藏在计算延迟中，在线程并行的情况下，当一个线程完成小矩阵的计算时，整个矩阵的计算也就完成了。

一个具体的例子

设矩阵大小为 $2048 \times 2048$，分块参数为 $b_m = 128,\quad b_n = 128,\quad b_k = 8,\quad r_m = 8,\quad r_n = 8$，其它衍生参数如下。

• block 数为 $(2048/128) \times (2048/128)=256$
• 每个 block 负责 $128 \times 128=16384$ 个 $\mathbf{C}$ 元素
• 每个 block 线程数为 $(128/8)\times(128/8)=256$
• 每个线程负责 $8\times8=64$ 个输出元素
• 每个 block 在 $K$ 方向有 $2048/8=256$ 次大迭代
• 每个 thread 在 $b_k$ 方向有 $8 / 1 = 8$ 次小迭代
• 每次小迭代会产生 $8\times8=64$ 条 FFMA 指令
• 对于子矩阵，每次大迭代需要将 $\mathbf{A}$ 的 $128 \times 8$ 个元素和 $\mathbf{B}$ 的 $8 \times 128$ 个元素从 global memory 搬运到 shared memory 中。

kernel实现此处不作详细讨论，不过概括来说，核心就是此处介绍的原理，用双层 buffer 把 global/shared/register 三段数据流串成流水线，沿着 global -> shared -> register -> compute 这条路径减少重复搬运、提高局部性，并用更多片上资源去换取更少的等待时间。

参考资料

深入浅出 GPU 优化系列：GEMM 优化（一）