[論文レビュー] Equality Saturation for Tensor Graph Superoptimization

Yichen Yang, Phitchaya Mangpo Phothilimtha+4 more|arXiv (Cornell University)|Jan 5, 2021

Parallel Computing and Optimization Techniques参考文献 29被引用数 31

ひとこと要約

本論文は Tensat を紹介する。等式飽和を用いたテンソルグラフの超最適化器で、e-グラフを介して全てのリライトを同時適用。実行時の速度向上は最大16%、従来の逐次法より最大48倍速く最適化。

ABSTRACT

One of the major optimizations employed in deep learning frameworks is graph rewriting. Production frameworks rely on heuristics to decide if rewrite rules should be applied and in which order. Prior research has shown that one can discover more optimal tensor computation graphs if we search for a better sequence of substitutions instead of relying on heuristics. However, we observe that existing approaches for tensor graph superoptimization both in production and research frameworks apply substitutions in a sequential manner. Such sequential search methods are sensitive to the order in which the substitutions are applied and often only explore a small fragment of the exponential space of equivalent graphs. This paper presents a novel technique for tensor graph superoptimization that employs equality saturation to apply all possible substitutions at once. We show that our approach can find optimized graphs with up to 16% speedup over state-of-the-art, while spending on average 48x less time optimizing.

研究の動機と目的

• テンソルグラフ最適化における逐次リライト戦略の最適性の欠如を動機づける。
• 等式飽和とe-グラフを導入し、同等グラフのより大きな空間を探索する。
• テンソルグラフに対する非局所的・マルチパターンリライトを扱えるよう、等式飽和を拡張する。
• 標準的なMLモデルで実用的な速度アップを示し、最適化時間を分析する。
• ILPを用いたスケーラブルな抽出技術と、探索を妥当な範囲に保つサイクルフィルタリング戦略を提供する。

提案手法

• 等式飽和に拡張した演算子ベースのDAGでテンソルグラフを表現する。
• 破壊的な置換を伴わず全ての同等グラフを生成するためにe-グラフを介してリライト規則を適用する。
• 複数パターンリライト処理を持つcanonical S式一致アプローチでe-グラフを飽和させる探索を実行する。
• グラフの根のe-クラスから最低コストのグラフを、貪欲法またはトポロジ的制約を伴うILPベースの定式化で抽出する。
• ハードウェア（GPU）上の演算子の実行時間に基づくコストモデルを抽出の指針とする。
• 抽出時に任意でILPを解いて最適ノードを選択し、トポロジカル順序によるサイクル回避機構を用いる。

実験結果

リサーチクエスチョン

• RQ1 等式飽和はテンソルグラフ書換の位相順序付けを克服して、全体最適なグラフを見つけられるか？
• RQ2 e-グラフをどのように非局所的・マルチパターンのテンソル書換を取り込むように効果的に拡張できるか？
• RQ3 テンソルグラフの実用的な抽出方法（貪欲法 vs ILP）はランタイムコストを最小にするには？
• RQ4 マルチパターンリライトを適用する際の性能トレードオフ（速度アップ vs 最適化時間）はどうなる？
• RQ5 Tensat は多様なモデルで最先端の逐次探索（TASO）とどのように比較されるか？

主な発見

モデル	検索時間 (s) TASO	実行時速度up (%) TASO	検索時間 (s) Tensat	実行時速度up (%) Tensat
NasRNN	177.3	0.5	45.4	68.9
BERT	13.6	1.4	8.5	9.2
ResNeXt-50	25.3	0.7	5.5	8.8
NasNet-A	1226	10.6	1.9	7.3
SqueezeNet	16.4	0.3	6.7	24.5
VGG-19	8.9	0.4	8.9	8.9
Inception-v3	68.6	5.1	6.3	10.0

• Tensat は評価済みモデルで最先端 TASO に対して最大16%の実行時間短縮を達成。
• 最適化時間は大幅に削減され、同等グラフをコンパクトに表現することで最大300倍高速化。
• 等式飽和は逐次バックトラッキングよりはるかに大きな探索空間をカバーでき、ほとんどのベンチマークでより良い速度アップを達成。
• ケースによってはマルチパターンの反復を増やすことで速度アップをさらに改善するが、e-グラフの爆発とILPのタイムアウトのリスクがある。サイクルフィルタリングがこれを緩和。
• Tensat は速度アップと最適化時間の間で実用的なトレードオフを提供し、コンパイルワークフローへの実用的な統合が可能。
• NasRNN、SqueezeNet、Inception-v3などのモデルでの評価は広範な適用性を示す。