导入 ONNX 定义并支持操作

目录

1. 概述
2. 添加操作
3. 定制操作
4. 构建
5. 版本详情

   概述

   ONNX-MLIR 定义了 ONNX 语言，用于表示 ONNX 指定的操作。ONNX 语言是使用 MLIR table gen 工具创建的。每个操作的定义通过 Python 脚本 utils/gen_onnx_mlir.py 自动从 ONNX 中转移过来。此脚本从 ONNX 包中检索操作定义，为 dialect table gen 生成 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 语言生成操作定义。

定制

添加接口和 Trait

• 默认情况下，所有操作都具有形状推断接口和 Pure trait。
• 如果一个操作具有 ResultTypeInferenceOpInterface，请将其添加到字典 OpsWithResultTypeInference 中。此接口推断结果张量的类型，而不是形状。
• 如果一个操作具有子图，它将具有接口 HasOnnxSubgraphOpInterface。此属性是从 ONNX 操作定义中推断出来的。
• 您可以使用字典 OpsWithHelpers 为操作定义辅助函数。

默认情况下，所有操作都具有形状推断接口和 Pure trait。如果一个操作具有 ResultTypeInferenceOpInterface，请使用字典 OpsWithResultTypeInference。此接口推断结果张量的类型，而不是形状。如果一个操作具有子图，它将具有接口 HasOnnxSubgraphOpInterface。

添加规范化接口

如果应在多个 pass 中将本地转换应用于某个操作，则可以使用规范化接口进行此转换。要为此操作启用规范化，请将此操作的名称添加到 OpsWithCanonicalizer 列表中，然后该操作将在其定义中具有 hasCanonicalizer = 1;。

定制构建器

操作的默认构建器需要结果的类型作为参数。但是，结果的类型是可以推断的。定制构建器可能有助于简化代码。根据推断的类型，有两种构建器：unranked 类型和 broadcast 类型。要为操作启用特殊构建器，您可以分别将其名称添加到 custom_builder_unranked_ops_list 和 custom_builder_broadcast_ops_list 中。

请注意，通过使用 returnType，可以避免在重写规则中使用特殊构建器。请参阅 MLIR 文档 或 ONNX-MLIR 中的示例。将此类类型推断代码移至 ONNXOpHelper.cpp 并取消定制构建器可能是更好的解决方案。

请注意，通过使用 returnType，可以避免在重写规则中使用特殊构建器。将此类类型推断代码移至 ONNXOpHelper.cpp 并取消定制构建器可能是更好的解决方案。

定制验证器

操作的操作描述列出了每个输入/输出和属性的允许类型。table gen 将生成一个默认验证器来检查 IR 的允许类型。如果操作具有额外约束，则应定义定制验证器以增强错误检测。例如，操作的两个输入可能需要相同的元素类型或相同的 rank。此类信息可在 ONNX 操作定义中找到，但无法通过 dialect 定义表达。检查这些约束的最佳方法是在验证器中进行。要将定制验证器接口添加到操作，请在 gen_onnx_mlir.py 中找到下面的数组并将您的操作添加进去。

OpsWithVerifier = ['AveragePool', 'Conv', 'InstanceNormalization', 'Mod']

然后，您将在 ONNXOps.td.inc 的操作定义中找到以下行

let verifier = [{ return ::verify(*this); }];

当新操作声明为使用定制验证器时，您需要在 src/Dialect/ONNX/ONNXOps.cpp 中添加实现代码。最好的方法是查看其他操作以了解通用模式，例如搜索 static LogicalResult verify(ONNXInstanceNormalizationOp op)。请注意，每次创建此类操作时，验证器都会执行。因此，您需要确保它可以处理 tensor 和 MemRef，以及可能的未 rank tensor。因此，在适当的情况下进行每个测试。例如，一旦 tensor 被 rank 化，您可以验证 rank 是否在批准的范围内（如果存在此类约束）；在它未 rank 化之前，请勿执行此测试。

提示

• 使用 operandAdaptor 对象获取输入（必须使用 operandAdaptor 获取输入的当前值），使用 op 对象获取属性（可以使用 op，因为属性通常是不可变的）。
• 使用 hasShapeAndRank(X) 测试输入 X 当前是否具有形状和 rank。如果没有，则返回成功，因为稍后我们将有机会使用此信息测试操作。请注意，某些输入也可能是标量，在这种情况下，它们可能被编码为形状类型，也可能不被编码为形状类型。
• 然后，您可以使用 MLIR 调用 mlir::cast<ShapedType>(X.getType()) 获取形状类型，从而获取 rank 和维度。目前，我们只检查运行时已知值的维度有效性。未知维度编码为负数。请仅在确定不会断言（即类型确实是 ShapedType）时使用 cast。
• 当您发现错误时，请使用 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

此命令将生成这两个文件（src/Dialect/ONNX/ONNXOps.td.inc 和 OpBuilderTable.inc），并将其复制到 src 目录中的正确位置。如果您修改了 gen_onnx_mlir.py，也需要提交这两个生成的文件。在 ONNX-MLIR 构建中，它们被视为源文件，因此 ONNX-MLIR 用户无需安装特定版本的 ONNX。请勿直接修改这些文件。您也可以直接运行脚本，并使用 utils 目录中生成的文件。 python ../utils/gen_onnx_mlir.py。

更新文档

添加新操作版本或更改 ONNX 版本时，我们也希望在支持操作的 ONNX 文档中体现这些更改。虽然最新的 ONNX 规范 始终可用，但我们支持的规范通常会落后一些，而且我们还支持上节所述的旧版本，使用版本化的名称。

有一个方便的命令可以更新 ONNX 和 Krnl 两种语言，如下所示。

make onnx-mlir-docs

上述命令在常规的 build 目录中运行，它将把新的 dialect md 文件直接安装到 docs/Dialects 目录中。

添加 Krnl 语言的操作/进行更改时，应使用相同的命令。

操作版本

ONNX-MLIR 项目始于 ONNX 版本 1.7.0，并且不打算向后兼容。我们依赖 onnx/converter 将模型转换为 ONNX-MLIR 支持的版本。随着 ONNX 版本的演进，ONNX-MLIR 尝试跟进，但可能落后于最新版本。

操作的版本

如前所述，我们尝试支持最新版本的 ONNX 操作。当前支持的每个操作的版本记录在 utils/gen_onnx_mlir.py 中。这种机制提供了一些版本稳定性。要检查版本更改，请使用标志 “--check-version” 运行 gen_onnx_mlir.py，将报告更改。要升级到新版本，请手动更新脚本中的版本字典。

支持多个版本

要支持操作的多个版本，应将选定的版本添加到 utils/gen_onnx_mlir.py 中的版本字典中。例如，对于 ReduceSum，支持 11 和 13 两个版本（opset）。version_dict 中的相应条目为 'ReduceSum': [13, 11]。

在 ONNX 语言中，最高版本操作的名称中没有版本号，而其他版本名称后跟 ‘V’ 和版本号。例如，opset 13 的 ReduceSum 将是 ONNXReduceSumOp，而 opset 11 的 ReduceSum 是 ‘ONNXReduceSumV11Op`。由于大多数 ONNX 操作在升级到更高版本时是兼容的，我们可以保留语言中的操作名称，只需更新 gen_onnx_mlir.py 中的 version_dict，而无需修改 ONNX-MLIR 中的代码。

导入模型时，使用不高于下一个可用版本的最高版本。以 ReduceSum 为例，如果 opset 是 12，则选择 ONNXReduceSumV11Op。

迁移

要迁移到新的 ONNX 版本，首先应升级 third_part/onnx 和您安装的 ONNX。然后，您可以运行 gen_onnx_mlir.py 并带有标志 --check_operation_version。所有操作的最高版本将作为新的 version_dict 输出。如果操作的接口保持不变（根据 ONNX 的变更文档），您可以直接使用新版本。如果接口确实发生了变化，您可以将新版本插入到版本列表的第一位。对于现有代码，所有相应的代码都必须更改。例如，当 ReduceSum 从版本 11 迁移到 13 时，首先将 ONNXReduceSumOp 替换为 ONNXReduceSumOpV11。然后版本 13 的代码将使用 ONNXReduceSumOp。这种设计的原因是大多数 ONNX 更改不会改变接口。我们不想给开发人员增加负担去记住使用了哪个版本的操作，除非绝对必要。并非总是需要保留旧版本的代码，因为它可以重写为新操作。因此，我们只需要 dialect 定义，而不需要推理或降低的代码。