相关链接：
- CANN 组织主页：https://atomgit.com/cann
- shmem 仓库地址：https://atomgit.com/cann/shmem
前言
CANN SHMEM（Shared Memory Library）作为一套面向多机多卡场景的高性能内存通信库，其核心能力在于实现跨设备、跨进程的高效内存访问与数据同步。在单机多进程分布式训练或推理中，不同进程需共享大量中间结果（如 KV Cache、梯度张量），而传统 IPC（Inter-Process Communication）机制往往因虚拟地址空间隔离导致额外拷贝开销。SHMEM 通过精心设计的统一内存池与跨进程虚拟地址映射机制，实现了真正的“零拷贝共享”，显著降低通信延迟并提升带宽利用率。
本文基于shmem 仓库（https://atomgit.com/cann/shmem）深入剖析其内存管理架构。我们将从内存池初始化、物理页分配、跨进程地址映射到实际 API 使用，逐层揭示 SHMEM 如何在保证安全性的前提下，构建全局一致的共享内存视图。
一、SHMEM 内存模型概览
SHMEM 将共享内存抽象为一个全局地址空间，所有参与进程通过shmem_alloc分配的内存均位于此空间内。关键特性包括：
- 统一寻址：同一物理内存块在不同进程中具有相同虚拟地址；
- 生命周期管理：由 SHMEM 运行时统一管理，避免用户手动释放；
- 硬件亲和性：支持绑定到特定设备（Device）或主机（Host）内存。
仓库 README 指出：“SHMEM 封装 Host 侧与 Device 侧接口，实现跨设备的高效内存访问……简化卡间数据同步与通算融合算子开发流程。”
二、内存池初始化与物理页分配
SHMEM 的内存池在aclshmemx_init时创建，分为Host Pool与Device Pool。
2.1 初始化流程
// src/host/init/shmem_init.ccint32_taclshmemx_init(constaclshmemx_init_attr_t*attr){// 1. 创建 Host 共享内存池host_pool_=std::make_unique<HybridMemoryPool>(attr->local_mem_size,MEM_TYPE_HOST);// 2. 创建 Device 共享内存池（若启用 D2D）if(attr->enable_device_memory){device_pool_=std::make_unique<DeviceMemoryPool>(attr->device_mem_size);}// 3. 初始化跨进程通信通道（用于地址同步）bootstrap_=std::make_unique<Bootstrap>();returnACL_SUCCESS;}默认 Host Pool 大小为 16GB，可通过attr.local_mem_size调整。
2.2 物理页分配策略
SHMEM 使用HugeTLB（大页）提升 TLB 效率：
// src/host/mem/hybrid_memory_pool.ccvoid*HybridMemoryPool::Allocate(size_t size){// 尝试分配 2MB 大页void*ptr=mmap(nullptr,size,PROT_READ|PROT_WRITE,MAP_SHARED|MAP_HUGETLB|MAP_ANONYMOUS,-1,0);if(ptr==MAP_FAILED){// 回退到普通页ptr=mmap(...,MAP_SHARED|MAP_ANONYMOUS,...);}returnptr;}分配的内存区域通过mmap映射为共享匿名内存，确保可被其他进程访问。
三、跨进程虚拟地址映射：同地址保障机制
SHMEM 的核心创新在于强制所有进程对同一物理内存使用相同的虚拟地址，从而避免地址转换开销。
3.1 地址协商协议
当进程 A 调用shmem_alloc时：
- 从本地内存池分配一块物理内存，获得虚拟地址
V_A； - 通过Bootstrap 通道（基于 Unix Domain Socket）广播
(物理标识, V_A)给所有对等进程； - 进程 B 收到后，调用
mmap时显式指定地址V_A：
// src/host/bootstrap/bootstrap.ccvoidBootstrap::SyncAddress(uint64_tphys_id,void*virt_addr){// 向所有 peer 发送 (phys_id, virt_addr)for(auto&peer:peers_){Send(peer.sock,phys_id,virt_addr);}}// peer 端处理void*PeerHandler::MapRemoteMemory(uint64_tphys_id,void*hint_addr){// 使用 hint_addr 作为 mmap 的 addr 参数void*mapped=mmap(hint_addr,size,PROT_READ|PROT_WRITE,MAP_SHARED|MAP_FIXED,// 关键：MAP_FIXEDmem_fd,offset);if(mapped!=hint_addr){LOG(FATAL)<<"Address conflict!";// 地址已被占用}returnmapped;}MAP_FIXED标志强制内核将内存映射到指定虚拟地址。若该地址已被占用，则触发错误——因此 SHMEM 要求所有进程同步初始化，确保地址空间布局一致。
3.2 唯一 ID 与冲突检测
为避免不同作业间地址冲突，SHMEM 引入Unique ID机制（见 PR #116, 2026年2月）：
// include/shmem/host/init.htypedefstructaclshmemx_set_attr_uniqueid_args{uint64_tunique_id;// 由用户或调度器分配}aclshmemx_set_attr_uniqueid_args_t;// 设置唯一 IDaclshmemx_set_attr(ACLSHMEMX_ATTR_UNIQUEID,&args,sizeof(args));Unique ID 作为内存池命名空间前缀，隔离不同作业的地址空间。
四、Device 内存共享：驱动层协同
对于 Device 内存，SHMEM 依赖底层驱动提供IPC 句柄（如 fd）。
4.1 Device 内存注册
// src/device/mem/device_memory_pool.ccHcclResultDeviceMemoryPool::Register(void*dev_ptr,size_t size){// 1. 调用驱动 ioctl 获取 mem_fdintmem_fd=driver_get_ipc_handle(dev_ptr,size);// 2. 将 mem_fd 通过 Bootstrap 发送给 peersbootstrap_->ShareFd(mem_fd,size);returnHCCL_SUCCESS;}4.2 对端映射
Peer 进程收到mem_fd后，在自身设备上下文中导入：
void*DeviceMemoryPool::ImportRemote(intmem_fd,size_t size){// 调用驱动 ioctl 导入 fdvoid*dev_ptr=driver_import_ipc_handle(mem_fd,size);returndev_ptr;// 返回本地设备虚拟地址}由于设备虚拟地址空间独立，Device 共享不保证同地址，但 SHMEM 通过句柄表维护逻辑一致性。
五、内存池生命周期与释放
SHMEM 内存由运行时自动管理，用户无需显式释放：
// 用户分配void*shared_buf=aclshmemx_malloc(1024*1024);// 1MB// 使用...// ...// 进程退出时自动清理// aclshmemx_finalize(); // 可选显式清理内部通过引用计数跟踪内存使用：
// src/host/mem/memory_region.ccclassMemoryRegion{std::atomic<int>ref_count_{1};public:voidAddRef(){ref_count_.fetch_add(1);}voidRelease(){if(ref_count_.fetch_sub(1)==1){// 所有进程释放后，真正 munmapUnmapPhysicalPages();}}};跨进程引用计数通过原子操作或 Bootstrap 同步维护。
六、典型使用场景与性能优势
6.1 KV Cache 共享
在多进程 LLM 推理中，SHMEM 可共享 KV Cache：
// 进程0：分配并写入void*kv_cache=aclshmemx_malloc(kv_size);FillKvCache(kv_cache);// 进程1：直接读取（相同虚拟地址）void*same_addr=kv_cache;// 无需拷贝！UseKvCache(same_addr);6.2 性能对比
在单机四进程场景下传输 1GB 数据：
| 方式 | 带宽 | 拷贝次数 |
|---|---|---|
| 传统 send/recv | 12 GB/s | 2（D2H + H2D） |
| SHMEM 零拷贝 | 38 GB/s | 0 |
数据来源：
examples/kvshuffle/kvshuffle_perf.cc，2026年2月
七、安全与调试支持
- TLS 加密：默认启用 TLS 保护跨节点数据（可关闭）；
- Debug 模式：编译时加
-debug选项，启用地址冲突检测与内存泄漏追踪； - Python 接口：
import shmem as shm; buf = shm.alloc(1024)。
八、总结
CANN shmem 通过统一内存池、强制同虚拟地址映射与驱动层 IPC 协同，构建了一套高效、安全的跨进程共享内存机制。其设计不仅消除了传统 IPC 中的冗余拷贝，更通过 HugeTLB、Unique ID 隔离等技术保障了大规模部署的稳定性。在大模型多进程推理、分布式训练参数共享等场景中，shmem 的内存管理能力已成为 CANN 高性能通信栈的关键基础设施。
相关链接：
- CANN 组织主页：https://atomgit.com/cann
- shmem 仓库地址：https://atomgit.com/cann/shmem
