ONNX-MLIR 运行时 API 文档

简介

ONNX-MLIR 项目带有一个可执行文件 onnx-mlir，能够将 ONNX 模型编译成共享库。在本教程中，我们将演示如何使用 ONNX-MLIR 的运行时 API 以编程方式与编译后的共享库进行交互。

C 运行时 API

数据结构

OMTensor 是用于描述与张量输入或输出关联的运行时信息（秩、形状、数据类型等）的数据结构。

OMTensorList 是用于容纳指向 OMTensor 的指针列表的数据结构，以便它们可以作为输入和输出进出编译后的模型。

OMEntryPoint 是用于返回模型中所有入口点名称的数据结构。这些入口点名称是模型中推理函数的符号。

OMSignature 是用于将给定入口点的输出签名以 JSON 字符串形式返回的数据结构。

模型入口点签名

所有编译后的模型都将具有完全相同的 C 函数签名，相当于

OMTensorList* run_main_graph(OMTensorList*);

直观地说，模型接受一个张量列表作为输入，并返回一个张量列表作为输出。

使用 C 运行时调用模型

API

我们将演示如何使用 API 函数来运行一个由加法操作组成的简单 ONNX 模型。要创建这样的 ONNX 模型，请使用此 Python 脚本

要编译上述模型，请运行 onnx-mlir add.onnx，此时应该会出现一个名为 "add.so" 的二进制库。我们可以使用以下 C 代码来调用编译后的函数，该函数计算两个输入的总和。

#include <OnnxMlirRuntime.h>
#include <stdio.h>
 
OMTensorList *run_main_graph(OMTensorList *);
 
OMTensorList *create_input_list() {
  // 共享形状和秩。
int64_t shape[] = {3, 2};
int64_t num_elements = shape[0] * shape[1];
int64_t rank = 2;
 
  // 构建填充有 1 或 2 的浮点数组。
  float *x1Data = (float *)malloc(sizeof(float) * num_elements);
  for (int i = 0; i < num_elements; i++)
x1Data[i] = 1.0;
  float *x2Data = (float *)malloc(sizeof(float) * num_elements);
  for (int i = 0; i < num_elements; i++)
x2Data[i] = 2.0;
 
  // 使用 omTensorCreateWithOwnership "true"，以便在销毁张量时自动
  // 释放浮点数组。
OMTensor *x1 = omTensorCreateWithOwnership(x1Data, shape, rank, ONNX_TYPE_FLOAT, true);
OMTensor *x2 = omTensorCreateWithOwnership(x2Data, shape, rank, ONNX_TYPE_FLOAT, true);
 
  // 使用张量构建 TensorList
OMTensor *list[2] = {x1, x2};
  return omTensorListCreate(list, 2);
}
 
int main() {
  // 生成输入 TensorList
OMTensorList *input_list = create_input_list();
 
  // 调用编译后的 ONNX 模型函数。
OMTensorList *output_list = run_main_graph(input_list);
  if (!output_list) {
    // 可以检查 errno 来获取有关错误的更多信息。
    return 1;
  }
 
  // 从输出列表中获取第一个张量。
OMTensor *y = omTensorListGetOmtByIndex(output_list, 0);
  float *outputPtr = (float *) omTensorGetDataPtr(y);
 
  // 打印其内容，应全部为 3。
  for (int i = 0; i < 6; i++)
printf("%f ", outputPtr[i]);
printf("\n");
 
  // 销毁列表及其中的张量。
  // 仅当只想销毁列表本身时，才使用 omTensorListDestroyShallow。
omTensorListDestroy(input_list);
omTensorListDestroy(output_list);
  return 0;
}

使用 gcc main.c add.so -o add 编译，您应该会看到一个名为 add 的可执行文件。运行它，输出应为

3.000000 3.000000 3.000000 3.000000 3.000000 3.000000

正是它应该有的样子。

释放张量内存

一般来说，如果调用者创建了一个张量对象 (omTensorCreate)，则在销毁张量后，他们需要单独释放数据缓冲区。如果 onnx-mlir 创建了张量 (run_main_graph)，那么张量对象拥有数据缓冲区，它会在张量销毁时自动释放。

此默认行为可以更改。创建张量时，用户可以使用 omTensorCreateWithOwnership 来显式设置数据缓冲区的拥有权。此外，在创建张量后，可以使用 omTensorSetOwning 来更改拥有权设置。

当调用 omTensorDestroy 时，如果拥有权标志设置为 "true"，则销毁张量还将释放任何相关的数据缓冲区内存。如果拥有权标志设置为 "false"，那么用户负责在销毁张量后释放数据缓冲区内存。

对于张量列表对象，当调用 omTensorListDestory 时，会在列表包含的所有张量上调用 omTensorDestory。每个张量的数据缓冲区将根据每个张量的拥有权设置进行释放。

要销毁一个 TensorList 而不自动销毁它包含的张量，请使用 omTensorListDestroyShallow。

参考

有关可用 C 运行时 API 的完整参考，请参阅 include/onnx-mlir/Runtime/OMTensor.h 和 include/onnx-mlir/Runtime/OMTensorList.h。