onnx-mlir
MLIRコンパイラインフラストラクチャにおけるONNXモデルの表現と参照ローワーリング
GitHub でプロジェクトを見る onnx/onnx-mlir
ハウツー
参考文献
開発
ツール
ツール
このプロジェクトは onnx によってメンテナンスされています。
GitHub Pages でホスト — テーマ by orderedlist
ONNX定義のインポートとオペレーションのサポート
目次
- 概要
- オペレーションの追加
- オペレーションのカスタマイズ
- ビルド
- バージョンに関する詳細
概要
ONNX-MLIRは、ONNXによって指定されたオペレーションを表すONNXダイアレクトを定義します。ONNXダイアレクトはMLIRテーブル生成ツールを使用して作成されます。各オペレーションの定義は、Pythonスクリプトutils/gen_onnx_mlir.pyを使用してONNXから自動的に転送されます。このスクリプトはONNXパッケージからオペレーション定義を取得し、ダイアレクトテーブル生成用のONNXOps.td.incとONNX-MLIRのONNXモデルインポータ用のOpBuilderTable.incを生成します。以降のセクションでは、gen_onnx_mlir.pyを使用してONNX-MLIRのONNXダイアレクトにオペレーションを追加する方法、およびオペレーションの定義を改良する方法について説明します。
オペレーションの追加
ONNXダイアレクトのオペレーションを生成するには、このオペレーションをgen_onnx_mlir.pyのディクショナリ「version_dict」に追加します。このディクショナリのキーはオペレーション名、値はオペレーションのopsetのリストです。通常、このオペレーションのトップバージョンのopset（onnx-mlir/third_party/onnx内）のみがサポートされます。バージョン管理の詳細については、バージョンセクションを参照してください。このエントリを使用すると、スクリプトはONNXダイアレクトのオペレーション定義を生成します。
カスタマイズ
インターフェースとトレイトの追加
- デフォルトでは、すべてのオペレーションはシェイプ推論インターフェースと
Pureトレイトを持っています。 - オペレーションが
ResultTypeInferenceOpInterfaceを持つ場合は、ディクショナリOpsWithResultTypeInferenceに追加します。このインターフェースは、シェイプではなく、結果テンソルの型を推論します。 - オペレーションがサブグラフを持つ場合、
HasOnnxSubgraphOpInterfaceインターフェースを持ちます。この属性はONNXオペレーション定義から推論されます。 - ディクショナリ
OpsWithHelpersを使用して、オペレーションのヘルパー関数を定義できます。
デフォルトでは、すべてのオペレーションはシェイプ推論インターフェースとPureトレイトを持っています。ResultTypeInferenceOpInterfaceを持つオペレーションの場合は、ディクショナリOpsWithResultTypeInferenceを使用します。このインターフェースは、シェイプではなく、結果テンソルの型を推論します。オペレーションにサブグラフがある場合、HasOnnxSubgraphOpInterfaceインターフェースを持ちます。
正規化インターフェースの追加
変換をパス全体でオペレーションにローカルに適用する必要がある場合、正規化インターフェースを使用できます。オペレーションの正規化を有効にするには、このオペレーションの名前をOpsWithCanonicalizerのリストに追加します。その後、オペレーションの定義にhasCanonicalizer = 1;が追加されます。
ビルダーのカスタマイズ
オペレーションのデフォルトビルダーは、パラメータとして結果の型を必要とします。しかし、結果の型は推論できます。コードを簡素化するために、カスタマイズされたビルダーが役立つ場合があります。型の推論に基づいて、アンランク型とブロードキャスト型の2種類のビルダーがあります。オペレーションの特別なビルダーを有効にするには、それぞれcustom_builder_unranked_ops_listとcustom_builder_broadcast_ops_listにその名前を追加できます。
returnTypeを使用することで、書き換えルールにおける特別なビルダーの必要性を回避できることに注意してください。MLIRドキュメントまたはONNX-MLIRの例を参照してください。そのような型推論コードをONNXOpHelper.cppに移動し、カスタマイズされたビルダーを削除する方が良い解決策かもしれません。
returnTypeを使用することで、書き換えルールにおける特別なビルダーの必要性を回避できることに注意してください。そのような型推論コードをONNXOpHelper.cppに移動し、カスタマイズされたビルダーを削除する方が良い解決策かもしれません。
検証者のカスタマイズ
オペレーションの説明では、各入力/出力と属性の許容される型を列挙しています。テーブル生成は、許容される型についてIRをチェックするデフォルトの検証者を生成します。オペレーションに追加の制約がある場合、エラー検出を強化するためにカスタマイズされた検証者を定義する必要があります。たとえば、オペレーションの2つの入力は、同じ要素型または同じランクを必要とする場合があります。そのような情報はONNXオペレーション定義に見られますが、ダイアレクト定義では表現できません。これらの制約をテストする最良の方法は検証者です。カスタマイズされた検証者のインターフェースをオペレーションに追加するには、gen_onnx_mlir.pyで以下の配列を探し、オペレーションを追加します。
OpsWithVerifier = ['AveragePool', 'Conv', 'InstanceNormalization', 'Mod']
次に、ONNXOps.td.incのオペレーション定義で次の行が見つかります。
let verifier = [{ return ::verify(*this); }];
新しいopがカスタマイズされた検証者を使用するように宣言された場合、src/Dialect/ONNX/ONNXOps.cppに実装コードを追加する必要があります。たとえば、static LogicalResult verify(ONNXInstanceNormalizationOp op)を検索して、一般的なパターンを確認することをお勧めします。検証者は、そのようなopが作成されるたびに実行されることに注意してください。そのため、テンソルとMemRefs、そしておそらくアンランクテンソルでも機能するようにする必要があります。したがって、各テストを適切な状況に合わせます。たとえば、テンソルがランク付けされたら、ランクが承認された範囲内にあることを検証できます（そのような制約がある場合）。ランク付けされる前はこのテストを実行しないでください。
ヒント
- 入力の現在の値を取得するには
operandAdaptorオブジェクト（入力の現在の値を取得するにはoperandAdaptorを使用する必要があります）、属性を取得するにはopオブジェクト（属性は通常不変であるためopを使用できます）を使用します。 X入力が現在シェイプとランクを持っているかどうかをテストするにはhasShapeAndRank(X)を使用します。そうでない場合は、後でこの情報を使用してオペレーションをテストする機会があるため、成功を返します。一部の入力はスカラーである場合もあり、その場合、シェイプ型としてエンコードされている場合とされていない場合があります。- 次に、MLIR呼び出し
mlir::cast<ShapedType>(X.getType())を使用してシェイプ型を取得できます。これにより、ランクと次元を取得できます。現時点では、実行時にわかっている値についてのみ次元の有効性をチェックしています。不明な次元は負の数としてエンコードされます。アサートしないことが確実な場合、つまり型が実際にShapedTypeである場合にのみ、キャストを使用してください。 - エラーが見つかった場合は、
op->emitError(msg)を使用してフレンドリーなエラーメッセージで報告します。
インポータのカスタマイズ
special_op_handler：frontend_dialect_transformer.cppに特別なインポート関数を生成します。現在、特別なハンドラーは、操作引数を持つオペレーションに使用されています。
任意の追加定義
オペレーションの定義に上記以外に追加のコードが必要な場合は、ディクショナリcustom_definition_miscにコードを配置できます。キーはオペレーション名、値はコードです。
インポータのカスタマイズ
special_op_handler：frontend_dialect_transformer.cppに特別なインポート関数を生成します。現在、特別なハンドラーは、操作引数を持つオペレーションに使用されています。
任意の追加定義
オペレーションの定義に上記以外に追加のコードが必要な場合は、ディクショナリcustom_definition_miscにコードを配置できます。キーはオペレーション名、値はコードです。
ビルド
gen_onnx_mlir.pyを実行するには、ONNXをインストールする必要があります。READMEを参照してください。ビルドディレクトリで、次のコマンドを実行します。
make OMONNXOpsIncTranslation
このコマンドは、これらの2つのファイル（src/Dialect/ONNX/ONNXOps.td.incとOpBuilderTable.inc）を生成し、srcディレクトリの正しい場所にコピーします。gen_onnx_mlir.pyを変更した場合は、生成された2つのファイルもチェックインする必要があります。これらはONNX-MLIRビルドでソースファイルとして扱われるため、ONNX-MLIRのユーザーは特定のバージョンのONNXをインストールする必要がありません。これらのファイルを直接変更しないでください。utilsディレクトリに生成されたファイルを使用してスクリプトを直接実行することもできます。python ../utils/gen_onnx_mlir.py。
ドキュメントの更新
新しいopバージョンを追加する場合やONNXバージョンに変更を加える場合、サポートされているオペレーションのONNXドキュメントにもこれらの変更を反映させたいと考えています。最新のONNX仕様は常に利用できますが、私たちがサポートしている仕様はしばしば少し遅れており、前のセクションで述べたように、バージョン付きの名前で古いバージョンもサポートしています。
ONNXとKrnlダイアレクトの両方を更新するための便利なコマンドがあります。以下に示します。
make onnx-mlir-docs
上記のコマンドは通常のbuildディレクトリで実行され、新しいダイアレクトmdファイルをdocs/Dialectsディレクトリに直接インストールします。
オペレーションを追加する場合やKrnlダイアレクトに変更を加える場合にも、同じコマンドを使用する必要があります。
オペレーションバージョン
ONNX-MLIRプロジェクトはONNXがバージョン1.7.0であったときに開始され、下位互換性を意図していません。ONNX-MLIRがサポートするバージョンにモデルを変換するには、onnx/converterに依存しています。ONNXバージョンは進化しているため、ONNX-MLIRはそれに追従しようとしますが、最新バージョンより遅れる可能性があります。
オペレーションのバージョン
前述の通り、ONNXオペレーションの最新バージョンをサポートしようと努めています。現在サポートされている各オペレーションのバージョンは、utils/gen_onnx_mlir.pyに記録されています。このメカニズムにより、バージョンにおいてある程度の安定性が確保されます。バージョンの変更を確認するには、gen_onnx_mlir.pyを「--check-version」フラグ付きで実行すると、変更点が報告されます。新しいバージョンに移行するには、スクリプト内のバージョン辞書を手動で更新します。
複数バージョンのサポート
オペレーションの複数バージョンをサポートするには、選択したバージョンをutils/gen_onnx_mlir.pyのバージョン辞書に追加する必要があります。例えば、ReduceSumにはサポートされているバージョン（opset）として11と13の2つがあります。version_dicに対応するエントリは'ReduceSum': [13, 11]となります。
ONNXダイアレクトでは、最新バージョンのオペレーション名はバージョンを含みません。一方、それ以外のバージョン名は「V」とバージョン番号が後に続きます。例えば、opset 13のReduceSumはONNXReduceSumOpとなりますが、opset 11のReduceSumは'ONNXReduceSumV11Op'となります。ほとんどのONNXオペレーションは上位バージョンにアップグレードしても互換性があるため、ダイアレクト内のオペレーション名を維持し、ONNX-MLIRのコードに触れることなく、gen_onnx_mlir.py内のversion_dictを更新するだけで済みます。
モデルのインポート時に、利用可能な次のバージョンよりも高くなく、最も高いバージョンが使用されます。ReduceSumの例では、opsetが12の場合、ONNXReduceSumV11Opが選択されます。
移行
新しいバージョンのONNXに移行するには、まずthird_part/onnxをアップグレードし、ONNXのインストールを更新する必要があります。その後、--check_operation_versionフラグ付きでgen_onnx_mlir.pyを実行できます。すべてのオペレーションの最新バージョンが新しいversion_dictとして出力されます。オペレーションのインターフェースが同じままであれば（ONNXの変更ドキュメントを参照）、新しいバージョンを使用できます。インターフェースが変更されている場合は、新しいバージョンをバージョンリストの先頭に挿入します。既存のコードについては、対応するコードをすべて変更する必要があります。例えば、ReduceSumがバージョン11から13に移行された場合、まずONNXReduceSumOpをONNXReduceSumOpV11に置き換えます。その後、バージョン13のコードはONNXReduceSumOpを使用します。このような設計になっている理由は、ONNXの変更のほとんどはインターフェースを変更しないためです。絶対に必要な場合を除き、開発者がどのバージョンのオペレーションを使用しているかを覚える負担を減らしたいと考えています。古いバージョンのコードを常に保持する必要はなく、新しいオペレーションに書き直すことができます。したがって、推論やローワーリングのためのコードではなく、ダイアレクトの定義のみが必要となります。