Anatomy Of High Performance Matrix Multiplication 高性能矩阵乘法剖析

下图左边是一个非常简单的多层内存模型，只有寄存器/ cache/RAM。 

在这种简单模型结构下考虑优化GEBP，Cmc,n+=Amc,kcBkc,nCmc,n+=Amc,kcBkc,n，其中

• A是block，A∈Rmc×kc
• B是panel，B∈Rkc×n
• C也是panel，C∈Rmc×n

3个假设

• Assumptions (a)：mc和kc都足够小[都为1]，使得cache能够容纳矩阵AA，矩阵BB的nrnr列Bj和矩阵C的nr列Cj
• Assumptions (b)：如果A，Bj和Cj都在cache里，那么Cj:=ABj+Cj能全速利用CPU
• Assumptions (c)：A 可以一直保留在cache里不被切换出去

基于以上三点假设，上图中GEBP的RAM和cache之间的数据搬移开销为mckc+kcn+2mcnmemops

• 把A从RAM加载到cache，mckc
• 把Bj从RAM加载到cache，kcn
• 把Cj从RAM加载到cache，mcn
• 计算完，把Cj结果从cache加载到RAM，mcn

而Cj:=ABj+Cj的计算量为2mckcnflops，那么计算量和数据搬移的比例。

问题变成，如何选择矩阵AA的尺寸mcmc和kckc，使得矩阵尺寸满足Assumptions (a)-(c)，而且2mckckc+2mc2mckckc+2mc值最大。

• 最大化kcmckcmc，尽可能选择最大尺寸的矩阵AA，使得A,Bj,CjA,Bj,Cj能够放进cache
• 在mckc≤Kmckc≤K前提下，求2mckckc+mc2mckckc+mc最大值问题，数学上转换成几何问题，一个长方形边长分别是2mc2mc和kckc，在长方形面积受限mckc≤Kmckc≤K条件下，如何最大化长方形周长2(kc+2mc)2(kc+2mc)，这个问题在数学上的最优解就是kc=2mckc=2mc，不应该是方阵?

现实中kckc的选择还受一些其他因素制约，我们将在6.3章节看到。 类似地可以分析GEPB和GEDOT操作。

4.2 Refinements

考虑更加实际的应用场景，假设矩阵是column-major order按列存储的。

4.2.1 Choosing the cache layer

Fig.6的右边描述了更准确的内存结构，实际上cache也是分成多层的。 首先问题是，尺寸为kc×mc的矩阵A应该保留在哪一层cache？答案肯定是，在满足Assumptions (a)-(c)的前提下，越快越好，最接近寄存器最快的那一层cache，也就是L1 cache。但是L1 cache天然地非常小，那么可以把矩阵AA应该保留在L2 cache里，同时让尺寸kc×mc更大一些? Rcomp表示CPU能处理的浮点运算操作速度，RloadRload表示CPU能从L2 cache加载到寄存器的速度，假设

• 矩阵AA保留在L2 cache，而矩阵BjBj和CjCj保留在L1 cache
• L1 cache和寄存器之间带宽足够，也就是忽略矩阵BjBj和CjCj的加载时间

那么Cj:=ABj+Cj的计算量为2mckcnr flops，时间开销为2mckcnrRcomp，带宽开销为从L2 cache加载矩阵AA到寄存器的数据量mckc，时间开销为mckcRload。为了克服矩阵AA的加载时间，必须保证 \(\frac{2m_ck_cn_r}{R_{comp}} \geq \frac{m_ck_c}{R_{load}} \\ n_r \geq \frac{R_{comp}}{2R_{load}}\)

Bj和Cj保留在L1 cache，A保留在L2 cache，希望矩阵计算时间大于L2 cache的加载时间，panel矩阵B,CB,C按nrnr行或nrnr列拆分成Bj,Cj矩阵时，拆分尺寸nr必须足够大，nr的最小值与CPU浮点运算操作速度Rcomp和L2 cache加载到寄存器的速度Rload相关。

4.2.2 TLB considerations

第二个考量与系统页管理有关。通常我们的系统使用虚拟内存，因此可用内存大小不受实际物理内存大小限制，内存是分页的。有一张page table来映射虚拟地址和物理地址，保持跟踪这一页实在内存还是硬盘。问题是这个表本身可能很大，比如若干Mbytes，妨碍虚拟内存到物理内存的快速转换。为了克服这个问题，一张更小的表Translation Look-aside Buffer (TLB)，用于存储最近用到的表信息。当一个虚拟地址能在TLB里找到时，转换速度是很快的。但没有找到时，称为TLB miss，需要重新访问page table。将信息从page table里拷贝到TLB里。TLB可以看做是page table的cache。最近更有些架构中出现类似L2 cache的L2 TLB。 TLB的存在意味着我们需要满足额外的假设

• Assumptions (d)：为了保证计算Cj:=ABj+Cj时不会出现TLB miss，mc,kc必须足够小使得，block矩阵A，panel矩阵B,C按nrnr行或nrnr列拆分成Bj,Cj矩阵可以同时被TLB访问
• Assumptions (e)：在直到计算完成之前，block矩阵A必须一直能被TLB访问。

4.2.3 Packing

通常矩阵A是一个大矩阵分解之后的小矩阵，因此矩阵AA的数据在内存中不是连续的， 不连续意味着可能的TLB miss。解决方法是将矩阵AA pack成连续数组A~。适当选择参数mc,kc使得A~,Bj,Cj可以容纳在L2 cache里，并且可以同时被TLB访问到。

Case 1: The TLB is the limiting factor

假设系统有TT个TLB entries，A~,Bj,CjA~,Bj,Cj对应的TLB entries数目是TA~,TBj,TCj，那么有 \(T_{\tilde{A}}+2(T_{B_j}+T_{C_j}) \leq T\)

TBj,TCj前面有系数2是因为当计算Cj:=A~Bj+Cj时，TLB需要同时准备好Bj+1,Cj+1 相比于将A加载到L2 cache的开销，将A pack成A~的操作的开销，不会太大。理由是，packing可以被安排，使得A~加载到L2 cache和能被TLB访问后，立即能为随后计算使用。前提是，A的访问开销不会明显大于将A加载到L2 cache，而这一点即使不做A没有做pack也是必然的。 GEPP或者GEPM会被分解成GEBP，在之前的case里，block矩阵BB会在多个GEBP里重用。这就意味着将BB拷贝成连续数组B~是有价值的，Bj对应的TLB entries数目是TBj也会对应下降到TB~j。

Case 2: The size of the L2 cache is the limiting factor.

可以得到case 1类似的结论。

4.2.4 Accessing data contiguously.

为了寄存器更高效的搬移数据，pre-fetch？连续CPU指令需要读取的数据在内存中必须连续。这不仅需要pack，还需要重排。第6章详细介绍。

4.2.5 Implementation of GEPB and GEDOT

类似GEBP

5. PRACTICAL ALGORITHMS

介绍Fig. 4中的6种实现

5.1 Implementing gepp with gebp

在gebp_opt1方法中，我们希望A一直保留在L2 cache里，B,C反复更新在L1 cache，因此需要确定A的尺寸最大可以多大，A的尺寸上限决定于可用TLB entry数目A~

• 将B packing成B~，Bj和Bj+1各需要一个TLB entry
• 用Caux尺寸是表示着对于这个buffer只需要用到一个TLB entry，Cj尺寸是mc×nr，如果mc很大的话，每个TLB entry包含一个mc，最多需要nr个TLB entry
• 那么剩余留给A~的TLB数目是TA~=T−(nr+3)

Cj是否连续不是太大的问题，因为这部分数据在GEPP计算时没有重用，只会访问一次。 一旦BB和AA变成连续的B~和A~后，gebp_opt1内的循环计算可以达到浮点运算速度峰值。

• packing B到B~的拷贝，是内存到内存的拷贝，开销和kc×n成正比，分摊到2m×kc×n的计算量里，每次拷贝需要2m次计算，这种packing操作打乱了之前的TLB上下文。
• packing A到A~的拷贝，如果合理编排的话，可以是内存到L2 cache的拷贝，开销和kc×mckc×mc成正比，分摊到2m×kc×n的计算量里，每次拷贝需要2n次计算。实际中，这种拷贝不是很耗时。

这种方法适合，m,n很大而k不大的GEMM类型，如GEPP的定义。

5.2 Implementing gepm with gebp

和GEBP不同的是，GEBP中的C只需访问一次，而GEPM里的C需要反复更新，因此值得将结果加到CC之前先累加C~=AB，而不是没计算一次C~=AB就加到CC上。$\check{B}p没有被重用，因此没有pack。此时没有被重用，因此没有pack。此时B_j最多需要最多需要n_r个TLBentry，个TLBentry，B{temp}, C_j, C_{j+1}各需要一个，那么剩余留给各需要一个，那么剩余留给\tilde{A}的TLB数目还是的TLB数目还是T_{\tilde{A}}=T-(n_r+3)$

5.3 Implementing gepp with gepb

A做了pack和转置来保证连续访问，在GEPB里B做pack而且保存在L2 cache，因此这里我们需要最大化TB~，同理TA~的上限是T−(nr+3)

5.4 Implementing gemp with gepb

用一个临时的C~来累加C~=(AB)T，B~保留在L2 cache里。同样因此这里我们需要最大化TB~，TA~的上限是T−(nr+3)

5.5 Implementing gepm and gemp with gepdot

C保留在L2 cache里，每次乘完A的一些列和B的一些行后累加到C上。下面会介绍这种方法不优。

5.6 Discussion

这里我们比较各种实现方法的优劣，结论是gebp实现gepp在列主序情况下最优。 首先排除gepdot实现。考量L2 cache带宽，gepdot实现中将C保留在L2 cache里，然后从内存里加载A和B，L2 cache和寄存器时间的带宽是gebp实现的两倍，因此gepdot实现是最差的一种实现。这个结论前提假设是，A和B的分片Aj和Bj都太大了，无法保留在L2 cache里。 比较gebp实现gepp Fig.8和gebp实现gepp Fig.9，主要区别在于前者pack B而从内存加载C，后者从内存加载B，将计算之后的临时C~放在buffer，再unpack C~并加到CC里。

• gebp实现gepp方法，隐藏从内存读写C的开销，而暴露了pack B的开销
• gebp实现gepp方法，隐藏从内存读B的开销，而暴露unpack C的开销

预期unpack C是比pack B更复杂的操作，因此gebp实现gepp比gebp实现gepp更优。同理gepb实现gepp也比gepb实现gemp更优。

最后比较gebp实现gepp和gepb实现gepp，如果矩阵是按照列主序排列的，那么更合适按列分解矩阵，基于此那么gebp实现gepp优于gepb实现gepp。

6. MORE DETAILS YET

我们只关注最优的gebp实现gepp方法。 Cmc,n+=Amc,kcBkc,nCmc,n+=Amc,kcBkc,n，其中

• A是block，A∈Rmc×kc
• B是panel，B∈Rkc×n
• C也是panel，C∈Rmc×n 

  

6.1 Register blocking

考虑Fig. 8中的Caux:=A~B，其中A~和B分别在L2和L1 cache。如下图所示，Caux分解成mr×nr的子矩阵加载到寄存器里计算。 

这意味着，计算Cj的时候不需要子矩阵在L1甚至L2 cache，基于mrnr memops 内存操作执行2mrnrkc flops计算量，这里kckc选的相对较大。 更详细地描述如何将A pack成A~，在我们的实现里，A的尺寸是mc×kc，进一步分解成内存连续的mr×kc子矩阵，每个子矩阵自身是列主序排列的，那么Caux:=A~B计算的时候，访问A~时就是内存连续的了。还有一种实现是将AA的转置保存成A~，这种做法复杂度稍微高一些。

6.2 Choosing mr×nrmr×nr

选择 mr×nr时有如下考量

• 一般可用寄存器的一半用于存储C分解成的mr×nrmr×nr子矩阵，留剩余寄存器来prefetching A~和B~
• mr≈nr时，加载的开销平摊到计算量上时，即计算密度最优。
• 如4.2.1章节提到，从L2 cache prefetching A~到寄存器的开销，不应该比之前的计算开销更长，最理想的是nr≥Rcomp2Rload。Rcomp=/fracflopscycle是CPU浮点运算操作速度，Rload是L2 cache加载到寄存器的速度

寄存器数目不够会限制gebp_opt1的性能。

6.3 Choosing kckc

每加载一个mr×nr的CC子矩阵，都需要和mr×kc的A子矩阵相乘，为了最大化平摊掉加载开销，kc必须尽量大。但同时kc受以下因素制约

• Bj会重用很多次，因此最好保留在L1 cache里。同时set associativity and cache replacement policy限制了BjBj能够占用多少L1 cache。一般，kcnr个浮点数应该占用少于一般的L1 cache，才能在加载A~和Caux时将Bj逐出cache
• A~的尺寸是mc×kc应该占据一定比例的L2 cache

经验做法，kc个双精度浮点数占据半页，这样得到的值在各种平台都满足。

6.4 Choosing mcmc

前面已经讨论了尺寸是mc×kc的A~A应该

• 能被索引TLB
• 小于L2 cache

事实上还有来自于set associativity and cache replacement policy的限制。 经验做法，mc选择前面限制条件的一半