矩阵乘法(GEMM)是使用最广泛的基本矩阵操作函数之一. 它不仅在很多计算应用中被频繁调用, 而且是基础代数库BLAS中三级函数的基础. 很多其他BLAS三级函数的性能提高也依赖于GEMM的优化. 因此, 对很多科学和工程应用来说, 提高稠密矩阵乘法的性能, 对加快整体计算速度有重要作用. 而GEMM函数的优化, 又严重依赖于硬件平台, 函数实现必须充分利用处理器特性才能发挥出计算平台的计算能力. 因此, 各处理器制造厂商都在致力于提供针对自己生产的处理器进行过优化的BLAS函数库, 尤其是进行了深度优化的GEMM函数. 例如, Intel公司提供优化的数学库MKL^[1]. AMD公司开发了数学库软件AOCL^[2], 面向AMD处理器进行优化. 在GPU上, NVIDIA公司提供深度优化的数学库软件cuBLAS^[3], AMD 开发了AMD GPU上的高性能库rocblas^[4], 并且开放到开源社区. 此外, GotoBLAS^{[5, 6]}, OpenBLAS^[7], KBLAS^[8]等开源基础函数库, 也都对GEMM进行了针对不同硬件平台的优化.

申威1621处理器是我国自主研发的多核高性能处理器, 采用自主sw指令集. 由于申威处理器的架构, 存储层次及所使用的指令集与国外商用CPU有很大不同, 因此在申威处理器上无法使用专为国外商用处理器优化的BLAS库. 而对开源库进行简单适配的版本性能又远远低于处理器所提供的性能. 基于此种状况, 研制面向申威CPU的BLAS库成为一个亟待实现的任务. 刘昊等人对早期申威处理器(申威1600)上的GEMM进行了优化^[9], 解庆春等人进行了向量化等方面的探索^[10]. 但是, 申威1621芯片与早期的1600芯片在指令发射方式, 缓存大小等方面, 有较大不同. 因此, 在申威1621平台上需要对GEMM函数的实现进行重新设计并进一步优化, 以充分发挥其硬件特性.

我们在申威1621平台上对GEMM从访问优化, 指令排布以及减少不必要计算和访存等方面, 进行了深度优化, 并开发了子函数代码生成器, 为不同规模的矩阵生成优化的代码并进行参数空间搜索. 相比于GotoBLAS, 单线程DGEMM实现7.39倍加速, 达到了系统峰值性能的85%, 16线程性能达到了系统峰值的80%.

1 研究现状及背景介绍

稠密矩阵乘法GEMM的计算公式是:

$ C=\alpha AB+ \beta C $

(1)

公式(1)中, 矩阵A, B, C大小分别是M×K, K×N, M×N, α和β是标量, 矩阵A和B可取原矩阵或其转置, 矩阵A和矩阵B均非转置的稠密矩阵乘法朴素算法如算法1所示.

算法1. GEMM朴素算法.

1. 　for i = 0 : M– 1 step 1

2. 　　for j= 0 : N– 1 step 1

3. 　　　 C_ij = βC_ij

4. 　　　for l = 0 : K– 1 step 1

5. 　　　　　 C_ij = αA_il × B_lj +C_ij

6. 　　　end for

7. 　　end for

8. 　end for

GEMM的浮点运算量是2MNK, 计算复杂度是O(N³); 访存量是MK+KN+2MN, 访存复杂度是O(N²), 是典型的计算密集型问题. 理论上讲, GEMM浮点操作性能应逼近硬件浮点峰值. 而在GEMM函数实现中, 除了O(N³)的运算操作, 访存操作也要占用大量时间, 属于不可避免的额外开销. 因此, 设计高效的GEMM实现, 不仅要充分优化计算指令流, 还要考虑如何减少或掩盖访存用时.

目前GEMM的优化主要有3个方面: (1)针对特定的平台, 在充分考虑硬件架构特征的基础上设计分块方案和单线程/多线程算法; (2)对核心代码进行精细手工优化, 例如, GotoBLAS就提供了几种不同CPU上的汇编核心代码; (3)基于自动调优(auto-tuning), 通过在不同配置下运行程序收集性能表现数据, 自动选择算法和参数. 该方法的代表是ATLAS数学库. 在库函数的开发当中, 往往需要将3种优化思路结合使用, 以期在提高开发效率的同时, 充分挖掘平台性能潜力.

各经典BLAS库中的GEMM函数及各方研究人员开发的GEMM函数中, 用到了多种体系结构相关的优化. 在广泛使用的开源软件GotoBLAS中, Goto等人考虑到硬件上缓存容量和TLB表项容量的限制, 给出了一系列分块的实际原则以及分块参数选择方法, 具有普适的指导性^{[5, 11]}. Lim等人在Intel KNL 架构上通过使用矩阵分块, 数据预取, 循环展开, 采用AVX512向量指令等方法来优化GEMM 函数, 最终性能达到了MKL DGEMM 性能的99%^[12]. Smith等人对多核处理器上GEMM的并行和优化方式进行了深入分析并提出了比较系统的优化方案^[13]. 近年来, 国产处理器迅速发展, 基于国产芯片的GEMM优化方法也应运而生. 张明针对龙芯平台, 围绕常用基础数学库函数的访存和计算关键问题展开研究^[14]. 对GEMM函数, 他提出一套适用于龙芯处理器的异步计算访存优化方法——通过流水化和分组处理计算任务, 充分掩盖了访存用时. 刘昊等人基于国产申威1600平台, 使用乘加指令, 向量化指令, 循环展开与软件流水, 数据预取与数据重排等优化技术, 给出了一种GEMM的高性能实现方法^[9]. Jiang等人在神威太湖之光超级计算机上实现了高性能的GEMM, 针对申威26010众核平台, 在不同存储层次上设计了分块方案, 并使用申威众核平台提供的异步DMA等功能充分掩盖了访存用时^[15].

ATLAS数学库采用了自动调优的方法, 在不同配置下运行程序, 收集性能数据, 选择性能表现最优的参数组合作为调优结果^[16]. 近些年出现的TVM深度学习编译器也借鉴了这一思路, 在一些GPU平台上, 对部分GEMM输入规模加速效果明显^{[17, 18]}.

国产申威1621平台^[19]有16个计算核心, 具有三级缓存结构. 每个核有独立的一级缓存和二级缓存, 所有核共享三级缓存. 申威1621处理器向量长度为32字节, 即可以同时操作4个双精度浮点数. 此外, 指令集中还包含预取指令, 可以将某个内存地址所属缓存行的数据从主存读入缓存. 针对申威1621多核平台的体系结构特点和访存特性, 我们设计并实现了GEMM单线程算法和多线程并行算法, 并进行了深入优化. 本文以矩阵A和矩阵B均非转置, 数据类型为双精度的DGEMM为例来说明实现以及优化技术, 其他输入情况和精度可以使用类似的优化方法.

2 面向申威1621处理器的DGEMM优化

本节介绍单线程DGEMM在申威1621平台上的实现方案及优化方法. 我们将文献[9]中基于申威1600处理器的GEMM实现与优化方法作为基础版本. 在此基础上, 我们针对1621处理器的特点, 从分块方式, 打包方式, 及核心计算优化等各方面, 进行了系统的分析与优化, 提出一套GEMM高效实现方案及基于此方案的一系列优化技术.

2.1 基础实现方案

本节介绍文献[9]所采用的DGEMM方案及实现, 即前文所述申威1621平台上进行DGEMM优化的基础版本.

2.1.1 分块方案及实现

为有效利用现代处理器上的分层缓存结构, GEMM采用分块算法, 每个维度上的循环迭代步长是此维度的分块大小, 最内层循环计算一个分块的矩阵乘法. 在循环顺序上, 我们采用了与GotoBLAS2相同的N-K-M的循环顺序. 鉴于传统BLAS库中矩阵使用列优先存储, 矩阵A和C都在M方向上连续, 这样的循环顺序可以保证计算时尽可能多的访问矩阵中连续元素, 提高缓存命中率, 从而减少访存开销. 程序中使用打包函数packA和packB将矩阵中不连续的数据放置在连续缓冲区中. 这样一来, 执行计算任务的KERNEL函数就可以从缓冲区中读取或写回数据, 避免了大量的不连续访存.

算法2展示了单线程DGEMM的分块算法. 矩阵A, B, C分别被分为BM×BK, BK×BN, BM×BN大小的子矩阵块(第2, 3, 5行), 记为§A_il, §B_lj,和§C_ij. 在循环开始之前, 将C矩阵全部乘以β.在计算过程中, 子矩阵块§A_il和§B_lj中的元素经过打包被放入缓冲区 $\tilde{A}\mathrm{和}\tilde{B}$ (第4, 6行)后, 函数KERNEL完成子矩阵块乘法的计算(第7行). 在文献[9]所使用的原始版本中, 对C矩阵未使用缓冲区, 而是在KERNEL中直接写入. 我们认为, 不对C矩阵使用缓冲区, 是因为对C的解包要在最内层循环每次迭代进行, 与packA的调用频率相同, 开销过大, 性能反而会受损. 我们也通过实验验证了上述推断.

算法2. GEMM分块算法.

1. 　 C = βC

2. 　for j = 0 : N – 1 step BN

3. 　　for l = 0 : K – 1 step BK

4. 　　　pack §B_lj into $\tilde{B}$

5.　　　 for i = 0 : M – 1 step BM

6. 　　　　pack §A_il into $\tilde{A}$

7. 　　　　调用计算函数KERNEL(), 计算C_ij += $\alpha \tilde{A}$ $\times \tilde{B}$

8. 　　　end for

9. 　　end for

10. 　end for

2.1.2 汇编KERNEL实现

函数KERNEL()完成M, N, K这3个维度大小分别为BM, BN, BK的矩阵乘法操作. 函数使用汇编语言开发, 使用向量化指令和乘加指令, 以充分发挥申威多核平台的浮点运算能力, 具体实现见算法3. 函数主体是N-M-K三层循环. 在最内层循环中, 矩阵A中的元素被装载到寄存器A_r中, B中元素被装载到寄存器B_r中(第4行), 对A_r和B_r进行乘法运算, 结果放入中间变量C_t中(第5行). 待K方向循环结束后, 进行写回操作, 即将C矩阵中元素装载到C_r, C_r和中间变量C_t累加后, 结果被写回C矩阵在内存中的地址(第7, 8, 9行). KERNEL()的计算过程如图1.

算法3. 计算函数KERNEL(A, B, C).

1. 　for j = 0 : BN– 1 step RN

2. 　　for i = 0 : BM– 1 step RM

3. 　　　for l = 0 : BK– 1 step RK

4. 　　　　load A, B to A_r, B_r

5. 　　　　 C_t = C_t + A_r × B_r

6. 　　　end for

7. 　　　load C to C_r

8. 　　　 C_r = C_r + C_t

9. 　　　store C_r to C

10. 　　end for

11.　 end for

图1所示, 在K方向上一次迭代中, 使用向量化指令来装载矩阵A一列上的RM个元素, 使用标量指令装载矩阵B一行中的RN个元素并将每个元素复制扩展为向量. RM个A矩阵向量与RN个B矩阵向量相乘后得到RM×RN个中间结果C_t. K方向RK次迭代后, 相加并写回到矩阵C. 考虑到申威1621处理器向量化长度是4, 所以K方向上一次迭代中, 矩阵A中元素需使用VM=RM/4个寄存器, 矩阵B中元素需使用RN个寄存器. 此外, 在K方向上循环展开2次, 为了充分利用指令级并行, 我们对展开的两个最内层循环使用不同的寄存器, 即矩阵A和矩阵B中元素分别使用VM×2和RN×2个寄存器. 存储C矩阵中间结果需要使用VM×RN个寄存器. 因此总共需要的寄存器数为

$ {{reg\_used}} = {{VM}}\times 2 + {{RN}}\times 2 + {{VM}}\times {{RN}} $

(2)

循环最内层共有VM×RN条计算指令, VM+ RN条访存指令, 为使得计算指令与访存指令的比例 $ \dfrac{{VM}\times {RN}}{{VM}+{RN}}$ 最大, 应使VM=RN. 根据公式(1), 当VM=RN=4时, 共需要32个寄存器. 而申威1621平台共有31个可用浮点寄存器. 为了充分利用寄存器资源, 我们在不影响性能的前提下, K方向上两次迭代中重用了一个寄存器, 此时计算指令比与访存指令比例达到2:1. 如果有更多的寄存器被重用, 可以进一步增大计算访存比. 在K方向展开的两次迭代中, 按照最大可能的重用度, A和B各只有一个寄存器不重用, 那么只需要VM+1+RN+1+VM×RN个寄存器. 基于这种方案, 可将现有的最内层分块增大, 假设在M方向增加1, 那么将会用到VM+1+RN+1+VM×RN=31个寄存器, 也就是全部浮点寄存器. 这种做法, 虽然能够在一定程度上提高计算访存比, 但是由于每个panel对A和B只有一个非重用的寄存器, 这就意味着最后的3条读取指令与倒数第4、3、2条乘加指令存在寄存器重用. 而神威6A处理器采用2条浮点乘加指令和2条访存指令同时发射的方式, 即浮点操作与访存双发射. 上述情况下, 由于浮点计算指令与读取指令所用的寄存器有重叠, 因此在最后两条指令上无法实现双发射, 从而影响性能. 我们对这种分块方式的KERNEL进行了测试, 性能略低于VM=RN=4的分块. 因此, 本文延续使用了VM=RN=4的分块方式.

为了提高数据访问的效率, KERNEL中插入了针对3个矩阵的预取操作. 由于cache line长度为128字节, 每次预取操作能够读取16个双精度浮点数. 因此, 对A的预取发生在K方向上的每次迭代中. 而对B的读取是每4次迭代处理16个双精度浮点数, 因此其预取频率是A的1/4. C矩阵的预取放置在K方向迭代之前, 即算法4第3行之前. 由于从主存中读取数据至cache的时间较长, 预取操作不能直接预取下一次迭代所需数据, 而是要设置一定的预取步长(fetch_stride). 也就是说, 每次预取操作要将fetch_stride个字节之后的128字节数据载入缓存中.

2.2 单线程GEMM优化方案

在现代计算机体系结构上实现高性能矩阵乘法, 分块方法的优化是必不可少的^{[11, 20]}. 尽管算法2对分块算法做了优化, 使用汇编代码实现计算函数KERNEL, 并在KERNEL中使用了预取指令来优化访存, 但是其性能仍然差强人意. 通过对各部分进行计时和分析, 我们发现, 在申威1621平台上, GEMM函数还有很大的性能提升空间. 首先, 打包函数时间过长, 这种额外开销降低了函数的整体效率. 此外, 我们将KERNEL中的访存指令删去, 其性能几乎可以达到峰值性能. 加入访存指令后, 性能大幅下降, 可见访存效率还有待提高. 1621处理器L1缓存仅有32 KB, 基础版本代码的分块大小A子矩阵就已经超过了L1缓存大小. 为提高函数整体性能, 我们考虑以下因素, 并提出有针对性的优化方法.

首先, 算法2中的分块算法无法兼顾打包开销的减少和KERNEL性能的优化. 一方面, 较大分块意味着更多的缓存缺失率, 影响KERNEL访存效率; 另一方面, 如果分块较小, 又会产生大量的重复打包操作, 增加打包操作引起的访存开销. 虽然通过调节分块大小能在一定程度上改善程序性能, 但是要想在打包和KERNEL两个环节都有性能提升, 需要对循环结构进行进一步优化以减少这两种开销. 其次, pack函数的运行时间是不可忽略的, 故对此类内存拷贝函数, 也需要进行精细的分析和优化. 此外, KERNEL的实现中也有一些可改进之处.

基于上述分析, 我们提出了兼顾减少打包开销和提高KERNEL效率的二层分块方案, 并对打包函数及KERNEL函数的实现都进行了优化. 基于优化之后的程序, 我们进一步对KERNEL的实现在减少访存量和提高数据预取质量方面进行了优化. 下面详细介绍各项优化技术.

2.2.1 二层分块及分块参数确定

算法2中, 函数packA()执行次数是 ${{N/BN \times K/BK \times M/BM}} $ . 单次执行访问的数据量是 ${{BM \times BK \times \textit{sizeof}(double)}}$ . 设机器带宽是bandwidth, 那么单次用时是 ${{BM \times BK\times \textit{sizeof}(double) / bandwidth}}$ , 所以函数packA()总用时为 ${{N \times K \times \textit{sizeof}(double) / bandwidth}}$ . 同理, 函数packB()总用时为 ${{N \times K \times \textit{sizeof}(double) / bandwidth}}$ .

当输入规模固定时, 函数packB()用时只取决于机器带宽; 函数packA()用时反比于分块参数BN, BN越大函数packA()用时越少. 但另一方面, 如果BN过大, 函数KERNEL()访问的子块数据量就会超出一级缓存容量, 导致缓存换入换出频繁, 增加不必要的开销. 为充分利用BN较大时函数packA()用时更少, BN较小时KERNEL()中缓存效率更高的特性, 在调用计算函数时加入了第2层分块, 将程序改为6层循环, 如算法4所示. 外层使用较大分块, 在外层分块中调用函数packA()和函数packB(), 将A, B矩阵中的子块打包放入大小分别为BM×BK和BK×BN的缓冲区 $\tilde{A}\mathrm{和}\tilde{B}$ 中(第3行和第5行). 内层使用较小分块, 3个维度分别为BBM, BBN和BBK. 函数KERNEL()处理BBM×BBN大小的子块, 进行BBK次累加. 因此, packA()和packB()将 $\tilde{A}, \tilde{B}$ 按BBM×BBK和BBK×BBN大小的子块排列, 也就是每个BBM×BBK大小的A子矩阵连续存放. 这样一来, 只要BBM, BBK和BBN足够小, 那么每次调用KERNEL时, A矩阵和B矩阵在缓存中的数据可以得到比较充分的重用. 在我们的实现中, BBM取4个向量的长度, 即KERNEL中一次循环处理的分块M方向大小. A子块驻留于L1缓存中, B子块驻留于L2缓存. 在KERNEL内部, 还有2层循环, 外层为N方向循环, 内层为B方向循环. 由于 KRENEL内部N方向分块大小仅为4 (RN=4), 在KERNEL内部可以保证A子矩阵驻留于L1. 当BBN足够小时, 能保证B子块驻留于L2缓存.

算法4. 二层分块算法.

1. 　 C = βC

2. 　for j = 0 : N – 1 step BN

3. 　　forl= 0 : K – 1 step BK

4. 　　　pack §B_lj into $ \tilde{B} $

5. 　　　for i = 0 : M – 1 step BM

6. 　　　　pack §A_il into $ \tilde{A} $

7.　　　　 for v = 0 : BN – 1 step BBN

8.　　　　　 for w = 0 : BK – 1 step BBK

9. 　　　　　　for u = 0 : BM – 1 step BBM

10. 　　　　　　　调用计算函数KERNEL(), 计算C_ij += $ \alpha \tilde{A} $ _uw $ \times \tilde{B} $ _wv,

11. 　　　　　　end for

12. 　　　　　end for

13. 　　　　end for

14.　　　 end for

15.　　 end for

16.　 end for

2.2.2 设置C矩阵缓冲区

如前所述, 在基础版本中, 对C矩阵未使用缓冲区, 而是在KERNEL中直接写入. 这是因为频繁的解包操作开销过大. 但通过分析和实验, 我们发现, 如果使用更大的缓冲区, 那么就可以将C矩阵的解包放到M方向循环之外, 这样就能保证只对C矩阵解包一次. 这种情况下, 使用缓冲区的性能收益超过C矩阵的解包操作开销. 因此, 我们在第2.1.2节优化的基础上, 对C矩阵也使用了缓冲区. 要保证只对C矩阵解包一次, 需要将解包操作放在K方向最外层循环(算法4第14行和第15行之间), 而M方向的两层循环都在K方向外层循环的内部, 这就意味着C缓冲区 $\tilde{C}$ 的大小应为M×BN. 也就是说, 缓冲区 $\tilde{C}$ 需要在内存中占据的空间达到整个矩阵规模的BN/N. 在M维度很大时, 采用减小BN的方法来避免内存分配失败.

2.2.3 打包函数优化

通过简单的计时统计, 我们发现, 打包操作的时间是不可忽略的, 因此, 对packA和packB操作, 我们对其实现进行了深入分析和优化.

由于KERNEL函数处理的是BBM×BBK大小的A矩阵子块, 因此packA要将矩阵中每个BBM×BBK大小的子块连续存放, 如图1所示, $\tilde{A}$ 中元素按数据块存储, 每一块大小是RM×RK, 块与块之间在行方向上连续, 块内在列方向上连续. 原始的打包操作, 对每个子块分别调用packA, 也就是每次读取矩阵A同一列上的BBM个元素, 并放入缓冲区. 然后再读入下一列的BBM个元素, 存入缓冲区. 这种做法每次读取的连续数据仅有BBM个, 容易造成cache缺失, 在M较大时还会有较高的TLB缺失. 而packA本身要被调用多次, 因此这种cache缺失和TLB缺失造成的开销会被进一步放大. 为了提高打包操作的性能, 我们对改进的循环进行分析, 发现二层分块给了我们改进packA操作的机会. 在这种循环模式下, 我们可以通过使用更大的缓冲区, 将打包函数改为集中式操作. 也就是说, 在分配缓冲区 $\tilde{A}$ 时, 同时分配BM/BBM个缓冲区. 先读取M方向上第1个列向量, 存入缓冲区 $\tilde{A}$ 中相应的位置. 然后读取M方向上下一个列向量. 依次类推, 直到全部BM个元素都已存入缓冲区, 然后开始处理下一列, 如图2所示. 这样, packA()函数可以实现M方向上的连续读取. 此外, 由于读取与写入都是列主序, 也就是读出和写入都是连续数据, 因此可以使用向量化指令完成读写, 进一步提高效率.

类似的, 函数packB()也要将B中数据打包成BBK×BBN大小的子块. 而且, 由于KERNEL中每次迭代处理的是N方向上的RN个元素, 打包时需要将N方向上每RK×RN个元素打包. 在基础版本程序中, RK=2, $\tilde{B}$ 中元素按行优先存储, 打包操作以4×4的小块为单位, 如图1所示. 这种存储方式使得在KERNEL中 $\tilde{B}$ 的元素读取都是连续的, 但它使packB函数的输入和输出使用了不同的存储方式, 因此无法使用向量化操作. 为了提高packB函数的运行效率, 我们将KERNEL内的循环从原来的2次展开改为4次展开, 如图3所示, 不但减少了分支判断语句的开销, 而且在KERNEL中将会连续访问K方向上的4个B元素. 这种访问方式允许我们对 $\tilde{B}$ 重新布局, 将其改为列优先存储. 也就是说, 一列当中(K方向上)的连续4个B元素在打包之后依然连续存储. 采用这种布局, 就可以使用向量化指令来完成packB()的操作.

2.2.4 减少对C矩阵的不必要读写

在本节及第2.2.5节中, 我们介绍对计算KERNEL的优化, 包括减少对C矩阵的不必要读写, 及针对预取操作进行的优化.

在KERNEL计算中, K循环结束时, 需要将矩阵C中的数据读出, 并与A和B的乘积结果相加, 再将加和的结果写回C. 在我们的改进方案中, 由于C矩阵使用了缓冲区, 我们可以在缓冲区 $\tilde{C}$ 中只保存中间结果, 而将与原始C矩阵块加和的操作推迟. 基于这种设计, 我们开发了3个KERNEL函数, 分别用在K方向的不同迭代阶段, 从而尽可能减少对C矩阵的读写.

在K方向第1次迭代中, 缓冲区 $\tilde{C}$ 中的初始值全部是0, 因此KERNEL不需要读取C的值, 而是直接将乘积写入缓冲区 $\tilde{C}$ , 此为第1个版本KERNEL_INIT. 之后的迭代中, C矩阵的读写都发生在缓冲区 $\tilde{C}$ , 即将缓冲区 $\tilde{C}$ 中的数据读出, 加上本次迭代中A与B的乘积, 然后写回 $\tilde{C}$ 缓冲区. 这些迭代中的KERNEL操作即为原始的KERNEL实现, 此处我们命名为KERNEL_GENERIC. 当到达K方向最后一次迭代时, 我们开发了KERNEL_UPDATE, 在此函数中把unpackC融合进来, 读取原始矩阵中C矩阵的值, 与本次迭代的乘积及 $\tilde{C}$ 中的值加和, 并写回原始C矩阵, 也就是完成C:= α $\tilde{A}$ × $\tilde{B}$ + β $\tilde{C}$ 的操作. 使用3个KERNEL后, 算法4第9行计算函数调用逻辑如算法5所示. 此方案中, C与β相乘的操作也被融合到KERNEL_UPDATE, 进一步减少了对C的读写.

算法5. 在K方向不同阶段使用3个版本的KERNEL函数.

1. 　if $\tilde{C}$ 上第1次K方向迭代

2. 　　KERNEL_INIT();

3. 　else if $ \tilde{C} $ 上最后一次K方向迭代

4. 　　KERNEL_UPDATE();

5. 　else

6. 　　KERNEL_GENERIC ();

7. 　end if

2.2.5 针对预取操作的优化

原始版本的KERNEL虽然对寄存器使用等进行了精细调优^[9], 但我们发现, 访存操作仍有可优化空间. 此外, 二层分块的结构也使KERNEL的访存行为有了变化. 基于对KERNEL访存特征的分析, 我们对KERNEL进行了以下两方面的改进.

首先, 原有KERNEL中对C矩阵中每一小块数据(RM×RN个浮点数)的读取都发生在一次迭代的写回阶段, 即K方向全部迭代之后(算法3第7行), 与前面循环体中的计算没有任何重叠. 实际上, 我们可以在不破坏寄存器使用依赖关系的前提下, 对C中的数据提前读入, 以充分利用指令级并行, 减少写回操作时对C数据的等待. 为实现此方案, 我们将K方向最后一次迭代从循环中剥离, 也就是将其放置于循环体之外, 然后在其中插入若干个对C矩阵元素的读取指令. 为了保证寄存器重用的安全性, 我们选取了9个C中的元素, 使它们的读取与K方向最后一次迭代的计算重合, 并在插入对矩阵元素C_i的读取指令时, 确保C_i使用的寄存器R_i在此读取指令和写回部分之间不被使用.

其次, 在原始程序中, 由于对A的预取发生在算法3第3到6行的循环中, 因此在每次迭代中对A的预取都使用同样的指令(即预取距离当前A指针同样步长的数据). 实际上, 每次K方向循环开始时, A都要被从头重新读取. 所以在K方向的最后几次迭代读取的是无效的数据, 而K方向起始的几次迭代中所使用的矩阵A数据却没有被预取. 为了解决这个问题, 我们通过插入新的预取指令对A的访存做了进一步优化. 如上一段所述, 我们已将K方向最后一次迭代置于循环体之外. 在这段代码中, 我们删除了原有的对矩阵A在K方向下一块数据的预取指令, 改为预取当前A子块(算法4第9行中的 $ \tilde{A}_{uw} $ )起始部分的数据.

矩阵B的预取也有类似的问题, 在每次KERNEL调用的末尾, 使用同样的预取指令, 读取的是下一块矩阵B的值. 而实际上, 只要不是N方向最后一次迭代, 就无需预取下一B子矩阵中的数据. 基于这种考虑, 我们将每个KERNEL进一步分为两个. 第1个KERNEL处理FBN列, 第2个KERNEL处理LBN=BBN–FBN列. 在第2个KERNEL中, 插入对当前B子矩阵起始部分数据的预取指令. 预取指令放在K方向循环体中. 如前所述, 改进的KERNEL中K方向是4次展开, 也就是每4次迭代, 进行一次矩阵B的预取. 两次KERNEL调用使用同样的矩阵B预取频率. FBN和LBN的大小决定了预取指令的间隔. 当LBN减少时, 预取距离不够; 当LBN增加时, 预取的数据可能在二级缓存中再次被换出, 都会使性能都有所降低. 我们通过自动调优选取最佳的KERNEL切割点, 即FBN和LBN的值.

2.3 多线程优化技术

对多线程矩阵乘法, 我们采用了与GotoBLAS相同的并行方案^[11], 即在M方向上划分任务分配给各个线程, 如图4所示. 原问题在N方向上分块后, 对C矩阵在N方向上每一块(即每个列块), 进一步在M方向上进行划分, 分配给各个线程并行计算. 这样划分任务后, 每个线程计算子数据块时, 各自使用无重叠区域的A矩阵子块, 线程间没有数据依赖; 多个线程使用相同的B矩阵子块, B矩阵子块的打包任务被划分到多线程共同执行, 打包后的缓冲区中B子矩阵由多个线程共同使用.

单线程GEMM的优化技术均适用于多线程DGEMM. 区别在于, 多线程GEMM在部署分块方案时, N方向需要在二层分块的基础上, 增加一层线程级分块.

3 子函数代码生成器及自动调优

在实现过程中, 我们开发了KERNEL和打包函数的代码生成器, 使用Python语言编写代码生成脚本. 通过调整参数, 该脚本可以生成不同版本的KERNEL函数和打包函数代码.

KERNEL函数的代码生成器将预取步长, 循环展开次数等作为参数传递, 可以生成使用不同优化策略和参数的KERNEL汇编代码. 对A, B, C这3个矩阵, 我们分别设置了预取步长FS_A, FS_B, 和FS_C. 由于C矩阵在K方向循环末尾使用, 因此FS_C设为0, 即K方向循环开始之前便预取本次M, N方向所需C子矩阵块. FS_A和FS_B通过遍历搜索选取. 理论上讲, FS_A的值应根据指令数与访存带宽计算. 设从主存装载数据到L1 cache需要W个时钟周期, K方向一次循环执行M条指令, 每时钟周期发射指令数为IPC, 那么预取应该在W×P/M迭代之后, 因此FS_A应等于W×IPC/M×16×sizeof(double). 我们根据以上推导, 确定出FS_A的大致初始值范围, 参数搜索将在初始值的周围进行遍历. FS_B的初始值可用类似方式确定. 打包函数使用C语言实现, 并采用内嵌汇编实现向量化的数据读写.

为了选取最优参数, 我们编写了自动调优脚本, 对参数空间进行遍历搜索. 需要调优的参数包括第1层分块和第2层分块的分块大小(BM, BN, BK, BBM, BBN, BBK), KERNEL中的预取步长(FS_A, FS_B), 及第2.2.5节所述的N方向KERNEL切分位置(FBN)等.

经过实验遍历搜索后得出性能最佳的分块和预取参数组合如表1和表2所示. 第2层分块参数BBM的值和寄存器层分块参数RM相同, 即KERNEL内部在M方向完全展开, 不做循环处理. 由于本文所做优化主要面向较大规模的矩阵, 因此搜索参数空间测试采用的是N=16384的矩阵. 当矩阵规模较小, 或者呈现极不规则的形状时, 由于GEMM操作的计算/访存占比相较大规模矩阵的情形已经出现了较为明显的变化^[21], 本文提出的策略不再适用于此类问题的优化. 此类小规模及不规则形状矩阵乘法的优化不在本文讨论的范围之内.

对多线程GEMM, 我们采用类似的方法, 对分块参数进行遍历搜索, 得出的多线程GEMM最优分块参数如表3所示. 由于L1和L2缓存为各核心私有, 故KERNEL内部最优预取步长与单线程相同. 而L3缓存是各核心共用的, 在外层循环中, 多线程环境下每次迭代中同时处理的数据量增加, 因此容易产生更多的L3缓存缺失及更多的缺页. 所以经遍历选得的最优分块大小小于单线程外层分块大小.

4 实验结果及分析

本节介绍我们在申威1621处理器上对双精度浮点实数矩阵乘法进行的测试, 包括单线程DGEMM使用各项优化技术后的性能, 以及优化的多线程DGEMM程序性能.

4.1 实验设置

实验采用的申威1621处理器有16个计算核心, 主频1.6 GHz, 内存128 GB. 单核双精度浮点运算峰值25.6 Gflops, 16核双精度浮点运算峰值409.6 Gflops. 三级缓存的大小分别是32 KB, 512 KB, 512 MB.

对于单线程DGEMM, 我们选择了7个版本进行了测试. 第1个版本Gemm_Goto使用稍加改动的GotoBLAS代码(GotoBLAS原始代码在申威处理器上不能直接运行). Gemm_baseline采用文献[9]中所述方案及实现, 作为基础版本. Gemm_op1–Gemm_opt5是5个优化版本, 每一版本都在上一个版本的基础上增加一种新的优化技术, 详细描述如后文表4所示. 其中, Gemm_opt1在Gemm_baseline的基础上加入了二层分块, Gemm_opt2引入了C矩阵缓冲区. Gemm_opt3使用改进的打包函数, Gemm_opt4使用了不同的KERNEL, 而Gemm_opt5包含了前文提到的所有优化技术. 实验中, 我们选取了4种输入规模, 均为方阵, 即M, N, K这3个维度相等, 下文为叙述方便用一个维度数字代表所有维度的大小. 在实验中, 我们发现矩阵的主维大小(即矩阵第1列长度)会影响内存到缓存映射关系, 进而影响访存性能, 尤其是矩阵C的主维LDC. 若矩阵C的LDC为256的倍数, 会造成较多的cache miss, 性能会有所下降. 因此我们将LDC的取值设为M加上一定的偏移量(16的倍数).

4.2 实验结果

本节我们将展示在使用了上文所述的优化方法后, 单线程和多线程GEMM的性能提升.

4.2.1 优化方法对单线程GEMM性能的影响

单线程DGEMM实验结果见图5. 与基础版本相比, 使用二层分块带来8.6%的性能提升. 对C矩阵使用缓冲区使性能又提高了10%. 对A和B打包函数的改进将性能提升了13%. Gemm_opt4和Gemm_opt5进一步增加了对KERNEL函数的改进, 分别得到了3.1%和5.3%的性能提升. 与GotoBLAS相比, 我们的最优版本实现了平均7.39倍的加速, 在规模较大(M=16384)时, 性能达到单核浮点计算峰值的85%.

4.2.2 多线程实验结果

多线程DGEMM实验在两组数据规模上进行, 都采用M, N, K相同的方阵, 大小为32768, 分别在1, 2, 4, 8, 16个线程上进行测试. 结果见表5. 从表中可以看出, 当线程数目小于等于8时, 扩展性效率超过95%, CPU效率超过82%; 线程数目为16时, 可扩展性和CPU效率稍有下降. 我们分析, 16线程时效率下降有两个可能的原因. 首先, 多线程DGEMM包含线程间同步的开销, 子函数计时结果(表6)数据说明, 当线程数目达到16时, 同步开销已不能完全忽略. 其次, 使用16线程时, 线程间对访存带宽及三级缓存资源的使用会产生更多竞争, 导致KERNEL函数用时比同样规模在单线程上的运行时间略长.

5 结　论

本文面向国产申威多核平台对稠密矩阵乘DGEMM单线程算法和多线程版的并行算法进行了重构与优化. 通过改变分块方案和打包方式等优化手段, 使程序运行最大化地减少额外开销, 同时又充分利用了数据局部性. 对KERNEL函数, 我们对函数在K方向迭代的不同阶段所读写的数据进行分析, 为每个阶段定制KERNEL, 并对迭代中每个部分进行精准的数据预取. 使用我们的优化方法之后, 实验数据显示, 单线程版性能达到浮点峰值性能的85%. 多线程版性能达到浮点峰值性能的80%. 这些优化技术对于典型分层存储架构和现代微结构上的应用优化都有参考价值. 此外, 我们开发了代码生成器和自动调优脚本, 为优化参数搜索和高效调优提供了实用工具. 在未来工作中, 我们准备将对GEMM函数的优化应用到其他BLAS函数中, 并对小规模矩阵乘法优化进行研究. 另外, 我们将对各种优化方法进行提炼, 进一步完善代码生成器, 增加可移植性, 形成较完整的优化工具链.