为何我们需要关心这个?
有做过 CPU 矩阵乘优化的同学可能知道,仅仅调整循环顺序就已经能够带来显著的性能差异了。有许多分析都是从局部性的角度进行分析的。即使用向量外积的方案可以利用到循环遍历的局部性,将一些重复访存使用寄存器缓存而避免无意义访存。例如我们补充一下采用向量外积方案关于寄存器的细节。
float a[M*N];float b[N*K];float c[M*N];for k in range(K): regB[0:N] = b[k*N:(k+1)*N] for i in range(M): regA = a[i*K+k]; for j in range(N): c[i*N+j]+=regA*regB[j];
其中 regA 和 regB 均为寄存器。其中我们不难发现,对于每一次循环 j ,使用的都是完全相同的 A 矩阵里的元素,因此可以用一个寄存器来缓存该值;对于每一次循环 k,使用的都是完全相同的一行 B 矩阵中的值,因此我们可以用 N 个寄存器缓存该值。于是将原本 次访存(底下两层循环需要访问一次 A 矩阵和一次 B 矩阵),通过使用 个寄存器缓存(B 使用 N 个,A 使用一个),优化为 N+M 次访存。同时我们也注意到, M 和 N 越大的情况下,提升效果越发显著,这也是为什么我们希望每一个线程负责的分块大一点比较好。但同时 M 和 N 越大,每一个线程多使用的寄存器就越多,而在 GPU 的语境下,更高的寄存器占用意味着更低的 Occupancy。因此当 M 和 N 大到 shared memory 带宽不是性能瓶颈即可。更详细的分析可以看李少侠的分析。
而我则从循环展开的角度解释一下为什么我们需要了解这个优化方案,同时解释一下为什么该优化方案在 GPU 上并不如 CPU 上那么有效。从循环展开的角度来看,第二种循环体构造与第一种循环最大的区别就在于它能在不展开 k 的情况下通过展开 m 和 n 处的循环就能自动的识别到重复访存,并使用相应的寄存器来避免重复访存。例如我们假定,那么展开 m 和 n 处循环的结果如下。
M=N=2;float a[M*N];float b[N*K];float c[M*N];for k in range(K): c[0*N+0]+=a[0*K+k]*b[k*N+0] c[0*N+1]+=a[0*K+k]*b[k*N+1] c[1*N+0]+=a[1*K+k]*b[k*N+0] c[1*N+1]+=a[1*K+k]*b[k*N+1]
只要是稍微现代一点的编译器,都能一眼看出这四条指令的 8 次访存,有 4 次是可以合并的。同时现代一点的编译器也能在一定程度上根据生成的汇编交叉排列计算和访存达到延迟覆盖的目的。而向量内积的方案需要把整个 k 维度展开才能看到这些潜在的访存合并机会。在 CPU 矩阵乘的语境下,一般计算 kernel 的 都比较大(好几百),而 都很小(一般取 6x16,根据架构来做具体确定),寄存器数量又非常少,因此基本上无法在 K 维上将循环完全展开并做优化。因为展开一个超长的循环不仅会带来额外的寄存器占用、优化难度,还会带来更多的汇编指令,使得最终的二进制文件臃肿不堪。但在 GPU 上,情况却恰恰相反。对于已知循环次数的小循环,即便你没有指定 #pragma unroll,nvcc 也会自动的展开这些循环。而对于一个 thread 所负责的小型矩阵乘,这三层循环的值均为 8,符合 nvcc 自动展开循环的条件。而在展开完成后,nvcc 会对所有的访存以及计算指令重排得到一个不错的汇编指令排列。
那么这就引出了下一个问题,我们为何还需要关心他究竟是向量内积还是向量外积?
答案就是 double buffer。如果我们能够提前知道一个循环需要什么数据,我们就能提前预取该循环第一次所需的数据,同时在该循环进行运算的时候预取下一次计算所需的数据。而显然这在向量内积的情况下是无法做到的。同时由于 double buffer 需要额外的寄存器保存从 global memory 转移到 shared memory 的数据,所以当一开始循环展开使用的寄存器过多时,尽管后续能优化到较少的寄存器,但编译器依然无法正确的在限定寄存器数量下实现 double buffer。这一点在优化 sgemm 的时候并不是那么重要(因为多使用一点寄存器也就从每个 SM 跑两个 block 变为一个 block),但是在优化 int8 矩阵乘时需要额外的关注(因为本身它就只能在一个 SM 上跑一个 block,如果实现不得当将会完全失去 double buffer)。
那么此时朴素的利用到向量外积和 shared memory 的代码:https://github.com/AyakaGEMM/Hands-on-GEMM/blob/main/src/cuda/shared_mem_gemm.cu
Double Buffer
由于 GPU 没有 prefetch 这种指令,同时我们又有 shared memory 这种可编程的 L1 cache,因此需要手动实现 prefetch 功能,而在 GPU 语境下一般被称作 double buffer。double buffer 的好处自不必多说,即它可以实现数据读取与计算在时间上重叠,利用 FFMA 单元与 Load/Store 单元可以并行执行指令的特点,达到覆盖延迟的目的。而尽管 GPU 可以在一个 warp 有延迟的情况下,通过切换去运行另一个 warp 达到延迟覆盖到目的,但由于可供 warp 调度器能切换到线程数量的限制,过于长的延迟并不能通过这种方式覆盖掉。这里引用一下李少侠更详细的分析:
若每 SM 有 4 个调度器,若每个调度器只有 4 个可调度 warp,当指令平均间隔超过 4 cycle 后就无法靠 warp 调度掩盖延迟了。考虑到 GEMM 中涉及 smem 读写的过程需要同步 thread block,进一步限制了 warp 调度空间,所以很难靠 warp 并行掩盖延迟。
而本文最终实现的 kernel occupancy 只有 50%,即每个 SM 只能调度 512 个 threads(16 warps),加上图灵架构每 SM 有 4 个 warp 调度器,最终结果与李少侠分析的一致。因此 double buffer 从指令角度提供的延迟覆盖方法最终还是会有效的。
但值得一提的是,在你自己动手实践时,尽可能的考虑在其他优化已经加无可加的情况下再加入这个优化。这是由于这个优化会大幅修改数据读取部分的代码,而且还会产生重复代码,不利于代码维护。同时在我自己的实践中发现,如果在一开始 kernel 写的比较垃圾,加了 double buffer 也没有什么卵用,还会让后续的优化不太好加上去。当然,这只是我的个人建议,如果你想实际看看 double buffer 的效果也可以一开始就加上去。
首先我们看一下每个 thread 的运行流程。
那么能实现 double buffer 的机会有两个地方:Global Memory to Shared Memory 与 Shared Memory to Register。即在每一次 FFMA 开始之前我们读取 Global Memory 的数据到寄存器中,在 FFMA 之后将该寄存器中的值写到 shared memory 中。由于在读取数据后 load from shared memory 以及 FFMA 两个流程中我们并不依赖于该寄存器中的数据,因此可以覆盖 Global Memory 的读取延迟。而同时在计算每一次 FFMA 之前,我们可以用寄存器提前取下一次 FFMA 需要的数据,也就能做到覆盖 shared memory 的延迟。
大概就是这样!我们在每一次运算之前提前将第一次循环所需的数据移动到寄存器中,这样我们就可以实现数据运算和数据存取指令级并行的功能了。
Warp 级优化
在做了不少铺垫之后,接下来的优化终于是可以带来一些看得见的性能提升的了。首先回顾我们之前的代码,可以看到每个 thread 负责的部分完全没有考虑到它们之间可能的协作关系,即同一个 warp 内的 thread 此时在同一块硬件上同时执行——它们共享同一个 register file,这表明它们可以通过寄存器快速共享数据(即 shared memory 的 broadcast 机制);它们会同时访存,这表明如何安排每一个 warp 内的 thread 访存是至关重要的。
Warp Tiling
已知我们指定一个 Block 计算 128x128 的矩阵,一个 Block 有 8 个 warp,一个 warp 有 32 个 thread,每个 thread 需要负责 8x8 的小型矩阵乘,那么我们沿用李少侠的定义:
一个 warp 由 个线程组成,可以是 ,我们把这些线程对应的 thread tile 拼在一起的区域称为 warp tile,尺寸为 ,如下图所示。
这里的图给的是 的排列方式。由于同一个 warp 在访问 shared memory 时有 broadcast 机制(即同一个 warp 在访问同一个内存地址内的值时只会实质发生一次数据读取),因此这一个 warp 计算时只会实际读取 个 float。与之相对的,这个 warp 会进行 次 FFMA。不难看出,在 固定为 32 的情况下,与 越相近,计算访存比就越大,因此取 最为合适。
而在确定了 warp tiling 后,如何读取和存储数据的细节还需要细扣,接下来我将会按照 GPU 的硬件特性讲解读写数据的细节。但这一部分的大致思路基本已经介绍完毕了,动手能力强的同学现在就可以自己试试如何写一个高效矩阵乘了!
向量化访存
向量化访存即是一条指令同时请求多个 float 数据,目前 CUDA 最高支持 128 bit 的向量化访存,即一条指令请求 4 个 float 数据。向量化访问主要的好处在于可以用更少的指令读取更多的数据。由于在访问全局内存时是以 32 Byte 为粒度进行访问的,因此如果同一个 warp 内的 thread 请求了一段连续内存的数据,每一个 thread 都请求两次 4 Byte 的数据(小于 GPU 全局访存的最小单位),那么 GPU 会在硬件处将 64 次数据请求按照 32 Byte 进行合并,最终形成 8 次 32 Byte 内存访问。
而如果每一个 thread 请求 8 Byte 数据,那么 GPU 会在硬件处同样将 32 次数据请求按照 32 Byte 进行合并,最终形成也形成 8 次 32 Byte 内存访问。
那么我们可以看出,对于访问同一数据量的数据,请求的指令越多,GPU 的聚合访问的压力就会越大。在极端情况下,尽管带宽足够,但大量的访存请求会塞满访问队列导致 stall。这在 Nsight Compute 中显示为 MIO Throttle 和 LG Throttle,即对应 shared memory 和 global memory。因此采用向量化访存能在一定程度上缓解 GPU 硬件层面的聚合访存压力(因为我们显式的用指令告诉 GPU 某些数据请求不需要聚合,直接用一个 sector 来处理就好了)。
但使用向量化访存——即用 float4 读写数据——也不是完美的。它的一个严重缺陷在于使用 float4 访存要求请求的数据地址要按照 float4 对齐,因此当 M、N、K 不为 4 的倍数时将会报 missaligned address 错误(因为第二行开始就不能按照 float4 对齐了)。
这么干对输入矩阵形状有一定要求,写出来的矩阵乘没有特别好的通用性。同时 sgemm 受聚合访存的影响也并不是那么大,因此在实操中往往并不会选择使用 float4 读写全局内存,而只会使用 float4 读写 shared memory。但由于我一开始学 CUDA 的时候对这一块理解也不深,然后发现许多人(李少侠除外)都很暴力的直接用 float4 读写全局内存,于是我也用了 float4 读写全局内存。
而我们这里对比李少侠的 kernel profile 和我们最终的成品发现,在 global memory 读取处是否使用向量化读取其实并不会对性能有多少影响。可以看到最终 profile 出来的 Stall LG Throttle 和 Stall MIO Throttle 占比都不高。
上图为李少侠的 kernel 下图为我最终写的 kernel。这两个 kernel 在数据读取方面的区别就是李少侠是以 4B 为单位访存的,而我是以 16B 为单位做访存的。这进一步印证了 sgemm 其实并不是非常关心读取 global memory 时是以怎样的粒度读取的。而向量化访存对于 shared memory 的影响就留给读者自行验证了。同时值得注意的是,在把数据读取方式从向量化访存修改为一个一个访存时需要注意 bank conflict 的问题。因为一个 warp 在执行 128-bit load 和 32-bit load 时的调度并不相同(这点会在后面提到)。
还有一个值得注意的是在 Global Memory 访存时,并不能直接将原先的向量化存取代码直接改成一个一个的读取。因为这样访存从原来一个 warp 并行访问一段连续的内存变成一个 warp 分成四次访问不连续的内存。虽然有 L2 cache 来平滑这种不规则的访存,但最终会带来 10% 左右的性能下降。代码如下:
// Original CodepreA = *reinterpret_cast<const float4 *>(baseA + i + rowA * K + colA); // Modified CodepreA.x = baseA[rowA * K + i + colA];preA.y = baseA[rowA * K + i + colA + 1];preA.z = baseA[rowA * K + i + colA + 2];preA.w = baseA[rowA * K + i + colA + 3];
可以看到这种简单的更改其实并不可取,更优的写法是每一条指令都是在 warp 视图下的连续访存。