onnx-mlir

Logo

ONNX 模型在 MLIR 编译器基础设施中的表示与参考 Lowering

在 GitHub 上查看项目 onnx/onnx-mlir

操作指南

使用 Python 进行推理
使用 C/C++ 进行推理
使用 Java 进行推理

参考资料

ONNX Dialect
OMTensor C99 运行时 API
OMTensorList C99 运行时 API
OMTensor Java 运行时 API
OMTensorList Java 运行时 API
生成 ONNX Dialect
关于文档

开发

添加操作
测试指南
错误处理
命令行选项
Instrumentation
常量传播
添加加速器

工具

工具

RunONNXModel.py
DocCheck

该项目由 onnx 维护

托管于 GitHub Pages — 主题由 orderedlist 提供

测试

在 onnx-mlir 中,有三种类型的测试来确保实现的正确性

  1. ONNX 后端测试
  2. LLVM FileCheck 测试
  3. 数值测试
  4. 使用 gdb
  5. ONNX 模型库

ONNX 后端测试

后端测试是基于 ONNX 节点和模型测试的 onnx-mlir 端到端测试。它们可用于测试 C/C++ .so 库和 JNI .jar 归档文件。对于每个 C/C++ 测试目标,添加 -jni 后缀可得到相应的 JNI 测试目标。要调用测试,请使用以下命令

cmake --build . --config Release --target check-onnx-backend[-jni]

需要安装诸如 third_party/onnx 等包才能运行后端测试。您可以使用命令 pip install your-onnx-mlir/third_party/onnx 安装您自己的 onnx 包。JNI 测试需要 jsoniter jar,如果在系统上找不到已安装的版本,则默认从其 maven 仓库下载。如果用户在构建 ONNX-MLIR 时开启了 cmake 选项 ONNX_MLIR_BUILD_JSONITER,则将从其 github 仓库克隆的源代码在本地构建 jsoniter jar。请注意,在本地构建 jsoniter jar 需要安装 maven 构建工具。

onnx 包提供的所有测试用例都列在文件 test/backend/all_test_names.txt 中。check-onnx-backend 将选择性地运行其中一些。check-onnx-backend 将运行的 onnx 中的节点和模型测试由 test/backend/test.py 中的变量 test_to_enable 定义。用户可以通过环境变量 TEST_CASE_BY_USER 测试一个测试用例。例如,

TEST_CASE_BY_USER=selected_test_name cmake --build . --config Release --target check-onnx-backend[-jni]

指定 TEST_CASE_BY_USER 后,中间结果、.onnx 文件和 .so 文件会保留在 build/test/backend 中以便调试。如果您需要检查生成的共享库是否包含某个特定指令,请将环境变量 TEST_INSTRUCTION_CHECK 设置为 true,并在测试名称后添加指令名称,例如 TEST_CASE_BY_USER=selected_test_name,instruction_name。请注意在 onnx 测试名称后添加后缀 _cpu

ONNX 支持的测试用例

文件 test/backend/all_test_names.txt 包含 ONNX 包提供的所有测试用例。您可以通过将其添加到 test/backend/inference_backend.py 中来启用测试用例。all_test_names.txt 使用命令 “make check-onnx-backend-case” 自动生成。只有在 ONNX 包升级时才需要更新。

将 ONNX 支持的测试用例添加到当前的后端测试集

添加操作符的 ONNX 到 Krnl 转换后,应将该操作符相应的后端测试添加到 test.py 中。可用的测试用例可以在 third_party/onnx/onnx/backend/test/case/node 中找到。您可以通过在 test/backend/all_test_names.txt 中查找新的操作符来识别新的测试。找到新的测试后,您可以将新的测试添加到 test/backend/inference_backend.py. 中。请注意在 onnx 测试名称后添加后缀 _cpu。与测试相关联,您可以定义如何运行新操作符的测试。例如

        "test_and2d_cpu": {STATIC_SHAPE:{}, DYNAMIC_SHAPE:{-1:{-1}}, CONSTANT_INPUT:{-1}},

表明测试 test_and2d_cpu 可以运行 (1) 使用静态形状,(2) 强制所有输入为动态形状,或 (3) 强制所有输入为定义的常量。这是大多数操作符的推荐设置。但是,有些操作符对某些参数不容忍动态形状;对于这些操作符,可以明确决定函数哪个参数可以是动态形状。这用表达式 {-1:{-1}} 指定。test/backend/inference_backend.py. 文件包含关于如何指定哪些参数和/或参数维度可以设置为动态的明确说明。

带有未知维度的测试

测试动态张量大小最简单的方法是使用以下命令,我们的检查器也使用此命令。

cmake --build . --config Release --target check-onnx-backend-dynamic[-jni]

ONNX 节点测试通常具有已知的输入张量维度大小。因此,为了测试带有未知维度的张量,模型导入器 (Build/FrontendONNXTransformer.cpp) 提供了一个功能来生成此类情况。设置环境变量 IMPORTER_FORCE_DYNAMIC 后,前端导入会将模型所有输入张量的所有维度(默认情况下)变为 -1。例如,

IMPORTER_FORCE_DYNAMIC='-1:-1' all dimensions of all the inputs will be changed
IMPORTER_FORCE_DYNAMIC='0:-1' all dimensions of the first input will be changed
IMPORTER_FORCE_DYNAMIC='0:-1|1:0,1' all dimensions of the first input and the 1st and 2nd dimensions of the second input will be changed

IMPORTER_FORCE_DYNAMIC 的巴科斯-纳尔范式 (BNF) 如下。

<ImportForceDynamicExpr> :== `'` <expr> `'`
                  <expr> ::= <inputString> | <inputString> `|` <expr>
            <inputString ::= <inputIndex> `:` <dimString>
             <dimString> ::= <dimIndex> | <dimIndex> `,` <dimString>
            <inputIndex> ::= <index>
              <dimIndex> ::= <index>
                 <index> ::= -1 | <number>
                <number> ::= <digit> | <digit><number>
                 <digit> ::= 0 | 1 | 2 | 3 | 4 | 5 | 6 | 7 | 8 | 9

-1 语义上表示所有输入或所有维度,并且具有最高优先级。例如 '0: -1, 0' 表示第一个输入的所有维度将被更改。输入和维度索引从 0 开始。

例如,test_add_cpu 的默认模型是

func @main_graph(%arg0: tensor<3x4x5xf32>, %arg1: tensor<3x4x5xf32>) -> tensor<3x4x5xf32>

使用 IMPORTER_FORCE_DYNAMIC='-1:-1',结果是

func @main_graph(%arg0: tensor<?x?x?xf32>, %arg1: tensor<?x?x?xf32>) -> tensor<?x?x?xf32>

使用 IMPORTER_FORCE_DYNAMIC='0:-1',结果是

func @main_graph(%arg0: tensor<?x?x?xf32>, %arg1: tensor<3x4x5xf32>) -> tensor<3x4x5xf32>

使用 IMPORTER_FORCE_DYNAMIC='0:0,2|1:1',结果是

func @main_graph(%arg0: tensor<?x4x?xf32>, %arg1: tensor<3x?x5xf32>) -> tensor<3x4x5xf32>

这是使用现有节点测试进行动态张量测试的一种方法。由于并非所有测试用例都可以使用动态张量通过,test/backend/test.py 中有一个列表 test_not_for_dynamic,用于指定定义了 IMPORTER_FORCE_DYNAMIC 时哪些测试无法通过。

带有常量输入的测试

因为 ONNX 节点测试在运行时接受输入张量,所以在编译 ONNX 模型时输入不是常量。然而,在实践中,输入可以是常量,我们希望测试这种情况。

测试常量输入最简单的方法是使用以下命令,我们的检查器也使用此命令。

cmake --build . --config Release --target check-onnx-backend-constant[-jni]

要测试单个 ONNX 节点,例如 test_add_cpu,使用两个环境变量 TEST_CONSTANTIMPORTER_FORCE_CONSTANT,例如

TEST_CONSTANT=true IMPORTER_FORCE_CONSTANT="0" TEST_CASE_BY_USER=test_add_cpu make check-onnx-backend[-jni]

这会将第一个输入(索引 0)转换为常量,因此模型现在只有一个输入而不是两个。

环境变量 IMPORTER_FORCE_CONSTANT 是一个索引列表,用 , 分隔(从 0 开始,或使用 -1 表示所有输入索引),例如 0, 2, 3-1

输入签名测试

使用以下命令测试 ONNX 模型具有各种数据类型的输入签名,我们的检查器也使用此命令。

cmake --build . --config Release --target check-onnx-backend-signature

启用 SIMD 指令

在支持的平台(目前是 s390x z14 及以上、x86 和 arm)上,后端测试可以为编译的模型生成 SIMD 指令。要启用 SIMD,请设置 TEST_MARCH 环境变量,例如,

TEST_MARCH=z16 cmake --build . --config Release --target check-onnx-backend[-jni]

后端测试的执行

定义在 utils/RunONNXLib.cpp 中的工具可用于轻松执行其 .so 模型中的文件,例如使用 TEST_CASE_BY_USER=selected_test_name make check-onnx-backend 命令生成的文件。通过将 onnx-mlir/src/Compiler/CompilerUtils.cpp 文件中的 overridePreserveFiles 值设置为 KeepFilesOfType::All,例如,也可以在以其他方式构建模型时保留模型。

当 ONNX 模型版本比 onnx-mlir 当前支持的版本旧时,可以通过设置环境变量 INVOKECONVERTER 为 true 来调用 ONNX 版本转换器。例如,对于 INVOKECONVERTER=true make check-onnx-backend,将对所有测试用例调用转换器。在 test.py 中,有一个名为 test_need_converter 的列表,供您在单个用例上调用转换器。

该工具直接扫描模型提供的签名,用随机值初始化所需的输入,然后调用模型中的函数。该程序可以与 gdblldbvalgrind 等其他工具结合使用。要列出工具选项,只需在运行时使用 -h--help 标志。

我们首先需要编译该工具,这可以通过两种模式之一完成。在第一种模式下,该工具使用静态链接的模型进行编译。此模式在编译期间除了包含 .so 文件外,还需要 -D LOAD_MODEL_STATICALLY=0 选项。最好使用 onnx-mlir/utils 目录中的 build-run-onnx-lib.sh 脚本,使用其模型编译该工具,模型作为参数传递给脚本。为避免 Mac 上的库路径问题,请在构建模型的目录中运行编译后的工具。

# Compile tool with model.
cd onnx-mlir/build
sh ../utils/build-run-onnx-lib.sh test/backend/test_add/test_add.so
# Run the tool to run the model (substitute `Release` for `Debug` for the release version).
Debug/bin/run-onnx-lib
# or, on Mac, run the tool in the directory where the model was built
(cd test/backend; ../../Debug/bin/run-onnx-lib)
# if test_add.so was built in `test/backend`:
cd test/backend; ../../Debug/bin/onnx-mlir --EmitLib test_add/test_add.onnx

(在 Mac 上可以使用 otool -L test_add.so 查看库的路径。)

在第二种模式下,该工具在编译时不包含模型,模型将在运行时传递。要启用此选项,只需使用 -D LOAD_MODEL_STATICALLY=1 选项编译该工具。您可以使用与上面相同的脚本,但不带参数。然后,只要在运行时将 .so 模型文件传递给该工具,就可以从任何目录运行该工具。

# Compile tool without a model.
cd onnx-mlir/build
sh ../utils/build-run-onnx-lib.sh
# Run the tool with an argument pointing to the model.
Debug/bin/run-onnx-lib test/backend/test_add/test_add.so

LLVM FileCheck 测试

我们可以通过提供中间表示作为输入,并使用 LLVM FileCheck 工具检查输出 IR 来测试某个 Pass 的功能。例如,我们有一个用于形状推理的测试用例 test.mlir。

func @test_default_transpose(%arg0 : tensor<5x5x1x32xf32>) -> tensor<*xf32> {
  %0 = "onnx.Transpose"(%arg0) : (tensor<5x5x1x32xf32>) -> tensor<*xf32>
  "std.return"(%0) : (tensor<*xf32>) -> ()

您可以在此测试用例上运行形状推理 Pass,并获得以下输出

module  {
  func @test_default_transpose(%arg0: tensor<5x5x1x32xf32>) -> tensor<32x1x5x5xf32> {
    %0 = "onnx.Transpose"(%arg0) {perm = [3, 2, 1, 0]} : (tensor<5x5x1x32xf32>) -> tensor<32x1x5x5xf32>
    return %0 : tensor<32x1x5x5xf32>
  }
}

手动检查输出是否正确。如果输出正确,将输出转换为将来可以自动检查的内容。使用命令

Debug/bin/onnx-mlir-opt --shape-inference test.mlir | python ../utils/mlir2FileCheck.py 

您将获得以下内容

// mlir2FileCheck.py
// CHECK-LABEL:  func @test_default_transpose
// CHECK-SAME:   ([[PARAM_0_:%.+]]: tensor<5x5x1x32xf32>) -> tensor<32x1x5x5xf32> {
// CHECK:           [[VAR_0_:%.+]] = "onnx.Transpose"([[PARAM_0_]]) {perm = [3, 2, 1, 0]} : (tensor<5x5x1x32xf32>) -> tensor<32x1x5x5xf32>
// CHECK:           return [[VAR_0_]] : tensor<32x1x5x5xf32>
// CHECK:         }

将源和检查代码组合起来,并添加到适当的测试用例中。ONNX Dialect 的所有测试用例都收集在 test/mlir/onnx 目录下。这些测试用例可以通过 make check-onnx-lit 调用。此目标是构建的基本要求。

数值测试

数值测试除了 ONNX 包提供的测试外,还用于测试数值正确性。目标是提供基于大量数值的单元测试;这对于确保优化转换的有效性和正确性非常重要:随着我们针对特定架构参数(如向量宽度)进行专业化,会出现更多的边缘情况。数值测试基于被测操作的简单、朴素(且极其缓慢)的实现,生成大量基于数值的单元测试,用于验证我们的操作 Lowering 和优化的正确性。

数值测试应按以下方式组织,使得以下两个组件是独立且分离的

这样做的动机是我们有两种方式生成测试用例参数

  // RapidCheck test case generation.
  bool success = rc::check("convolution implementation correctness", []() {
    const auto N = *rc::gen::inRange(1, 10);
    const auto C = *rc::gen::inRange(1, 20);
    const auto H = *rc::gen::inRange(5, 20);
    const auto W = *rc::gen::inRange(5, 20);

    const auto kH = *rc::gen::inRange(1, 15);
    const auto kW = *rc::gen::inRange(1, 15);

    // We don't want an entire window of padding.
    const auto pHBegin = *rc::gen::inRange(0, kH - 1);
    const auto pHEnd = *rc::gen::inRange(0, kH - 1);
    const auto pWBegin = *rc::gen::inRange(0, kW - 1);
    const auto pWEnd = *rc::gen::inRange(0, kW - 1);

    // Make sure we have at least 1 output per dimension.
    RC_PRE((H >= kH) && (W > kW));

    RC_ASSERT(isOMConvTheSameAsNaiveImplFor(
        N, C, H, W, kH, kW, pHBegin, pHEnd, pWBegin, pWEnd));
  });
  assert(success && "error while performing RapidCheck tests");

有时能够看到与数值测试相关的 mlir 文件很方便。为此,最简单的方法是将 src/Compiler/CompilerUtils.cpp 文件中的 overridePreserveFiles 变量设置为您想要保留的文件类型(例如 KeepFilesOfType::All)。然后,无论您如何编译模型,输入和输出 mlir 文件以及未优化和优化的字节码文件以及一些额外的二进制文件都将保留下来。

在失败的情况下,RapidCheck(用于数值测试的基础设施)和 ONNX 模型都允许用户使用相同的值重新运行测试。运行测试时,您可能会看到以下输出。

Model will use the random number generator seed provided by "TEST_SEED=1440995966"
RapidCheck Matrix-Vector test case generation.
Using configuration: seed=4778673019411245358

通过记录以下两个环境变量中的种子值

export RC_PARAMS="seed=4778673019411245358"
export TEST_SEED=1440995966

您可以分别强制 RapidCheck 中使用的随机种子和用于填充 ONNX 输入向量的随机种子相同。仅设置第一个(RC_PARAMS),您将看到运行相同的测试配置但输入值不同。同时设置两者,您将看到相同的配置和相同的输入,从而获得完全相同的运行结果。

如果需要更改用于精度检查的 ATOL 和 RTOL,请将环境变量 TEST_ATOLTEST_RTOL 设置为新值。

启用 SIMD 指令

在支持的平台(目前是 s390x z14 及以上、x86 和 arm)上,数值测试可以为编译的模型生成 SIMD 指令。要启用 SIMD,请设置 TEST_ARGS 环境变量,例如,

TEST_ARGS="-march=z16" CTEST_PARALLEL_LEVEL=$(nproc) cmake --build . --config Release --target check-onnx-numerical

特定加速器的测试

目前我们提供了加速器 NNPA 的测试。它在此处 here 进行了描述。

使用 gdb

获取 ONNX 模型源代码

编译 ONNX 模型时,添加选项 --preserveMLIR。将创建一个 MLIR 格式的模型源代码文件,命名为 your_model_name.input.mlir。操作的行信息将附加并一直传播到二进制文件。在 gdb 中运行编译后的库时,您可以在模型中停止并按照 ONNX 操作进行单步调试。这是模型 test_add.onnx 的示例

$Debug/bin/onnx-mlir --preserveMLIR test_add.onnx
$. ../utils/build-run-onnx-lib.sh
$gdb Debug/bin/run-onnx-lib
(gdb) b run_main_graph
(gdb) run ./test_add.so
(gdb) list
1	builtin.module  {
2	  builtin.func @main_graph(%arg0: tensor<3x4x5xf32>, %arg1: tensor<3x4x5xf32>) -> tensor<3x4x5xf32> {
3	    %0 = "onnx.Add"(%arg0, %arg1) : (tensor<3x4x5xf32>, tensor<3x4x5xf32>) -> tensor<3x4x5xf32>
4	    return %0 : tensor<3x4x5xf32>
5	  }
(gdb) b 3
Breakpoint 2 at 0x3fffdf01778: file /home/chentong/onnx-mlir/build/test_add.input.mlir, line 3.
(gdb) c
Continuing.

Breakpoint 2, main_graph () at /home/chentong/onnx-mlir/build/test_add.input.mlir:3
3	    %0 = "onnx.Add"(%arg0, %arg1) : (tensor<3x4x5xf32>, tensor<3x4x5xf32>) -> tensor<3x4x5xf32>
(gdb) n
[Detaching after vfork from child process 2333437]
#  0) before op=     Add VMem:  6804
[Detaching after vfork from child process 2333470]
#  1) after  op=     Add VMem:  6804
4	    return %0 : tensor<3x4x5xf32>
(gdb)

请注意,Instrumentation 的输出显示 gdb 在 ONNX 操作级别正确地进行了单步调试。您需要为 onnx-mlir 运行 Instrumentation 添加额外标志,而 gdb 不需要。源文件是 test_add.input.mlir。未来的工作之一是在 gdb 中支持 ONNX 级别的符号。如果可以在 gdb 中打印出张量,那将非常有用。

使用 LLVM 调试支持

在 LLVM 和 MLIR 项目中添加跟踪代码的标准方法是使用 LLVM_DEBUG 宏。LLVM 的官方文档在此处 here

要在调试控制下插入单个“打印输出”,可以使用以下模板。

#include "llvm/Support/Debug.h"

#define DEBUG_TYPE "my_opt_name_here"
...
LLVM_DEBUG(llvm::dbgs() << "debug msg here" <<  obj_to_print << "\n");

要触发调试跟踪,只需使用 –debug-only=my_opt_name_here 选项调用编译器。

在源文件可能只有一个跟踪消息的情况下,可以使用另一个名为 DEBUG_WITH_TYPE 的宏。在这种情况下,您可以放弃定义 DEBUG_TYPE,改用以下内容。

DEBUG_WITH_TYPE("my_debug_msg", llvm::dbgs() << "my trace msg here\n");

要保护更大段代码,可以使用此模板。

LLVM_DEBUG({
  for(i...) {
    llvm::dbgs() << "emit trace for a: " << a << "\n";
    compute b;  // should be side effects free
    llvm::dbgs() << "emit trace for 'b':" << b << "\n";
    ...
});

项目中使用此支持的一些示例位于这些文件中。

同样,可以通过向 onnx-mlironnx-mlir-opt 添加 --debug-only=my_opt_name_here 选项来激活这些调试语句。

ONNX 模型库

我们提供了一个 Python 脚本 RunONNXModelZoo.py,用于检查 ONNX 模型库中模型的推理精度。 RunONNXModelZoo.py 要求 RunONNXModel.py 在同一文件夹中。例如,要检查 mnist-8 的推理精度

$ mkdir test && cd test
$ ln -s /onnx-mlir/utils/RunONNXModel.py
$ ln -s /onnx-mlir/utils/RunONNXModelZoo.py
$ ONNX_MLIR_HOME=/onnx-mlir/build/Release/ python RunONNXModelZoo.py -m mnist-8 -c="-O3"

运行脚本时使用 -h 可以查看所有选项。除了用于指定模型的 -m 标志和用于指定编译选项的 -c 标志外,有用的选项还有 -k 标志(将 onnx 模型保留在当前目录中作为 .tgz 文件)和 -l debug 标志(打印大量调试信息)。

要找出哪些模型可用,请运行脚本时使用 -p 打印可用模型列表;或者使用 -m 后跟不完整的名称,脚本将建议确切的名称。

如果不使用 -m 指定模型,脚本将检查 ONNX 模型库中的所有模型。

ONNX 模型库性能分析

如果您想收集有关模型库(或任何模型)的性能信息,最简单的方法是在编译时请求所需统计信息(使用 -profile-ir 标志),将输出统计信息重定向到文件,然后使用 make-report.py 进行分析。例如

> ONNX_MLIR_INSTRUMENT_FILE=run.log RunONNXModelZoo.py -c "-O3 --march=arm64 --profile-ir=Onnx" -m bertsquad-10
...
> make-report.py -r run.log
...
Statistics start (all ops).
  onnx.Add, 112, 0.0130570
  onnx.Cast, 105, 0.0001860
  onnx.Concat, 55, 0.0001290
  onnx.Constant, 473, 0.0008220

运行时性能分析信息也可以与特定的编译时统计信息结合使用。假设我们对 SIMD 统计信息感兴趣。我们使用 -opt-report 选项通知编译器发出编译时统计信息,并使用 --log-to-file 选项通知 RunONNXModelZoo.py 我们希望保留编译器输出。例如

> ONNX_MLIR_INSTRUMENT_FILE=run.log RunONNXModelZoo.py -c "-O3 --march=arm64 -opt-report=Simd --profile-ir=Onnx" -m bertsquad-10 --log-to-file compile.log
...
> make-report.py -c compile.log -r run.log
...
Statistics start (all ops).
  onnx.Add-simd, 112, 0.0130570
  onnx.Cast, 23, 0.0000650
  onnx.Gemm, 1, 0.0003570
  onnx.Gemm-simd, 72, 0.8109330

在上面的列表中,被向量化的操作被单独汇总,其各自的操作名称后附加了 -simd 后缀。

相同的选项和环境变量对于 RunONNXModel.pyRunONNXModelZoo.py 同样有效。