[論文レビュー] 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

ひとこと要約

本論文は、深層学習ワークロードにおける行列乗算（GEMM）のコンパイラレベル最適化のための2つの新規技術—グリーディー最良優先探索（G-BFS）およびネイバーヒュッド・アクトロープラティスティック・アドバンテージ・クリティック（N-A2C）—を提案する。マルコフ意思決定過程における近接構成関係を活用することで、両手法ともXGBoostおよびRNNベースのチューナーよりもそれぞれ24%および40%低いGEMM計算時間（1024×1024行列において）を達成したが、探索する構成空間は0.1%にとどまった。

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は次のステップを優先するためのヒューリスティックスコアを用いる。
• 両アプローチは、1024×1024の行列を対象としたNVIDIA Titan XP GPU上で、TVMコンパイラーフレームワーク内での評価が行われた。

実験結果

リサーチクエスチョン

• RQ1 G-BFSのようなヒューリスティックベースの探索手法は、最小限の探索コストで、高パフォーマンスなカーネルを効率的に探索できるか？
• RQ2 N-A2Cのような強化学習アプローチは、解の質および探索効率の両面で、既存の最先端のチューニング手法を上回れるか？
• RQ3 G-BFSおよびN-A2Cの性能は、XGBoostおよびRNNベースのチューナーと比較して、計算時間の短縮率および探索する構成空間の割合においてどう異なるか？
• RQ4 これらの手法は、2048×2048のようなより大きな行列サイズにおいて、探索効率および解の質の観点でどの程度スケーラブルか？

主な発見

• N-A2C手法は、1024×1024行列において、RNNベースのチューナーよりも40%低いGEMM計算時間を達成したが、探索する構成空間は0.1%にとどまった。
• G-BFS手法は、同じ0.1%の探索空間制約下で、XGBoostチューナーよりも24%低い計算時間を達成した。
• 両手法とも、XGBoostおよびRNN手法よりも最適構成をより速く発見したが、N-A2Cは多段階計画による大規模な行列において優れた性能を示した。
• 10回の試行において、両手法とも性能のばらつきが低く、より安定的かつ信頼性の高いチューニング行動を示した。
• N-A2C手法は、2048×2048のような大規模な行列サイズにおいてG-BFSを上回り、複雑な探索空間における逐次的意思決定の利点を示した。
• 結果から、G-BFSおよびN-A2Cの両手法が、大規模な構成空間を持つテンソルオペレータのコンパイラーレベル最適化において有効かつ汎用的であることが確認された。