CANN 系列底层篇：基于 shmem 实现 NPU 设备内存的高效共享

CANN 系列底层篇：基于 shmem 实现 NPU 设备内存的高效共享

cann组织链接：https://atomgit.com/cann

ops-nn仓库链接：https://atomgit.com/cann/ops-nn

在高性能 AI 推理服务、分布式训练或边缘协同计算场景中，多个进程或线程往往需要访问同一块 NPU 设备内存。例如：

• 一个预处理进程将图像数据拷入 NPU 显存；
• 多个推理线程并发读取该数据进行模型推理；
• 后处理进程从显存读取结果并返回给客户端。

若缺乏统一的内存管理机制，极易出现：

• 内存重复分配 → 显存浪费
• 访问冲突 → 数据错乱
• 生命周期不一致 → 段错误或悬空指针

为此，CANN 提供了 shmem 模块，专门解决 NPU 设备内存的跨进程/线程共享与生命周期管理问题。

📌 项目地址：https://gitcode.com/cann/shmem

---

一、shmem 的核心功能

shmem 并非简单的内存分配器，而是一个带引用计数、命名空间隔离、跨进程同步的设备内存管理器。其关键特性包括：

特性	说明
命名内存块（Named Buffer）	通过字符串名称标识一块设备内存，便于跨进程查找
引用计数（Reference Counting）	自动管理内存生命周期，最后一个使用者释放内存
POSIX 共享内存后端	利用 /dev/shm 或 shm_open 实现进程间元数据同步
线程安全	所有接口加锁保护，支持多线程并发访问
零拷贝共享	多个进程映射同一物理设备地址，避免数据复制

💡 类比：shmem 相当于 NPU 上的 "mmap + shared_ptr" 组合。

---

二、shmem 的典型使用场景

1. 多进程推理服务

   • 主进程加载模型权重到共享显存；
   • 工作进程直接使用，无需重复加载。

2. 生产者-消费者流水线

   • 生产者写入 input_buffer；
   • 消费者（推理引擎）读取并计算；
   • 结果写入 output_buffer，由另一消费者读取。

3. 多卡数据广播

   • 在 device 0 分配共享内存；
   • 通过 hcll 将其映射到 device 1~N，实现高效广播。

---

三、实战示例：多进程共享 NPU 输入缓冲区

✅ 场景描述

• 主进程（Producer） ：分配一块 4MB 的 NPU 设备内存，命名为 "input_tensor"，并写入测试数据。
• 子进程（Consumer） ：通过名称打开同一块内存，将其作为 ge 图的输入执行推理。

注意：为简化，本例使用 fork() 模拟多进程（实际部署可用多进程服务框架如 Gunicorn + NPU）。

---

✅ 主进程代码（producer.cpp）

#include <sys/wait.h>
#include <unistd.h>
#include <iostream>
#include "shmem.h"
#include "hcll.h"

int main() {
    const size_t size = 4 * 1024 * 1024; // 4MB
    const char* name = "input_tensor";

    hcllInit();
    hcllSetDevice(0);

    // Step 1: 创建命名共享设备内存
    ShmemHandle handle;
    void* dev_ptr = shmem_create(name, size, &handle);
    if (!dev_ptr) {
        std::cerr << "Failed to create shared memory!" << std::endl;
        return -1;
    }

    // Step 2: 准备 Host 数据
    std::vector<float> host_data(size / sizeof(float));
    for (size_t i = 0; i < host_data.size(); ++i) {
        host_data[i] = static_cast<float>(i % 100);
    }

    // Step 3: 拷贝到共享设备内存
    hcllMemcpy(dev_ptr, host_data.data(), size, HCLL_MEMCPY_HOST_TO_DEVICE);

    std::cout << "Producer: Shared memory '" << name 
              << "' created and filled." << std::endl;

    // Step 4: 启动消费者进程
    pid_t pid = fork();
    if (pid == 0) {
        // 子进程：执行 consumer
        execl("./consumer", "consumer", nullptr);
    } else {
        wait(nullptr); // 等待子进程结束
        // 主进程退出时自动释放（引用计数归零）
    }

    shmem_close(handle);
    hcllFinalize();
    return 0;
}

---

✅ 子进程代码（consumer.cpp）

#include <iostream>
#include "shmem.h"
#include "ge_api.h"
#include "hcll.h"

int main() {
    const char* name = "input_tensor";
    const size_t size = 4 * 1024 * 1024;

    hcllInit();
    hcllSetDevice(0);

    // Step 1: 打开已存在的共享内存
    ShmemHandle handle;
    void* dev_ptr = shmem_open(name, &handle);
    if (!dev_ptr) {
        std::cerr << "Consumer: Failed to open shared memory!" << std::endl;
        return -1;
    }

    std::cout << "Consumer: Opened shared memory '" << name << "'" << std::endl;

    // Step 2: 构建简单图（例如：output = input * 2）
    ge::Graph graph("scale_graph");
    ge::TensorDesc desc(ge::Dims({static_cast<int64_t>(size / sizeof(float))}), 
                        ge::FORMAT_ND, ge::DT_FLOAT);
    auto input = graph.AddInput("x", desc);
    
    auto mul_op = ge::OperatorFactory::CreateOperator("Mul", "Mul");
    mul_op.SetInput("x1", input)
          .SetAttr("x2", static_cast<float>(2.0f)); // 标量乘法
    graph.SetOutput(mul_op.GetOutput("y"), "output");

    auto session = ge::CreateSession(graph, {});
    session->BindInput("x", dev_ptr);

    void* output_ptr;
    hcllMalloc(&output_ptr, size);
    session->BindOutput("output", output_ptr);

    session->Run();

    // 拷回部分结果验证
    std::vector<float> result(5);
    hcllMemcpy(result.data(), output_ptr, 5 * sizeof(float), HCLL_MEMCPY_DEVICE_TO_HOST);
    std::cout << "Consumer result (first 5): ";
    for (float v : result) std::cout << v << " ";
    std::cout << std::endl;

    // 清理
    hcllFree(output_ptr);
    shmem_close(handle); // 引用计数减1
    ge::DestroySession(session);
    hcllFinalize();
    return 0;
}

---

🔧 编译命令

# 编译 producer
g++ -O2 -std=c++17 producer.cpp -o producer \
    -I/path/to/shmem/include -I/path/to/hcll/include -I/path/to/ge/include \
    -L/path/to/libs -lshmem -lhcll -lge -lnpu_runtime

# 编译 consumer
g++ -O2 -std=c++17 consumer.cpp -o consumer \
    -I/path/to/shmem/include -I/path/to/hcll/include -I/path/to/ge/include \
    -L/path/to/libs -lshmem -lhcll -lge -lnpu_runtime

---

四、shmem 的内部机制简析

1. 元数据存储

   共享内存的名称、大小、设备ID、引用计数等信息存储在 POSIX 共享内存段（如 /dev/shm/shmem_input_tensor.meta）。

2. 设备内存句柄映射
   shmem_create 调用 hcllMalloc 分配真实设备内存，并将物理地址通过驱动注册为可共享句柄。

3. 引用计数同步

   每次 shmem_open 增加计数，shmem_close 减少计数；当计数归零且无活跃映射时，自动调用 hcllFree。

4. 权限与安全

   支持设置访问权限（如只读/读写），防止非法修改。

---

五、性能与可靠性优势

指标	传统方式	使用 shmem
内存占用	每进程独立拷贝 → N×显存	单份共享 → 1×显存
启动延迟	每次加载模型/数据	一次加载，多次复用
数据一致性	需手动同步	自动保证
容错性	进程崩溃易导致内存泄漏	引用计数+内核清理保障

---

六、结语：shmem ------ 高效协同的基石

虽然 shmem 不直接参与计算，但它是构建高并发、低资源消耗 AI 服务 的关键基础设施。它让多个计算单元（进程/线程/设备）能够安全、高效地共享同一份数据，从而最大化 NPU 的利用率。

至此，我们已经覆盖了 CANN 软件栈的主要开源模块：

• 计算层 ：ops-math, tbe
• 通信层 ：hcll
• 调度层 ：ge
• 资源管理层 ：shmem

它们共同构成了一个完整、高效、可扩展的国产 AI 软件生态。

---

🔗 探索全部项目：https://gitcode.com/cann

📌 建议结合 shmem/samples/ 中的多进程示例深入实践。

是否希望下一篇对 CANN 整体架构做一次全景式总结？或者深入某个特定应用场景（如大模型推理、CV/NLP 定制 pipeline）？欢迎告诉我方向！