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.h
struct GlobalSyncArea {
    // 屏障控制块（每个 Team 一个）
    BarrierControlBlock barriers[MAX_TEAMS];
    
    // 信号量数组（用于 p2p 同步）
    volatile uint64_t signals[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.cpp
void ShmemBarrier::BarrierAll() {
    int my_pe = shmem_my_pe();
    int npes = shmem_n_pes();
    
    // 阶段1：标记到达
    barrier_area_->arrival_flags[my_pe] = 1;
    
    // 内存屏障，确保写入全局可见
    __sync_synchronize();
    
    // 阶段2：轮询等待所有到达
    while (true) {
        bool all_arrived = true;
        for (int i = 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.cpp
void MTETransport::PostBarrierNotify(int dst_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.h
void shmem_signal_set(shmem_ctx_t ctx, int pe, uint64_t value);
void shmem_signal_wait_until(shmem_ctx_t ctx, int pe, uint64_t cmp_value, int cmp_op);

3.2 信号量内存布局

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

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

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

3.3 等待循环优化

shmem_signal_wait_until 实现高效等待：

// shmem/src/sync/signal.cpp
void ShmemSignal::WaitUntil(int pe, uint64_t cmp_val, int op) {
    volatile uint64_t* target = &global_signals[pe];
    
    while (true) {
        uint64_t current = *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.h
long shmem_long_atomic_fetch_add(long *dest, long value, int pe);
double shmem_double_atomic_compare_swap(double *dest, double cond, double value, int pe);

4.1 原子操作实现

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

// shmem/src/atomic/atomic_x86.cpp
long ShmemAtomic::FetchAddX86(volatile long* addr, long val) {
    return __sync_fetch_and_add(addr, val);
}

// shmem/src/atomic/atomic_arm.cpp
long ShmemAtomic::FetchAddARM(volatile long* addr, long val) {
    long old, 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
    ));
    return old;
}

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__ void MatmulAllreduceKernel(...) {
    // ... 计算 ...
    
    // 设备侧屏障
    aclshmemx_device_barrier();
    
    // ... 通信 ...
}

5.2 Host-Device 事件同步

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

// shmem/src/sync/event_sync.cpp
int32_t ShmemEventSync::WaitForEvent(int event_id) {
    // 等待 Host 设置事件完成标志
    while (!event_table_[event_id].completed) {
        usleep(1);
    }
    return SUCCESS;
}

// Device Kernel 中
__global__ void Kernel() {
    // 触发事件完成
    aclshmemx_device_set_event(event_id);
}

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

---

6. 跨机同步与安全通信

6.1 跨机初始化

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

// shmem/src/init/bootstrap.cpp
void Bootstrap::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.cpp
void TLSManager::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