相关链接：

• 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)
namespace shmem {
// 对远程 PE (Processing Element) 上的目标地址执行原子加法
__device__ void int32_add_atomic(
    int32_t* dest, // 远程目标地址
    int32_t value,
    int pe         // 目标PE的ID
);

// 原子比较并交换 (Compare-and-Swap)
__device__ int32_t int32_cswap_atomic(
    int32_t* dest,
    int32_t cond,  // 比较值
    int32_t value, // 新值
    int pe
);

// 获取并原子加 (Fetch-and-Add)
__device__ int32_t int32_fadd_atomic(
    int32_t* dest,
    int32_t value,
    int pe
);
}

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

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__ void mte_int32_add_atomic(int32_t* remote_addr, int32_t val, int pe) {
    // 1. 将远程虚拟地址转换为硬件可识别的物理/逻辑地址
    uint64_t hw_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 操作
    // 这是一个非阻塞的写操作，硬件会自动完成原子性保证
    *(volatile MteAtomicCmd*)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__ void example_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__ void centralized_barrier(int team_size, int my_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