基于 Apache TVM 的一类矩阵乘法例程自动生成器
计算与语言
2023-11-01 v1
摘要
我们探索利用 Apache TVM 开源框架自动生成一类算法，这些算法沿用 GotoBLAS2、BLIS 和 OpenBLAS 等流行线性代数库所采取的方法，以获得通用矩阵乘法（GEMM）的高性能分块形式。此外，我们还通过借助 Apache TVM 框架推导 GEMM 的完整多种处理器专用微内核，将生成过程完全自动化。这与高性能库中的惯例形成对比，后者使用汇编代码为每个架构手工编码单一微内核。总体而言，我们将 TVM 生成的分块算法与 GEMM 微内核相结合：1) 提升了可移植性、可维护性，并从全局上简化了软件生命周期；2) 提供高度灵活性，可轻松定制并优化针对不同数据类型、处理器架构和矩阵操作数形状的方案，取得与手工调优库相当（在特定矩阵形状下甚至更优）的性能；3) 具有较小的内存占用。
引用
@article{arxiv.2310.20347,
title = {Automatic Generators for a Family of Matrix Multiplication Routines with Apache TVM},
author = {Guillermo Alaejos and Adrián Castelló and Pedro Alonso-Jordá and Francisco D. Igual and Héctor Martínez and Enrique S. Quintana-Ortí},
journal= {arXiv preprint arXiv:2310.20347},
year = {2023}
}
备注
35 pages, 22 figures. Submitted to ACM TOMS