Tackling the Matrix Multiplication Micro-kernel Generation with EXO
Tackling the Matrix Multiplication Micro-kernel Generation with EXO
- 文章标题:利用 EXO 解决矩阵乘法微内核生成问题
- 作者/机构:Adrian Castelló (瓦伦西亚理工大学), Julian Bellavita (康奈尔大学), Grace Dinh (加州大学伯克利分校), Yuka Ikarashi (麻省理工学院 CSAIL), Hector Martínez (科尔多瓦大学)
A1 主要贡献
本文旨在解决现有线性代数库(如 BLIS, OpenBLAS)中通用矩阵乘法(GEMM)实现所面临的若干限制。这些库的性能高度依赖于针对特定硬件手动编写和优化的微内核(micro-kernel),导致开发工作量大、可移植性差、难以维护,并且可能无法高效处理深度学习(DL)等领域中出现的非方形矩阵等特殊情况。
为应对这些挑战,本文展示了如何使用 EXO 编译器自动生成一系列针对 GEMM 的微内核,并提出以下核心优势与贡献:
- 替代方案:通过 EXO 生成的一系列 C 代码微内核集合,来替代传统库中单一、庞大的微内核。每个生成的微内核可以专门处理一个特定的边界情况(edge case)。
- 提升可移植性与可维护性:通过使用适当的后端,代码生成和优化过程可以轻松地针对不同的数据类型进行特化,从而增强了解决方案的可移植性和可维护性。
- 性能超越:通过将微内核的尺寸与具体问题相匹配,生成的代码能够超越现有广受认可的高性能 GEMM 实现。
- 简化优化流程:针对每个问题的优化过程被大大简化,归结为评估少量生成的微内核。
- 开源工具:本文提供了一个公开可用的 ARM Neon 微内核生成器。
- 展示 EXO 的可用性:本文证明了自动生成的 C 代码在性能上可以媲美手动编写的解决方案,特别是通过为不同问题尺寸生成不同代码的方式。
- 对 EXO 的贡献:本研究工作为 EXO 编译器贡献了对两种 Neon 内在指令(intrinsic instructions)和 ARM 16位浮点数据类型的支持。
A3 背景知识/关键Observation/设计原则
A. BLIS 的 GEMM 算法
GEMM 操作与宏内核结构。考虑 GEMM 操作 C = C + AB,其中矩阵 A, B, C 的维度分别为 m x k, k x n 和 m x n。BLIS 框架及其他线性代数库遵循 GotoBLAS【1, K. Goto and R. A. van de Geijn, “Anatomy of a high-performance matrix multiplication,” ACM Trans. Math. Softw., vol. 34, no. 3, pp. 12:1–12:25, May 2008.】的方法,将此操作编码为围绕两个数据打包(packing)例程的三层嵌套循环。这种结构被称为宏内核(macro-kernel),其代码通常在不同硬件架构间共享。在 BLIS 中,宏内核被进一步分解为围绕一个微内核的两个额外循环,而微内核则包含一个单循环,每次迭代执行一次外积操作。图1和图2展示了 BLIS 的基础 GEMM 算法,包括六层循环、两个打包例程和微内核。
缓存层次结构优化。算法的最外三层循环遍历问题的 n、k 和 m 维度,通过对矩阵操作数进行分块,使其适应处理器的缓存层次结构。根据目标处理器架构调整的缓存配置参数 nc、kc 和 mc【9, T. M. Low, F. D. Igual, T. M. Smith, and E. S. Quintana-Ort´ı, “Analytical modeling is enough for high-performance BLIS,” ACM Trans. Math. Softw., vol. 43, no. 2, pp. 12:1–12:18, Aug. 2016.】,使得在微内核执行期间,Bc 保持在 L3 缓存,Ac 保持在 L2 缓存,而 C 则从主内存流式传输到处理器寄存器中,如图2所示。
数据打包的作用。此外,宏内核的数据打包例程确保了微内核在访问 Ac 和 Bc 中的数据时能够以单位步长(unit stride)进行,从而提高访存效率。
微内核的特性与局限。微内核是面向硬件的代码片段,通常使用汇编语言编码。在微内核内部,一个 mr x nr 大小的 C 矩阵块在 kc 循环内被更新。用于更新的操作数是 Ac 和 Bc 缓冲区的一部分,即维度分别为 mr x kc 和 kc x nr 的 Ar 和 Br。mr 和 nr 的值通常被用来命名微内核。BLIS 为 AMD、Intel、ARM 和 IBM 的许多处理器提供了专门的微内核。然而,为每个新架构开发一个新的面向硬件的微内核是一项昂贵的开发工作,因此 BLIS 对每种架构采用单一化的方法,即提供一个微内核。
B. 编译与调优框架
自动代码生成的动机。为不同架构开发优化的微内核之所以成本高昂,主要是因为需要执行架构特定的优化以获得最高性能,并生成平台特定的指令(如硬件 intrinsics)。为了解决这个问题,自动代码生成方法被应用于实现跨现有和新兴架构(包括通用处理器或专用加速器【10, T. Moreau, T. Chen, Z. Jiang, L. Ceze, C. Guestrin, and A. Krishnamurthy, “VTA: an open hardware-software stack for deep learning,” CoRR, vol. abs/1807.04188, 2018. http://arxiv.org/abs/1807.04188】 【11, H. Genc, S. Kim, A. Amid, A. Haj-Ali, V. Iyer, P. Prakash, J. Zhao, D. Grubb, H. Liew, H. Mao, A. Ou, C. Schmidt, S. Steffl, J. Wright, I. Stoica, J. Ragan-Kelley, K. Asanovic, B. Nikolic, and Y. S. Shao, “Gemmini: Enabling systematic deep-learning architecture evaluation via full-stack integration,” in Proceedings of the 58th Annual Design Automation Conference (DAC), 2021.】)的性能可移植性,且只需最少的程序员干预【12, M. Li, Y. Liu, X. Liu, Q. Sun, X. You, H. Yang, Z. Luan, and D. Qian, “The deep learning compiler: A comprehensive survey,” CoRR, vol. abs/2002.03794, 2020. https://arxiv.org/abs/2002.03794】 。
用户可调度语言与编译器框架。用户可调度语言和编译器框架,如 Halide【13, J. Ragan-Kelley, C. Barnes, A. Adams, S. Paris, F. Durand, and S. Amarasinghe, “Halide: A language and compiler for optimizing parallelism, locality, and recomputation in image processing pipelines,” in Proceedings of the 34th ACM SIGPLAN Conference on Programming Language Design and Implementation, ser. PLDI ’13. New York, NY, USA: Association for Computing Machinery, 2013, p. 519–530. https://doi.org/10.1145/2491956.2462176】、MLIR 【14, C. Lattner, J. A. Pienaar, M. Amini, U. Bondhugula, R. Riddle, A. Cohen, T. Shpeisman, A. Davis, N. Vasilache, and O. Zinenko, “MLIR: A compiler infrastructure for the end of moore’s law,” CoRR, vol. abs/2002.11054, 2020. https://arxiv.org/abs/2002.11054】、TVM 【15, T. Chen, T. Moreau, Z. Jiang, H. Shen, E. Q. Yan, L. Wang, Y. Hu, L. Ceze, C. Guestrin, and A. Krishnamurthy, “TVM: end-to-end optimization stack for deep learning,” CoRR, vol. abs/1802.04799, 2018. http://arxiv.org/abs/1802.04799】 或 EXO【8, Y. Ikarashi, G. L. Bernstein, A. Reinking, H. Genc, and J. Ragan-Kelley, “Exocompilation for productive programming of hardware accelerators,” in Proceedings of the 43rd ACM SIGPLAN International Conference on Programming Language Design and Implementation, ser. PLDI 2022. New York, NY, USA: Association for Computing Machinery, 2022, p. 703–718. https://doi.org/10.1145/3519939.3523446】,提出了一种清晰的关注点分离思想,即分离操作的定义(如矩阵乘法)和应用于该操作的调度(或优化集)。这使得通用优化(如针对高层内存层次结构的循环分块)可以从架构特定的优化(如向量化)中抽象出来 。
硬件规范的集成方式。编译器框架还必须获得描述硬件指令集(包括向量 intrinsics)的硬件规范。传统编译器如 Halide、TVM 和 LLVM 将硬件规范紧密集成到编译器中,需要手动指定代码生成过程以支持新硬件。因此,有时生成工具会绑定到特定的编译器,导致用户无法控制一些代码决策。例如,TVM 和 LLVM 是绑定的,用户无需编译代码,但同时也对某些方面不知情。
EXO 的独特之处。另一方面,EXO【8, Y. Ikarashi, G. L. Bernstein, A. Reinking, H. Genc, and J. Ragan-Kelley, “Exocompilation for productive programming of hardware accelerators,” in Proceedings of the 43rd ACM SIGPLAN International Conference on Programming Language Design and Implementation, ser. PLDI 2022. New York, NY, USA: Association for Computing Machinery, 2022, p. 703–718. https://doi.org/10.1145/3519939.3523446】通过将硬 件 intrinsics 的定义外部化,以用户定义的库作为输入。例如,图3提供了 ARM Neon 内在函数 vstlq_f32 和 vfmaq_laneq_f32 的语义等效 Python 定义;编译器使用这些定义来自动生成调用。一个完整的硬件规范由每个硬件 intrinsic 的类似定义,以及内存层次结构(本例中为 DRAM 和 Neon 寄存器)和数据类型(本例中为 f32 浮点数)的定义组成。具体来说,这些定义将通过检查 intrinsic 替换与预期模式是否匹配,来确保用户方法不会改变原始代码的行为。没有这种安全定义,用户可能会用任何 intrinsic 指令替换任何循环,从而产生不同的代码。EXO 的另一个关键方面是它独立于任何编译器/优化器工具。具体而言,EXO “仅”生成带有 intrinsic 指令的 C 代码,该代码需要被编译,因此用户可以尝试不同的硬件/编译器组合,以获得生成代码的最大性能。
框架依赖性问题。其中一些框架的一个严重缺点是依赖现有库。例如,TVM 生成的代码使用一组 TVM 对象,这可能迫使用户重写大部分(或全部)软件栈,此外还会因将数据类型转换为这些库支持的类型而产生运行时开销。然而,像 EXO 这样的工具生成的是纯 C 代码,可以用于所有现有的面向性能的库中。
C. 自动调优与优化
自动调优的挑战。通过将调度视为输入,用户可调度语言促进了通过自动调优技术【16, R. T. Mullapudi, A. Adams, D. Sharlet, J. Ragan-Kelley, and K. Fatahalian, “Automatically scheduling halide image processing pipelines,” ACM Trans. Graph., vol. 35, no. 4, jul 2016. https://doi.org/10.1145/2897824.2925952】 【17, L. Zheng, C. Jia, M. Sun, Z. Wu, C. H. Yu, A. Haj-Ali, Y. Wang, J. Yang, D. Zhuo, K. Sen, J. E. Gonzalez, and I. Stoica, “Ansor: Generating High-Performance Tensor Programs for Deep Learning,” arXiv, Tech. Rep., arXiv:2006.06762 [cs, stat] type: article. http://arxiv.org/abs/2006.06762】对调度空间的自动探索。例如,作 为 TVM 一部分的 AutoTVM【18, T. Chen, L. Zheng, E. Q. Yan, Z. Jiang, T. Moreau, L. Ceze, C. Guestrin, and A. Krishnamurthy, “Learning to optimize tensor programs,” CoRR, vol. abs/1805.08166, 2018. http://arxiv.org/abs/1805.08166】对用户指定模板定义的搜索空间进行全面探索。虽然这类搜索方法在实践中表现出强大的性能,但它们必须应对一个计算成本高昂的可能调优参数搜索空间,该空间随设计空间维度的增加呈指数级增长,并且在泛化(对不同问题规模和不同硬件目标)和对开发者的可解释性方面存在困难 。
基于模型的优化方法。最近,作为 BLIS 项目一部分的工作【9, T. M. Low, F. D. Igual, T. M. Smith, and E. S. Quintana-Ort´ı, “Analytical modeling is enough for high-performance BLIS,” ACM Trans. Math. Softw., vol. 43, no. 2, pp. 12:1–12:18, Aug. 2016.】表明,使用分析模型选择最优配置参数是实现高性能的有效方法,无需进行自动调优。用基于模型的解决方案替代自动调优方案也已在【2, Z. Xianyi, W. Qian, and Z. Yunquan, “Model-driven level 3 BLAS performance optimization on Loongson 3A processor,” in 2012 IEEE 18th International Conference on Parallel and Distributed Systems (ICPADS), 2012.】、【19, K. Yotov, X. Li, M. J. Garzaran, D. Padua, K. Pingali, and P. Stodghill, ´ “Is search really necessary to generate high-performance BLAS?” Proceedings of the IEEE, special issue on “Program Generation, Optimization, and Adaptation”, vol. 93, no. 2, 2005.】、【20, A. Olivry, G. Iooss, N. Tollenaere, A. Rountev, P. Sadayappan, and F. Rastello, “IOOpt: automatic derivation of i/o complexity bounds for affine programs,” in Proceedings of the 42nd ACM SIGPLAN International Conference on Programming Language Design and Implementation. ACM, jun 2021. https://doi.org/10.1145/3453483.3454103】和 【21, Q. Huang, M. Kang, G. Dinh, T. Norell, A. Kalaiah, J. Demmel, J. Wawrzynek, and Y. S. Shao, “Cosa: Scheduling by constrained optimization for spatial accelerators,” in 2021 ACM/IEEE 48th Annual International Symposium on Computer Architecture (ISCA). IEEE, 2021, pp. 554–566.】中成功探索。
本研究的定位。本工作专注于微内核生成,与使用 MLIR 描述整个 GEMM 算法早期经验的【22, U. Bondhugula, “High performance code generation in MLIR: an early case study with GEMM,” CoRR, vol. abs/2003.00532, 2020. https://arxiv.org/abs/2003.00532】、提出该原语高级自动调优方案的 【23, Y. Zhang, Parallel solution of integral equation-based EM problems in the frequency domain. IEEE Press, 2009.】以及使用 Python 脚本和 C 宏构建微内核的【24, G. Alaejos, A. Castello, H. Mart´ınez, P. Alonso-Jorda, F. D. Igual, and ´ E. S. Quintana-Ort´ı, “Micro-kernels for portable and efficient matrix multiplication in deep learning,” The Journal of Supercomputing, pp. 1–24, 2022.】不同。具体来说,我们使用 EXO 来扩展和进一步分析 GEMM 微内核的手动生成过程,并将生成的代码集成到类 BLIS 的 GEMM 算法中。
A2 方法细节
本节将分步解释如何从零开始使用 EXO 为 ARMv8 架构构建一个高性能计算(HPC)微内核代码。微内核的维度将是 8 × 12,与 BLIS 库中针对该特定架构的微内核相同。对于每个构建阶段,我们首先介绍所使用的 EXO 指令,并解释生成的中间代码。我们还将展示如何扩展代码生成以处理微内核的其他特性。分步生成的代码可在 https://github.com/adcastel/EXO_ukr_generator 获取。
A. 微内核生成
初始微内核代码与调整。图4展示了一个基于外积(如 BLIS 中)的列主序微内核的初始代码。为了满足 BLIS GEMM 算法的特性,我们对初始微内核代码进行了如下修改:1) 由于 C 语言是行主序分配器,我们转置了 C 矩阵的维度;2) BLIS 对微内核的 A 和 B 操作数使用数据打包,这保证了数据的单位步长访问。因此,为了在 Ac 操作数上实现这一点,我们也交换了 Ac 的维度。Bc 操作数已经是单位步长访问,因此无需更改;3) 我们将内部循环重新排列为 k, j, i 的顺序,以匹配期望的结构。注意,Ac 和 Bc 操作数在调用微内核的算法中是在一维结构中分配的,因此即使我们转置了数据缓冲区 C 和 Ac,也不会有访问模式的风险。
初始 EXO 代码详解。首先解释图4中的代码,第一行的 @proc 告诉 EXO 编译器下一个函数是可调度的,并将因此产生 C 语言代码。函数的参数也展示了 EXO 语言的一些方面。首先,参数名后跟“:”和类型,可以是 size、scalar、vector 或 matrix。对于需要内存分配的数据,我们应该指定其位置。在此示例中,我们使用 @DRAM 注解将 Ac、Bc 和 C 缓冲区映射到 RAM。此版本涵盖了 alpha 和 beta 值的所有组合。位于第8行和第9行的缓冲区 Cb 和 Ba 用于计算 C * beta 和 Bc * alpha。具体来说,第12-14行计算 Cb 的结果,第17-19行计算 Ba 的结果。第22-25行通过执行 Cb = Cb + Ac * Ba 来计算微内核的结果。最后,第28-30行将计算结果返回到 C 矩阵。
简化版微内核。为简化起见,从现在开始,我们将优化一个特定版本的微内核。具体来说,我们将对图5所示的代码进行分步转换,该代码对应于 alpha 和 beta 值均为1时的微内核。对初始代码的优化将涉及更多针对 Cb 和 Ba 循环(第11-14、17-19和27-30行)的调度函数,这些函数与从此处开始展示的函数等效。
a) 基础微内核
特化微内核维度。首先,我们生成要调度的函数,并通过指定我们希望为 MR 和 NR 参数分别使用值8和12来特化生成的代码。这种转换通过 partial_eval 函数完成,该函数替换了 MR 和 NR 的值。图6显示了用户代码和生成的表示。第3行和第4行获取微内核的初始版本,并分别设置 MR 和 NR 变量。生成的代码已将变量更改为其值(例如,第二个循环的迭代现在从0到12,而不是 NR)。
b) 循环结构
划分循环以匹配向量长度。在图7的第2行和第3行,我们分割了 i 和 j 循环,以匹配架构的向量长度。由于基于 ARM 的 NVIDIA Carmel 处理器使用128位向量寄存器,并且使用 float32 数据类型,向量长度为4。此外,本工作中使用的类 BLIS GEMM 算法确保 Ac 和 Bc 缓冲区大小分别是 MR 和 NR 值的倍数。这一操作导致了位于第14-17行的嵌套循环 it、itt、jt 和 jtt。此外,EXO 自动对 C、Ac 和 Bc 数据的访问进行了分块,以匹配新的循环结构(第18-20行)。
c) C 矩阵
将 C 矩阵绑定到向量寄存器。图8显示了微内核生成中最复杂的步骤之一,即将 C 矩阵绑定到向量寄存器,这包括:声明、加载和存储结果。具体过程如下:
- 映射内存到寄存器:使用
stage_mem函数将 C 矩阵的内存与一个向量寄存器(C_reg)绑定(第3、4行),这反映在第51行。这使得 EXO 可以设置 C 操作数的数据移动。 - 调整寄存器维度:使用
expand_dim方法将向量寄存器调整为一个3D结构,每个维度对应于每个循环的迭代次数。第一次调用(第7行)将声明大小调整为向量长度(本例中为4);第8行完成了MR维度的大小;第9行绑定到 C 的NR维度。这些行的结果可以在第34行看到,其中出现了最终的分配。 - 提升分配位置:
lift_alloc语句将 C 的寄存器声明移动到生成代码的顶部。 - 分离加载和存储:第15-18行使用
autofission函数将 C 矩阵的加载和存储操作移出计算循环(分别为第37-43行和第56-62行)。 - 替换为内在函数:第21行和第22行使用
replace函数将加载和存储的itt循环替换为内在指令。第40行和第59行突出了这些替换。 - 设置寄存器类型:第25行使用
set_memory函数将 C 寄存器变量的类型设置为 Neon。
d) Ac 和 Bc 操作数
生成 Ac 和 Bc 的加载代码。图9列出了从 Ac 和 Bc 加载到寄存器的操作。注意,为简化起见,代码使用名称 Xc,因为这些操作必须对两个操作数都执行。每个操作数的过程如下:
- 映射内存到寄存器:将
Xc矩阵映射到一个向量寄存器(第3行)。 - 调整寄存器维度:将向量寄存器调整为一个2D结构,其中第一个维度是向量长度,第二个维度是最外层循环(第6、7行)。这些行的结果可以在第32行和第33行看到,其中出现了最终的分配。
- 提升分配位置:
lift_alloc语句将寄存器的声明移到 k-循环之外。 - 移动加载操作:第13行和第14行将操作数的加载移动到 k-循环内。
- 替换为向量加载指令:第17行将
xtt循环替换为 neon 向量加载指令。 - 设置寄存器类型:第20行将寄存器变量设置为 Neon 类型。
e) GEMM 操作
优化计算核心。图10说明了这一步骤。我们重新排序了计算的 jtt 和 it 循环(第2行),以便对 B 寄存器值的访问是顺序的。第3行将最内层循环替换为 fmla 语句,如第24-26行所示。
f) 循环展开
手动循环展开。虽然这是某些编译器默认使用的技术,但在 EXO 中也可以手动完成。图11显示了一个将 Ac 和 Bc 操作数加载到寄存器的循环展开示例。我们通过第2行和第3行的 unroll_loop 语句来展示它,结果如第21-25行所示。
g) 生成的 C 代码
验证生成的汇编代码。为了确保生成的代码不仅在 C 语言层面得到优化,而且到汇编的编译也正确完成,我们用 gcc-10 -S 命令编译了 c 代码,k-循环的最终汇编代码如图12所示。这个输出与 BLIS 中手工实现的汇编代码一样优化。
B. 边界情况(Edge Cases)
单一微内核方法的性能问题。HPC 库采用的每个架构一个微内核的方法存在一个问题,即当微内核的维度与优化维度不匹配时,性能会下降。这种情况被称为边界情况。像 BLIS 这样的解决方案对这些情况使用非专用版本的微内核,因为对于大的问题规模,边界情况没有性能影响。然而,像 DL 这样的新兴 HPC 场景充满了这些边界情况。一个明显的例子是 ResNet50-v1.5 卷积模型第一层的尺寸,在应用 IM2ROW 方法后,得到的 GEMM 维度 M、N 和 K 分别为 12544、64 和 147。
EXO 的解决方案。使用 EXO,并假设 BLIS 的打包存在于 GEMM 算法中,我们所需要做的就是使用图6中所示的代码,并更改 MR 和 NR 的值以匹配边界情况。然后,运行所有步骤将生成一个与这些值匹配的新微内核。
无数据打包的场景。然而,可能我们不需要打包,因为数据已经打包好了,或者问题规模足够小,打包的成本不值得。在这种情况下,我们应该调整微内核生成的过程(例如,不对 A 进行打包)如下:
- 图7中的循环 i 不应被分割。
Ac和A_reg之间的映射将改变,后者的维度将匹配MR的值。- 在 k-循环内部,
A_reg将用Ac的值进行广播。 - 计算将使用
neon_vfmadd_4xf32_4xf32,它计算广播的A_reg与整个B_reg的乘积。
C. 架构可移植性
简化架构移植。EXO 只需要更改用户代码中 replace 语句的第三个参数即可生成所需的代码。如果新架构提供了具有相同功能的 API,这就是所需的更改。
处理 API 差异。然而,本示例中使用的某个语句可能在其他硬件的内在函数应用程序编程接口(API)中不存在(例如,ARM Neon vfmaq_laneq_f32)。这时,可以采用与数据非打包可用时类似的方法。
移植示例。举一个简单的例子,将图8中第21行的 replace(p, 'for itt in _: _', neon_vld_4xf32) 改为 replace(p, 'for itt in _: _', _mm512_loadu_ps),将把加载内在函数从 ARM Neon 更改为 Intel AVX512。
D. 数据类型
支持不同数据类型。为不同数据类型生成微内核非常简单,只需对代码中的每个内存分配和寄存器使用 set_precision 函数即可。例如,set_precision(p, A_reg, "f16") 将为 A_reg 变量使用16位浮点寄存器。此外,set_memory() 语句中的 Neon 参数必须更改为 Neon8f。
A4 实验环境
- 硬件配置:实验在 NVIDIA Jetson AGX Xavier 板上进行,使用其嵌入的 NVIDIA Carmel 处理器(ARM v8.2)的单个核心。
- 软件配置:
- 对比了三种微内核实现:
- Neon:一个使用 Neon intrinsics 手动开发的微内核。
- BLIS:BLIS v.0.9 库中的微内核(汇编实现)。
- EXO:本文第三节介绍的自动生成的代码。
- 使用了类 BLIS 的 GEMM 算法框架来调用这些微内核。
- 对比了三种微内核实现:
- 实验参数:所有实验均使用 IEEE 32 位浮点算术(FP32)。
- 实验场景:
- Solo Mode:单独运行微内核,评估其峰值性能及在不同维度(边界情况)下的表现。
- Squared Matrices:评估在大型方阵上的完整 GEMM 算法性能。
- Rectangular Matrices:评估在高度“矩形”问题上的性能,这些问题来自于深度学习模型(ResNet50 v1.5, VGG16)的卷积层经过
IM2ROW转换后的 GEMM 维度。
A4 实验结果
A. Solo 模式
- 实验内容:直接运行微内核5秒钟并计算 GFLOPS,以评估其峰值性能。对于 8x12 的标准情况,Kc 设置为512。同时测试了多种其他维度(如 4x4, 4x8 等)作为边界情况。
- 实验结果(图13):
- 在优化的 8x12 维度下,三种实现性能相近。NEON(基于 intrinsics)略慢于 BLIS(基于汇编)。EXO 性能稍好,因其代码仅为 8x12 情况特化,不包含处理边界情况的逻辑。
- 在其他边界情况下,NEON 和 BLIS 使用相同的微内核,性能下降明显。而 EXO 为每种尺寸都生成了专用的微内核,因此在所有边界情况下都获得了最佳性能。
- 分析结论:EXO 通过生成专用微内核的方式,能有效解决传统单一微内核在边界情况下的性能瓶颈问题。
B. 方形矩阵
- 实验内容:在不同尺寸(1000x1000 到 5000x5000)的方形矩阵上执行完整的 GEMM 算法。
ALG+前缀表示使用类 BLIS 的算法框架,BLIS列表示直接调用官方 BLIS 库函数。 - 实验结果(图14):
- 官方 BLIS 库的性能最佳,因为它在微内核内部实现了数据预取,而
ALG+BLIS的组合方式未使用该特性。 ALG+EXO的性能优于其他ALG+组合。
- 官方 BLIS 库的性能最佳,因为它在微内核内部实现了数据预取,而
- 分析结论:考虑到所有
ALG+组合使用了相同的打包策略(基于模型),性能差异主要源于微内核。EXO 通过为不同问题尺寸选择不同的特化微内核(例如,为1000和4000尺寸使用8x4微内核,为2000和5000尺寸使用8x8微内核),证明了使用微内核集合优于单一微内核的策略。
C. 矩形矩阵(深度学习模型)
- 实验内容:使用 ResNet50 v1.5 和 VGG16 模型(批大小为1)的卷积层经 IM2ROW 转换后产生的 GEMM 维度进行测试。具体维度见表I和表II。
- 实验结果:
- ResNet50 v1.5(图15):在20个不同的层中,
ALG+EXO在9个层上性能最佳,而带预取的官方 BLIS 在6个层上最佳。ALG+EXO在此过程中使用了 8x12, 8x4, 4x4, 4x8, 4x12, 1x8, 和 1x12 等多种微内核。 - ResNet50 v1.5 聚合时间(图16):在整个模型的总执行时间上,
ALG+EXO表现最好,略微领先于 BLIS。 - VGG16(图17):在9个不同的层中,EXO 在3个层上表现最好,带预取的 BLIS 在4个层上领先。
- VGG16 聚合时间(图18):在总执行时间上,
ALG+EXO和 BLIS 的性能非常接近。
- ResNet50 v1.5(图15):在20个不同的层中,
- 分析结论:实验结果再次证实了为边界情况生成特设微内核的优势。
ALG+EXO能够在多种实际的深度学习负载中取得与高度优化的库(如 BLIS)相媲美甚至更好的性能。
TABLE I: ResNet50 v1.5 DNN 模型(批大小为1)的各层应用 IM2ROW 转换后得到的 GEMM 维度。
TABLE II: VGG16 DNN 模型(批大小为1)的各层应用 IM2ROW 转换后得到的 GEMM 维度。
A5 结论
本文成功解决了线性代数(LA)库采用的单一化微内核方法所带来的问题,通过使用 EXO 工具生成了特定的、面向硬件的微内核 C 代码。研究中详细描述了一个从零开始构建微内核生成器的分步过程,该生成器产生的 C 代码性能接近甚至优于 BLIS 库中广为人知的代码。同时,文章还给出了如何调整微内核生成器以满足不同软件需求(如不同数据类型)的提示。
通过在三种不同场景(微内核独立执行、大型方形矩阵乘法、以及由深度学习卷积模型生成的矩形 GEMM)下的性能分析,本文验证了 EXO 生成的代码与基于 Neon intrinsics 和汇编的微内核的竞争力。此外,这项工作还为 EXO 工具本身贡献了对某些 ARM 特性的支持。
作为未来工作,作者计划将微内核生成器工具适配并分析到其他架构上,如 Intel、RISC-V 或矩阵引擎。此外,还计划解决为线性代数库生成其他代码片段或特定领域代码(如卷积代码)的问题。
A6 附录
A. 摘要
制品文档说明。本附录记录了 EXO ukr generator 软件制品及其安装和执行流程,旨在复现论文中展示的结果。该软件包旨在复现论文的所有章节内容,包括分步的微内核设计和实验结果。此软件为 ARMv8 处理器设计和配置,因为这是论文中使用的硬件。不过,它也可以在支持 Neon intrinsics 指令的 ARM 处理器上执行。该制品可在 https://github.com/adcastel/CGO_paper44_artifact 公开获取,并包含一系列用于构建环境和执行实验的脚本。
B. 制品清单(元信息)
- 算法:高性能矩阵乘法
- 编译:Exo 编译器
- 硬件:ARM Neon (v8)
- 指标:GFLOPS
- 实验:独立的微内核和深度神经网络 GEMM
- 准备工作流所需时间(大约):20分钟
- 完成实验所需时间(大约):10分钟
- 是否公开可用:是,地址为 https://github.com/adcastel/CGO_paper44_artifact
- 代码许可证(如果公开可用):无
C. 描述
- 交付方式:此制品通过 https://github.com/adcastel/CGO_paper44_artifact Github 仓库提供,无任何许可证。此特性在
REQUIREMENTS.txt文件中指明。 - 硬件依赖:该制品适用于 ARM v8(支持 Neon 指令),并已在 NVIDIA Xavier、Orin 和 Nano 平台上测试。此特性在
REQUIREMENTS.txt文件中指明。 - 软件依赖:此制品包含 Python 和 C 代码,因此必须安装 Python3.9 和 gcc-10(或更高版本)。可能存在一些当前未安装的 Python 包。此可能情况已在配置脚本中处理,但仍可能发生包安装失败。Exo 和 Blis 库的软件也包含在包中。此特性在
REQUIREMENTS.txt文件中指明。对于绘图过程,此制品使用gnuplot工具。 - 数据集:实验所需的数据集已包含在仓库代码中,因此没有额外的要求。
D. 安装
- 克隆仓库:首先,我们需要用以下命令克隆仓库:
git clone https://github.com/adcastel/CGO_paper44_artifact.git - 进入目录:然后我们需要进入目录:
- 执行构建脚本:并像这样执行
build.sh脚本:source build.sh。此脚本将检查现有的编译器,并构建和安装 Blis 和 Exo 软件,以及设置环境变量。
E. 实验工作流
微内核生成。对于微内核生成,我们只需执行以下脚本:./microkernel_generator.sh。此脚本将生成论文第三节中解释的微内核,并显示过程中每一步生成的代码,如前述章节所示。
F. 评估和预期结果
复现评估。为了复现评估,我们需要使用两个不同的脚本,对应论文中的不同评估。
1. solo-mode 评估:要评估微内核在 solo-mode 下的性能,我们需要执行相应的脚本:./execute_ukernel_solo.sh。此脚本将执行图13中显示的实验。
2. 算法实验复现:为了复现图14-18的实验,我们应该执行以下脚本:./execute_algorithm.sh。由于算法、微内核和 GEMM 尺寸(包括方形和 DL 模型)的不同组合,这项评估耗时最长。
3. 生成图表:要生成图表,我们应该执行绘图脚本:./do_plots.sh。论文的图表将位于 plots 文件夹中。
请注意,用户可以使用一个脚本来构建和执行制品:./build_and_execute_all.sh。此脚本将构建环境并自行执行所有实验。
G. 实验定制
定制化选项。此制品可以通过多种方式进行定制。用户能够修改生成的微内核或 GEMM 的评估。
修改微内核生成。在微内核生成的情况下,用户应修改 EXO_ukr_generator 目录下的 generator.py 文件。在那里,用户可以设置 MR 和 NR 的值(即微内核的尺寸),以及数据类型。请注意,在该文件中,第一个调用是论文第三节中使用的调用,而其他注释掉的行可用于生成不同尺寸的微内核。
修改实验设置。如果用户想更改算法+微内核评估的实验设置,需要更改 CGO_paper44_artifact/gemm_blis_family/cnn_models/ 中的文件,或者生成一个新的输入文件,然后将相应的行添加到 execute_algorithm.sh 文件中。请注意,对不同微内核(非论文中使用的微内核)的整体评估将涉及每个测试/配置中的不同更改。因此,不建议在没有制品开发者指导的情况下对非通常的微内核尺寸进行“随机”尝试。
💬 评论讨论
欢迎在这里分享您的想法和见解!