[论文解读] Compiler-Level Matrix Multiplication Optimization for Deep Learning

Huaqing Zhang, Xiaolin Cheng+2 more|arXiv (Cornell University)|Sep 23, 2019

Parallel Computing and Optimization Techniques参考文献 34被引用 6

一句话总结

本文提出两种新颖的基于编译器的优化技术——贪心最佳优先搜索（G-BFS）与邻域智能体优势评论（N-A2C），用于深度学习工作负载中的矩阵乘法（GEMM）优化。通过在马尔可夫决策过程（MDP）中利用邻域配置关系，两种方法在1024×1024矩阵上仅探索0.1%的配置空间，分别比XGBoost和RNN-based调优器实现24%和40%更低的GEMM计算时间。

ABSTRACT

An important linear algebra routine, GEneral Matrix Multiplication (GEMM), is a fundamental operator in deep learning. Compilers need to translate these routines into low-level code optimized for specific hardware. Compiler-level optimization of GEMM has significant performance impact on training and executing deep learning models. However, most deep learning frameworks rely on hardware-specific operator libraries in which GEMM optimization has been mostly achieved by manual tuning, which restricts the performance on different target hardware. In this paper, we propose two novel algorithms for GEMM optimization based on the TVM framework, a lightweight Greedy Best First Search (G-BFS) method based on heuristic search, and a Neighborhood Actor Advantage Critic (N-A2C) method based on reinforcement learning. Experimental results show significant performance improvement of the proposed methods, in both the optimality of the solution and the cost of search in terms of time and fraction of the search space explored. Specifically, the proposed methods achieve 24% and 40% savings in GEMM computation time over state-of-the-art XGBoost and RNN methods, respectively, while exploring only 0.1% of the search space. The proposed approaches have potential to be applied to other operator-level optimizations.

研究动机与目标

• 解决因在多样化硬件上生成次优GEMM内核而导致的深度学习推理与训练性能瓶颈。
• 减少深度学习框架中对GEMM算子进行手动、硬件特定调优的依赖。
• 实现在不同硬件平台上高效、自动化的高性能GEMM配置搜索。
• 提升在大规模、复杂配置空间中张量算子优化的搜索效率与解决方案质量。

提出的方法

• G-BFS方法采用启发式引导的贪心搜索策略，在GEMM配置空间中探索邻近配置，优先选择预期性能提升较大的配置。
• N-A2C方法采用基于邻域动作空间的演员-评论家强化学习框架，迭代优化配置选择。
• 两种方法均将GEMM优化建模为马尔可夫决策过程（MDP），基于分块与布局变换定义状态转移。
• 配置空间由多级分块参数定义，状态表示为矩阵A、B和C的分块大小元组。
• N-A2C方法使用实测内核执行时间作为奖励信号以训练策略网络，而G-BFS则使用启发式评分来优先选择下一步。
• 两种方法均在NVIDIA Titan XP GPU上的TVM编译器框架中，针对不同大小的GEMM工作负载进行了评估。

实验结果

研究问题

• RQ1 基于启发式的搜索方法如G-BFS能否以极低的搜索成本高效探索GEMM配置空间，从而找到高性能内核？
• RQ2 如N-A2C这样的强化学习方法是否能在解决方案质量与搜索效率两方面均超越现有最先进调优方法？
• RQ3 在计算时间减少与配置空间探索比例方面，G-BFS与N-A2C相较于XGBoost和RNN-based调优器表现如何？
• RQ4 在搜索效率与解决方案质量方面，这些方法在更大矩阵尺寸（如2048×2048）上扩展到何种程度？

主要发现

• 在1024×1024矩阵上，N-A2C方法实现比RNN-based调优器低40%的GEMM计算时间，且仅探索0.1%的配置空间。
• 在相同0.1%的配置空间探索约束下，G-BFS方法相比XGBoost调优器将计算时间降低了24%。
• 两种方法发现最优配置的速度均快于XGBoost与RNN方法，其中N-A2C因具备多步规划能力，在大矩阵上表现更优。
• 两种方法在10次试验中表现出更低的性能方差，表明其调优行为更具稳定性和可靠性。
• 在更大矩阵尺寸（如2048×2048）上，N-A2C优于G-BFS，证明了在复杂搜索空间中顺序决策的优势。
• 结果证实，G-BFS与N-A2C在具有大规模配置空间的张量算子编译器级优化中均具有效性与泛化能力。