跳转至

SHMEM Symmetric Memory

导言

我最初把 HCCL 理解成集合通信,把 HiXL 理解成单点通信。这个分法可以作为起点,却会让 SHMEM 无处安放:它既能一对一 Put/Get,也能做原子操作与同步,为什么还需要“对称内存”这套约束?

关键在于,HCCL、HiXL 和 SHMEM 并不只是在争夺同一种通信 API。HCCL 更接近“我要完成什么群体操作”,HiXL 更接近“我要把哪段数据传到哪里”,SHMEM 则把问题改写成“我要访问哪个 PE 上的哪个内存坐标”。它用跨 PE 的布局约束,换取设备侧可以低开销地定位和操作远端数据;但当各 PE 的容量需求严重不均时,这份约束也确实可能造成浪费。

先换一个分类轴

“集合通信、点到点通信、共享内存通信”很容易被理解成三个并列的通信拓扑。更准确的区分是:程序向通信系统表达的意图不同。

抽象 程序主要表达什么 最自然的心智模型
HCCL 一组 Rank 共同得到某个结果 “帮我完成一次 AllReduce、AllGather 或 AllToAll”
HiXL 把一段显式地址的数据传给另一个端点 “把这里的字节送到那里”
SHMEM 访问某个 PE 上、某个对称对象中的位置 “写入 <PE, object, offset> 这个坐标”

因此,“HCCL 主要用于集合通信、HiXL 用于单点通信”并没有错,只是还不够。这里的“主要”也很重要:HCCL 还可以提供 Send/Recv 一类接口,不能只用参与 PE 的数量来划线。更稳定的区分是,HCCL 的核心是操作语义,HiXL 的核心是显式传输,SHMEM 的核心则是分布式内存语义。SHMEM 不是又发明了一种链路;它提供的是一套让远端内存能够被稳定命名、访问和同步的编程模型。

不是互斥的三选一

同一个系统完全可以让规则的大块交换走 HCCL,让显式的 Host Buffer 传输走 HiXL,再让 Kernel 内细粒度、不规则、需要计算通信流水的访问走 SHMEM。它们可以共享底层搬运能力,却对上层承诺不同的语义。

对称的是坐标系

SHMEM 中每个 PE(Processing Element) 都拥有自己的本地内存。所谓“对称”,不是所有 PE 真的共享同一片物理内存,也不是其中的数据必须相同,而是它们建立了可互相翻译的对象布局

假设所有 PE 都按同样的顺序申请:

A = symmetric_alloc(4 elements)
B = symmetric_alloc(8 elements)
flag = symmetric_alloc(1 element)

那么每个 PE 都拥有自己的 ABflag。它们的物理地址甚至本地虚拟地址可以不同,但运行时能够把“我这里的 B + 0”翻译成“PE 5 那里的 B + 0”。远端对象因此可以写成一个二元坐标:

\[ \text{remote object}=\langle \text{target PE},\ \text{symmetric address}\rangle \]

把对象再拆开,也可以写成:

\[ \text{remote object}=\langle \text{target PE},\ \text{object},\ \text{offset}\rangle \]

SHMEM 对称地址与非对称 Put

自绘技术图:两个 PE 保持相同的对象布局,因此 `A[2]` 和 `B[0]` 都是可跨 PE 翻译的坐标;具体操作却可以从 PE 0 的 `A[2]` 写到 PE 5 的 `B[0]`,并不要求源和目标位置相同。

这份契约约束的是可寻址性:参与同一对称内存域的 PE 需要遵循一致的分配与生命周期规则,作为远端目标的对象必须能够在对端被同一坐标识别。它不要求:

  • 各 PE 的对象保存相同数据;
  • 各 PE 同时读写同一个位置;
  • 源地址与目标地址具有相同偏移;
  • 所有业务内存都放进对称堆;
  • 每个 PE 执行相同的计算分支。

一句话概括:对称内存相当于大家使用同一份“街道与门牌规则”,不是每栋房子必须摆放同样的家具。

操作不需要对称

因此,下面这种操作在语义上完全成立:PE 0 从自己的第 3 个位置取数,写入 PE 5 的第 1 个位置。用教学化伪代码表示:

A = symmetric_alloc(4 * sizeof(int));
B = symmetric_alloc(8 * sizeof(int));

if (my_pe == 0) {
    A[2] = 42;
    shmem_put(
        target = &B[0],   // 以本地对称指针表达 PE 5 上的 B[0]
        source = &A[2],   // 本地源可以是另一个位置
        count = 1,
        pe = 5
    );
    shmem_quiet();        // 在发布 ready flag 前确认此前远端更新完成
}

这里最容易误解的是 &B[0]:代码虽然在 PE 0 上执行,但它不是要求把数据写回 PE 0 的 B[0]。它提供了一个本地可计算的对称坐标,再由 pe = 5 指定这个坐标属于哪个远端 PE。

PE 5 不需要同时执行一条镜像的 Put,也不需要从 B[0] 再复制回 PE 0。单边通信的含义正是:发起方可以主动读写远端内存,而目标方不必以一条匹配的收发调用参与。目标方只需在消费数据前使用正确的 wait、signal、quiet、fence 或 barrier 等同步边界。

对于“计算引用”也要区分两件事:

  • 本地计算是任意的。每个 PE 可以按自己的控制流读写本地对称对象,不需要与其他 PE 镜像执行。
  • SHMEM 原语不是远程函数调用。Put/Get 搬运数据,AMO 对远端位置执行受定义的原子更新;它们不会自动让远端 PE 执行一段任意计算。远端计算仍要由目标 PE 的 Kernel 或后续任务消费数据。

所以真正需要共同遵守的不是“大家做同一操作”,而是地址契约、并发规则和完成语义。操作可以不对称,竞态却不能不管。

OpenSHMEM 为什么规则严格

如果操作本来可以任意,为什么 OpenSHMEM Specification 还要定义那么多分配、同步和顺序规则?因为“任意访问”只有建立在稳定契约上才可组合。

  1. 统一命名。如果 PE 0 和 PE 5 的第二次分配分别代表完全不同的对象,那么 B + 0 就无法作为跨 PE 坐标。
  2. 统一生命周期。如果某些 PE 已经释放或复用了对象,另一些 PE 仍把旧指针当成远端坐标,地址翻译就失去意义。
  3. 统一完成语义。发起方需要知道何时可以复用源 Buffer,目标方需要知道何时可以消费数据;否则“写调用已经返回”会被误当成“远端数据已经按需要可见”。
  4. 统一可移植性。同一份程序可能运行在不同网络、DMA 引擎和 Runtime 上。标准规定 PE、Team、对称对象、RMA、AMO、顺序与同步的语义,具体实现再选择怎样搬运字节。

因此,OpenSHMEM 的价值不只是统一函数名。它把一组原本容易依赖具体硬件的技巧,收敛成可以推理和移植的内存模型契约。规则严格,恰恰是为了让规则之上的单边操作足够自由。

设备侧直驱为何重要

CANN SHMEM 相对显式传输 API 的一个重要吸引力,确实是它能够把常用的数据面原语放到 Device/AICore 侧:Kernel 可以根据运行时数据直接发起 Put/Get、原子操作或同步,而不必为每个细粒度步骤重新回到 Host 编排。这对不规则通信尤其重要,因为目标 PE、偏移和数据量可能直到 Kernel 执行时才确定。

但“SHMEM 的操作都是 AIV 直驱”仍然过于绝对。完整路径至少要分成三层:

Host 控制面:初始化、建链、注册/分配内存、建立 PE/Team、启动 Kernel
Device 控制面:AIV/AICore 执行算法,准备参数、提交操作、检查完成
数据搬运面:MTE、SDMA、RDMA、UDMA 等引擎和互联真正搬运字节

AIV 直驱中的“直”主要指提交和推进控制路径更直接,不是 AIV 亲自替代所有 DMA 或网络硬件。某个具体调用是否真的走设备直驱,还要同时确认 CANN 版本、芯片、所用 API、目标拓扑、搬运引擎和完成路径,不能只看库名。

这类能力最适合以下模式:

  • 计算依赖的数据交换。Kernel 根据中间结果动态决定目标 PE 与偏移。
  • 细粒度生产者—消费者。写入数据后用 signal/flag 通知对端继续处理。
  • 不规则稀疏更新。不同 PE 收到的更新数量和位置并不一致。
  • MoE、图计算或分布式工作队列。通信目标离散,需要把数据准备、传输和消费组成流水。

如果数据交换规则、规模大且所有 Rank 本就共同参与,一次 HCCL 集合操作通常更自然;如果只是少量明确的 Buffer 搬运,显式传输接口也可能更简单。设备直驱是选择 SHMEM 的一个强理由,但不是绕过语义设计的性能开关。更完整的控制路径说明可参考 AIV Direct Drive

对称内存会浪费吗

会,而且这是一个真实的工程代价,但需要先区分数据量对称容量对称。SHMEM 不要求每个 PE 当前保存同样多的有效数据;它要求对称对象在每个 PE 上都具有可识别的容量与布局。若为了覆盖最坏情况,把每个 PE 的窗口都按最大需求分配,偏斜就会转化成闲置容量。

设有 \(P\) 个 PE,第 \(i\) 个 PE 的实际需求为 \(s_i\),而统一对称容量按最大值 \(S=\max_i s_i\) 规划,则:

\[ \text{capacity utilization}=\frac{\sum_{i=0}^{P-1}s_i}{P\cdot S} \]

例如 4 个 PE 的实际需求是 [1, 1, 1, 10] GiB,若每个 PE 都为这个对象准备 10 GiB,则总容量为 40 GiB,实际只使用 13 GiB,利用率只有 32.5%

浪费来自偏斜,不是来自“地址一样”

对称设计真正昂贵的情况是:各 PE 峰值差异大、对象长期驻留,而且容量按全局最大值静态配置。若各 PE 工作集接近、窗口很小或能够跨阶段复用,对称布局的额外容量可能并不显著。至于“保留容量”是否立刻全部形成物理内存占用,还取决于具体 Runtime 的分配实现;做容量规划时至少要按最坏预算检查。

降低浪费的关键,不是放弃 SHMEM,而是缩小必须对称的边界

  1. 只对称通信窗口。模型参数、长期状态和普通临时张量继续使用本地分配;只有远端要直接访问的 staging buffer、flag 和元数据进入对称域。
  2. 按 Tile 或 Ring Buffer 流水。窗口覆盖一个可复用工作集,而不是一次容纳完整张量或全量 Token。
  3. 容量对称、长度可变。每个 PE 保持相同上限,同时用 length、offset table 或 signal 表达本轮实际有效范围。
  4. 跨阶段复用。不同通信阶段共享同一组窗口,避免为每个算子永久保留峰值容量。
  5. 混合使用通信库。规则的大块 Collective 交给 HCCL,明确的 Buffer 传输交给 HiXL,只有需要远端坐标、细粒度同步和设备侧动态发起的部分使用 SHMEM。

三种抽象如何选择

问题 更偏向 HCCL 更偏向 HiXL 更偏向 SHMEM
程序最想表达什么 群体运算结果 一次显式传输 远端内存访问
参与关系 规则、集体参与 明确端点 可不规则、单边发起
地址管理 主要由库隐藏 应用显式给出 Buffer 对称对象 + PE + offset
控制位置 Host 或 Device 集合算法 取决于具体 API 路径 适合 Device Kernel 内动态推进
内存代价 Collective workspace Buffer 与注册资源 对称堆、窗口和同步状态
典型优势 算法成熟、表达简洁 传输意图直接、容量灵活 细粒度 RMA、AMO、同步与通算流水

最终判断可以压缩成三句话:

  1. SHMEM 对称的是内存布局和地址解释,不是复制、计算或控制流。PE 0 的 A[2] 写到 PE 5 的 B[0] 完全合理。
  2. CANN SHMEM 的设备侧数据面是重要优势,但不能把它简化成“所有操作都由 AIV 搬运”。Host 仍有控制面,实际字节由具体引擎与互联搬运。
  3. 对称容量在负载偏斜时会浪费。工程上应把对称域限制为小而可复用的通信工作集,并与 HCCL、HiXL 组合,而不是把全部业务内存都对称化。

这也解释了 SHMEM 看似矛盾的设计:底层先用严格规则建立一个稳定坐标系,上层才获得不对称、单边、细粒度操作的自由。

相关资料

  1. CANN SHMEM
  2. OpenSHMEM Specification
  3. Ascend Communication API Evolution
  4. Memory Semantics vs RDMA

评论