GEMM优化原理
整理 GPU GEMM 优化的一般分块思路：global memory 到 shared memory、shared memory 到 register、register 分块和 prefetch
针对 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 计算会不断进行下面的流程
- 从 global memory 读取下一组 \(\mathbf{A}/\mathbf{B}\) tile 到 shared memory。
- 从 shared memory 读取当前小片段到 register。
- 用 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 这条路径减少重复搬运、提高局部性，并用更多片上资源去换取更少的等待时间。