基于纯编译器分层数据重组与内建指令降级的快速矩阵乘法
分布式、并行与集群计算
2023-05-30 v1 性能
摘要
机器学习的复兴增加了对高性能基础线性代数子程序(BLAS)的需求,这类程序长期依赖库来在商用硬件上达到峰值性能。高性能 BLAS 实现依赖于一种分层方法,该方法由用于数据(重)组织的分块与打包层,以及执行实际计算的微内核组成。高性能微内核的创建需要大量开发工作,为每种架构编写定制的汇编代码。IBM 的 POWER10 MMA、Intel AMX 和 Arm ME 近期引入矩阵引擎以提供高性能矩阵运算,使这一手工优化任务更加复杂。本文提出了一种替代高性能库的纯编译器方案,据我们所知且属首次,将分层方法的自动生成集成到生产级编译器 LLVM 中。该算法的模块化设计,例如使用 LLVM 的矩阵乘内建指令作为分块打包层与微内核之间的清晰接口,使得代码生成易于重定向到多种加速器。内建指令的使用使得全面的性能研究成为可能。在无硬件矩阵引擎的处理器中,对于小矩阵,分块与打包相较广泛使用的多面体优化器 PLuTo 在 Intel 上实现高达 22 倍的性能提升,在 POWER9 上对于大矩阵实现超过 6 倍的加速。其性能也接近高性能库,仅比 OpenBLAS 慢 34%,且对于大矩阵与 Eigen 相当。在 POWER10 的 MMA 下,对于大矩阵,该方案比向量扩展方案快逾 2.6 倍,与 Eigen 性能持平,并达到 BLAS 峰值性能的 96%。
引用
@article{arxiv.2305.18236,
title = {Fast Matrix Multiplication via Compiler-only Layered Data Reorganization and Intrinsic Lowering},
author = {Braedy Kuzma and Ivan Korostelev and João P. L. de Carvalho and José E. Moreira and Christopher Barton and Guido Araujo and José Nelson Amaral},
journal= {arXiv preprint arXiv:2305.18236},
year = {2023}
}