CPU Case Study - Optimizing DGEMM

Yet Another GEMM Study.

两年前我按 Ref. 1 的页面(以前还没有 GitHub Repo 和 Markdown pages 呢)做过一次 DGEMM Optimization,当时做的效果其实不是很好。去年叶老师给我看了一下 BLIS 这个项目,说里面分块和分级 Cache 的思路值得一看。前两天一搜,居然出了 Ref. 2 这个 Repo,有如此详细的指导和参考代码,不自己造一次轮子简直说不过去了。我最后撸出来的代码在 这里

测试平台

测试平台有两个:

  1. 我自己的笔记本:Core i7-4702MQ (Haswell) @ 2.2GHz, AVX2, 每核心理论双精度峰值是 2.2 GHz (256bits / 64bits) 2 SIMD Units * 2 FMA = 35.2 GFlops
  2. 实验室的服务器:Xeon E5-2620v2 (Sandy Bridge) @ 2.1 GHz, AVX, 每核心理论双精度峰值是 2.1 GHz (256bits / 64bits) (1 Add SIMD Unit + 1 Mul SIMD Unit) = 16.8 GFlops
    两个平台均使用 CentOS 7.3 x64 + GCC 4.8.5-11, 我的笔记本宿主系统是 Windows 10.0.15063 x64,在 VMware Workstation 12 Pro 里跑的 CentOS 虚拟机。
    为了方便,我只做了单线程的 DGEMM,测试也只用单线程进行测试。使用 OpenBLAS v0.2.19 作为性能参考。

Kernel 1~3 Naive ways

Kernel 1、2、3 都是非常 Naive 的方法:原始的 ijk 循环、改变顺序的 ikj 循环和将 M、N、K 填充到 8 的倍数再去计算。Kernel 3 没有比 Kernel 2 慢多少,大概是因为矩阵尺寸都还不是很大,填充操作不是很耗时。

Kernel 4: Registers for C

Kernel 4 每次计算一个 4x4 的 C 子块,并且使用寄存器来保存这 16 个 C 的元素。Ref. 1 给出的结果显示,这样做应该是可以提高一点点速度的;按 Ref. 1 做的 MMult_4x4_6 (列优先存储)比起不分块不使用寄存器的版本也的确可以有较大的性能提高。然而,我自己重新写的这个却没有变快,甚至还变慢了。这与我使用行优先进行存储应该没有关系,因为无论是行优先还是列优先,要么访问 A 的一列会造成访存跳行,要么访问 B 的一行会造成访存跳行。这让我感到比较疑惑。

Kernel 5: 假装做了分块

Kernel 5 是按 Ref. 1 的 MMult_4x4_11 的思路来做的:每次算 C 的一个若干大小的块,算这个块的时候调用 Kernel 4。按 Ref. 1 的说法和我以前做的结果,这样可以在矩阵尺寸变大的时候维持性能。然而 Kernel 5 并没有这个效果,我依旧无法解释这个结果——见了鬼了。

BLISLab_dgemm_kernel 1 & 2

这两个 Kernel 就是对下面这个 BLIS Framework 的原样照搬:
BLIS Framework

这张框架图可以分成两个部分来看:外三层(3-rd, 4-th, 5-th)循环是对 A、B、C 进行整体分块的,然后调用里面的 Marco Kernel 去处理。Marco Kernel 的外两层(2-nd, 1-st)进一步分块,分成若干个 cache size awared 的块状内积,即 micro-kernel 的那一条东西。最内层的 micro-kernel 将块状内积拆分成了若干个 rank-1 update,以便操作数都可以放在寄存器里。

blislab_dgemm_kernel1 其实只实现了外三层,对 A B 分块的打包只是简单地 memcpy 出来,没有做重排。测试表明,这个 kernel 虽然性能比前面几个的要好一点,但依旧很糟糕。毕竟只进行简单的分块只能保证每次 Marco Kernel 要算的数据不会掉出 L3 cache,但是并没有很好地利用 L1, L2 cache。

blislab_dgemm_kernel2 对应 Ref. 2 的 step3/dgemm/my_dgemm.c 以及 step3/kernels/bl_dgemm_int_8x4.c,完整地按部就班地(几乎可以说是公式翻译器一样)实现了上面这个框架。其中,进行 rank-1 update 的两种方法我都实现了出来:最直观的读 4 个 B 的元素广播一个 A 的元素(广播法),以及一次读 A B 向量然后通过 shuffle 和 permutate 改变 B 的排列来 update 不同 C 元素(混洗法,或者叫蝴蝶法)。广播法的图我直接用 Ref. 2 教程里的图了(Ref. 2 用的列优先,所以图里的广播对象和我的相反):
bcast_4x4

混洗法的图我觉得 Ref. 2 教程里的图有些问题,自己画了一个更容易理解一点的:
shuffle_4x4

实验表明,广播法在我的笔记本上只能跑到 13~14 GFlops 的速度,但是混洗法可以跑到 20 GFlops,与 Ref. 2 GitHub Repo 里的代码速度相同(然而我的代码写得简单得多)。然而,在实验室服务器上,混洗法和广播法的速度竟然区别不大。因此我在实验室的服务器上开了一下 VTune 对比了一下这两个版本的区别。先看广播法的:
bcast_generalexploration
bcast_hotspots

呃……我很怀疑进行 general-exploration 测试时, VTune 是否采集到了有效的数据,要不然为什么 Port Utilization 一片都是 0.0。对比看一下混洗法的结果:

shuffle_generalexploration
shuffle_hotspots

这次 general-exploration 里的 Core Bound 消失了,全都变成了 Memory Bound,还是看不出什么信息。对比 hotspot 图,可以看出混洗操作几乎没有使用什么时间,而广播 A 还是需要不少时间的,这里已经省下了一部分时间。

Performance Compare

下面是所有的性能测试结果。其中,Ref. 2 的代码和我实现的代码均使用 M*N*K = 128*512*512 的分块大小,MR, NR 的值均为 8 和 4。为了对比,我还将 Ref. 1 里最后一个 Kernel MMult_4x4_15 用 AVX 指令集改写了一下,一并加入测试。
下面先看我的笔记本上的结果:
i7-4702mq_perf

OpenBLAS 基本能跑到峰值性能。我按部就班写出来的代码(Kernel 7)和 Ref. 2 材料里的代码跑的速度基本一样,都没能接近峰值性能,可能需要上汇编才行。MMult_4x4_15 的结果也不太坏,但是那种分块方式显然没有上面那张图的那么好。其他 Kernel 基本就是自己跟自己玩了。

再看一下实验室服务器上的结果:
e5-2620v2_perf

这次 OpenBLAS 距离峰值还差一点。很意外地,Ref. 2 材料里的代码居然能够逼近 OpenBLAS 的性能,并且没有使用汇编代码。我按部就班写的代码(Kernel 7)则显得力不从心了,可能是我的代码只在同一个迭代内 prefetch 了一次,没有连续 prefetch。

Update: 我将 Kernel 7 里的代码改成了交替连续 prefetch 的,但是仍然没有用 __asm__ 汇编指令。这一改进使得程序在实验室服务器上可以跑到 12.48 GFlops 了,交替连续 prefetch 还是可以用计算掩盖一下数据传输的。

小结

以下方法有助于提高 CPU 程序的性能:

  • 将分支调整到外层循环,避免在内层发生分支;
  • 将内层循环调整为访存连续的 SIMD 操作;
  • 使用向量化 pragma、intrinsics 来进行向量化 SIMD 操作;
  • 分层次打包重排某些需要反复使用的数据,使之可以留在各级 cache;
  • 展开循环,增大指令吞吐量;
  • 适当使用寄存器变量保存频繁使用的变量;
  • 适当调整指令顺序,以计算指令掩盖长延时的取数据指令。

References

  1. How To Optimize GEMM
  2. BLISLab: A Sandbox for Optimizing GEMM