导入 ONNX 定义和支持的操作

目录

1. 概述
2. 添加操作
3. 自定义操作
4. 构建
5. 版本详细信息

   概述

   ONNX-MLIR 定义了一个 ONNX 方言来表示 ONNX 指定的操作。ONNX 方言是使用 MLIR table gen 工具创建的。每个操作的定义都通过 Python 脚本 utils/gen_onnx_mlir.py 从 ONNX 自动转移而来。该脚本从 ONNX 包中检索操作定义，以生成用于方言 table gen 的 ONNXOps.td.inc 和用于 ONNX-MLIR 中 ONNX 模型导入器的 OpBuilderTable.inc。以下各节将描述如何使用 gen_onnx_mlir.py 将操作添加到 ONNX-MLIR 中的 ONNX 方言，以及如何细化操作的定义。

添加操作

要为 ONNX 方言生成操作，请将其添加到 gen_onnx_mlir.py 中的字典 ‘version_dict’ 中。此字典的键是操作名称，值是此操作的版本列表。通常只支持此操作的最高版本（在 onnx-mlir/third_party/onnx 中）。有关版本化的详细信息，请参阅 版本部分。有了这个条目，脚本将生成 ONNX 方言的操作定义。

自定义

添加接口和 Trait

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

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

添加规范化接口

如果一个转换应该在各个 Pass 中局部应用于一个操作，则可以使用规范化接口来实现此转换。要为此操作启用规范化，请将该操作的名称添加到 OpsWithCanonicalizer 列表中，然后在其定义中将该操作设置为 hasCanonicalizer = 1;。

自定义构建器

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

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

请注意，使用 returnType 可以避免在重写规则中使用特殊构建器的需求。将此类类型推断代码移至 ONNXOpHelper.cpp 并移除自定义构建器可能是一个更好的解决方案。

自定义验证器

操作的操作描述列出了每个输入/输出和属性的允许类型。Table gen 将生成一个默认验证器来检查 IR 的允许类型。如果操作有额外的约束，则应定义自定义验证器以增强错误检测。例如，操作的两个输入可能需要相同的元素类型或相同的秩。此类信息可以在 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，因为属性通常是不可变的）。
• 使用 hasShapeAndRank(X) 来测试输入 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

此命令将生成这两个文件（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。

更新文档

添加新 op 版本或对 ONNX 版本进行更改时，我们也希望将这些更改反映在我们支持的操作的 ONNX 文档中。虽然最新的 ONNX 规范始终可用，但我们支持的规范通常有点滞后，此外，我们如前一节所述，以版本化名称支持旧版本。

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

make onnx-mlir-docs

上述命令在通常的 build 目录中运行，它将把新的方言 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 有两个版本（opset），11 和 13，都得到支持。version_dic 中的相应条目是 '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 安装。然后，您可以运行带标志 --check_operation_version 的 gen_onnx_mlir.py。所有操作的最高版本将作为新的 version_dict 输出。如果操作的接口保持不变（根据 ONNX 的更改文档），则只需使用新版本。如果接口确实发生了更改，您可以将新版本插入版本列表中的第一个。对于现有代码，所有相应的代码都必须更改。例如，当 ReduceSum 从版本 11 迁移到版本 13 时，ONNXReduceSumOp 将首先被替换为 ONNXReduceSumOpV11。然后版本 13 的代码将使用 ONNXReduceSumOp。这种设计的原因是，大多数 ONNX 更改都不会改变接口。我们不想给开发人员带来记住使用哪个操作版本的负担，除非绝对必要。保留旧版本的代码并不总是需要，它可以重写为新操作。因此，我们只需要方言定义，而不需要推断或降低代码。