RKNN Runtime C API V2.5.0 — 张量、内存分配与同步

平台参考与来源

本文依据 RKNN V2.5.0 原厂 PDF 整理。适用于原文列出的芯片和接口；ArmSoM 板卡须先核对 SoC、NPU 驱动、Toolkit/Runtime、SDK/BSP 和内核版本。尚未进行 ArmSoM 硬件验证，原文示例及测量值不是 ArmSoM 实测。

系列导航

• 平台支持、头文件与编译链接
• 内存所有权与通用、零拷贝流程
• 设备、上下文与初始化标志
• 模型、张量与动态形状查询
• 输入、执行、输出与动态形状设置
• 张量、内存分配与同步
• 矩阵乘法接口与动态形状
• 自定义算子注册与属性
• 状态码、张量、内存与扩展结构
• 矩阵乘法数据结构
• 自定义算子结构与废弃 API 迁移

5.4. 张量和内存

tensor 内存分配、释放、同步及用戶自定义内存设置。

rknn_create_mem_from_phys

当用戶需要自己分配内存让NPU使用时,通过rknn_create_mem_from_phys函数可以创建一个rknn_tensor_mem结构体并得到它的指针,该函数通过传入物理地址、虚拟地址以及大小,外部内存相关的信息会赋值给rknn_tensor_mem结构体。

示例代码如下:

//suppose we have got buffer information as input_phys, input_virt and size
rknn_tensor_mem* input_mems [1];
input_mems[0] = rknn_create_mem_from_phys(ctx, input_phys, input_virt, size);

rknn_create_mem_from_fd

当用戶要自己分配内存让NPU使用时,rknn_create_mem_from_fd函数可以创建一个rknn_tensor_mem结构体并得到它的指针,该函数通过传入文件描述符fd、偏移、虚拟地址以及大小,外部内存相关的信息会赋值给rknn_tensor_mem结构体。

原稿差异说明：中文 PDF 第 25 页将该参数写成 void virt_addr，星号 * 位于说明句末；英文 PDF 第 26 页写为 void *virt_addr。上表保留对应语言原稿写法。这是原稿中的类型标注差异，不是已验证的接口定义；使用时请核对实际 V2.5.0 SDK 的头文件声明。

示例代码如下:

//suppose we have got buffer information as input_fd, input_virt and size
rknn_tensor_mem* input_mems [1];
input_mems[0] = rknn_create_mem_from_fd(ctx, input_fd, input_virt, size, 0);

rknn_create_mem

当用戶要NPU内部分配内存时,rknn_create_mem函数可以分配用戶指定的内存大小,并返回一个rknn_tensor_mem结构体。

示例代码如下:

//suppose we have got buffer size
rknn_tensor_mem* input_mems [1];
input_mems[0] = rknn_create_mem(ctx, size);

rknn_create_mem2

当用戶要NPU内部分配内存时,rknn_create_mem2函数可以分配用戶指定的内存大小及内存类型,并返回一个rknn_tensor_mem结构体。

rknn_create_mem2与rknn_create_mem的主要区别是rknn_create_mem2带了一个alloc_flags,可以指定分配的内存是否cacheable的,而rknn_create_mem不能指定,默认就是cacheable。

示例代码如下:

//suppose we have got buffer size
rknn_tensor_mem* input_mems [1];
input_mems[0] = rknn_create_mem2(ctx, size, RKNN_FLAG_MEMORY_NON_CACHEABLE);

rknn_destroy_mem

rknn_destroy_mem函数会销毁rknn_tensor_mem结构体,用戶分配的内存需要自行释放。

示例代码如下:

rknn_tensor_mem* input_mems [1];
int ret = rknn_destroy_mem(ctx, input_mems[0]);

rknn_set_weight_mem

如果用戶自己为网络权重分配内存,初始化相应的rknn_tensor_mem结构体后,在调用rknn_run前,通过rknn_set_weight_mem函数可以让NPU使用该内存。

示例代码如下:

rknn_tensor_mem* weight_mems [1];
int ret = rknn_set_weight_mem(ctx, weight_mems[0]);

rknn_set_internal_mem

如果用戶自己为网络中间tensor分配内存,初始化相应的rknn_tensor_mem结构体后,在调用rknn_run前,通过rknn_set_internal_mem函数可以让NPU使用该内存。

示例代码如下:

rknn_tensor_mem* internal_tensor_mems [1];
int ret = rknn_set_internal_mem(ctx, internal_tensor_mems[0]);

rknn_set_io_mem

如果用户自己为网络输入/输出 tensor 分配内存，初始化相应的 rknn_tensor_mem 结构体后，在调用 rknn_run 前，通过 rknn_set_io_mem 函数可以让 NPU 使用该内存。该接口属于零拷贝 API，输入/输出 tensor 的内存可以是运行时分配的（通过 rknn_create_mem），也可以是外部分配的（通过 rknn_create_mem_from_fd / rknn_create_mem_from_phys）。与 rknn_init 类似，rknn_set_io_mem 只需在开始调用一次，之后反复调用 rknn_run 即可。

使用说明:

1. 传入的 rknn_tensor_attr 应使用 rknn_query 接口查询得到。零拷贝 API 的输入/输出属性应使用 RKNN_QUERY_NATIVE_INPUT_ATTR / RKNN_QUERY_NATIVE_OUTPUT_ATTR 查询（而非 RKNN_QUERY_INPUT_ATTR / RKNN_QUERY_OUTPUT_ATTR），以获得 NPU 硬件效率最优的 layout 和 type。

1. 创建输入/输出内存时,若无需修改查询得到的layout和type,直接按size_with_stride分配内存大小;若修改了layout和type,则需按修改后的size分配(例如输出数据类型由int8修改为float32时,size需改为原size的四倍,即n_elems *sizeof(float))。

1. 输入数据需按照查询或修改后的layout排列,例如输入layout为NHWC时按NHWC排列;当pass_through=1时,width可能需按w_stride做对⻬。

1. 推理完成后从rknn_tensor_mem结构体读取输出内存信息;外部分配的内存还需调用与分配内存类型关联的接口刷新输出缓存。

示例代码如下:

// 零拷⻉API:查询原生属性 -> 修改类型/布局 -> 分配内存 -> 设置输入/输出内存
rknn_tensor_attr input_attrs[1], output_attrs[1];
rknn_tensor_mem* input_mems[1];
rknn_tensor_mem* output_mems[1];

memset(input_attrs, 0, sizeof(input_attrs));
input_attrs[0].index = 0;
rknn_query(ctx, RKNN_QUERY_NATIVE_INPUT_ATTR, &(input_attrs[0]), sizeof(rknn_tensor_attr));
// 默认输入类型为 int8(归一化和量化需在外部计算);
// 设置为 uint8 时,会将归一化和量化融合到 NPU 上执行
input_attrs[0].type = RKNN_TENSOR_UINT8;
// 默认 fmt 为 NHWC,零拷⻉模式下 NPU 仅支持 NHWC 输入
input_attrs[0].fmt = RKNN_TENSOR_NHWC;
input_mems[0] = rknn_create_mem(ctx, input_attrs[0].size_with_stride);
// 将输入数据拷⻉到 input_mems[0]->virt_addr ...

memset(output_attrs, 0, sizeof(output_attrs));
output_attrs[0].index = 0;
rknn_query(ctx, RKNN_QUERY_NATIVE_OUTPUT_ATTR, &(output_attrs[0]), sizeof(rknn_tensor_attr));
// 默认输出类型取决于模型,这里需要 float32 用于计算 top5
output_attrs[0].type = RKNN_TENSOR_FLOAT32;
output_mems[0] = rknn_create_mem(ctx, output_attrs[0].n_elems * sizeof(float));

// 设置输入/输出内存
ret = rknn_set_io_mem(ctx, input_mems[0], &input_attrs[0]);
ret = rknn_set_io_mem(ctx, output_mems[0], &output_attrs[0]);

// 推理
ret = rknn_run(ctx, NULL);
// 从 output_mems[0]->virt_addr 读取输出结果 ...

rknn_mem_sync

rknn_create_mem函数创建的内存默认是带cacheable标志的,对于带cacheable标志创建的内存,在被CPU和NPU同时使用时,由于cache行为会导致数据一致性问题。该接口用于同步一块带cacheable标志创建的内存,保证CPU和NPU访问这块内存的数据是一致的。

示例代码如下:

ret =rknn_mem_sync(ctx, &outputs[0].mem,
RKNN_MEMORY_SYNC_FROM_DEVICE);
if (ret < 0) {
   fprintf(stderr, " rknn_mem_sync error! ret=%d\n", ret);
   return -1;
}

来源与校对说明

原始技术章节按用途拆分，代码、标识符和参数保留。中文采用中文原稿；英文、俄文、德文采用对应英文稿。无图章节不添加虚构配图。

原稿保留事项：章节编号与示例输出按原文保留，示例日志可能显示早期 SDK 版本，不代表本文当前软件版本。Runtime 原文直接由第 3 章跳到第 5 章；查询小节序号存在重复。rknn_init 参数表称 extend 当前禁用，但后文列出 user_key/共享权重等字段；查询原生动态输出的说明写 RKNN_QUERY_CURRENT_OUTPUT_ATTR，示例使用 RKNN_QUERY_CURRENT_NATIVE_OUTPUT_ATTR。原文对部分 native NHWC/NC1HWC2 查询描述也存在输入/输出混称。请以实际 V2.5.0 头文件及目标 SDK 核对，不把这些说明当作已验证行为。RK2118 引用文件在原文写为 RKNN3_RK2118_Quick_Start.pdf，当前目录对应 RKNN2 快速入门。

RKNN-Toolkit2 · ArmSoM documentation