CUDA Case Study - SGEMM on Pascal

最近自己重新学 CUDA （以前上过课，长时间不用又忘记了），找些经典的 case study 自己照猫画虎弄一次加深一点认识。HPC 领域里一个绕不开的例子就是 xGEMM，即稠密矩阵-矩阵乘法。网上关于 CUDA 如何实现高性能 xGEMM 的介绍不多，而且很多都是 Fermi 时代的资料，面对 Kepler 有详细介绍的只有 Ref No.2 那个网页。我以 Ref No.1, No.2 两个网页的资料和代码为蓝本，一并参考了其他一些文章，自己动手实践了一下，代码在 这里。

2018-01-13 Update notes: 上个学期我的 HPC 课期末作业选了写 GPU SGEMM, 所以我重新写了一次这个代码，以及之前的代码似乎在尺寸不为 128 倍数的时候有 Bug, 在我的 GTX 1070 上面也跑得不是很快。

测试平台

我用我座机上的 EVGA GTX 1070 进行测试：1920 SP @ 2002 MHz, 8GB GDDR @ 4004 MHz, 理论峰值单精度性能 7.68 TFlops, 理论峰值带宽 256BG/s. 编译平台是 Ubuntu 16.04 LTS + CUDA 8.0 + GCC 5.1。以 CUBLAS 作为性能和计算结果的参考。CUBLAS 录得最好性能为 7.2 TFlops, 为理论峰值单精度性能的 93.75%。

Kernel 1: Naive

最基本的思想：每个 Block 计算 C 的一小块（这一点一直延续到后面所有的 kernel 中），每个 thread 直接计算 C 的一个元素的值，直接从显存里读数据进来。当然，性能也是要多差有多差了：只有 186 GFlops。

Kernel 2, 3： Tiling

分块乘法是所有平台进行 xGEMM 都绕不开的步骤，因为这可以有效提高计算-访存比，充分利用高速缓存。

我使用了 32 * 32 的分块大小。然而，shared memory 里的分块矩阵存放顺序如果搞错了会引起严重的 bank conflict 问题：kernel 2 只有 202 GFlops, kernel 3 则有 1001 GFlops。这两个 kernel 对应 Ref No.1 的第 2，3 步。

Kernel 4：4x Workload Per Thread

Kernel 4 中，每一个 thread 负责计算 C 中的 4 个元素。为了保持 coalesced memory access, 每半个 warp 中的 threads (即 16 个 threads) 应该访问连续的元素。记 shm_C 为一个 32x32 的 C 的子块，分配给一个 16x16 的thread block, 并分别记 threadIdx.{x, y} 为tid{x, y}. 则每个元素需要负责计算 shm_C[tidy][tidx], shm_C[tidy+16][tidx], shm_C[tidy][tidx+16], shm_C[tidy+16][tidx+16] 这四个元素。这样每个 thread 的工作量更大，指令并行度（ILP）更高，性能也会更好。Kernel 4 的性能可以达到 1801 GFlops

Kernel 5: 8x Work Per Thread, Zero Padding

Kernel 5 对应 Ref No.2 里的 kernel 3。这个 kernel 的思路相对简单一点：每一个 thread 读入 A, B 同一列四个相距为 8 的元素，计算 C 的同一行四个相距为 8 的元素，thread block size 是 (32, 8)。Ref No.2 里的实现是使用列优先 (Column-Major) 格式的，我仍旧使用行优先格式。同时，使用 #pragma unroll 来指示编译器展开最内层循环，加大 ILP。

Kernel 1~4 中我没有显式为 A, B, C 矩阵填充 0 以满足 32x32 的分块大小，而是在 kernel 里进行判断。这样写起来比较麻烦，性能也会受到影响。因此我在 kernel 5 里先进行了填0，使得 A, B 的尺寸都是 32 的倍数，然后再计算，最后再把 C 周围填充的 0 给去掉。

Kernel 5 的性能为 2375 GFlops.

Kernel 6 & 7: More^2 Work Per Thread & 2D Register

更进一步！这次每个线程不止管 C 的一列了，而是管 C 的一个子块：每一列管 C 的 8 * 4 个元素，和 kernel 4 中的思路一样，每个线程管的元素在同一行中两两间隔有 15 个元素。这个思路其实就是 Ref No.3 论文里的 Fig.2，只是论文里的是一个线程管 16 个元素而不是 32 个。每个线程每次从 shared memory 的 A, B 分块中取出一列 4 个元素和一行 8 个元素，对自己寄存器中的 C 子块进行秩1修正（一个列向量乘一个行向量），这样可以最大限度提高读入到寄存器中的数据的使用率（秩1修正读 $2n$ 个元素计算 $2n^2$ 次，前提是计算的结果必须能存在 fast mem 里）。Kernel 7 和 kernel 6 类似，不过存在 shared memory 里的 A 和 B 不再是正方形，而是长方形。因此从 global memory 读入到 shared memory 的时候会有一些变化。每次秩一修正也变成 8x8 的子块。

Kernel 6 和 7 的性能很接近，分别是 3102 GFlops 和 3280 GFlops。

Kernel 8: Vector Read

这个 Kernel 对应 Ref No.2 里的 kernel 7，计算上和上面的 kernel 7 一致，只是使用了 float4 类型来进行读入。这样，GPU 可以使用 LD.128 这样的向量载入指令，每次读取更多数据，减少总的指令数，使得全局访问的速度更快（参见这个）。然而读到寄存器里的 float4 再拆开写到 shared mem 里的时候免不了要出现一次 bank conflict。这个 Kernel 的速度可以达到 4807 GFlops。如果忽略掉填充 0 所需的时间，单纯是计算乘法的性能可以达到 5280 GFlops。

数据汇总

Kernel 8 与 CUBLAS sgemm 在不同尺寸下的对比：

再来个横向对比，看看各个 kernel 的差别：

小结

以下方法有助于提高 CUDA 程序的性能：

• 避免分支，或者将分支转变为只有在某个 warp 中有不同值的情况；
• 每个 warp 访问的内存地址连续，避免 shared memory bank conflict 和 global memory 不合并的访问；
• 为每个 thread 分配更多的工作量，提高 SP 利用率；
• 对循环进行多路展开以减少循环判断次数和增大指令吞吐量；
• 利用 shared memory 来保存一个 thread block 共用的数据，利用寄存器来保存每个 thread 各自的计算结果和数据；
• 适当调整指令顺序，以计算指令掩盖长延时的取数据指令。

References

1. 5KK73 CUDA GEMM
2. OpenCL SGEMM Tutorial
3. An Improved MAGMA GEMM for Fermi GPUs
4. Performance Upper Bound Analysis and Optimization of SGEMM on Fermi and Kepler GPUs
5. Fast Implementation of DGEMM on Fermi GPU （sci-hub 链接已失效）
6. nvprof & nvvp user guide
7. nvprof metrics & events meaning