4. 张量内存分配与数据拷贝

任务目标

在NPU设备上为输入/输出张量分配显存，将CPU端数据拷贝至NPU，完成一次推理后将结果拷回CPU，并校验数据一致性。本案例以华为Ascend NPU + CANN接口为例，使用最小可复现实现完成闭环。

环境依赖

• Atlas 300/800 系列推理卡，已安装CANN Toolkit（≥5.1）
• 编译链接：-lascendcl -lacl_op_compiler
• 头文件：#include "acl/acl.h"

关键完整代码

#include <cstdio>
#include <cstdint>
#include <cstring>
#include <vector>
#include "acl/acl.h"

#define CHECK(expr) do { \
    aclError err = (expr); \
    if (err != ACL_SUCCESS) { \
        printf("FAIL at %s:%d, ret=%d\n", __FILE__, __LINE__, err); \
        return -1; \
    } \
} while(0)

// 最小NPU推理：输入加1（用Add算子模拟）
static const char* MODEL_PATH = "./add_model.om";

int main() {
    // 1. 初始化
    CHECK(aclInit(nullptr));
    CHECK(aclrtSetDevice(0));
    aclrtContext ctx = nullptr;
    CHECK(aclrtCreateContext(&ctx, 0));

    // 2. 加载模型
    aclmdlDesc* modelDesc = aclmdlCreateDesc();
    CHECK(aclmdlLoadFromFile(MODEL_PATH, &modelId));
    CHECK(aclmdlGetDesc(modelDesc, modelId));

    // 3. 获取输入/输出尺寸
    size_t inputSize = aclmdlGetInputSizeByIndex(modelDesc, 0);
    size_t outputSize = aclmdlGetOutputSizeByIndex(modelDesc, 0);
    printf("inputSize=%zu, outputSize=%zu\n", inputSize, outputSize);

    // 4. 分配CPU与NPU内存
    void* cpuInput = nullptr;
    void* cpuOutput = nullptr;
    void* npuInput = nullptr;
    void* npuOutput = nullptr;
    CHECK(aclrtMallocHost(&cpuInput, inputSize));
    CHECK(aclrtMallocHost(&cpuOutput, outputSize));
    CHECK(aclrtMalloc(&npuInput, inputSize, ACL_MEM_MALLOC_HUGE_FIRST));
    CHECK(aclrtMalloc(&npuOutput, outputSize, ACL_MEM_MALLOC_HUGE_FIRST));

    // 5. 准备CPU输入数据（4个float）
    float* inData = static_cast<float*>(cpuInput);
    inData[0] = 1.0f; inData[1] = 2.0f; inData[2] = 3.0f; inData[3] = 4.0f;

    // 6. CPU → NPU 拷贝
    CHECK(aclrtMemcpy(npuInput, inputSize, cpuInput, inputSize, ACL_MEMCPY_DEVICE_TO_DEVICE));

    // 7. 创建输出描述并执行推理
    aclmdlDataset* inputSet = aclmdlCreateDataset();
    aclmdlDataset* outputSet = aclmdlCreateDataset();
    aclDataBuffer* inBuf = aclCreateDataBuffer(npuInput, inputSize);
    aclDataBuffer* outBuf = aclCreateDataBuffer(npuOutput, outputSize);
    aclmdlAddDatasetBuffer(inputSet, inBuf);
    aclmdlAddDatasetBuffer(outputSet, outBuf);
    CHECK(aclmdlExecute(modelId, inputSet, outputSet));

    // 8. NPU → CPU 拷贝
    CHECK(aclrtMemcpy(cpuOutput, outputSize, npuOutput, outputSize, ACL_MEMCPY_DEVICE_TO_DEVICE));

    // 9. 数据一致性校验（预期：每个元素+1）
    float* outData = static_cast<float*>(cpuOutput);
    bool pass = true;
    for (int i = 0; i < 4; i++) {
        if (outData[i] != inData[i] + 1.0f) {
            printf("Mismatch at %d: got %.1f, expect %.1f\n", i, outData[i], inData[i]+1);
            pass = false;
        }
    }
    printf(pass ? "PASS: data consistent\n" : "FAIL: data mismatch\n");

    // 10. 清理
    aclDestroyDataBuffer(inBuf);
    aclDestroyDataBuffer(outBuf);
    aclmdlDestroyDataset(inputSet);
    aclmdlDestroyDataset(outputSet);
    aclrtFree(npuInput);
    aclrtFree(npuOutput);
    aclrtFreeHost(cpuInput);
    aclrtFreeHost(cpuOutput);
    aclmdlDestroyDesc(modelDesc);
    aclmdlUnload(modelId);
    aclrtDestroyContext(ctx);
    aclrtResetDevice(0);
    aclFinalize();
    return pass ? 0 : -1;
}

编译与运行

# 编译（假设模型文件add_model.om已生成）
g++ -o npu_copy npu_copy.cpp -lascendcl -lacl_op_compiler

# 运行
./npu_copy

验证方法与判定标准

• 程序输出 inputSize=16, outputSize=16（4个float各占4字节）
• 最终输出 PASS: data consistent 即通过
• 若输出 Mismatch 或返回非0，说明拷贝或推理链路存在数据错误

关键说明

• 第6步使用 ACL_MEMCPY_DEVICE_TO_DEVICE 是CANN中CPU↔NPU拷贝的标准模式（实际走PCIe/DMA通道）
• 模型文件需提前用ATC工具将onnx/pb转换为om格式，本案例不展开模型生成细节
• 若需验证更大张量，修改第5步的循环赋值与第9步的校验范围即可