[論文レビュー] DISC: A Dynamic Shape Compiler for Machine Learning Workloads

Kai Zhu, Wenyi Zhao+9 more|arXiv (Cornell University)|Mar 9, 2021

Parallel Computing and Optimization Techniques参考文献 9被引用数 4

ひとこと要約

DISC は、MLIR を拡張して完全に動的形状の IR（DHLO）を導入し、実行時解釈のオーバーヘッドを回避するためのコンパイル時制御フローを生成することで、マシンラーニングワークロード向けの動的形状コンパイラである。また、形状伝搬と制約収集を通じてホスト・デバイスの共同最適化を可能にし、Transformer などの動的形状ワークロードで Tensorflow/PyTorch より最大 3.3× の高速化を達成している。

ABSTRACT

Many recent machine learning models show dynamic shape characteristics. However, existing AI compiler optimization systems suffer a lot from problems brought by dynamic shape models, including compilation overhead, memory usage, optimization pipeline and deployment complexity. This paper provides a compiler system to natively support optimization for dynamic shape workloads, named DISC. DISC enriches a set of IR to form a fully dynamic shape representation. It generates the runtime flow at compile time to support processing dynamic shape based logic, which avoids the interpretation overhead at runtime and enlarges the opportunity of host-device co-optimization. It addresses the kernel fusion problem of dynamic shapes with shape propagation and constraints collecting methods. This is the first work to demonstrate how to build an end-to-end dynamic shape compiler based on MLIR infrastructure. Experiments show that DISC achieves up to 3.3x speedup than TensorFlow/PyTorch, and 1.8x than Nimble.

研究の動機と目的

• 動的形状マシンラーニングワークロードにおけるパフォーマンスボトルネックに対処する。既存のコンパイラは、高いコンパイルオーバーヘッド、メモリの肥大化、および融合の機会が限られている。
• XLA や同様のシステムにおける静的形状最適化の限界を克服する。これらのシステムは動的形状を効率的に処理できず、しばしばこのようなワークロードで融合を無効化する。
• 実行時解釈やパディングベースの回避策に依存せずに、動的形状計算をネイティブにサポートするソリューションを提供する。
• コンパイル時に制御フローのコードを生成することで、動的形状の状況下でも効率的なカーネル統合とホスト・デバイスの共同最適化を実現する。
• MLIR のインfra を活用して、複数のフレームワークをサポートするエンドツーエンドの動的形状コンパイルを実現し、可能な場合は静的最適化に自動的にフォールバックする。

提案手法

• MLIR の HLO ディアレクトを拡張して、DHLO（Dynamic HLO）を構築。DHLO は、新規 IR を再設計する必要なく、動的形状計算をネイティブにサポートする完全に動的形状の表現である。
• 実行時制御フロー（形状推論、バッファ管理、ホスト側ロジックなど）をすべてコンパイル時に生成することで、VM 解釈のオーバーヘッドを回避し、ホスト・デバイスの共同最適化を可能にする。
• 生産者と消費者間の形状伝搬を活用して、互いに互換性のある動的形状を持つ隣接演算子の融合候補を特定する。
• DHLO に下げる段階で形状制約を収集することで、コンパイル時に完全な形状を把握していない状況下でも、形状互換性を考慮した大規模な統合を可能にする。
• 静的および動的最適化の混合をサポートし、サブグラフの形状が既知の静的形状である場合には自動的に静的コンパイルにフォールバックする。
• DHLO ハブを介して TensorFlow および PyTorch と統合し、動的形状モデルのエンドツーエンドコンパイルを可能にする。

実験結果

リサーチクエスチョン

• RQ1 実行時解釈に依存せずに、マシンラーニングコンパイラが動的形状計算を効率的に表現・最適化するにはどうすればよいか？
• RQ2 コンパイル時に完全な形状情報が入手できない状況下で、動的形状ワークロードにおけるカーネル統合を実現するにはどのような技術が有効か？
• RQ3 動的形状実行における実行時オーバーヘッドを低減するために、ホスト・デバイスの制御フローをどのように共同最適化できるか？
• RQ4 動的形状コンパイラが静的コンパイラと同等のパフォーマンスを達成できる範囲はどの程度か？
• RQ5 MLIR のインfra を基盤とするコンパイラシステムは、最小限の実行時コストでエンドツーエンドの動的形状コンパイルをネイティブにサポートできるか？

主な発見

• DISC は、動的形状ワークロードにおいて TensorFlow や PyTorch より平均 2.27× の高速化を達成し、特定のモデルでは最大 3.3× の高速化を実現している。
• Transformer のメモリ集約的演算において、DISC は Nimble より 2.61× の高速化を達成している。これは、形状ヒントと制約によって有効に統合が可能になったためである。
• 実行時制御フローの CPU 時間を Nimble の 36.6% まで削減している。これは、わずかなカーネル数の削減とは対照的に、効率的なコード生成と共同最適化のおかげである。
• DISC は、平均的に Nimble より 1.8× の高速化を達成しており、VM による解釈とは対照的に、コンパイル時フロー生成の利点が顕著に表れている。
• 静的最適化へのフォールバックを無効化した場合、DISC は平均で静的コンパイルの 85% のパフォーマンスを達成し、74.5% から 91.4% の範囲にとどまる。これは、完全な形状情報が得られない状況下で融合の機会が限られることが原因のパフォーマンスギャップを示している。
• Transformer モデルにおいて、カーネル数は Nimble の 13,924 個から DISC の 10,734 個に削減され、明示的なカーネル数削減を目的としないにもかかわらず、統合効率の向上が示されている。