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

---

在超大规模分布式深度学习训练中，节点间数据交换的效率是决定模型扩展性的关键瓶颈。SHMEM 库作为 CANN 计算平台的核心通信基础，通过引入分区全局地址空间（PGAS）模型和对底层硬件的直接控制，实现了 PE-to-PE 的零拷贝、微秒级延迟通信。它为从模型并行到数据并行的高级通信原语提供了最底层的加速通道，是支撑万亿级参数模型高效训练的基石。

1. SHMEM 的理论基石：PGAS 模型与内存架构

PGAS（Partitioned Global Address Space，分区全局地址空间）模型是 SHMEM 区别于传统 MPI 消息传递机制的核心所在。它在逻辑上构建了一个统一的全局地址空间，但在物理上保留了数据的局部性。这种设计既提供了共享内存编程的便利性，又具备了分布式内存系统的可扩展性。

1.1 统一视图下的内存区域划分

在 SHMEM 的视角下，集群内的内存被严格划分为两个截然不同的区域，分别承担计算与通信的职责：

• 私有内存（Private Memory）：这是每个 PE（Processing Element）独占的本地存储区域。它存放着局部的栈变量、非共享的静态数据以及计算过程中的临时张量。私有内存的操作仅限于本地 PE，拥有最高的访问速度，是计算密集型任务发生的核心场所。
• 对称堆（Symmetric Heap）：这是 SHMEM 实现全局通信的关键。系统在所有参与通信的 PE 上预先分配一段物理地址对齐、大小一致的内存段。通过这一机制，PE $i$ 上的变量 $X$ 与 PE $j$ 上的变量 $X$ 在逻辑上是对称的。任何 PE 都可以通过特定的全局地址句柄，直接索引到远程 PE 对称堆中的对象，从而实现跨设备的直接内存访问。

1.2 物理地址映射与句柄解析

为了在硬件层面实现高效的远程访问，SHMEM 维护了一套复杂的地址映射表。当用户发起远程访问请求时，Runtime 层不需要进行昂贵的操作系统系统调用，而是直接查询内部维护的 Rank 映射表。该表将逻辑上的 PE ID 和偏移量（Offset）迅速转换为目标设备在互联总线上的物理地址。这种用户态的地址解析机制，彻底消除了内核态上下文切换的开销。

1.3 静态分配策略与内存稳定性

与动态频繁申请内存不同，SHMEM 极度推崇静态规划。在通信初始化阶段，建议一次性申请足够大的对称堆空间。这种策略不仅避免了运行时内存碎片化的问题，更重要的是保证了物理连续性。物理连续的内存块是启用硬件 Burst 传输模式的前提，能够大幅提升总线带宽的利用率，确保在大数据块传输时的吞吐量达到理论峰值。

2. 单边通信范式：突破双边协议的延迟屏障

传统的双边通信（Two-sided Communication，如 Send/Recv）要求发送方和接收方必须在时间点上进行复杂的握手同步，任何一方的延迟都会导致整个链路的阻塞。SHMEM 采用的单边通信（One-sided Communication）彻底打破了这一限制。

2.1 硬件直通与 CPU 旁路机制

SHMEM 的核心原语直接映射到底层 NPU 芯片内部的 DMA（Direct Memory Access）控制器。

1. 指令下发：Host CPU 或控制核心仅需向 DMA 引擎配置源地址、目的地址和传输长度。
2. 硬件接管：一旦指令下发，数据搬运的全过程完全由 DMA 硬件独立完成。
3. 计算并行：在此期间，计算核心（AI Core）可以立即释放去执行其他计算任务，实现了真正的“计算与通信重叠”。

2.2 Put/Get 原语的异步特性

SHMEM 提供的 shmem_put（写）和 shmem_get（读）操作本质上是异步的。当函数调用返回时，仅代表传输请求已成功提交到硬件队列，并不意味着数据搬运已经完成。这种非阻塞的设计允许开发者在一个时间窗口内连续发射成百上千条通信指令，充分填满互联总线的流水线，从而掩盖单次通信的固有延迟（Latency Hiding）。

2.3 零拷贝（Zero-Copy）的实现

在传统的网络协议栈中，数据往往需要在用户缓冲区、内核缓冲区和网卡缓冲区之间多次拷贝。SHMEM 利用 RDMA（Remote Direct Memory Access）类技术思想，允许本地 PE 直接将数据写入远程 PE 的显存中。整个路径中不存在中间缓冲区的周转，既节省了宝贵的显存带宽，也降低了传输延迟至微秒级别。

3. 核心 API 实战：显式控制与算子开发

SHMEM 提供了一套类 C++ 的接口标准，专门针对异构计算环境进行了深度的定制和优化。开发者可以在自定义算子（Kernel）代码中直接嵌入这些原语。

3.1 基础数据传输接口

最常用的接口用于块数据的搬运。为了适配不同的数据类型和对齐要求，API 提供了多种变体，例如针对连续内存块的 shmem_putmem 和针对特定类型步长的传输接口。

3.2 算子开发中的通信实现示例

下面的代码展示了如何在一个自定义的计算核函数中，利用 SHMEM 从远程节点拉取数据，并在本地进行计算处理。

/* 
 * SHMEM 远程数据拉取与计算示例 
 * 场景：当前 PE 需要从相邻 PE 获取权重分片进行计算
 */

__aicore__ void RemoteFetchAndCompute(int target_pe_id, uint64_t remote_offset, uint64_t data_len) {
    // 1. 地址解析：获取远程 PE 对称堆中对象的全局句柄
    // ShmemGetGlobalAddr 会根据 target_pe_id 计算出在互联总线上的可访问地址
    void* remote_addr_ptr = ShmemGetGlobalAddr(target_pe_id, remote_offset);
  
    // 2. 资源分配：在本地 Unified Buffer (UB) 中申请接收空间
    // UB 是 AI Core 内部的高速缓冲，离计算单元最近
    LocalTensor<float> input_tensor = ub_queues.AllocTensor<float>();
    uint64_t local_phy_addr = input_tensor.GetPhyAddr();
  
    // 3. 发起异步 GET 请求 (One-sided Read)
    // 该指令仅配置 DMA 描述符，耗时极短，立即返回
    ShmemGetAsync(local_phy_addr, remote_addr_ptr, data_len);
  
    // -----------------------------------------------------------
    // 在此处可以插入不依赖 input_tensor 的其他计算任务，实现掩盖
    // PreComputeIndependentWork();
    // -----------------------------------------------------------

    // 4. 同步屏障：等待数据传输完成
    // ShmemQuiet 确保此前发起的所有 GET 操作数据已完全落盘到 UB
    ShmemQuiet(); 
  
    // 5. 执行核心计算：此时数据已安全位于本地 UB
    // 使用 Vector 单元对数据进行处理
    VectorAdd(input_tensor, input_tensor, 1.0f, data_len / sizeof(float));
  
    // 6. 释放资源
    ub_queues.FreeTensor(input_tensor);
}

3.3 原子操作与分布式协作

除了基本的数据搬运，SHMEM 还提供了 shmem_atomic_add、shmem_atomic_inc、shmem_atomic_cas（Compare-and-Swap）等原子操作接口。这些接口允许 PE 对远程内存地址执行不可分割的读-修改-写操作。这在实现分布式梯度累加、全局计数器或分布式锁时至关重要，且由于是硬件原子指令，无需软件层面的 Mutex 介入。

4. 内存一致性模型：数据可见性的生命线

在单边通信模型中，由于接收方 CPU/NPU 不参与数据接收过程，如何确保“数据何时到达”成为了核心难题。SHMEM 定义了严格的内存一致性模型来解决这一问题。

4.1 shmem_quiet 的全局栅栏作用

shmem_quiet() 是最强力度的同步原语。它能够保证当前 PE 发起的所有待处理远程写操作（Put）和非阻塞读操作（Get）全部完成。对于写操作，它保证数据已到达远程内存并对远程 PE 可见；对于读操作，它保证数据已完全拷贝到本地内存。在大多数计算依赖通信结果的场景下，这是必须调用的指令。

4.2 shmem_fence 的细粒度排序

与 quiet 等待所有传输完成不同，shmem_fence() 仅保证指令发射的顺序性。它确保在 fence 之前发出的所有写操作，一定会在 fence 之后发出的写操作之前被远程观测到。

• 应用场景：当你需要先写入数据块，再写入一个“完成标志位（Flag）”时，必须在两者之间插入 fence。否则，由于网络乱序，远程节点可能先看到 Flag 变更为完成，却读取到旧的数据。

4.3 点对点同步与信号机制

除了基于内存的屏障，SHMEM 还提供了基于信号量的点对点同步机制。例如 shmem_wait_until 可以让当前 PE 挂起，直到特定的内存地址通过远程原子操作被修改为指定的值。这种机制常用于构建复杂的生产者-消费者流水线，避免了轮询带来的总线带宽浪费。

5. 层次化通信栈：从 SHMEM 到 HCCL 的协同

SHMEM 并非孤立存在，它是 CANN 通信栈的底座，向上支撑着 HCCL（Huawei Collective Communication Library）等高级集合通信库，向下对接 Runtime 和硬件驱动。

5.1 HCCL 的加速引擎

HCCL 提供的 AllReduce、Broadcast 等高级集合通信算子，在底层执行时会根据物理拓扑智能选择路径。

• 片内/片间互联：在单机多卡场景下，HCCL 会优先编译为一系列 SHMEM 的 Put/Get 指令序列，利用 HCCS 高速互联直接在设备显存间交换数据，完全绕过主机内存。
• 拓扑感知：SHMEM 的性能极度依赖物理距离。HCCL 的算法调度器会感知当前的 Rank 拓扑，尽量在物理相邻的 PE 间通过 SHMEM 交换大量数据，减少跨 Server 的以太网/InfiniBand 流量。

5.2 Runtime 的环境初始化

Runtime 负责整个通信环境的建立，是 SHMEM 能够运行的前提。

• Rank Table 构建：Runtime 会扫描集群硬件，为每个物理设备分配逻辑 Rank ID，并建立全量的物理地址映射表。
• 资源预留：它负责初始化 SDMA 通道，分配 notify 信号量池，确保 SHMEM 调用时硬件资源已就绪。

5.3 Graph Engine (GE) 的调度融合

在静态图执行模式下，GE 负责将通信算子与计算算子进行统一编排。GE 能够识别计算图中的数据依赖关系，自动在计算节点之间插入基于 SHMEM 的通信节点，并利用 Streams（流）机制实现通信任务的异步下发，最大化硬件单元的并行度。

6. 性能调优指南：挖掘硬件极限带宽

要充分发挥 SHMEM 的性能，单纯调用 API 是不够的，必须遵循特定的优化范式。

6.1 对称堆的对齐与 Burst 传输

硬件总线在处理对齐的大块数据传输时效率最高。

• 对齐约束：建议所有用于通信的缓冲区首地址至少按照 64 字节对齐，最佳实践是 128 字节或 256 字节对齐。
• 大小规整：传输的数据长度若为 Burst Length 的整数倍，可以减少尾部零散数据的传输开销。如果数据未对齐，驱动层可能需要进行额外的 padding 或拆分，导致性能严重下降。

6.2 通信掩盖与双缓冲技术（Double Buffering）

这是提升分布式训练效率最有效的手段。

• 原理：将数据切分为两个缓冲区 A 和 B。当 AI Core 在计算缓冲区 A 的数据时，DMA 引擎并行地将下一批数据拉取到缓冲区 B。
• 流水线：通过这种“乒乓操作”，只要计算时间大于等于通信时间，通信延迟就可以被完全掩盖，对外体现为零通信开销。

6.3 锁争用规避与归约树

在高并发的分布式场景下，原子操作虽然方便，但如果所有 PE 同时对同一个远程地址执行 atomic_add，会造成严重的总线热点（Hot Spot）和锁争用。

• 树状归约：应采用树状或蝶形（Butterfly）交换算法，将全局归约分解为多轮的点对点 SHMEM 操作，最后仅在局部进行少量原子更新。
• 本地聚合：尽可能在私有内存中完成部分结果的聚合，减少对全局对称堆的原子访问频率。