[論文レビュー] Ansor: Generating High-Performance Tensor Programs for Deep Learning

Lianmin Zheng, Chengfan Jia+10 more|arXiv (Cornell University)|Jun 11, 2020

Parallel Computing and Optimization Techniques参考文献 53被引用数 82

ひとこと要約

Ansorは、自動的に大規模な階層的探索空間を構築し、完全なプログラムをサンプルし、学習済みコストモデルを用いた進化的探索でDNNサブグラフをCPUとGPU上で最適化するテンソルプログラム生成フレームワークである。最先端手法と比較して顕著なスピードアップを達成する。

ABSTRACT

High-performance tensor programs are crucial to guarantee efficient execution of deep neural networks. However, obtaining performant tensor programs for different operators on various hardware platforms is notoriously challenging. Currently, deep learning systems rely on vendor-provided kernel libraries or various search strategies to get performant tensor programs. These approaches either require significant engineering effort to develop platform-specific optimization code or fall short of finding high-performance programs due to restricted search space and ineffective exploration strategy. We present Ansor, a tensor program generation framework for deep learning applications. Compared with existing search strategies, Ansor explores many more optimization combinations by sampling programs from a hierarchical representation of the search space. Ansor then fine-tunes the sampled programs with evolutionary search and a learned cost model to identify the best programs. Ansor can find high-performance programs that are outside the search space of existing state-of-the-art approaches. In addition, Ansor utilizes a task scheduler to simultaneously optimize multiple subgraphs in deep neural networks. We show that Ansor improves the execution performance of deep neural networks relative to the state-of-the-art on the Intel CPU, ARM CPU, and NVIDIA GPU by up to $3.8\times$, $2.6\times$, and $1.7\times$, respectively.

研究の動機と目的

• DNNのためにさまざまなハードウェアプラットフォーム上で高性能なテンソルプログラムが必要である理由を動機づける。
• 手動のテンプレート設計を排除し、テンソルプログラムの探索空間を自動的に生成・探索する。
• スケジューリングとリソース意識的選択を通じてDNNサブグラフのエンドツーエンド最適化を可能にする。
• AnsorがCPUおよびGPUターゲットで最先端のアプローチと比較して優れた性能を達成することを実証する。

提案手法

• 高レベル構造（スケッチ）を低レベルの詳細（注釈）から分離する階層的探索空間を導入する。
• 計算グラフを表すDAGに対する derivationベースの列挙を通じてスケッチを自動生成する。
• スケッチにランダムに注釈を付けて、評価と調整のための完全なプログラムを取得する。
• 学習済みコストモデルを用いた進化的探索を適用してサンプルされたプログラムを微調整する。
• 勾配降下法ベースのタスクスケジューラを用いて、エンドツーエンドのDNN性能を最も向上させるサブグラフを優先する。
• 勾配ブースティング木で訓練されたコストモデルを活用し、プログラムのスループットを予測して探索を導く。

実験結果

リサーチクエスチョン

• RQ1 Ansorは手動テンプレートなしで大規模・総合的なテンソルプログラム探索空間を自動生成できるか。
• RQ2 学習済みコストモデルを用いた進化的探索は、CPUおよびGPUバックエンド全体で高性能プログラムを効果的に識別できるか。
• RQ3 スケジューラによるサブグラフの優先順位付けは、演算子ごとの最適化と比較してエンドツーエンドのDNN性能を改善できるか。
• RQ4 Ansorの性能は複数のハードウェアプラットフォーム上で最先端の探索ベースのフレームワークと比較してどうか。

主な発見

• AnsorはIntel CPU、ARM CPU、およびNVIDIA GPU上のDNN実行性能を、それぞれ最大で3.8×、2.6×、1.7×向上させる。
• Ansorが見つける最も高性能なプログラムの多くは、既存の探索空間の外に位置している。
• Ansorは探索空間が大きいにもかかわらず、より効率的に探索を行い、短時間で高性能プログラムを見つける。
• Ansorは探索時間を桁違いに削減して、最先端の性能に匹敵する。
• Ansorは数十倍の探索時間を要することなく、新しい演算子への自動拡張を可能にするには、手動テンプレートではなく数学的定義のみを必要とする。