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"，以便浮点数数组在 Tensor 销毁时自动释放
  //
OMTensor *x1 = omTensorCreateWithOwnership(x1Data, shape, rank, ONNX_TYPE_FLOAT, true);
OMTensor *x2 = omTensorCreateWithOwnership(x2Data, shape, rank, ONNX_TYPE_FLOAT, true);
 
  // 使用 Tensors 构造一个 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。