[論文レビュー] Compiling ONNX Neural Network Models Using MLIR

Jin Tian, Gheorghe-Teodor Bercea+9 more|arXiv (Cornell University)|Aug 19, 2020

Parallel Computing and Optimization Techniques参考文献 3被引用数 39

ひとこと要約

この論文は、ONNXモデルをMLIRを用いてネイティブコードへ翻訳するコンパイラである onnx-mlir を提示し、ONNX および krnl ダイアレクトといくつかの最適化パスを導入します。複数のアーキテクチャに跨る MNIST および ResNet50 の予備結果を示します。

ABSTRACT

Deep neural network models are becoming increasingly popular and have been used in various tasks such as computer vision, speech recognition, and natural language processing. Machine learning models are commonly trained in a resource-rich environment and then deployed in a distinct environment such as high availability machines or edge devices. To assist the portability of models, the open-source community has proposed the Open Neural Network Exchange (ONNX) standard. In this paper, we present a high-level, preliminary report on our onnx-mlir compiler, which generates code for the inference of deep neural network models described in the ONNX format. Onnx-mlir is an open-source compiler implemented using the Multi-Level Intermediate Representation (MLIR) infrastructure recently integrated in the LLVM project. Onnx-mlir relies on the MLIR concept of dialects to implement its functionality. We propose here two new dialects: (1) an ONNX specific dialect that encodes the ONNX standard semantics, and (2) a loop-based dialect to provide for a common lowering point for all ONNX dialect operations. Each intermediate representation facilitates its own characteristic set of graph-level and loop-based optimizations respectively. We illustrate our approach by following several models through the proposed representations and we include some early optimization work and performance results.

研究の動機と目的

• モデルのポータビリティを促進するため、ONNX モデルをターゲット環境で効率的に動作させられるようにする。
• onnx-mlir コンパイラのアーキテクチャと設計を説明する。
• ONNX の意味論を低レベルのコードへ橋渡しするため、2つの新しい MLIR ダイアレクト（onnx および krnl）を導入する。
• ニューラルネットワークのコンパイルに対する最適化パスとその影響を示す。

提案手法

• ONNX グラフを LLVM IR およびネイティブコードへ落とす多ダイアレクトのコンパイラパイプラインを MLIR を用いて構築する。
• ONNX の意味論を符号化する ONNX ダイアレクトと、ループベースの降下のための krnl ダイアレクトの2つを導入する。
• グラフ書換え、形状推論、定数伝搬などの最適化パスを適用する。
• ONNX モデルのインポートから推論用の動的エントリポイント関数までのエンドツーエンドのフローを提供する。
• POWER システム上で MNIST および ResNet50 モデルをコンパイルして実行することでパイプラインを実証する。

実験結果

リサーチクエスチョン

• RQ1 ONNX モデルを専用ダイアレクトを用いた MLIR でネイティブコードへ効率的にコンパイルできるか。
• RQ2 ONNX を LLVM/Native コードへ落とすのに必要なダイアレクトとパスは何か。
• RQ3 初期の最適化が代表的なモデルのコンパイルと推論性能にどのように影響するか。

主な発見

• onnx-mlir コンパイラは ONNX モデルをエントリ関数 _dyn_entry_point_main_graph を持つネイティブコードライブラリへ変換できる。
• ONNX および krnl ダイアレクトは構造化された降下とループ最適化を可能にし、その後 affine、std、llvm ダイアレクトへの変換が行われる。
• グラフ書換えと定数伝搬は MLIR ベースのパイプラインにおける効果的な早期最適化である。
• MNIST は POWER9 上で迅速にコンパイル可能（0.237 s コンパイル、0.001 s 推論）、一方 ResNet50 は POWER9 上で 7.661 s のコンパイル、7.540 s の推論を示し、大規模モデルの実用性を示唆する。
• 現状の実装は正確性と基礎的な性能に焦点を当てており、ポリヘドラル、SIMD、ループ融合の最適化を追加して速度を向上させる計画がある。
• システムはアーキテクチャ間の課題（例：ビッグエンディアンのシステム）に直面しており、ONNX と MLIR へのパッチが必要だったが、現在は x86、Power、Z、Windows のサポートへ拡張されている。