关于添加新自定义加速器的指南

通常，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 运行与选定加速器相关的通道。

加速器定义的通道可以通过 onnx-mlir-opt 命令使用选项 --maccel 来运行或测试，这类似于 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_<accel_name>
2. INSTRUMENTSTAGE_CL_ENUM_<accel_name>
3. PROFILEIR_CL_ENUM_<accel_name>
4. OPTREPORT_ENUM_<accel_name>
5. OPTREPORT_CL_ENUM_<accel_name>

将 <accel_name> 替换为加速器的名称，例如，如果您的加速器名为 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 both onnx-mlir and onnx-mlir-opt drivers
//===--------------------------------------------------------------------===//

/// Configure passes for the accelerator.
virtual void configurePasses() 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 中有许多通道，但我们只为 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 加速器，我们定义 ZHigh 语言 用于 onnx-to-krnl，定义 ZLow 语言 用于 krnl-to-llvm。

3. 测试

加速器的测试应放在文件夹 test 中。特别是，

• LIT 测试放在 mlir/accelerators 下新创建的文件夹中
• 其他测试放在 accelerators 下新创建的文件夹中