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);
 
  // 使用张量构造一个张量列表
OMTensor *list[2] = {x1, x2};
  return omTensorListCreate(list, 2);
}
 
int main() {
  // 生成输入张量列表
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。