新しいカスタムアクセラレータの追加に関するガイドライン

一般的に、onnx-mlir ではカスタムアクセラレータを、onnx-mlir のビルドとモデルのコンパイル時にオン/オフできるプラグインとして処理します。処理は主に cmake を介して行われ、その手順をこのドキュメントの概要として説明します。

このドキュメントに加え、NNPA アクセラレータ は onnx-mlir にデプロイされている例として使用できます。

1. コードフォルダー

onnx-mlir では、アクセラレータのすべてのコードは src/Accelerators の下の別のフォルダ内に配置する必要があります。そのため、アクセラレータをサポートするための最初のステップは、src/Accelerators 内にそれ専用のフォルダを作成することです。

フォルダ名は、onnx-mlir でのアクセラレータ名として使用されます。特に、次の目的で使用されます。

1. cmake にアクセラレータフォルダ内のコードのビルドを指示する
2. onnx-mlir コマンドを使用してアクセラレータのモデルをコンパイルする
3. onnx-mlir-opt コマンドを使用してアクセラレータに関連するパスを有効にする

フォルダの内容はアクセラレータごとに柔軟です。ただし、可能な限り onnx-mlir のルートフォルダと同じ構造に従うことをお勧めします。これにより、プロジェクト全体に一貫性が維持されます。

1.1 onnx-mlir でアクセラレータをビルドする

onnx-mlir でアクセラレータをビルドするには、onnx-mlir をビルドするときに cmake 変数 ONNX_MLIR_ACCELERATORS を使用します。ONNX_MLIR_ACCELERATORS は、アクセラレータ名のセミコロン区切りのリストを受け付けます。たとえば次のようになります。

$ cd build
$ cmake .. -DONNX_MLIR_ACCELERATORS=accel1;accel2

1.2 選択したアクセラレータで実行するようにモデルをコンパイルする。

コンパイラコマンド onnx-mlir には、選択したアクセラレータ用にモデルをコンパイルするためのオプション --maccel があります。各アクセラレータについて --maccel=accel_name エントリを追加します。たとえば次のようになります。

$ onnx-mlir --maccel=accel1 --maccel=accel2 model.onnx

--maccel と一緒に使用できるのは、ビルドされたアクセラレータのみです。

1.3 選択したアクセラレータに関連するパスを実行する。

アクセラレータによって定義されたパスは、オプション --maccel を使用して onnx-mlir-opt コマンドで実行またはテストできます。これは onnx-mlir の --maccel に似ています（セクション 1.2 を参照）。たとえば、アクセラレータ accel1 によって定義されたパス --optimize-data-layout を呼び出すには、次のようになります。

$ onnx-mlir-opt --maccel=accel1 --optimize-data-layout model.mlir

--maccel と一緒に使用できるのは、ビルドされたアクセラレータのみです。

2. コード統合

2.1 マクロ

各アクセラレータでは、いくつかのマクロを定義する必要があります。これらは onnx_mlir::accel::Accelerator に含める必要があります。これらのマクロは次のとおりです。

1. INSTRUMENTSTAGE_ENUM_<アクセラレータ名>
2. INSTRUMENTSTAGE_CL_ENUM_<アクセラレータ名>
3. PROFILEIR_CL_ENUM_<アクセラレータ名>
4. OPTREPORT_ENUM_<アクセラレータ名>
5. OPTREPORT_CL_ENUM_<アクセラレータ名>

<アクセラレータ名> をアクセラレータ名に置き換えてください。たとえば、アクセラレータの名前が ACCEL1 の場合、次を使用します。

#define INSTRUMENTSTAGE_ENUM_ACCEL1
#define INSTRUMENTSTAGE_CL_ENUM_ACCEL1
#define PROFILEIR_CL_ENUM_ACCEL1
#define OPTREPORT_ENUM_ACCEL1
#define OPTREPORT_CL_ENUM_ACCEL1

2.2 方言とパス

MLIR でコードを記述するには、通常、方言とパスの設計が必要です。アクセラレータのサポートでも同様です。したがって、アクセラレータコードを onnx-mlir に統合するには、onnx-mlir に方言とパスを登録します。

基底クラス onnx_mlir::accel::Accelerator を提供します。ユーザーはこの基底クラスから継承したクラスを定義し、方言とパスの登録に対するフックを記述できます。

//===--------------------------------------------------------------------===//
// Hooks for onnx-mlir driver
//===--------------------------------------------------------------------===//

/// Add the transformations necessary to support the accelerator.
virtual void addPasses(mlir::OwningOpRef<mlir::ModuleOp> &module,
    mlir::PassManager &pm,
    onnx_mlir::EmissionTargetType &emissionTarget) const = 0;

//===--------------------------------------------------------------------===//
// Hooks for onnx-mlir-opt driver
//===--------------------------------------------------------------------===//

/// Register the MLIR dialects required to support an accelerator.
virtual void registerDialects(mlir::DialectRegistry &registry) const = 0;

/// Register accelerator transformation passes to make available as
/// command line options.
virtual void registerPasses(int optLevel) const = 0;

//===--------------------------------------------------------------------===//
// Hooks for onnx-to-krnl pass
//===--------------------------------------------------------------------===//

/// Convert TensorType to MemRefType.
/// Acccelators may have special versions of TensorType. If not, override this
/// method and return nullptr.
virtual mlir::MemRefType convertTensorTypeToMemRefType(
    const mlir::TensorType tensorType) const = 0;

/// Define conversion target to be used with ONNXToKrnl.
virtual void conversionTargetONNXToKrnl(
    mlir::ConversionTarget &target) const = 0;

/// Define rewrite patterns to be used with ONNXToKrnl.
virtual void rewritePatternONNXToKrnl(mlir::RewritePatternSet &patterns,
    mlir::TypeConverter &typeConverter, mlir::MLIRContext *ctx) const = 0;

//===--------------------------------------------------------------------===//
// Hooks for krnl-to-llvm pass
//===--------------------------------------------------------------------===//

/// Define conversion target to be used with KrnlToLLVM.
virtual void conversionTargetKrnlToLLVM(
    mlir::ConversionTarget &target) const = 0;

/// Define rewrite patterns to be used with KrnlToLLVM.
virtual void rewritePatternKrnlToLLVM(mlir::RewritePatternSet &patterns,
    mlir::LLVMTypeConverter &typeConverter, mlir::MLIRContext *ctx) const = 0;

onnx-mlir には多数のパスがありますが、2 つのパス onnx-to-krnl と krnl-to-llvm についてのみフックを提供します。その理由は、原則としてこれらが onnx-mlir における最初と最後のパスであるためです。パス onnx-to-krnl は、どの ONNX 演算子をホストで実行するか (Krnl 方言にローワーすることで）、またはアクセラレータで実行するか (アクセラレータ用に定義された方言にローワーすることで）を決定できる場所です。パス krnl-to-llvm は、Krnl 演算子とアクセラレータ演算子を LLVM 方言にローワーする場所です（アセンブリコードを生成したり、アクセラレータ用の外部 API を呼び出すことができます）。アクセラレータ用に onnx-to-krnl と krnl-to-llvm の間に存在する方言とパスは任意です。

たとえば、NNPA アクセラレータの場合、onnx-to-krnl で使用する ZHigh 方言 と、krnl-to-llvm で使用する ZLow 方言 を定義します。

3. テスト

アクセラレータのテストは、フォルダー test 内に配置する必要があります。特に

• LIT テストは mlir/accelerators 配下の新しく作成されたフォルダー内に配置されます
• その他のテストは accelerators 配下の新しく作成されたフォルダー内に配置されます