借助 mperf 进行矩阵乘法极致优化
作者：旷视 MegEngine 架构师 洪超
前言
单精度矩阵乘法（SGEMM）是非常典型的计算密集型算子，对 SGEMM 的优化也经常被当作算子优化从业人员的练手项目。本文将借助于 mperf，在 ARM A55 cpu 核心上对 SGEMM 的性能进行极致优化，过程中会展示 mperf 辅助性能优化工作的基本逻辑。希望本文的读者对计算机体系结构有一定了解，并且可以补充一些 TMA 的基础概念。另外，关于本文使用的 mperf 工具，更多详细介绍请参见之前的文章 mperf 介绍。
本文需要优化的计算目标是 C=AB，假设矩阵 A 大小为 MK，矩阵 B 大小为 KN，则得到矩阵 C 大小为 MN 。为了后文分块操作的方便，这里假设 M，N 是 4 的倍数，并选择 M=N=K 分别为 100，200，300，500，700，900 的矩阵尺寸，对不同优化版本进行性能测试。本文所有试验代码均在 mperf/apps/cpu_pmu_analysis 目录下。
##矩阵乘法优化
寄存器和 FPU 优化——Naive 实现到循环展开
矩阵乘法的 Naive 实现为三层循环计算：
此时，测试 M=N=K 分别为 100，200，300 时的 mperf 性能数据见 Naive_mperf , 关于这些测试数据的简单介绍请参见 mperf 测试数据。测试命令(后文其他优化版本的测试命令不再重复贴出)：
从上面链接的 mperf 数据可以看到 Naive 实现的 GFLOPS 甚至不到 0.5，而 mperf 在 ARM A55 平台实测的峰值浮点算力可以达到 14GFLOPs，显然我们还有很长的路要走。
分析 TMA Level-1 的数据，也即比较 Frontend_Bound、Bad_Speculation、Backend_Bound 和 Retiring 的占比，可以发现占比最高的为 Backend_Bound，约为 50%，据此首先可以判断 Naive 版本的性能瓶颈位于处理器后端。进一步看 TMA Level-2 及以上级别的数据，发现 Backend_Bound 中占比最大的是 Core_Bound ，而 Interlock_Bound 又占 Core_Bound 的 95% 以上，紧接着 Interlock_FPU 又几乎占 Interlock_Bound 的 100%。
由此可得，当前的瓶颈在于 Interlock_FPU，Interlock_FPU 占比大说明浮点计算触发了大量的 pipeline stall。除此之外，我们还观察到 Metric_Neon_Port_Util 为 0，这是因为 Naive 版本没有进行向量计算。
关于如何降低 Interlock_FPU，我们早期基于 mperf 做过验证，证明 unroll 可以有效降低 Interlock_FPU，过程如下:
在这个例子中，分别对不进行循环展开和以粒度为 2、5、10、20 的循环展开的 fmla 计算进行测试，结果如下：
通过这个简单的例子我们可以看到，随着循环展开粒度的增加，Interlock_FPU 会逐渐降低，并逐渐趋于稳定。这是因为没有循环展开或循环展开特别小的时候，上一次循环还没有执行完成，下一次循环又要读取相同的寄存器，进而造成寄存器依赖。要消除寄存器依赖，首先要保证循环体中不同条指令使用不同的寄存器，并且尽量保证循环体内指令的条数大于该指令的 latency 和该指令的 throughput 的乘积。以本文在 a55 平台测试 fmla 向量计算为例，fmla 的 latency 为 4 个 cycle，并且 a55 处理器后端只有一个 port 可以执行 fmla 向量计算，也即 fmla 的 throughput 为 1，所以要求循环体 fmla 指令数大于 4 * 1 才可以解除寄存器依赖，这也符合上述试验结果的变化情况。
理解上面信息后，我们考虑在 Naive 版本上进行循环展开，并依据 Naive 版本 Metric_Neon_Port_Util 为 0，同时施加 SIMD 优化。具体做法是，先对前面的行列采用 812 展开，对剩余的列采用 84 展开，再对剩余的行先采用 412 展开，最后用 44 展开，期间配合 ARM NEON 指令进行向量化操作。
这里选择 8*12 是为了更充分的利用寄存器资源，是因为 ARM A55 上面有 32 个 128 bit 向量寄存器，以及考虑到本文最后 matmul 汇编版本的最内层实现大体如下（注意这里使用的外积计算矩阵乘，具体参见外积算矩阵乘）：
从 A 矩阵读 8 个 float 到 2 个向量寄存器，此时应该是 8 行里面每行的第一个数。
从 B 矩阵读 12 个 float 到 3 个向量寄存器，应该是 12 列里每列的第一个数。
用 fmla 指令，让 B 的每个 float 分别乘以 A 的两个向量寄存器，产生 24 个向量结果，也全部存储在寄存器中( 24 个向量寄存器存储 8x12=96 个中间结果)。
采用 8*12 展开之后的代码结构如下：
关于循环展开，可以用一个图来直观地理解，这里 mr=8，nr=12，我们每次用 A 矩阵的 8 * K 的小块和 B 矩阵的 K * 12 小块来计算 C 的一个 8 * 12 小块：
用 mperf 分析 unroll 版本的 matmul 实现，具体数据见[ Unroll_mperf ] (https://www.megengine.org.cn/blog/mperf2#:~:text=Metric_GFLOPs_Use%20%3A%200.30557-,Unroll_mperf,-M%3DN%3DK)
对比 M，N，K 较小时(注: 在优化寄存器和 FPU 单元利用率的时候可以把问题规模先限制得比较小，减少访存相关 issue 的干扰)，Unroll 版本与 Naive 版本的 Interlock_FPU、Metric_Neon_Port_Util 和 GFLOPS：
可以看到循环展开后可以看到 Interlock_FPU 占比下降明显，同时 Metric_Neon_Port_Util 明显上升，而 GFLOPS 指标有了数量级的提高，说明充分利用寄存器资源和提高 FPU 利用率有效提高了程序性能。
###访存优化——分块和 PACK
上图是 Unroll 版本与 Naive 版本在不同矩阵尺寸下的性能对比，可以看到循环展开在矩阵尺寸比较小的情况下，性能提升还是很明显的，但是随着矩阵尺寸增大 unroll 的效果迅速下降。所以现在我们将注意力转移到，如何解决矩阵尺寸增大 GFLOPS 下降的问题。通过分析上节链接的 unroll 版本的 mperf 数据，可以看到，随着矩阵尺寸变大到一定程度，Memory_Bound 占比逐渐接近 50%，替代之前的 Interlock_FPU 成为新的性能瓶颈。而 Memory_Bound 占比高主要是 Load_Cache 造成的，这就指明了我们接下来需要进行访存相关的优化。
分块
之所以矩阵尺寸增加，unroll 版本的性能会下降，主要原因就是数据无法全部驻留在 Cache 中，导致数据频繁地在 Cache 和主存之间换入换出，而处理器对主存的访问是非常昂贵的。为了减少对主存的重复访存，首先我们能想到的就是分块( Unroll 版本中提到的分块是内层分块，目的是优化寄存器和 FPU 的利用率，请注意区分)，将分块之后的数据保存在 Cache 中，尽量使处理器发起的访存操作都能命中 cache 中的数据。
这里我们选择对 N，K 维分别进行 Nr 和 Kr 粒度的分块，结合循环展开部分的逻辑，整体的分块方式即为：外层选取 A 矩阵的 MKr 小块和 B 矩阵的 KrNr 小块，内层再对这两个小块分别进行 mr 行和 nr 列的划分，所以最终内层每次计算 mrKr 的 A 小块和 Krnr 的 B 小块，得到 C 矩阵 mr*nr 小块的部分中间结果。
关于如何确定 Nr 和 Kr 的大小，我们的目标就是使得计算时需要用到的分块可以根据访存频繁的程度保存在 CPU 的各级存储中，原则就是访问越频繁的分块存储在速度越快的存储上，以及保证优先用满速度快的存储资源之后再下溢。针对 matmul, 具体约束条件设定为：
将重复访存率最高的 mr*nr 大小的 C 小块保存在访存速度最快的寄存器上(unroll 版本就是这样假设的)。
将下图中红色部分(包括计算完一个 mrNr 的 C 行块需要重复访问次数最多的 mr Kr 的 A 行块，内层一次计算迭代需要用到的 Kr*nr 大小的 B 列块)都保存在 L1 中。
由于计算完每一个 mrNr 的 C 行块，都需要重复遍历一次整个 KrNr 大小的 B 块，因此希望将 KrNr 大小的 B 块存放在 L2 中，使得每次读取 Krnr 的 B 列块的时候，都是从 L2 中读取。
依据上面的分配策略，并结合 CPU 中的各级存储资源(寄存器数量，L1D 和 L2 )的尺寸，便可以确定最佳的 Nr，Kr 取值：
可以根据 CPU 处理器的寄存器数量得到 mr 和 nr 的具体大小，寄存器容量>mr*nr (unroll 版本就是遵循这个约束条件取的 mr=8,nr=12)
根据 L1D Cache 的大小结合 mr 和 nr 计算出 Kr，Kr=L1D/(mr+nr)
再根据 L2 的大小计算出 B 矩阵中的 Nr，Nr=(L2-L1D)/Kr
在 ARM A55 上，最终得到的 Nr 为 252，Kr 为 256。注意：这里计算得到 Nr 为 256，但是由于我们选择 nr 为 12，因此为了避免不必要的余数处理，选择 Nr 为 12 的倍数简化问题。
Kr/Nr 分块的代码逻辑如下：
添加 Kr/Nr 分块后的 mperf 性能数据见 Block_mperf
对比 Kr/Nr 分块前后以下变化明显的数据：
Kr/Nr 分块后可以看到在矩阵较大时进行快速矩阵乘法，GFLOPS 确实上升不少，Load_Cache 和 Metric_L3D_RD_Miss_Ratio 等均有明显下降。这说明 Kr/Nr 分块确实大幅减少了对 latency 非常大的系统主存的访问，优化了程序的访存性能。
数据 PACK
分析 Kr/Nr 分块后测得的 mperf 数据，可以看到随着矩阵尺寸变大，性能也能保持相对稳定，没有明显下降。但是此时 mperf 拿到的 TMA 指标显示， Backend_Bound 中的 Memory_Bound 占比依旧很高，性能瓶颈还是停留在访存部分。
进一步思考，在计算 mrnr 大小的 C 小块的时候，每一次迭代都需要读取 A 矩阵 mr1 的数据，而本文测试的矩阵数据都是行主序，即不同行相同列的数据是内存不连续的，访存不连续就意味着对 Cache 不友好。同样在不同次迭代中，需要读取矩阵 B 中不同行的 1*nr 小块的数据，自然也存在数据读取不连续的情况。考虑到在分块计算 matmul 的逻辑下，A 的所有行块和 B 中的所有列块将被读取多次，因此可以通过对 A 和 B 提前进行数据 PACK，这样只在第一次 PACK 时候对 Cache 不友好，而后面计算的时候多次访问数据均为连续访存，因此收益巨大。下图说明了对 A/B 矩阵分别 PACK 的过程：
PACK 后的 mperf 性能数据见 Pack_mperf
对比 PACK 前后的数据：
可以看到进行 PACK 优化后 GFLOPS 又有一定幅度的增长。此时可以观察到 L1D_TLB_Miss_Ratio 有所降低，并且在 M=N=K >= 700 时L2D_TLB_REFILL明显降低，说明 PACK 确实可以通过减少缺页的发生，减少 TLB miss，从而提升性能。
pipeline 优化——嵌入汇编
分析 PACK 之后的 mperf 数据，可以发现 Backend_Bound 占比 40% 以上，其中 Core_Bound 类别下的 Interlock_FPU 再次成为性能瓶颈。回想 unroll 版本减少 Interlock_FPU 的思路，是通过循环展开给编译器更大的优化空间，让编译器充分利用寄存器来减少流水线依赖，但编译器也是有局限的。一般来说，编译器考虑到通用性，是很难生成针对特定处理器架构特点的最优汇编实现的（ ARM in-order 架构的小核上更是如此）。因此下一步的想法就是优化编译器生成的汇编，通过把内层计算逻辑替换为嵌入式汇编，依据架构特性调整指令选择和指令排布，进一步减少 pipeline 上的依赖和冲突，从而达成降低 Interlock_FPU 的目的。
首先，我们对比了 PACK 版本的 mperf 数据中的 FPU_util 和纯算力测试情况下的 FPU_util，发现 PACK 版本的 FPU_util 相对低了很多。也就是说，PACK 版本的 matmul 对处理器 SIMD 单元的利用率还是有一定提升空间的（下这个判断的一个前提也是考虑到前文已经对 matmul 访存部分优化得很充分了，并且考虑到内层循环中没有分支判断等复杂逻辑，只是比较存粹的访存和计算指令的 interleave。换句话说，汇编优化一般是留到最后进行的）。
其次，我们注意到编译器生成的汇编代码中，数据加载使用的是 ldq 指令(ldq 指的是 armv8 isa 中 128bit load 操作)，但是结合上面 FPU_util 数值比理想情况低的观察（会关注 ldq 指令，是因为我们了解到 ARM A55 的访存能力弱，一个 cycle 最多 load 64bit，store 128bit），我们有理由怀疑 ldq 指令的选择可能会造成 pipeline stall。通过查询 ARM A55 trm 手册，我们发现 ldq 在 a55 上需要两个 cycle 才能 issue 出去，并且 ldq 跟 fmla 不能双发射(注： ARM A55 是双发射架构)，这就证实了 ldq 会造成计算和访存指令无法双发射，并造成 Metric_FPU_util 数值的下降。进一步，我们发现 ldr,ldx,ins 三种指令都能与 fmla 双发射，且发射周期都是 1 个 cycle，而这三条指令可以组合出 ldq 等价的操作。因此，我们使用 ldr,ldx,ins 指令组合来替换 ldq 指令，就可以提高流水线的满载程度，进而提高性能。
我们用一个小的测试例子来进行验证： 考虑下面两段代码： 优化前，使用 ldq 指令加载数据，代码如下：
优化后，将 ldq 拆分为 ldr，ldx，ins 指令加载数据，代码如下：
我们对优化前和优化后的版本进行测试，对分析 mperf 拿到的数据，看到主要的变化为：
可以看到Metric_Load_Port_Util 和 Metric_GFLOPs_Use 均提升明显，也就验证了前面的猜想。
因此，在嵌入的汇编代码中，我们应用上面提到的加载指令的替换，测得的 mperf 性能数据见：ASM_mperf
对比分析 PACK 和 ASM 两个版本的 mperf 数据，可以看到 Metric_Neon_Port_Util 和 Metric_Load_Port_Util 均显著上升。
注意到完成汇编优化后，在矩阵大小 M，N，K <= 300 时，GFLOPS 性能已经达到峰值算力的 90% 以上，基本上可以判定较小尺寸的情况下，matmul 在 ARM A55 平台上已经优化到位了。而当 M，N，K 比较大时，依旧有进一步的优化空间。但是考虑到本文的主要目的，是借助 matmul 优化来说明 mperf 可以给性能优化工作带来哪些助益，所以就不继续进一步的优化关工作了。
##总结
本文以 ARM A55 平台上的矩阵乘优化为例，详细介绍了如何用 mperf 分析当前实现的性能表现，找到性能瓶颈，进而确定下一步的优化方向，如此反复迭代，最终取得了接近硬件性能峰值的性能表现。本文各个版本 matmul 的性能对比，见下图：
附录
mperf 测试数据
这里贴了不同优化版本的 mperf 测试数据，每一个版本下的测试数据都包含两部分：一部分是非 Metric 前缀命名的，都是 TMA 范式下的指标；另一部分是 Metric 前缀命名的则是非 TMA 范式下，但是对性能优化有辅助作用的指标。这两类指标都是基于 PMU 原始 event 四则运算加工得到的，具体的计算公式请参见 ARM a55 tma 。
另，TMA 范式下的指标呈现多级展开的关系，由于当前 mperf 未提供可视化工具，请参考下图示意的层级关系进行对号入座。并且，请注意附录贴的 TMA 数值均为无量纲百分比。
Naive_mperf
Unroll_mperf
Block_mperf
Pack_mperf
####ASM_mperf
术语表
TMA: 也即 Top-down Microarchitecture Analysis 的缩写。TMA 是 intel 工程师为分析 intel X86 cpu 微架构的性能瓶颈而总结的方法论，可参见 intel x86 TMA 。一句话总结的话，TMA 就是一种在微架构层面自顶而下的瓶颈定位方法。mperf 在 ARM cpu 上实现了 TMA 的方法论，并进行了包括本文在内的实践验证。
Core_Bound: This metric represents fraction of slots where Core non-memory issues were of a bottleneck. Shortage in hardware compute resources; or dependencies in software's instructions are both categorized under Core Bound. 表示 CPU 处理器后端非访存相关的 stall 引起的性能损失，通常包括 ALU 单元的资源不足，以及流水线上的指令依赖引起的 stall 等。
Interlock_Bound: No operation issued due to the backend interlock. 表示 CPU 处理器后端因为发生了 pipeline interlock 而产生的 stall, pipeline interlock 就是 A 指令的执行依赖 B 指令某一个 pipeline stage 的结果，而 B 指令还没有执行到对应 stage，指令 A 就等 stall 等待依赖条件满足，比如 raw 依赖。考虑到本文中频繁出现关于处理器 pipeline 以及 pipeline interlock 相关的术语，如果读者对这部分概念比较模糊，推荐一个很好的教学文档 mips pipeline 。
Interlock_FPU: No operation issued due to the backend, interlock, FPU. This event counts every cycle that issue is stalled and there is an interlock that is due to an FPU/NEON instruction. 这个指标是上面 Interlock_Bound 的子类别，表征由 FPU/NEON 指令的 pipeline interlock 而产生流水线停顿的比例，比如前文降低 Interlock_FPU 的试验中，如果循环体中只有一个
fmla v0.4s, v0.4s, v0.4s指令，下一个循环 iter 的 fmla 指令需要读取 v0 的数值，这就依赖上一个 iter 的 fmla 指令更新完 v0，而 fmla 指令的 latency 是多个 cycle，也就造成了相邻 iter 见 fmla 指令的 interlock, 流水线中将产生大量 stall。Load_Cache: No operation issued due to the backend, load, Cache miss. 表征因为 Cache 读延迟导致的 stall。
Metric_Neon_Port_Util: 处理器后端 neon port 的利用率，计算公式是 ASE_SPEC/(CYCLES*NEON_PORTS)，拿实际执行的 neon 指令数除以处理器运行的 cycle 和处理器后端支持 neon 操作的 port 端口数的乘积，ARM a55 的 NEON_PORTS 是 1，在这里用于表征处理器 neon 单元的利用率。
Metric_L3D_RD_Miss_Ratio: L3 Cache 读的缺失率，计算公式是 L3D_CACHE_REFILL_RD/L3D_CACHE_RD。
Metric_Load_Port_Util: 处理器后端 load port 的利用率，计算公式是 LD_SPEC/(CYCLES*LD_PORTS)，拿实际执行的 load 指令数除以处理器运行的 cycle 数和处理器后端支持 load 操作的 port 端口数的乘积，ARM a55 的 LD_PORTS 是 1，在这里用于表征处理器 LD 单元的利用率。
L1D_TLB_Miss_Ratio: 表征的是 L1D 的 TLB miss 的比例，L1D TLB miss 会触发 L2D TLB 的 read。
L2D_TLB_REFILL: 表征的是 L2D 的 TLB miss 的数量，会触发昂贵的 page table walker。
附
更多 MegEngine 信息获取，您可以：查看文档和 GitHub 项目，欢迎参与 MegEngine 社区贡献，成为 Awesome MegEngineer，荣誉证书、定制礼品享不停。
工业级研发的开源深度学习框架-MegEngine 2022-10-21 加入
官网：https://www.megengine.org.cn/ MegEngine 技术交流 QQ 群：1029741705 框架使用相关交流或反馈，欢迎访问论坛：https://discuss.megengine.org.cn/ GitHub ：http://github.com/MegEngine/MegEngine
评论