CANN shmem 的原子操作语义与内存一致性保障

相关链接：

• CANN 组织主页：https://atomgit.com/cann
• shmem 仓库地址：https://atomgit.com/cann/shmem

前言

在 CANN（Compute Architecture for Neural Networks）生态中，SHMEM（Shared Memory）库作为面向多机多卡场景的高性能内存通信库，其核心价值在于为分布式 AI 应用提供了一种高效、低延迟的跨设备内存访问范式。然而，当多个计算单元（如不同设备上的 AI Core）并发地访问同一块远程共享内存时，如何保证操作的原子性（Atomicity）和内存一致性（Memory Consistency）就成为了正确性和性能的关键。CANN SHMEM 基于 OpenSHMEM 标准，并针对 AI 工作负载的特点，设计了一套精密的原子操作语义与内存一致性保障机制。

一、SHMEM 原子操作 API 体系

SHMEM 为开发者提供了一系列标准的原子操作接口，这些接口位于include/device/atomic.h中，涵盖了整数和浮点类型的基本运算。

1.1 核心原子操作接口

SHMEM 的原子操作遵循shmem_<type>_<op>_atomic的命名规范，其中<type>包括int,long,float等，<op>包括add,and,or,xor,max,min,swap等。

// include/device/atomic.h (Device-side API)namespaceshmem{// 对远程 PE (Processing Element) 上的目标地址执行原子加法__device__voidint32_add_atomic(int32_t*dest,// 远程目标地址int32_tvalue,intpe// 目标PE的ID);// 原子比较并交换 (Compare-and-Swap)__device__int32_tint32_cswap_atomic(int32_t*dest,int32_tcond,// 比较值int32_tvalue,// 新值intpe);// 获取并原子加 (Fetch-and-Add)__device__int32_tint32_fadd_atomic(int32_t*dest,int32_tvalue,intpe);}

这些接口的设计直接映射到硬件能力，并为上层应用（如自定义的同步屏障、分布式计数器、无锁数据结构）提供了构建模块。

1.2 Host 与 Device 接口的统一

值得注意的是，SHMEM 同时提供了 Host 侧和 Device 侧的原子操作接口。虽然调用位置不同，但它们操作的是同一块由shmem_alloc分配的全局共享内存。这种统一的地址空间视图极大地简化了编程模型。

二、底层实现：硬件原子指令的封装

SHMEM 的高性能原子操作能力，根植于对底层硬件原子指令的直接封装。根据仓库代码（src/device/atomic/），其实现主要依赖于两种硬件机制：**MTE **(Memory Tagging Extension) 和 **xDMA **(eXtended Direct Memory Access) 引擎。

2.1 MTE 引擎驱动的原子操作

对于支持 MTE 的硬件，SHMEM 会优先使用 MTE 提供的原子指令。在src/device/atomic/mte_atomic_impl.cu中，我们可以看到具体的内联汇编实现：

// src/device/atomic/mte_atomic_impl.cu__device__voidmte_int32_add_atomic(int32_t*remote_addr,int32_tval,intpe){// 1. 将远程虚拟地址转换为硬件可识别的物理/逻辑地址uint64_thw_addr=translate_to_hw_addr(remote_addr,pe);// 2. 构建 MTE 原子加法指令的操作码和参数MteAtomicCmd cmd;cmd.opcode=MTE_CMD_ATOMIC_ADD;cmd.addr=hw_addr;cmd.value=val;cmd.size=sizeof(int32_t);// 3. 通过专用寄存器或内存-mapped I/O 触发 MTE 操作// 这是一个非阻塞的写操作，硬件会自动完成原子性保证*(volatileMteAtomicCmd*)MTE_CMD_QUEUE_ADDR=cmd;// 4. 【可选】如果需要返回旧值（如 fadd），则需等待完成事件// wait_for_mte_completion(cmd.id);}

MTE 引擎在硬件层面保证了对指定内存地址的读-改-写（Read-Modify-Write）操作的原子性，完全绕过了传统的缓存一致性协议，从而实现了极低的延迟。

2.2 xDMA 引擎的后备方案

对于不支持 MTE 或特定操作的场景，SHMEM 会回退到 xDMA 引擎。xDMA 虽然主要用于大块数据传输，但也支持基本的原子操作。其实现位于src/device/atomic/xdma_atomic_impl.cu，其原理是通过发送一个特殊的 xDMA 原子请求包到目标设备，由目标设备的 DMA 控制器执行原子操作。

三、内存一致性模型：顺序一致性 **(SC)

在并发系统中，仅仅有原子操作是不够的，还需要明确定义内存一致性模型（Memory Consistency Model），即所有处理单元观察到的内存操作顺序应遵循何种规则。

CANN SHMEM 采用了最强的顺序一致性（Sequential Consistency, SC）模型。这意味着：

“任何一次执行的结果，都等同于所有处理器的操作以某种全局顺序执行的结果，且每个处理器的操作在其程序顺序中出现。”

3.1 SC 模型的实现保障

为了实现 SC，SHMEM 在以下层面进行了保障：

1. 原子操作本身的顺序性：所有针对同一内存地址的原子操作，在全局范围内是严格有序的。
2. 显式内存屏障（Memory Fence）：SHMEM 提供了shmem_fence和shmem_quiet等同步原语，用于强制刷新未完成的内存操作，确保操作的全局可见性。

// 使用示例：确保原子操作的全局可见性__device__voidexample_usage(){// 1. 执行一系列本地计算...// 2. 对全局计数器进行原子加shmem::int32_add_atomic(global_counter,1,0);// 3. 【关键】插入 fence，确保此原子操作在后续操作前对所有PE可见shmem_fence();// 4. 执行依赖于计数器更新的后续操作...}

3. 通信引擎的保序性：无论是 MTE 还是 xDMA，其硬件设计都保证了来自同一个源 PE 的请求在目标 PE 上按发送顺序被处理。

3.2 与 C++ 内存模型的映射

SHMEM 的 SC 模型可以看作是 C++11 内存模型中memory_order_seq_cst的分布式扩展。这使得熟悉 C++ 多线程编程的开发者能够很容易地理解和使用 SHMEM 的同步原语。

四、高级同步原语的构建

基于上述原子操作和内存一致性保障，SHMEM 在其sync模块（src/host/sync/和src/device/sync/）中构建了更高级的同步原语，如全局屏障（Barrier）。

4.1 基于原子操作的屏障实现

一个典型的中心化屏障（Centralized Barrier）可以利用原子操作来实现：

// src/device/sync/barrier_impl.cu (概念示意)__device__voidcentralized_barrier(intteam_size,intmy_pe){static__shared__int32_t*arrival_count;static__shared__int32_t*release_flag;if(my_pe==0){// PE 0 作为协调者*release_flag=0;// 重置释放标志}shmem_barrier_all();// 确保协调者已重置// 所有 PE 原子增加到达计数shmem::int32_fadd_atomic(arrival_count,1,0);// 协调者检查是否所有人都已到达if(my_pe==0){while(*arrival_count<team_size){// 忙等待或使用事件通知}// 通知所有人可以离开shmem::int32_swap_atomic(release_flag,1,0);}// 所有其他 PE 等待释放信号if(my_pe!=0){while(*release_flag==0){// 忙等待}}// 最后一个同步点shmem_barrier_all();}

这个例子清晰地展示了如何利用fadd_atomic和swap_atomic来构建一个功能完整的屏障，而其正确性完全依赖于底层原子操作的 SC 语义。

五、总结

CANN SHMEM 通过一套精心设计的原子操作 API、对底层硬件原子指令（MTE/xDMA）的深度集成，以及对顺序一致性（SC）内存模型的严格遵守，为多机多卡环境下的并发编程提供了坚实的基础。它不仅提供了add,cswap,fadd等基本原子原语，还以此为基础构建了全局屏障等高级同步工具。

这种从硬件指令到高级抽象的完整栈，使得开发者能够在享受极致通信性能的同时，无需过度担忧底层的并发正确性问题。SHMEM 的原子操作与内存一致性保障机制，是 CANN 生态能够高效支撑复杂分布式 AI 应用（如大规模强化学习、分布式优化器）的核心技术支柱之一。

相关链接：

• CANN 组织主页：https://atomgit.com/cann
• shmem 仓库地址：https://atomgit.com/cann/shmem