导入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”中。该字典的键是操作名称，值是此操作的ops。通常只支持此操作的最高版本opset（在onnx-mlir/third_party/onnx中）。有关版本控制的详细信息，请参阅版本部分。有了此条目，脚本将为ONNX方言生成操作定义。

自定义

添加接口和特性

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

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

添加规范化接口

如果需要在传递中局部应用于操作的转换，可以使用规范化接口进行此转换。要为操作启用规范化，请将此操作的名称添加到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); }];

当新操作被声明为使用自定义验证器时，您需要在src/Dialect/ONNX/ONNXOps.cpp中添加实现代码。最好的方法是查看其他操作以获取通用模式，例如通过搜索static LogicalResult verify(ONNXInstanceNormalizationOp 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。

更新文档

在添加新操作版本或更改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的大部分更改不会改变接口。我们不希望给开发人员带来负担，让他们记住使用哪个版本的操作，除非绝对必要。并不总是需要保留旧版本的代码，旧版本可能会被重写为新操作。因此，我们只需要方言定义，而不需要推断或降级代码。