CANN shmem 共享内存通信模型的进程间同步机制

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

前言

在多进程、多设备协同计算场景中，高效的进程间同步是确保数据一致性和执行正确性的关键。CANN（Compute Architecture for Neural Networks）开源项目中的shmem（Shared Memory Communication Library）仓库（https://atomgit.com/cann/shmem）基于 OpenSHMEM 标准，提供了一套面向高性能计算的共享内存通信模型。其核心不仅在于远程内存访问（RMA），更在于一套轻量级、低延迟、硬件亲和的进程间同步原语，包括屏障（Barrier）、信号量（Signal）、原子操作（Atomic）与事件等待（Event Wait）等。

1. shmem 同步模型的整体架构

shmem 的同步机制建立在全局共享内存池之上，所有进程通过映射同一物理内存区域实现状态共享。

1.1 内存布局

初始化时，shmem 分配一块全局共享内存（默认 16GB），其布局如下：

// shmem/include/shmem_internal.hstructGlobalSyncArea{// 屏障控制块（每个 Team 一个）BarrierControlBlock barriers[MAX_TEAMS];// 信号量数组（用于 p2p 同步）volatileuint64_tsignals[MAX_PROCESSES];// 原子操作缓冲区AtomicBuffer atomics;// 事件 ID 池（用于流水线同步）EventIdPool event_pool;};

该区域通过mmap或posix_memalign分配，并由所有进程共享。

1.2 同步原语分类

shmem 提供两类同步：

• 集体同步（Collective）：如shmem_barrier_all()；
• 点对点同步（P2P）：如shmem_signal_wait_until()。

2. 集体同步：屏障（Barrier）的实现

2.1 双阶段屏障算法

shmem 采用经典的双阶段（Two-phase） 屏障：

1. 到达阶段（Arrival）：进程写入自己的到达标志；
2. 离开阶段（Departure）：等待所有进程到达后继续。

// shmem/src/sync/barrier.cppvoidShmemBarrier::BarrierAll(){intmy_pe=shmem_my_pe();intnpes=shmem_n_pes();// 阶段1：标记到达barrier_area_->arrival_flags[my_pe]=1;// 内存屏障，确保写入全局可见__sync_synchronize();// 阶段2：轮询等待所有到达while(true){boolall_arrived=true;for(inti=0;i<npes;++i){if(barrier_area_->arrival_flags[i]!=1){all_arrived=false;break;}}if(all_arrived)break;// 降低 CPU 占用usleep(1);}// 重置标志（为下一次屏障准备）if(my_pe==0){memset(barrier_area_->arrival_flags,0,npes*sizeof(int));}}

⚠️优化点：实际实现中，shmem 使用硬件缓存一致性协议（如 MESI）避免显式刷新，仅依赖编译器屏障。

2.2 MTE 加速的屏障通知

对于支持 MTE 的平台，shmem 可通过硬件中断替代轮询：

// shmem/src/transport/mte_transport.cppvoidMTETransport::PostBarrierNotify(intdst_pe){// 构造 MTE 描述符，写入远端屏障标志MTE_DESC desc;desc.src_addr=&local_done_flag_;desc.dst_addr=remote_barrier_addr_[dst_pe];desc.size=sizeof(uint32_t);desc.notify=true;// 触发中断// 提交至 MTE 引擎mte_submit(&desc);}

接收方通过中断处理程序唤醒等待线程，延迟从 μs 级降至 ns 级。

3. 点对点同步：信号量与事件等待

3.1 信号量接口

shmem 提供shmem_signal*系列函数：

// shmem/include/shmem.hvoidshmem_signal_set(shmem_ctx_t ctx,intpe,uint64_tvalue);voidshmem_signal_wait_until(shmem_ctx_t ctx,intpe,uint64_tcmp_value,intcmp_op);

3.2 信号量内存布局

每个 PE 拥有一个 64 位信号量槽：

// 全局共享内存中volatileuint64_tglobal_signals[MAX_PES]__attribute__((aligned(64)));

对齐至 cache line（64 字节）避免 false sharing。

3.3 等待循环优化

shmem_signal_wait_until实现高效等待：

// shmem/src/sync/signal.cppvoidShmemSignal::WaitUntil(intpe,uint64_tcmp_val,intop){volatileuint64_t*target=&global_signals[pe];while(true){uint64_tcurrent=*target;if(Compare(current,cmp_val,op)){break;}// 使用 PAUSE 指令降低功耗_mm_pause();// 若支持，插入内存提示（如 CLFLUSHOPT）if(IsHardwareHintSupported()){__builtin_ia32_clflushopt(target);}}}

✅可移植性：_mm_pause()在 ARM 上替换为yield指令。

4. 原子操作与内存顺序

shmem 支持整数与浮点原子操作：

// shmem/include/shmem.hlongshmem_long_atomic_fetch_add(long*dest,longvalue,intpe);doubleshmem_double_atomic_compare_swap(double*dest,doublecond,doublevalue,intpe);

4.1 原子操作实现

底层通过CAS 循环或硬件原子指令实现：

// shmem/src/atomic/atomic_x86.cpplongShmemAtomic::FetchAddX86(volatilelong*addr,longval){return__sync_fetch_and_add(addr,val);}// shmem/src/atomic/atomic_arm.cpplongShmemAtomic::FetchAddARM(volatilelong*addr,longval){longold,new_val;do{old=__atomic_load_n(addr,__ATOMIC_RELAXED);new_val=old+val;}while(!__atomic_compare_exchange_n(addr,&old,new_val,false,__ATOMIC_ACQ_REL,__ATOMIC_RELAXED));returnold;}

4.2 内存顺序语义

shmem 遵循 C++11 内存模型，提供：

• __ATOMIC_ACQ_REL：用于 fetch_add；
• __ATOMIC_SEQ_CST：用于 compare_swap。

确保跨进程的 happens-before 关系。

5. 设备侧同步与 Host-Device 协同

shmem 的独特优势在于Host 与 Device 侧同步原语统一。

5.1 设备侧屏障调用

在 Device Kernel 中可直接调用：

// examples/matmul_allreduce/device_kernel.cu__global__voidMatmulAllreduceKernel(...){// ... 计算 ...// 设备侧屏障aclshmemx_device_barrier();// ... 通信 ...}

5.2 Host-Device 事件同步

shmem 支持通过事件 ID实现跨流水线等待：

// shmem/src/sync/event_sync.cppint32_tShmemEventSync::WaitForEvent(intevent_id){// 等待 Host 设置事件完成标志while(!event_table_[event_id].completed){usleep(1);}returnSUCCESS;}// Device Kernel 中__global__voidKernel(){// 触发事件完成aclshmemx_device_set_event(event_id);}

🔧最新特性（PR !113）：支持用户传入特定 EVENT ID 用于跨流水线等待，避免全局同步开销。

6. 跨机同步与安全通信

6.1 跨机初始化

shmem 通过bootstrap 通道（如 TCP）交换共享内存地址：

// shmem/src/init/bootstrap.cppvoidBootstrap::ExchangePeerInfo(){// 1. 通过 TCP 发送本地共享内存地址SendOverTCP(local_shmem_addr_);// 2. 接收远端地址并 mmap 到本地虚拟地址空间void*remote_addr=RecvFromTCP();remote_shmem_=mmap(...,remote_addr,...);}

6.2 TLS 加密同步数据

默认启用 TLS 加密保护同步元数据：

// shmem/src/security/tls_manager.cppvoidTLSManager::EncryptSyncData(void*data,size_t size){// 使用 OpenSSL AES-GCM 加密EVP_CIPHER_CTX*ctx=EVP_CIPHER_CTX_new();EVP_EncryptInit_ex(ctx,EVP_aes_256_gcm(),nullptr,key_,iv_);EVP_EncryptUpdate(ctx,encrypted_data,&len,data,size);EVP_EncryptFinal_ex(ctx,encrypted_data+len,&len);}

可通过aclshmemx_set_conf_store_tls(false, NULL, 0)关闭（需在 init 前调用）。

结语

CANN shmem 通过精心设计的共享内存布局与多层次同步原语，实现了高效、可靠的进程间同步。其不仅支持传统的屏障与信号量，更通过 MTE 硬件加速、设备侧原生调用及跨机安全通信，构建了一套完整的同步解决方案。作为 CANN 多设备协同计算的核心组件，shmem 的同步机制为通算融合算子、分布式训练等场景提供了低延迟、高吞吐的协同基础。随着对更多同步模式（如分段屏障、条件变量）的支持，shmem 的能力将持续演进，成为高性能计算领域的重要基础设施。

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