3. 输入张量内存对齐优化

任务目标

RK3588的NPU（RKNN）要求输入张量在内存地址和每个维度stride上满足特定对齐约束（通常为16字节对齐）。当输入数据来自非对齐的缓冲区（如普通malloc或图像解码结果）时，直接调用rknn_inputs_set会导致NPU内部执行额外的内存拷贝，降低推理效率。本案例实现一个最小对齐分配器，确保输入张量内存首地址及行stride均对齐到16字节，从而消除隐式拷贝。

核心实现

// align_alloc.h
#ifndef ALIGN_ALLOC_H
#define ALIGN_ALLOC_H

#include <cstdlib>
#include <cstdint>

// 分配对齐内存，返回对齐后的指针（调用free_align释放）
inline void* alloc_align(size_t size, size_t align = 16) {
    if (align & (align - 1)) return nullptr; // align必须是2的幂
    size_t offset = align - 1 + sizeof(void*);
    unsigned char* raw = (unsigned char*)malloc(size + offset);
    if (!raw) return nullptr;
    
    uintptr_t addr = (uintptr_t)raw + sizeof(void*);
    uintptr_t aligned = (addr + align - 1) & ~(uintptr_t)(align - 1);
    ((void**)aligned)[-1] = raw;  // 在对齐地址前保存原始指针
    return (void*)aligned;
}

inline void free_align(void* ptr) {
    if (!ptr) return;
    void* raw = ((void**)ptr)[-1];
    free(raw);
}

#endif // ALIGN_ALLOC_H
// main.cpp - 演示对齐分配与RKNN输入设置
#include <stdio.h>
#include <vector>
#include "rknn_api.h"
#include "align_alloc.h"

// 假设输入为 224x224 RGB 图像
#define IMG_H 224
#define IMG_W 224
#define IMG_C 3

int main() {
    // 1. 计算对齐后的行stride（16字节对齐）
    size_t row_size = IMG_W * IMG_C;  // 672字节
    size_t aligned_row = (row_size + 15) & ~(size_t)15;  // 对齐到16 -> 672本身已对齐
    
    // 2. 分配对齐内存（首地址16字节对齐）
    size_t total_size = aligned_row * IMG_H;
    unsigned char* input_data = (unsigned char*)alloc_align(total_size, 16);
    if (!input_data) { printf("alloc failed\n"); return -1; }
    
    // 填充测试数据（实际使用时从图像解码写入）
    for (int h = 0; h < IMG_H; h++) {
        unsigned char* row_ptr = input_data + h * aligned_row;
        for (int w = 0; w < IMG_W * IMG_C; w++) {
            row_ptr[w] = (unsigned char)(h + w);  // 任意测试值
        }
    }
    
    // 3. 设置RKNN输入（关键：传入对齐后的指针和stride）
    rknn_input inputs[1];
    memset(inputs, 0, sizeof(inputs));
    inputs[0].index = 0;
    inputs[0].type = RKNN_TENSOR_UINT8;
    inputs[0].size = total_size;
    inputs[0].fmt = RKNN_TENSOR_NHWC;
    inputs[0].buf = input_data;
    inputs[0].pass_through = 0;
    
    // 注意：rknn_inputs_set内部会检查buf地址和size是否满足对齐要求
    // 若满足则直接零拷贝引用，否则内部拷贝到对齐缓冲区
    // 实际调用（需已有ctx）：
    // rknn_inputs_set(ctx, 1, inputs);
    
    printf("Input buffer: addr=%p, aligned_row=%zu, total=%zu\n", 
           input_data, aligned_row, total_size);
    printf("Address 16-byte aligned: %s\n", 
           ((uintptr_t)input_data % 16 == 0) ? "yes" : "no");
    
    free_align(input_data);
    return 0;
}

编译与运行

# 编译（需RKNN Toolkit的include路径）
g++ -o align_demo main.cpp -I/path/to/rknn/include -L/path/to/rknn/lib -lrknnrt

# 运行
./align_demo

验证方法

  • 对齐检查：程序输出中Address 16-byte aligned: yes表明首地址满足要求；aligned_row值（672）为16的倍数，满足行stride对齐。
  • 零拷贝验证：在真实推理代码中，可在rknn_inputs_set前后对比rknn_query的RKNN_QUERY_INPUT_ATTR返回的stride值——若传入的stride与NPU要求一致，则内部不会产生额外拷贝。本案例中aligned_row=672恰好等于NPU默认的NHWC布局stride，因此可直接零拷贝。
  • 性能对比：使用非对齐malloc分配相同大小缓冲区并设置输入，对比两者单次推理耗时（需运行100次取平均）。对齐版本应无额外拷贝开销，耗时更短。

关键说明

  • RK3588 NPU要求输入缓冲区首地址16字节对齐，且每个维度的stride需满足特定对齐（通常也是16的倍数）。本案例的alloc_align通用分配器可满足首地址要求；行stride需根据实际输入尺寸计算并填充到aligned_row。
  • 当输入尺寸本身导致行字节数非16倍数时（如单通道灰度图宽为奇数），必须手动将aligned_row向上取整到16倍数，并在填充数据时按aligned_row作为行跨度写入，而非原始IMG_W*IMG_C。
  • 若使用RKNN Toolkit的Python接口，对应功能由rknn.config(force_float16=False)内部处理，C接口需自行保证对齐。
← 上一章 📖 返回目录 下一章 →