CANN SHMEM 中基于 ACLGraph 的 RDMA AllGather 同步机制与运行实战：rdma_aclgraph

CANN SHMEM 中基于 ACLGraph 的 RDMA AllGather 同步机制与运行实战：rdma_aclgraph_demo 全解析

【免费下载链接】shmemCANN SHMEM 是面向昇腾平台的多机多卡内存通信库，基于OpenSHMEM 标准协议，实现跨设备的高效内存访问与数据同步。项目地址: https://gitcode.com/cann/shmem

CANN SHMEM 是面向昇腾平台的多机多卡内存通信库，基于 OpenSHMEM 标准协议实现跨设备的高效内存访问与数据同步。本文以仓库中的 examples/rdma_aclgraph_demo 为对象，剖析在 ACLGraph（Ascend 图执行模型）中嵌入 RDMA AllGather 通信算子的关键技术——即通过aclshmemx_roce_barrier_all屏障接口解决"AllGather 仅发出 Write 操作、无法保证图执行到相同阶段"的同步问题。读完本文，你将掌握该样例的 ACLGraph 图结构设计、环境准备、单机/跨机运行方式、命令行参数含义，以及从源码层面理解其同步原理与精度校验逻辑。

样例场景：为什么 ACLGraph 中的 AllGather 需要显式同步

rdma_aclgraph_demo的目标是在一张 ACLGraph 图中编排add + allGather + allGather + add的计算序列，并验证在 RDMA（ROCE）传输下 AllGather 算子嵌入图执行后的功能正确性。

该场景的核心难点在于：RDMA AllGather 只发送 Write 操作。与传统的读-改-写通信方式不同，aclshmemx_roce_put_nbi这类单边 PUT 操作将数据直接写入远端对称内存，不会主动通知对端"写入已完成"。因此，如果没有额外同步，所有 PE 上的图执行进度可能不一致：某些 PE 已经开始消费数据，而另一些 PE 的数据尚未到达，导致 AllGather 结果错误，进而使 ACLGraph（model）的精度异常。

解决办法是在 rdma allgather 算子中显式加入 aclshmemx_roce_barrier_all 接口进行全量屏障同步，确保所有 PE 上的图均已执行到相应阶段后再进行数据交换，从而保证 ACLGraph 模型的精度。

ACLGraph 图结构与数据流

完整的图结构如下（其中allGather为 RDMA 后端实现）：

add → allGather → allGather → add

具体数据流（对应 rdma_aclgraph_demo_kernel.cpp 与 main.cpp）：

每个 PE 生成一组长度为TRANS_SIZE（16 个 int32）的本地输入input[i] = pe_id + 10，通过aclrtMemcpy写入对称内存ptr和b_ptr的本地分片；
第一次add：run_vector_add<int>对ptr与add_ptr做向量加，结果写回temp_ptr，再经aclrtMemcpyAsync拷回ptr；
两次allgather_demo：分别对ptr与b_ptr执行 RDMA AllGather，将各 PE 的数据片汇聚成完整数组；
第二次add：对两个 AllGather 的输出ptr与b_ptr求和，得到c_ptr；
通过aclmdlRICaptureBegin/End将上述整条序列捕获为 ACLGraph 模型model，之后每轮循环仅执行aclmdlRIExecuteAsync(model, stream)重放；
每一轮结束后，把c_ptr拷回 Host，逐元素与期望值(10 + i) * (3 + zz)（zz为循环轮次）比对，任一元素不等即报错退出。

循环共执行LOOP_TIMES = 3轮：第 0 轮负责图捕获，第 1、2 轮为图重放。通过每轮校验第二次add的输出，即可判断 AllGather 算子在 ACLGraph 中是否工作正常。

设备端内核的同步实现

设备端内核device_all_gather_test（见 rdma_aclgraph_demo_kernel.cpp）直观展示了"屏障—写入—屏障"的完整模式：

aclshmemx_roce_barrier_all(); // 图执行阶段同步：确保所有 PE 都到达该点 for (int i = 0; i < pe_size; i++) { if (i == my_rank) { continue; } aclshmemx_roce_put_nbi( gva + message_length * my_rank, gva + message_length * my_rank, (__ubuf__ uint8_t*)ubLocal.GetPhyAddr(), message_length, i, 0); } aclshmemx_roce_barrier_all(); // 数据写入完成同步：确保所有 PUT 均已被远端接收

第一次aclshmemx_roce_barrier_all()保证进入数据交换前，所有 PE 的图都已经执行到 AllGather 阶段，这是"精度正常"的前提；
中间的aclshmemx_roce_put_nbi将自己对称内存分片广播给其它 PE；
第二次aclshmemx_roce_barrier_all()保证所有 PUT 完成后才继续后续add，避免读到未完成写入的远端数据。

aclshmemx_roce_barrier_all在设备侧声明见 include/device/gm2gm/engine/shmem_device_rdma.h，其实现位于 src/device/gm2gm/engine/shmem_device_rdma.hpp，并同样被 rdma_demo、rdma_aggregate_demo、rdma_sync_barrier_demo 及多个单测内核（如 tests/unittest/device/sync/rdma_sync_barrier_kernel.cpp）复用，可见该屏障是 ROCE 后端示例与测试的公共原语。

环境要求

rdma_aclgraph_demo的环境要求与 rdma_demo 一致，主要包括：

RDMA 环境可用：RDMA 网卡及驱动已正确安装并配置。
Ascend950 平台需 CANN 9.1.0 包：RDMA 相关样例在 Ascend950 平台上要求 CANN 9.1.0，其它版本目前不支持。

检查 RDMA 环境

A2/A3 平台：使用hccn_tool检查网卡 IP 与网络健康状态（数字 7 按实际设备数调整）：

for i in {0..7}; do hccn_tool -i $i -ip -g; done for i in {0..7}; do hccn_tool -i $i -net_health -g; done

Ascend950 平台：使用ibv_devinfo检查 RDMA 设备：

XSCALE 网卡：ibv_devinfo | grep xscale
HNS 1825 网卡：ibv_devinfo | grep hrn

注意：HNS 1825 网卡同端口通信时，若交换机未开启端口桥（port bridge），RDMA 收发可能失败，相关配置可参考调试手册 - 同端口通信需开启端口桥。

IBV_EXTEND_DRIVERS 环境变量（Ascend950 必设）

Ascend950 平台运行前需将IBV_EXTEND_DRIVERS设置为对应网卡的用户态 Verbs provider 插件库路径：

XSCALE 网卡：

export IBV_EXTEND_DRIVERS=<path_to_libxscale_nda.so>

HNS 1825 网卡：

export IBV_EXTEND_DRIVERS=<path_to_libhrn5-rdmav34.so>

libxscale_nda.so与libhrn5-rdmav34.so由对应网卡驱动安装包提供（而非 SHMEM 项目构建产物）。驱动安装后可用find / -name "libhrn5-rdmav34.so"或find / -name "libxscale_nda.so"定位库路径并设置。IBV_EXTEND_DRIVERS是 libibverbs 的环境变量，用于加载默认搜索路径之外的 Verbs provider 插件。

编译构建

在shmem/根目录按平台选择编译参数（RDMA 后端参数详见编译与构建指南 - RDMA 参数使用说明）：

# A2/A3 平台 bash scripts/build.sh -enable_rdma -examples # Ascend950 + XSCALE 网卡 bash scripts/build.sh -soc_type Ascend950 -enable_rdma -rdma_backend XSCALE -examples # Ascend950 + HNS 1825 网卡 bash scripts/build.sh -soc_type Ascend950 -enable_rdma -rdma_backend HNS_1825 -examples

构建目标由 examples/rdma_aclgraph_demo/CMakeLists.txt 中的aclshmem_add_collective_example(rdma_aclgraph_demo)声明，产物为./build/bin/rdma_aclgraph_demo。

运行方式

方式一：直接运行脚本（单机双卡）

在shmem/examples/rdma_aclgraph_demo目录下执行：

bash run.sh # 单机双卡用例

run.sh 的行为要点：

自动将PROJECT_ROOT指向仓库根目录，并把${PROJECT_ROOT}/build/lib加入LD_LIBRARY_PATH；
清理旧输出目录后，通过msprof --application=...分别启动 PE 0 与 PE 1 两个进程（rdma_aclgraph_demo 2 0 tcp://127.0.0.1:8899 2 0 0与rdma_aclgraph_demo 2 1 tcp://127.0.0.1:8899 2 0 0），并等待全部进程结束；
任一进程返回非 0 时脚本以非 0 退出，便于在 CI 中判断样例是否通过。

注：Ascend950 平台必须预先设置IBV_EXTEND_DRIVERS环境变量（见上文）。

方式二：手动启动（跨机多卡）

样例支持任意规模扩展。以双机 8 卡为例：总 PE 数n_pes=8，每机启动 4 个 PE（g_npus=4）；第 1 机 PE 号为 0–3（f_pe=0），第 2 机 PE 号为 4–7（f_pe=4）。将run.sh中对应的启动命令替换为如下形式（注意第 2 机命令行中的f_pe从 0 改为 4，并将tcp://{IP}:{Port}统一填 PE0 所在主机地址）：

# Server 1 pids=() msprof --application="./build/bin/rdma_aclgraph_demo 8 0 tcp://{IP address of server 1}:{Port number} 4 0 0" --output=${PROJECT_ROOT}/examples/rdma_aclgraph_demo/output/ & # pe 0 pid=$! pids+=("$pid") msprof --application="./build/bin/rdma_aclgraph_demo 8 1 tcp://{IP address of server 1}:{Port number} 4 0 0" --output=${PROJECT_ROOT}/examples/rdma_aclgraph_demo/output/ & # pe 1 pid=$! pids+=("$pid") msprof --application="./build/bin/rdma_aclgraph_demo 8 2 tcp://{IP address of server 1}:{Port number} 4 0 0" --output=${PROJECT_ROOT}/examples/rdma_aclgraph_demo/output/ & # pe 2 pid=$! pids+=("$pid") msprof --application="./build/bin/rdma_aclgraph_demo 8 3 tcp://{IP address of server 1}:{Port number} 4 0 0" --output=${PROJECT_ROOT}/examples/rdma_aclgraph_demo/output/ & # pe 3 pid=$! pids+=("$pid") # Server 2 pids=() msprof --application="./build/bin/rdma_aclgraph_demo 8 4 tcp://{IP address of server 1}:{Port number} 4 4 0" --output=${PROJECT_ROOT}/examples/rdma_aclgraph_demo/output/ & # pe 4 pid=$! pids+=("$pid") msprof --application="./build/bin/rdma_aclgraph_demo 8 5 tcp://{IP address of server 1}:{Port number} 4 4 0" --output=${PROJECT_ROOT}/examples/rdma_aclgraph_demo/output/ & # pe 5 pid=$! pids+=("$pid") msprof --application="./build/bin/rdma_aclgraph_demo 8 6 tcp://{IP address of server 1}:{Port number} 4 4 0" --output=${PROJECT_ROOT}/examples/rdma_aclgraph_demo/output/ & # pe 6 pid=$! pids+=("$pid") msprof --application="./build/bin/rdma_aclgraph_demo 8 7 tcp://{IP address of server 1}:{Port number} 4 4 0" --output=${PROJECT_ROOT}/examples/rdma_aclgraph_demo/output/ & # pe 7 pid=$! pids+=("$pid")

提示：跨机测试时，tcp://中的 IP 必须为 PE0 所在主机的 IP；在容器中执行跨机测试，启动容器需指定--net=host模式。

命令行参数说明

程序入口 main.cpp 要求恰好 6 个参数，格式为：

./rdma_aclgraph_demo <n_pes> <pe_id> <ipport> <g_npus> <f_pe> <f_npu>

参数	含义	说明
`n_pes`	全局 PE 数量	参与本次通信的总进程数，如单机双卡为 2、双机 8 卡为 8
`pe_id`	当前 PE 号	全局唯一，取值`[0, n_pes)`
`ipport`	SHMEM 初始化地址	格式`tcp://<IP>:<端口号>`；跨机测试时 IP 须为 PE0 所在主机的 IP
`g_npus`	当前机器上启动的 NPU 数量	用于`pe_id % g_npus + f_npu`计算本进程使用的设备号
`f_pe`	当前机器使用的第一个 PE 号	第 1 机通常为 0，第 2 机为 4
`f_npu`	当前机器使用的第一个 NPU 卡号	物理设备起始偏移，通常为 0

以双机 8 卡为例：第 1 机的 4 个进程参数为8 0..3 tcp://ip1:port 4 0 0，第 2 机的 4 个进程参数为8 4..7 tcp://ip1:port 4 4 0。

初始化与内存分配要点

从源码可进一步确认样例内部的关键初始化步骤（见 main.cpp）：

ACL 环境初始化：aclInit→aclrtSetDevice(pe_id % g_npus + f_npu)→aclrtCreateStream，完成设备绑定与 Stream 创建；
SHMEM 初始化：填充aclshmemx_init_attr_t attributes（含pe_id、n_pes、local_mem_size = 1 GB、ipport），并设置attributes.option_attr.data_op_engine_type = ACLSHMEM_DATA_OP_ROCE以指定 ROCE 数据通路，随后调用aclshmemx_init_attr(ACLSHMEMX_INIT_WITH_DEFAULT, &attributes)；
对称内存分配：通过aclshmem_malloc(SYMMETRIC_MEM_SIZE)（1024 字节）分配 AllGather 所需的对称内存分片，远端进程可通过对称地址直接访问；
退出流程：aclshmem_free/aclrtFree释放资源 →aclmdlRIDestroy销毁模型 →aclshmem_finalize→aclrtDestroyStream/aclrtResetDevice/aclFinalize。

结果验证与退出码

每轮图执行后，Host 端比对y_host[TRANS_SIZE * i + j]与期望值(10 + i) * (3 + zz)，任一不一致即打印ERROR loop: ...并将status置为-1；
全部通过则打印check transport result success, relative pe=<pe_id>；
main以状态码返回：成功输出[SUCCESS] demo run success in relative pe <pe_id>，失败输出[FAILED] demo run failed in relative pe <pe_id>并以EXIT_FAILURE退出，供上层脚本判断。

CANN SHMEM 中基于 ACLGraph 的 RDMA AllGather 同步机制与运行实战：rdma_aclgraph_demo 全解析