来源：机器翻译学堂作者：张译单位：东北大学自然语言处理实验室
本文将从一个 cuda 初学者的角度来阐述如何优化一个形状较大的正方形乘正方形的 FP32 矩阵乘。
矩阵乘作为目前神经网络计算中占比最大的一个部分，其快慢会显著影响神经网络的训练与推断所消耗的时间。虽然现在市面上已经有非常多的矩阵乘的高效实现——如基于 cpu 的 mkl、基于 arm 设备的 ncnn 与 emll、基于 cuda 的 cublas ——掌握了矩阵乘优化的思路不仅能帮助你更好的理解编写高性能代码的一些基本原则，而且许多神经网络加速领域进阶的技巧如算子融合都是与矩阵乘交互从而达到更高的性能。
由于矩阵乘的性能优化与两个矩阵的形状有着非常密切的联系，因此，为了降低本文的撰写难度（以及辅助读者更好的理解矩阵乘优化），本文将从一个 cuda 初学者的角度来阐述如何优化一个形状较大的正方形乘正方形的 FP32 矩阵乘。同时本文按如下顺序讲解：
- Goals：本文的目标是什么？
- Performance：我们达到了多少性能？
- 朴素 GEMM 与前置知识：简单介绍一下我们的任务是什么，我们需要提前了解什么。
- Tiling：如何做矩阵分块？即如何将一个巨大的矩阵乘任务合理的分配到 GPU 的不同线程上。
- Thread 级优化：在 Thread 这个维度，我们能做什么优化？
- Warp 级优化：在 Warp 这个维度，我们能做什么优化？
- Block 级优化：在 Block 这个维度，我们能做什么优化？
- Epilogue：尾声。
Goals
首先明确一下本文的目标是：
• 实现一个比 cublas 更快的形状较大的正方形乘正方形的 FP32 矩阵乘。
• 从理论角度与硬件规格能够简单的推导矩阵分块与排布的方法。
• 可以大致清楚各个优化技术效果的阶段性的 benchmark。
• 如何使用 Nsight Compute 等性能分析工具分析潜在的性能瓶颈。
本文不含：
• 使用 Tensor Core 加速矩阵乘。（这也是为什么这篇文章叫传统 CUDA GEMM）
• 使用安培架构新提出的 async memcpy。
• CUDA 语法知识。
• 汇编。（主要是现在并没有官方支持汇编的操作，目前的汇编器几乎都是逆向的产物，不是很稳定。同时汇编带来的好处如消除寄存器的 bank conflict nvcc 也在不断的做相应的改进，因此就不介绍了）
开源地址：https://github.com/AyakaGEMM/Hands-on-GEMM
同时本文在相当程度上参考了李少侠的 GEMM 优化指南（写得非常！非常！非常！不错），本文的优势在于补了阶段性代码和在某些少侠一笔带过的地方做了一些扩展。
Performance
为了让大家更有动力阅读下去，这里先放出来性能效果！
测试平台
• 系统：Arch Linux
• 驱动：520.56.06
•CUDA：11.8
•GPU：Nvidia RTX 2080
测试结果
我们也可以注意到，在较大形状上手写的矩阵乘有着与 cublas 相近，甚至更优的性能。
从这张图我们可以看出，手写的矩阵乘能够达到硬件 95% 的峰值性能，效果还是很不错的。
朴素 GEMM 与前置知识
首先写一个朴素矩阵乘。
# 数组 A：M 行 K 列的行主序矩阵# 数组 B：K 行 N 列的行主序矩阵# 数组 C：M 行 N 列的行主序矩阵# alpha：一个标量# beta：一个标量# 计算方法：# c=alpha*A*B+beta*C; __global__ void matrixMul(const float *A, const float *B, float *C, int M, int N, int K, float alpha, float beta){ int tx = blockIdx.x * blockDim.x + threadIdx.x; int ty = blockIdx.y * blockDim.y + threadIdx.y; int baseX = blockIdx.x * blockDim.x; int baseY = blockIdx.y * blockDim.y; float c = 0; if (tx < M && ty < N) { for (int i = 0; i < K; i++) { c += A[tx * K + i] * B[i * N + ty]; } C[tx * N + ty] = beta * C[tx * N + ty] + alpha * c; // we multiply alpha here to reduce the alpha cal num. }}
这里 GPU 矩阵乘与 CPU 矩阵乘最大的区别就在于 GPU 可以为目标矩阵 C 中每一个元素分配一个 thread 进行计算。这也是可以切实的感知到 GPU 多线程编程的一点。但这个矩阵乘的朴素实现会非常慢，而分析性能瓶颈中最常见的两个指标即是带宽和延迟。这里借用 Nvidia 在 GTC 2018 上的分享来做说明。
这里以一个自动扶梯作为例子来讲解。
- 带宽：指这个自动扶梯每秒能够运送多少个人，以这张图为例子，这个扶梯每秒能运 0.5 个人，这就是这个自动扶梯的带宽。
- 延迟：指一个人踏上这个扶梯直到他被运到顶所需要的时间。同样以这张图为例，这个扶梯需要 40s 的延迟。
那么回到指令上来，每一个指令都有对应的延迟和带宽，而以朴素矩阵乘为例，每一个乘法运算需要读两次内存和一次 FFMA，假如没有其他额外的优化（如循环展开与指令重排），相当于是两个级联的自动扶梯，一个负责运送数据，一个负责做数学运算。假设数据运送扶梯的带宽与延迟与图中一致而不考虑 FFMA 的带宽与延迟，那么一次 FFMA 需要等待 40s（扶梯延迟）+ (1/0.5)s（第一个数据到达后第二个数据到达的时间）才能拿到所需的数据，这与扶梯的带宽 0.5s / 人的峰值性能相去甚远。那么此时这个 kernel 就完全被延迟卡住了，而无法发挥出应有的性能。
而对于带宽部分，这里我们引用李少侠的带宽分析：
对于 FP32 数据，如上图所示，一个 warp 一次做 32 次 FFMA，对应 64 OP，需读取 A 矩阵 1 个元素和 B 矩阵 32 个元素，共 132byte。通过寄存器累加，且忽略 C 矩阵写回开销，那么计算访存比为 64OP / 132byte = 0.48。虽然 dram 最小访问单位为一个 memory transaction，但考虑到 L1 cache 的存在也不会影响实际的计算访存比。通过 repo 中提供的 l2cache_bandwidth.cu 可测得 Titan V L2 cache 带宽约 1.9TB/s，那么最乐观的结果即使 L2 cache 100% 命中，此方案的理论上限也只有 1.9T * 0.48 = 0.912 tflops，远低于 14.9 tflops 的硬件算力。
由此我们可以看出，朴素的矩阵乘实现方法无论从延迟和带宽上都无法满足需要。
这里一个 warp（即 32 个线程）是指 GPU 调度线程的粒度，可以简单的理解为同一个 warp 内的线程总是同时运行、同时休眠的。当然这种说法并不完全准确，毕竟还有 warp divergent 问题，感兴趣的同学可以自行了解。但总之，思考 GPU 执行时总是从 warp 的角度思考是非常合理的。那么对于一个 warp 而言，我们可以根据李少侠的分析看出，就算我们假设延迟能够被完全覆盖，这种分配方案也并不能达到硬件的峰值性能。
这里我用自己的话总结一遍就是：在每个线程执行指令设计时，需要尽可能的覆盖掉每个指令的延迟；在性能分析时，则从带宽角度分析矩阵分块是否合理。
而在于延迟部分还有一顿免费的午餐。在实际应用中，编译器会自动的做一些优化，如循环展开与指令重排等。例如展开循环后可以将多个读取 A 矩阵的元素和读取 B 矩阵的元素排在一起，使得取数据的自动扶梯能够一次多上几个人，从而去覆盖掉扶梯的延迟。而且 GPU 与 CPU 还有一个非常不同的地方在于，GPU 的线程切换代价非常低，因此可以在等待延迟的时候转而去运行其他线程从而达到延迟覆盖的目的。还是以扶梯为例子，GPU 上有很多个扶梯，在等待第一个人到达扶梯末尾时，GPU 可以转到第二个扶梯送几个人上扶梯。理想情况下，在 GPU 送完第 N 个扶梯的人时，第一个扶梯的人刚好达到扶梯顶部，那么这个运送的延迟就被覆盖掉了。
Tiling
矩阵乘分块是为了将一个大问题化解为小问题求解，这里 CPU 与 GPU 分块的需求也不尽相同。CPU 是希望保持计算的局部性，可以充分利用 L1、L2 高速缓存来避免缓慢的内存访问。而 GPU 在此基础之上还需要将一个大问题合理拆分到不同的 thread 上，使得其能够充分利用 GPU 上的硬件资源。下面我将从局部性和合理拆分两个方面讲解如何做矩阵分块。
局部性原理
局部性原理，在我的理解中便是为了能够尽可能的使用高速缓存器内的数据进行运算所提出的一个程序设计理念。由于高速缓存器往往十分昂贵（或者需要很高的功耗），因此空间都不大。由此我们需要尽可能的将一些重复访存聚合起来，放到高速缓存器里面来加速数据访问，或者在进行访存的时候尽可能连续访存来使用 cache 加速访存。我们先还是让每一个 thread 负责一个目标矩阵元素的计算。虽然这种分配方式十分朴素、十分直接，同时各个 thread 之间也没有数据依赖关系，不需要做额外的同步之类的操作，但这种分配方式却是十分访存不友好的，因为每一个 thread 都是直接与内存做交互，而 GPU 的全局内存访问带宽完全不足以匹配上它的计算速度。
同时我们注意到处于同一行的 thread 总是会同样的读取 A 矩阵的同一行数据；同一列的 thread 总是会读取 B 矩阵的同一列数据。那么一个非常自然的想法则是对于每一个 Block，我们将数据移动到这个 Block 共享的一块高速存储区 shared memory 上，从而减少与全局内存交互的次数。同时我们考虑到 shared memory 的容量有限，因此可以一次只取一部分 k，然后通过循环迭代的方式完成这个 Block 所负责的矩阵乘区域。
值得一提的事，shared memory 虽然叫做 memory，但他却有着非常高的访存速度与极低的延迟。实际上，shared memory 可以被看作是一块可以显式控制的 L1 cache。从图灵架构开始，在硬件上 shared memory 与 GPU 上的 L1 cache 共享同一块区域，同时 shared memory 与 Load/Store 单元交互也是直连的（没有中间商赚差价）。
在将一个大型矩阵乘划分为一个个由 Block 负责的小型矩阵乘之后，我们接下来还需要把一个 Block 负责的矩阵乘分配给 Block 内部的 warp；分配到 warp 之后我们还需要把一个 warp 负责的矩阵乘分配给 warp 内部的 thread。经过这么一步一步的划分，我们便可以把一个巨大的矩阵乘任务高效的分配到各级速度不一的存储器上，最终尽可能打满硬件峰值性能，实现高效矩阵乘。有了前面划分 Block 的经验，我们也就可以依葫芦画瓢，实现大矩阵的拆分（Tiling），在此就不过多赘述了，最终整体流程图如下。
当然这只是一个较为粗糙的流程图，例如每一个 thread 负责的分块也并不是图中所示的连续一块矩阵乘，我们也将在后续一步一步完善细节，但这种分解的框架却是一种非常经典的思路。
如何确定分块大小？
在拥有分块的基本理念之后，我们还有一个问题没有解决。那便是每一个 Block 该负责多大的矩阵乘？每一个 thread 又应该负责多大的矩阵乘？为了让文字变得清晰起来，我们定义每一个 Block 负责的矩阵大小为
，每次迭代 
的 k 维数据，每一个 warp 负责的矩阵大小为
，每一个 thread 负责的大小为
。其中这些符号都在上图出现过，可以自行对照一下。
这里我们同样引用李少侠的计算访存分析：
假设我们不考虑 shared memory 的访存代价（因为可以做到覆盖掉shared memory 的访存延迟，而且其带宽能够满足 FFMA 单元的计算速度），只考虑全局内存的访问，可以看到选择在 K 上缩水（即不把整个 K 维度都放到 shared memory 里）还是比较合理的，因为 的大小其实并不影响计算访存比。而对计算访存比有决定性影响的是每一个 Block 计算的大小。如果取 为 64，带入 RTX 2080 的数据，可以得到 10.1 Tflops / 16 = 631.25 GB/s。即内存访问带宽达到 631.25 GB/s 就能避免内存访问瓶颈了。同样，我们取 L2 命中率为 20%（还是比较好达到的），加权内存访问带宽为：，即可避免内存访问瓶颈。
那是否我们只要取分块大小为 64x64 就万事大吉了呢？也不尽然。我们前面只分析了带宽，而在延迟无法被覆盖的情况下，整个 kernel 性能也不会太好。而更大的分块意味着每一个 thread 会计算更多的数据，可以使用一些手段实现更优的延迟覆盖。这一点会在后面讨论如何具体实现，大致思想也是局部性的原理，只不过这次是将数据从 shared memory 保存到寄存器，从而实现使用更高速的缓存计算的目的。
那是否我们取分块越大越好呢？那也不一定。更大的分块使用了更多的寄存器，从而使得同一个 SM 能够同时承载的线程数变少，这里 Nvidia 将之称为 Occupancy。如前文所述，当一个 warp 被卡住时，GPU 可以切换到另一个 warp 执行指令，Occupancy 越低，可供 GPU 切换的线程就越少。
而 Occupancy 也是和硬件强相关的。一个 GPU 由多个 SM 构成，每一个 SM 拥有有限的寄存器数量、 shared memory 和最大可调度线程数量。而 Occupancy 是指每个 SM 能够同时调度的线程数量除以一个 SM 的最大可调度线程数量。关于 Occupancy 的计算我们可以通过在编译时添加 --ptxas-options=-v 参数，使编译器在编译时输出每个 kernel 所花费的寄存器数量和 shared memory，然后通过随 cuda 提供的一个 excel 表格进行计算。（尽管这个 Excel 已经 deprecated 了，但他用起来确实挺方便的。）
例如我们每个 thread 需要 128 个寄存器，2048 bytes 的 shared memory，那么由于 RTX 2080 每个 SM 只有 65536 个寄存器，因此每个 SM 最多只能同时跑 512 个 threads。又因为每个 SM 最多能够承载 1024 个 threads，所以此时 Occupancy 为。
值得一提的是，虽然较高的 Occupancy 使得在一个线程卡住时，SM 能够马上切换到别的线程，通过将其他线程需要执行的指令填充到流水线中从而达到覆盖延迟的目的，但这并不代表高性能。例如，如果每一个线程本身就能够通过更多的寄存器占用从而达到延迟覆盖的目的，自然也就不需要 SM 来做这件事了，反倒是如果无脑的去提高 Occupancy 使得一些 thread 内的延迟甚至都无法被 SM 通过切换执行线程的方式覆盖，那属实是得不偿失了。
因此，我们能够做的就是在有一定理论分析的情况下确定好一些矩阵的分块大小的方案，然后要不就是经验性的去选择最终用哪个分块，要不就是跑一个 profile 来直接得到最快的分块。这里由于已经有非常多的先例证明了 128x128x8 是一个较优的选择，因此本文则遵从这个分块方案。那么，目前我们能够确定的分块如下表。
当然有些同学可能会问，既然最终还是需要用跑 profile 的方式来确定最优分块，那理论分析还有什么意义呢？答案就是如果提前通过理论分析，那么就能够在一定程度上缩小需要跑 profile 的分块数量。用算法上的语言来讲就是如果我们将需要搜索的所有分块作为搜索空间，那么理论分析便是搜索算法中的 A* 算法，你掌握了越多的理论分析知识那么这个搜索过程就会越高效。同时对 CUDA 底层越了解，在同一个分块策略下，你更容易写出能达到理论性能的 kernel。
Thread 级优化
对于一个 thread 能做的优化其实并不多，因为 GPU 是以一个 warp（即 32 个 thread）进行调度的，所以许多基于单线程的优化，如访存优化，其实并不能直接套到 GPU 上。而为数不多值得一提的优化手段便是单个线程在计算时应该采用向量内积还是向量外积以及 double buffer。但实质上向量外积严格意义上也不能算作是一个优化，因为这一步编译器就能在编译阶段帮忙做了。之所以提一句是还是为了给 double buffer 做铺垫，即我们应该怎么预取数据。
首先我们取了 128x128 的分块策略，一个 Block 有 256 个线程，那么每个线程需要负责一个 8x8 的矩阵乘运算。而一个线程完成一个小型矩阵乘有两种实现方法。
向量内积
向量内积的实现方法如图所示，即将 A 矩阵拆分为多个行向量、B 矩阵拆分为多个列向量，这些向量通过向量内积的方法求得最终答案。
用代码描述如下：
M=N=K=8;float a[M*N];float b[N*K];float c[M*N];for i in range(M): for j in range(N): for k in range(K): c[i*N+j]+=a[i*K+k]*b[k*N+j];
向量外积
向量外积的实现方法如图所示，即将 A 矩阵拆分为多个列向量、B 矩阵拆分为多个行向量，这些向量通过向量外积的方法求得最终答案。
用代码描述如下：
M=N=K=8;float a[M*N];float b[N*K];float c[M*N];for k in range(K): for i in range(M): for j in range(N): c[i*N+j]+=a[i*K+k]*b[k*N+j];
可以看到，向量内积和向量外积的区别在代码上仅仅体现在循环方式上。
