矩阵乘法优化：GEMM中如何将大矩阵切割成小矩阵
矩阵乘法优化：GEMM中如何将大矩阵切割成小矩阵
矩阵乘法是线性代数中最基本的运算之一，在高性能计算、机器学习和深度学习等领域有着广泛的应用。为了提高矩阵乘法的计算效率，研究人员提出了许多优化算法，其中GEMM（General Matrix Multiply）是最常用的一种。本文将详细介绍GEMM算法中如何将大矩阵切割成小矩阵以提高计算效率。
如何拆分
一个矩阵乘法有6种拆分方式，其中对row-major效率最高的是：
第一次拆分
先做第一次拆分，取A的kc列（PanelA）和B的kc行（PanelB），得到待累加的一部分结果；第二次拆分
第二次拆分，把PanelB按nc大小切分成多个BlockB，分别与PanelA相乘；第三次拆分
最后一次切分，把PanelA按mr大小切分，与BlockB乘（注意具体代码里的kernel4x4或者kernel8x8是指mr的大小）。
GEPB约束条件
我们考察最后一次GEPB的切分过程，如何使其性能最高？
以上是GEPB的参数，如果这个运算所需数据都加载到L1 cache里，那计算效率自然是最高的。此时运算ops为2mkcnc，同时做的数据搬运量为kcnc+m(kc+2nc)，我们期望做最小的io产生最大的计算，即
max(2mkcnc/(kcnc+m(kc+2nc)))
由于nc远小于m，上式约等于
maximize(2kcnc/(kc+2*nc))
可以得到以下约束/倾向：
- 约束1: maximize(2kcnc/(kc+2*nc))
- 约束2：kc*nc越大效果越好
Cache层级分配
由于L1Cache通常比较小，不太可能得到比较大的kcnc值，是否可能把BlockB放到L2Cache里，尽量让kcnc最大？已知GEPB计算的伪代码为
Load PanelB into cache
for j = 0...M-1
Load Aj into cache
Load Cj into cache
subkernel_calculation
Store Cj
endfor
循环里每次取PanelA的mr行做计算，每次计算的时间开销为
2mrkc*nc/cpu_frequency
假设PanelB放到了L2Cache，每次加载需要的时间为
kc*nc/l2cache_speed
计算时间要大于加载时间，才能防止CPU等待IO从而让CPU ALU跑满，代入上式可得到约束条件
- 约束3：mr≥cpu_frequency/(l2cache_speed*2)
因为计算用的数据都在L1Cache，所以每次从PanelA加载的nr行不宜超过L1Cache大小。同时计算结果C也需要空间，可以得到
- 约束4：mr*kc≤l1_cachesize/2
当数据全放到了L1Cache里计算时，每次是用kernel4x4还是kernel8x8，取决于CPU有多少寄存器。一般是尽可能利用所有reg。可以得到
- 约束5：nr kr的选择完全看寄存器情况
- 需要一半的寄存器存结果
- nr和kr接近，效果最佳
- armv8浮点一般是、8x8
来自TLB的约束
我们知道计算机存储结构长这个样子
上文说到要把BlockB放到了L2Cache，随着ALU的计算不断加载到L1Cache。实际上两个Cache中间还隔着个TLB，因此要考虑TLB寻址空间的限制。L1Cache miss可以用prefetch指令处理掉，TLB miss会真的stall CPU，影响ALU利用率。
Cache寻址空间（256K）一般小于L2Cache大小（2M），因此有
- 约束6：GEPB内循环一次计算所需数据总大小（的行的行BlockB+A的mr行+C的mr行）不能超过TLB寻址空间
- 约束7：计算数据要Packing，防止TLB miss
总结一下行主序GEMM分块的参数选择
mr、nr和kr如何取值
- 尽可能是个正方形，尽量用满寄存器
- 约束3限制了mr的最小值
kc的取值
- nc*kc尽可能大，但不能大过TLB寻址空间的一半
- kc不超过l1_cachesize/2/mr，其中mr已在上文选取
- 按经验是页表大小的一半，即pagesize/2
nc的取值
- nc=kc，参见约束1
mc的取值
- mc*k<l2_cachesize，我们期望A尽可能放到L2Cache里面
就已知的情况来看，fp32 GEMM极限可达到10-11 gflops左右（即CPU峰值的80+%，packA和packB开销跟着矩阵形状走大约是20%），结合具体的任务，例如2D卷积、全连接，可以做到90%左右。