[論文レビュー] A Learned Performance Model for Tensor Processing Units

Samuel J. Kaufman, Phitchaya Mangpo Phothilimthana+5 more|arXiv (Cornell University)|Aug 3, 2020

Parallel Computing and Optimization Techniques参考文献 26被引用数 33

ひとこと要約

この論文は、グラフニューラルネットワークを用いてTPU向けの学習済み性能モデルを提示します。タイルサイズ選択と演算子融合において分析モデルを上回り、ハードウェアアクセスが限られている場合の自動調整を支援します。

ABSTRACT

Accurate hardware performance models are critical to efficient code generation. They can be used by compilers to make heuristic decisions, by superoptimizers as a minimization objective, or by autotuners to find an optimal configuration for a specific program. However, they are difficult to develop because contemporary processors are complex, and the recent proliferation of deep learning accelerators has increased the development burden. We demonstrate a method of learning performance models from a corpus of tensor computation graph programs for Tensor Processing Unit (TPU) instances. We show that our learned model outperforms a heavily-optimized analytical performance model on two tasks -- tile-size selection and operator fusion -- and that it helps an autotuner discover faster programs in a setting where access to TPUs is limited or expensive.

研究の動機と目的

• TPU上のテンソルプログラムの正確な実行時間予測を提供し、実機ハードウェアへの依存を減らす。
• 重い特徴量設計を必要とせず、多様なテンソルプログラムと最適化タスクに対して一般化する。
• 異なるコンパイラ最適化タスクに対して再ターゲット可能な性能モデリングを可能にする。
• TPUアクセスが制限されているまたは高価な場合に、自動調整の効率を改善することを示す。

提案手法

• テンソルプログラムを操作ノードと張力エッジを持つデータフローグラフとして表現する。
• 局所近傍からノード埋め込みを計算するグラフニューラルネットワーク(GraphSAGE)を使用する。
• 長距離依存関係をとらえるため、ノード埋め込みの上にシーケンスモデル（LSTMまたはTransformer）をオプションで組み合わせる。
• オペコードの埋め込みとノード/カーネルごとの特徴を組み込んでモデルの入力ベクトルを形成する。
• タスク特化の目的で訓練する：タイルサイズ選択にはペアワイズランク損失、融合には対数変換ターゲットの平方誤差損失を用いて絶対実行時間を予測する。
• 高度に手作業で調整された分析モデルと比較し、学習済みモデルをXLA自動チューナに組み込み、設定探索を指導する。

実験結果

リサーチクエスチョン

• RQ1 学習済みモデルは、タイルサイズ選択と演算子融合を指導するのに十分な精度でTPUカーネルの実行時間を予測できるか？
• RQ2 最小限の特徴量設計でグラフベースの表現は、見たことのないテンソルプログラムや異なる世代のハードウェア（TPU v2/v3）に対して一般化するか？
• RQ3 学習済みモデルを自動チューナに組み込むと、高速な構成を見つけるために必要な実機評価の数を減らすか？

主な発見

• 学習済みモデルは、タイルサイズタスクと融合タスクで、それぞれ真の測定値に対して96.3%および95.5%の精度を達成する。
• 学習済みモデルは、タイルサイズと融合タスクで、最良の手作業で調整された分析モデルよりそれぞれ2.4%と26.6%高い精度をもたらす。
• TPU v3では、タイルサイズ評価で平均誤差3.8%、平均相関0.65とやや低めを達成。
• アブレーション研究は、エッジの方向、静的性能特徴、およびカーネル埋め込みの選択が、融合タスクの精度に大きく影響し、タイルサイズタスクの精度には比較的限定的に影響することを示す。
• 学習済みモデルをXLA自動チューナに組み込むと、ハードウェアアクセスが制限されているまたは高価な場合に効率的な構成の発見を高速化できる。