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 都按同样的顺序申请:
那么每个 PE 都拥有自己的 A、B 和 flag。它们的物理地址甚至本地虚拟地址可以不同,但运行时能够把“我这里的 B + 0”翻译成“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 还要定义那么多分配、同步和顺序规则?因为“任意访问”只有建立在稳定契约上才可组合。
- 统一命名。如果 PE 0 和 PE 5 的第二次分配分别代表完全不同的对象,那么
B + 0就无法作为跨 PE 坐标。 - 统一生命周期。如果某些 PE 已经释放或复用了对象,另一些 PE 仍把旧指针当成远端坐标,地址翻译就失去意义。
- 统一完成语义。发起方需要知道何时可以复用源 Buffer,目标方需要知道何时可以消费数据;否则“写调用已经返回”会被误当成“远端数据已经按需要可见”。
- 统一可移植性。同一份程序可能运行在不同网络、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\) 规划,则:
例如 4 个 PE 的实际需求是 [1, 1, 1, 10] GiB,若每个 PE 都为这个对象准备 10 GiB,则总容量为 40 GiB,实际只使用 13 GiB,利用率只有 32.5%。
浪费来自偏斜,不是来自“地址一样”
对称设计真正昂贵的情况是:各 PE 峰值差异大、对象长期驻留,而且容量按全局最大值静态配置。若各 PE 工作集接近、窗口很小或能够跨阶段复用,对称布局的额外容量可能并不显著。至于“保留容量”是否立刻全部形成物理内存占用,还取决于具体 Runtime 的分配实现;做容量规划时至少要按最坏预算检查。
降低浪费的关键,不是放弃 SHMEM,而是缩小必须对称的边界:
- 只对称通信窗口。模型参数、长期状态和普通临时张量继续使用本地分配;只有远端要直接访问的 staging buffer、flag 和元数据进入对称域。
- 按 Tile 或 Ring Buffer 流水。窗口覆盖一个可复用工作集,而不是一次容纳完整张量或全量 Token。
- 容量对称、长度可变。每个 PE 保持相同上限,同时用
length、offset table 或 signal 表达本轮实际有效范围。 - 跨阶段复用。不同通信阶段共享同一组窗口,避免为每个算子永久保留峰值容量。
- 混合使用通信库。规则的大块 Collective 交给 HCCL,明确的 Buffer 传输交给 HiXL,只有需要远端坐标、细粒度同步和设备侧动态发起的部分使用 SHMEM。
三种抽象如何选择¶
| 问题 | 更偏向 HCCL | 更偏向 HiXL | 更偏向 SHMEM |
|---|---|---|---|
| 程序最想表达什么 | 群体运算结果 | 一次显式传输 | 远端内存访问 |
| 参与关系 | 规则、集体参与 | 明确端点 | 可不规则、单边发起 |
| 地址管理 | 主要由库隐藏 | 应用显式给出 Buffer | 对称对象 + PE + offset |
| 控制位置 | Host 或 Device 集合算法 | 取决于具体 API 路径 | 适合 Device Kernel 内动态推进 |
| 内存代价 | Collective workspace | Buffer 与注册资源 | 对称堆、窗口和同步状态 |
| 典型优势 | 算法成熟、表达简洁 | 传输意图直接、容量灵活 | 细粒度 RMA、AMO、同步与通算流水 |
最终判断可以压缩成三句话:
- SHMEM 对称的是内存布局和地址解释,不是复制、计算或控制流。PE 0 的
A[2]写到 PE 5 的B[0]完全合理。 - CANN SHMEM 的设备侧数据面是重要优势,但不能把它简化成“所有操作都由 AIV 搬运”。Host 仍有控制面,实际字节由具体引擎与互联搬运。
- 对称容量在负载偏斜时会浪费。工程上应把对称域限制为小而可复用的通信工作集,并与 HCCL、HiXL 组合,而不是把全部业务内存都对称化。
这也解释了 SHMEM 看似矛盾的设计:底层先用严格规则建立一个稳定坐标系,上层才获得不对称、单边、细粒度操作的自由。
