ONNX

2020年10月、ONNX-MLIRという機械学習コンパイラ（ONNXモデルを実行可能バイナリへ変換するツール）に macOSのCI設定を追加するPRを送り、マージされました。 本記事では、このPRを作成してからマージされるまでの試行錯誤を紹介します。 経緯 OSS機械学習コンパイラの調査の一環でONNX-MLIRを見つけ、開発用PC(macOS; MacBook Pro 13")に導入していました。 リポジトリ内を見ていたところ、x86-Linux、s390-Linux、x86-WindowsのCIは走っていますが、どうやらmacOSのCIはないみたいでした。 Issueに「macOS向けCIパイプラインを追加しようと思ってるけどどう？」と聞いてみたところ、 メンテナの一人から是非やってやって欲しい旨のコメントがあり、作業に着手しました。 CIサービスの選定 GitHubをサポートしているCIサービスはいくつかありますが、メンテナからはTravis CIでやって欲しいとの要望があったため、 Travis CI向けの調査に着手しました。 ONNX-MLIRでは、ビルドにDockerを使っています。 単一のDockerイメージでビルドするのではなく、以下のようにビルドを2段階に分ける構成とすることで、 依存ライブラリのビルド成果物を再利用可能となり、ビルド時間の短縮を実現しています。 依存ライブラリをビルドするDockerイメージ llvm/mlirが一番大きな依存ライブラリで、ビルドに1時間程度かかる 依存ライブラリの更新時にビルドされる 上記でビルドした生成物を使ってONNX-MLIR本体をビルドするDockerイメージ このイメージはTravis-CI実行時に毎回ビルドされる Travis CIとDockerのmacOSサポート状況を調査したところ、 DockerはゲストOSとしてのmacOSをサポートしていない（多分） Travis CIのmacOS環境はDockerをサポートしていない（Using Docker in Builds） ため、Travis CIは使えないと判断しました。 ただ、実現したいことは「キャッシュ機構による高速化」であることは分かった、代替案の調査に着手しました。 別のCIサービスとしてGitHub Actionsを調査したところ、キャッシュ機構(actions/cache@v2)をサポートしており、これを提案しようと考えました。 （今考えると、Travis CIのキャッシュ機構を使えば同じことができますね。。。） しかし、GitHub Actionsのキャッシュ機構は以下の制約があるとわかりました。 7日間以上アクセスのないキャッシュエントリは削除される リポジトリ全体でキャッシュサイズの上限が5GB 上記の制約があるが、GitHub Actionsは使えそうとのことをメンテナに相談したところ、 1.の制約は問題なく、2.が大丈夫ならGitHub Actionsを使おうとのコメントをもらいました。 というわけで、キャッシュ対象のファイルサイズの調査をします。 actions/cache@v2調査 キャッシュ対象のファイルサイズの調査の前に、どのようにキャッシュされているか知るためにactions/cache@v2を調査します。 GitHub Actionsは、リポジトリ全体でキャッシュサイズの上限が5GBという制約があります。 実はこのサイズは、キャッシュ対象の圧縮後のサイズを指しています。（公式ドキュメントには記載がなく、ソースコードを読むとわかりました） 以下のようにコードを読んでいくと、キャッシュ対象はGzip or zstdで圧縮されることがわかります。 キャッシュ保存時に呼ばれるメソッドはsaveCache()を呼び出しています。 https://github....

ONNX MLIRに含まれるDialectとその変換パスを調査。 Dialectの変換はソースコードの以下に存在。 src/Conversion Dialect一覧 ONNX MLIRには以下のDialectが含まれている。 ONNX Dialect Krnl Dialect Handshake Dialect StaticLogic Dialect LLHD Dialect その他、MLIRでデフォルトで定義されているDialectとして以下がある。 Affine Dialect Standard Dialect LLVM Dialect Dialect変換パス ONNXモデルをLLVM IRを生成する流れは、 ONNX parserでONNXモデルからONNX Dialectへ変換し、 ONNX DialectからKrnl Dialectへ変換、 Krnl DialectからLLVM Dialectへ変換し、 最後にLLVM DialectからLLVM IRを生成する

ONNX MLIRに付属のdebug.pyをベースに作成。 ONNX MLIRでONNXモデルファイルをshared libraryに変換し、PyRuntimeで実行する。 以下はメモ。 ExecutionSession(shared_lib_path, "_dyn_entry_point_main_graph")は、第2引数にエントリポイントを指定している。 ExecutionSession.runでグラフを実行するときは、引数を辞書で与えなくていい ONNX Runtimeだと{name: value}の辞書形式で与える必要がある import os import sys import argparse import onnx import subprocess import numpy as np import tempfile from collections import OrderedDict ONNX_MLIR = os.path.join(os.environ['ONNX_MLIR_HOME'], "bin/onnx-mlir") RUNTIME_DIR = os.path.join(os.environ['ONNX_MLIR_HOME'], "lib") sys.path.append(RUNTIME_DIR) from PyRuntime import ExecutionSession def main(model_path): model = onnx.load(model_path) subprocess.run([ONNX_MLIR, model_path], stdout=subprocess.PIPE, check=True) shared_lib_path = os.path.splitext(os.path.basename(model_path))[0] + ".so" sess = ExecutionSession(shared_lib_path, "_dyn_entry_point_main_graph") # calculate input shape shape_proto = model....

ONNX MLIRに付属のdebug.pyを動かすで動かしたdebug.pyの中を見る。 debug.pyの処理内容 概要 入力として指定したONNXモデルを、ONNX MLIRでshared libraryとしてビルド、実行し、リファレンスバックエンドで実行した結果と比較する リファレンスバックエンドにはONNX Runtimeを使用 Operatorのoutputごとに比較 モデルの出力だけでなく、Operatorの実行結果単位で比較している PyRuntimeはおそらく、ONNX MLIRでビルドしたshared libraryを、Pythonから実行するためのPythonバインディング shared libraryの実行方法はPyRuntimeの実装を見る必要がありそう 詳細 import onnxruntimeでONNX Runtimeをインポートする。 import os import sys import argparse import onnx import subprocess import numpy as np import tempfile from collections import OrderedDict # Reference backend, use onnxruntime by default import onnxruntime prepare = onnxruntime.InferenceSession ONNX_MLIR_HOMEが設定されているか確認。 if (not os.environ.get('ONNX_MLIR_HOME', None)): raise RuntimeError( "Environment variable ONNX_MLIR_HOME is not set, please set it to the path to " "the HOME directory for onnx-mlir....

ONNX MLIRに付属のdebug.pyを動かす手順。 .soファイルの実行手順を解析するための事前準備。 前提条件 環境 macOS Catalina 10.15.6 Python 3.7.5 (pyenvで導入) ONNX MLIR commit id dbe0d734b5687e0aa7da911684912163cea07bd2 ONNX MLIRをビルド済み(ONNX MLIRをビルドする) ビルドディレクトリは$HOME/workspace/onnx-mlir/buildとする 作業ディレクトリは$HOME/workspace/debug_testとする 環境変数の設定 ONNX_MLIR_HOMEにONNX MLIRのビルドディレクトリを設定する。 export ONNX_MLIR_HOME=$HOME/workspace/onnx-mlir/build 準備 debug.pyを作業ディレクトリにコピーする。 ONNXモデルファイルの例としてmnistをダウンロードする。 cp $HOME/workspace/onnx-mlir/utils/debug.py $HOME/workspace/debug_test/ wget https://github.com/onnx/models/raw/master/vision/classification/mnist/model/mnist-8.onnx Pythonのonnxパッケージを導入する。 pip3 install onnx 実行 debug.pyを見る限り、ONNX MLIRでONNXモデルファイルを.soへ変換し、実行するようだ。 python3 debug.py mnist-8.onnx Temporary directory has been created at /var/folders/mg/v57g_1jj3s52wgwtd61rt8t80000gn/T/tmp1megvpue %8 = "krnl.getref"(%7, %c0_i64) : (memref<10368xi8>, i64) -> memref<1x8x18x18xf32> Verifying value of Parameter193_reshape1 Verifying value of Convolution28_Output_0 Verifying value of Plus30_Output_0 Verifying value of ReLU32_Output_0 Verifying value of Pooling66_Output_0 Verifying value of Convolution110_Output_0 Verifying value of Plus112_Output_0 Verifying value of ReLU114_Output_0 Verifying value of Pooling160_Output_0 Verifying value of Pooling160_Output_0_reshape0 Verifying value of Times212_Output_0 Verifying value of Plus214_Output_0 ヘルプ情報 python3 debug....

前回の記事でビルドしたONNX MLIRを使って、ONNX Model Zooで配布されている学習済みモデルを変換する。 環境変数パスの設定 ONNX MLIRのビルドディレクトリは~/workspace/onnx-mlir/buildとする。 ~/workspace/onnx-mlir/build/bin以下に次の2つの実行ファイルがあるはず。 onnx-mlir おそらくコンパイラのフロントエンドに相当？ onnx-mlir-opt コンパイラPassをテストするためのツール この実行ファイルをコマンドで実行できるように環境変数パスにディレクトリを追加する。 export PATH=$PATH:~/workspace/onnx-mlir/build/bin ONNXモデルを取得 例としてMNISTの学習済みモデルを取得する。 cd ~/workspace wget https://github.com/onnx/models/raw/master/vision/classification/mnist/model/mnist-8.onnx 変換実行 shared libraryへ変換 $ nnx-mlir --EmitLib mnist-8.onnx %22 = "krnl.getref"(%21, %c0_i64) : (memref<4096xi8>, i64) -> memref<1x1x32x32xf32> %20 = "krnl.getref"(%19, %c0_i64) : (memref<25088xi8>, i64) -> memref<1x8x28x28xf32> %18 = "krnl.getref"(%17, %c0_i64) : (memref<25088xi8>, i64) -> memref<1x8x28x28xf32> %16 = "krnl.getref"(%15, %c0_i64) : (memref<25088xi8>, i64) -> memref<1x8x28x28xf32> %14 = "krnl....

ONNX MLIRをビルドする手順。基本はGitHubのREADMEに記載の手順と同じ。 検証環境は Ryzen 5 1600 Ubuntu 20.04 on WSL2 ONNX MLIR commit id ac67900 ビルドツール導入 ビルドに必要なツール群を導入する。 Ubuntu 18.04の場合、libprotoc-devが古くてビルドできないため、別途手動で導入する必要がある。 sudo apt install protobuf-compiler build-essential cmake libprotoc-dev \ ninja-build libncurses python3 python3-dev python-is-python3 MLIRのビルド カレントディレクトリは~/workspaceとする。 LLVM内にあるMLIRを取得し、ONNX MLIRのビルドが確認できているブランチをチェックアウトする。 cd ~/workspace git clone https://github.com/llvm/llvm-project.git cd llvm-project && git checkout 1d01fc100bb5bef5f5eaf92520b2e52f64ee1d6e && cd .. ビルドする。Ryzen 5 1600で20分程度。 mkdir llvm-project/build cd llvm-project/build cmake -G Ninja ../llvm \ -DLLVM_ENABLE_PROJECTS=mlir \ -DLLVM_BUILD_EXAMPLES=ON \ -DLLVM_TARGETS_TO_BUILD="host" \ -DCMAKE_BUILD_TYPE=Release \ -DLLVM_ENABLE_ASSERTIONS=ON \ -DLLVM_ENABLE_RTTI=ON cmake --build ....