CANN SHMEM 的 ACLGraph 集成实践基于图捕获与重放的多卡 allGather 精度验证【免费下载链接】shmemCANN SHMEM 是面向昇腾平台的多机多卡内存通信库基于OpenSHMEM 标准协议实现跨设备的高效内存访问与数据同步。项目地址: https://gitcode.com/cann/shmem导读本文面向在昇腾平台上使用 CANN SHMEM 的开发者围绕仓库中的 aclgraph_demo 样例完整讲解如何将allGather集合通信算子编排进 ACLGraphAscendCL 图模式模型并通过「首次循环图捕获、后续循环图重放」的方式验证算子在图内逐次循环中的功能与精度。读完本文你将掌握 ACLGraph 图捕获/重放的 API 使用流程、SHMEM 集合通信算子入参的图适配改造方法magic参数从int值改为设备侧地址__gm__ int *、样例的编译运行方式与参数含义以及设备侧 allGather 内核基于信号量同步的实现原理。一、样例背景为什么要在 ACLGraph 中集成 SHMEM 集合通信ACLGraph对应aclmdlRI*系列接口是 AscendCL 提供的图执行模式先将一段由算子构成的 Host 侧调用序列「捕获」为图模型model后续循环只需重放该模型即可避免反复下发算子降低调度开销。而 CANN SHMEM 作为面向昇腾多机多卡的内存通信库其集合通信算子如allGather通常以设备侧 kernel 形式在 Stream 上异步执行天然可以与 ACLGraph 的算子编排模型结合。aclgraph_demo 的定位就是验证在 ACLGraph 图中混合编排普通的向量add算子与 SHMEM 的allGather集合通信算子并确认这种混合图在多次循环重放中功能与精度均正确。仓库中同样基于aclmdlRI接口的还有 rdma_aclgraph_demo可见「ACLGraph SHMEM 通信」是仓库着力验证的一类重要使用场景。二、ACLGraph 图结构与算子编排样例构造的 ACLGraph 模型包含四个算子节点按顺序执行为add → allGather → allGather → add具体数据流如下与 main.cpp 中捕获区间的调用顺序一一对应第一个add将输入input_ptr与input_b_ptr逐元素相加结果写入input_c_ptr通过aclrtMemcpyAsync将input_c_ptr的结果设备到设备拷回input_ptr作为第一个allGather的输入第一个allGather对input_ptr做全收集结果写入output_ptr第二个allGather对input_b_ptr做全收集结果写入output_b_ptr第二个add将两个allGather的输出相加结果写入output_c_ptr。在 Host 侧这四步被包在aclmdlRICaptureBegin/aclmdlRICaptureEnd之间完成图捕获main.cppif (zz 0) { aclmdlRICaptureBegin(stream, ACL_MODEL_RI_CAPTURE_MODE_RELAXED); run_vector_addT(trans_size, input_ptr, input_b_ptr, input_c_ptr, stream); aclrtMemcpyAsync(input_ptr, trans_size * sizeof(T), input_c_ptr, trans_size * sizeof(T), ACL_MEMCPY_DEVICE_TO_DEVICE, stream); allgather_demoT(BLOCK_NUM, stream, ffts_addr, (uint8_t*)input_ptr, (uint8_t*)output_ptr, (uint8_t*)ptr, trans_size, magic_ptr); allgather_demoT(BLOCK_NUM, stream, ffts_addr, (uint8_t*)input_b_ptr, (uint8_t*)output_b_ptr, (uint8_t*)ptr, trans_size, magic_ptr2); run_vector_addT(n_pes * trans_size, output_ptr, output_b_ptr, output_c_ptr, stream); aclmdlRICaptureEnd(stream, model); aclmdlRIExecuteAsync(model, stream); } else { aclmdlRIExecuteAsync(model, stream); }捕获完成后model保存了整条算子链后续每次循环只调用aclmdlRIExecuteAsync(model, stream)完成重放。样例默认loop_times 3main.cpp即 1 次捕获 2 次重放每次循环结束后都会执行一次精度校验。三、核心适配改造magic入参从int值改为设备地址为了适配 ACLGraphallGather算子用于同步控制的magic入参需要从「Host 侧传入的int值」改为「设备侧地址__gm__ int *」。这样做的原因非常直接被捕获进图内的 kernel 参数在重放时是固定的只有通过指针引用设备内存才能在每次重放前由 Host 修改该内存中的值从而让图内通信逻辑感知到循环次数的变化。从源码可以清楚看到这一改造的落点内核侧签名中magic已是设备地址__gm__ int* magic并在内核入口通过int magic_data *magic;读取aclgraph_demo_kernel.cppHost 侧在每次循环开始前先计算新的 magic 值并拷贝到设备内存再触发图执行main.cppfor (int zz 0; zz loop_times; zz) { magic; magic_value magic * MAGIC_MULTIPLIER; // MAGIC_MULTIPLIER 1024 magic_value_2 (magic loop_times 10) * MAGIC_MULTIPLIER; *magic_host_ptr magic_value; aclrtMemcpy(magic_ptr, sizeof(int), magic_host_ptr, sizeof(int), ACL_MEMCPY_HOST_TO_DEVICE); *magic_host_ptr magic_value_2; aclrtMemcpy(magic_ptr2, sizeof(int), magic_host_ptr, sizeof(int), ACL_MEMCPY_HOST_TO_DEVICE); ... }magic在图中承担同步角色设备侧内核用magic_data构造信号signal值通过aclshmemx_signal_op/aclshmem_signal_wait_until完成跨 PE 的同步握手详见下文第五节。由于每次循环magic单调递增可以有效区分不同轮次的同步信号避免重放时旧信号被误当作新一轮完成信号。这种「参数指针化 循环前更新」的模式正是 SHMEM 算子被安全嵌入 ACLGraph 图模型的关键技巧。四、编译与运行4.1 编译构建在仓库根目录shmem/下执行构建脚本并携带-examples选项即可编出所有示例含 aclgraph_demoA2 / A3 平台bash scripts/build.sh -examplesAscend950 平台bash scripts/build.sh -soc_type Ascend950 -examples构建完成后可执行文件输出到build/bin/aclgraph_demo动态库输出到build/librun.sh 中的LD_LIBRARY_PATH即指向这两个目录。4.2 运行 demo在examples/aclgraph_demo目录下执行# 完成 PEs 卡下的 ACLGraphadd allGather allGather add并在每次循环中验证精度。 # 支持 PEs : [2, 4, 8] bash run.sh -pes ${PEs}运行脚本会依次完成三件事run.sh生成 golden 数据调用python3 ./scripts/data_gen.py PE_SIZE TEST_TYPE生成每个 PE 的输入与期望输出设置运行环境导出SHMEM_UID_SESSION_ID127.0.0.1:8899并设置LD_LIBRARY_PATH指向构建产物拉起多进程按GNPU_NUM个进程数用msprof以--application方式启动build/bin/aclgraph_demo每个进程对应一个 PE。4.3 run.sh 可选参数run.sh除-pes外还支持以下参数默认值见 run.sh参数含义默认值说明-pesPE进程/rank数量2必须为正整数样例验证 PEs 2 / 4 / 8-ipport通信服务监听地址tcp://127.0.0.1:8766非空字符串-gnpus本机 NPU 数量进程数8必须为正整数若大于PE_SIZE会自动收敛为PE_SIZE-fnpu起始 NPU 编号0非负整数-type测试数据类型int当前仅支持int-tool性能采集工具msprof仅支持msprof脚本内置了严格的参数校验非法参数、空值、非整数、超出支持范围的值都会报错退出当GNPU_NUM PE_SIZE时自动把GNPU_NUM赋值为PE_SIZE并打印提示。4.4 直接运行二进制也可跳过脚本直接运行二进制参数顺序固定见 main.cpp./build/bin/aclgraph_demo n_pes pe_id ipport g_npus f_pe f_npu data_type loop_times位置参数含义样例值n_pesPE 总数2 / 4 / 8pe_id当前进程的 PE 编号0 ~ n_pes-1ipport通信地址tcp://127.0.0.1:8998main 默认值g_npus本机 NPU 数8f_pe/f_npu起始 PE / NPU 偏移0 / 0data_type数据类型仅支持intloop_times循环次数含捕获3进程内会按device_id pe_id % g_npus f_npu绑定 NPUmain.cpp随后依次执行aclInit、aclrtSetDevice、aclshmemx_init_attr初始化模式为ACLSHMEMX_INIT_WITH_DEFAULT并在结束时按aclshmem_finalize→aclrtResetDevice→aclFinalize的顺序清理。五、主机侧实现细节初始化、捕获与逐轮校验5.1 初始化与公共参数main.cpp开头定义了若干编译期常量main.cpp理解它们有助于把握样例的资源边界constexpr int64_t SYNC_FLAG_INTERVAL 16; // 同步 flag 间隔 constexpr int64_t UB_DMA_MAX_SIZE 190 * 1024; // UB 单次 DMA 搬运上限约 190 KB constexpr int64_t GVA_BUFF_MAX_SIZE 100 * 1024 * 1024; // 对称内存缓冲上限100 MB constexpr uint32_t MAGIC_MULTIPLIER 1024; // magic 信号值放大倍数 constexpr uint32_t DATA_SIZE_THRESHOLD 2097152; // 大数据/小数据内核分界2 MB constexpr uint32_t BLOCK_NUM_SMALL_DATA 8; // 小数据内核 AIV 核数 constexpr uint32_t BLOCK_NUM_LARGE_DATA 16; // 大数据内核 AIV 核数测试共分 2 个 case数据长度为262144 * (1 i)i 0, 1即 256K 与 512K 个int元素main.cpp。根据trans_size * sizeof(T)是否小于 2 MB 阈值选择 8 核或 16 核 AIV 的内核启动配置main.cpp。5.2 设备内存布局每个 case 内样例使用aclrtMalloc分配输入/输出缓冲并使用 SHMEM 对称堆接口分配同步与数据缓冲main.cppvoid *ptr aclshmem_malloc(aiv_num * SYNC_FLAG_INTERVAL * sizeof(T) GVA_BUFF_MAX_SIZE / sizeof(T));这段aclshmem_malloc分配的内存同时承载两块用途起始aiv_num * SYNC_FLAG_INTERVAL个int32用作同步 flag 区每核一个 flag 槽位后续最多GVA_BUFF_MAX_SIZE字节用作GVA 数据区即 allGather 各 PE 写入、读取的对称内存中转区。5.3 逐轮精度校验每次循环在aclrtSynchronizeStream(stream)之后将output_c_ptr拷回 Host与 golden 对比main.cppfor (int zzz 0; zzz n_pes * trans_size; zzz) { if (static_castint(output_host[zzz]) ! static_castint(golden_host[zzz]) * (zz 3)) { std::cout ERROR each loop: ... std::endl; status -1; break; } }期望值是golden * (zz 3)其中zz为循环序号0、1、2golden为所有 PE 输入拼接的全收集正确结果。该期望值与数据流严格自洽第一次add把input更新为2 × 原输入随后两个 allGather 分别贡献2 × golden与1 × golden第二次add得3 × golden此后输入随循环不断累加期望值随之递增为4 × golden、5 × golden。可见校验同时覆盖了两点allGather 图内功能正确以及每次重放时经设备地址更新的magic参数确实生效。每个 case 结束时还会调用aclshmemx_get_prof(nullptr, true)触发性能数据落盘并执行内存与模型销毁清理。六、设备侧 allGather 内核实现信号量同步与大小数据分路内核入口ShmemAllGather_int依据数据量分派两条实现路径aclgraph_demo_kernel.cpp小于 2 MB 走all_gather_small_data否则走all_gather_big_data。两条路径共享同一套「两步走」流水Step 1本地写远端每个 AIV 核把本 PE 输入input的一段数据通过aclshmemx_mte_put_nbi写入自己的对称内存区GVA 数据区随后通过aclshmemx_signal_op(..., ACLSHMEM_SIGNAL_SET, my_rank)发出递增信号Step 2远端读本地各核通过aclshmem_signal_wait_until(..., ACLSHMEM_CMP_EQ, magic_data)等待对端信号到达再通过aclshmemx_mte_get_nbi从对端对称内存区把数据拉回output的对应分片位置。小数据路径直接以elements / aiv_num切分数据配合aclshmem_quiet()与AscendC::SyncAll()保证本端写任务完成后再发信号aclgraph_demo_kernel.cpp大数据路径则将数据按GVA_BUFF_MAX_SIZE分块循环处理并把 AIV 核分成「发送半区」与「接收半区」两组接收组使用 ping-pong 双缓冲ping_buff/pong_buff轮转拉取数据同时用flags_ub1/flags_ub2两级 UB 标志位跟踪各数据块的就绪进度避免等待空转aclgraph_demo_kernel.cpp。可以推断这套内核的同步信号语义正是上文「magic 入参指针化」改造的直接受益者magic_data参与信号值的构造与比较Host 侧每轮循环更新设备内存中的magic即可让重放图内每一轮的信号匹配新值从而实现图内通信算子的多轮安全复用。内核启动封装allgather_demoT位于 aclgraph_demo_kernel.h以block_dimAIV 核数、stream、FFTS 配置地址、输入输出及magic设备地址为参数。图内的两个向量add算子由仓库内的 AscendC 算子实现add_custom.hadd_custom内核按块切分数据使用深度为 2 的TPipe流水队列执行 CopyIn →AscendC::Add→ CopyOut并单独处理尾块Host 封装run_vector_add依据platform_ascendc查询到的 UB 大小与 AIV 核数自动计算numBlocks、blockLength与tileSize从而让一个简单的 add 算子也能在全卡规模上并行。七、golden 数据生成与校验基准run.sh启动前会自动调用 scripts/data_gen.py 生成基准数据避免手工准备 bin 文件固定随机种子np.random.seed(42)保证多节点、多次运行的数据可复现每个 PE 的输入input_gm_i.bin用np.random.uniform(0.0, 10.0)生成全收集的正确结果golden.bin为所有 PE 输入的按序拼接目录按golden/allgather_data_len_pe_size/组织主程序运行时按同样的路径拼接规则加载input_gm_pe_id.bin与golden.binmain.cpptype_map预留了int32/float16/bfloat16等类型的映射但当前样例的run.sh与主程序均限定int类型。八、构建与运行的适用前提构建命令中的-soc_type Ascend950仅用于 Ascend950 平台A2 / A3 平台直接使用默认目标即可运行需要ASCEND_HOME_PATH环境变量指向已安装的 CANN 工具链run.sh用它拼接lib64路径多进程由msprof --application拉起每个进程绑定一块 NPU若PE_SIZE超过单机 NPU 数需调整-gnpus与-fnpu以适配实际机器拓扑样例输出目录为examples/aclgraph_demo/output/msprof 落盘位置可在其中检查采集到的性能与日志数据。总结aclgraph_demo 为「ACLGraph 图模式 SHMEM 集合通信算子」的融合使用提供了完整的参考实现通过将allGather的magic同步参数改造为设备侧地址__gm__ int *成功让带通信算子的图模型在捕获后支持多轮重放且每轮的信号同步与精度均正确。其主机侧覆盖了 ACL 初始化、图捕获/重放、golden 校验与资源回收的完整生命周期设备侧则给出了基于signal_op/signal_wait_until与 UB ping-pong 缓冲的大小数据双路 allGather 内核实现。开发者可直接复用本文的编译、运行与参数说明将该模式迁移到其他集合通信算子如 reduce、allreduce 等与 ACLGraph 的集成场景中。【免费下载链接】shmemCANN SHMEM 是面向昇腾平台的多机多卡内存通信库基于OpenSHMEM 标准协议实现跨设备的高效内存访问与数据同步。项目地址: https://gitcode.com/cann/shmem创作声明：本文部分内容由AI辅助生成（AIGC），仅供参考 CANN SHMEM 的 ACLGraph 集成实践：基于图捕获与重放的多卡 allGather 精度验证
CANN SHMEM 的 ACLGraph 集成实践基于图捕获与重放的多卡 allGather 精度验证【免费下载链接】shmemCANN SHMEM 是面向昇腾平台的多机多卡内存通信库基于OpenSHMEM 标准协议实现跨设备的高效内存访问与数据同步。项目地址: https://gitcode.com/cann/shmem导读本文面向在昇腾平台上使用 CANN SHMEM 的开发者围绕仓库中的 aclgraph_demo 样例完整讲解如何将allGather集合通信算子编排进 ACLGraphAscendCL 图模式模型并通过「首次循环图捕获、后续循环图重放」的方式验证算子在图内逐次循环中的功能与精度。读完本文你将掌握 ACLGraph 图捕获/重放的 API 使用流程、SHMEM 集合通信算子入参的图适配改造方法magic参数从int值改为设备侧地址__gm__ int *、样例的编译运行方式与参数含义以及设备侧 allGather 内核基于信号量同步的实现原理。一、样例背景为什么要在 ACLGraph 中集成 SHMEM 集合通信ACLGraph对应aclmdlRI*系列接口是 AscendCL 提供的图执行模式先将一段由算子构成的 Host 侧调用序列「捕获」为图模型model后续循环只需重放该模型即可避免反复下发算子降低调度开销。而 CANN SHMEM 作为面向昇腾多机多卡的内存通信库其集合通信算子如allGather通常以设备侧 kernel 形式在 Stream 上异步执行天然可以与 ACLGraph 的算子编排模型结合。aclgraph_demo 的定位就是验证在 ACLGraph 图中混合编排普通的向量add算子与 SHMEM 的allGather集合通信算子并确认这种混合图在多次循环重放中功能与精度均正确。仓库中同样基于aclmdlRI接口的还有 rdma_aclgraph_demo可见「ACLGraph SHMEM 通信」是仓库着力验证的一类重要使用场景。二、ACLGraph 图结构与算子编排样例构造的 ACLGraph 模型包含四个算子节点按顺序执行为add → allGather → allGather → add具体数据流如下与 main.cpp 中捕获区间的调用顺序一一对应第一个add将输入input_ptr与input_b_ptr逐元素相加结果写入input_c_ptr通过aclrtMemcpyAsync将input_c_ptr的结果设备到设备拷回input_ptr作为第一个allGather的输入第一个allGather对input_ptr做全收集结果写入output_ptr第二个allGather对input_b_ptr做全收集结果写入output_b_ptr第二个add将两个allGather的输出相加结果写入output_c_ptr。在 Host 侧这四步被包在aclmdlRICaptureBegin/aclmdlRICaptureEnd之间完成图捕获main.cppif (zz 0) { aclmdlRICaptureBegin(stream, ACL_MODEL_RI_CAPTURE_MODE_RELAXED); run_vector_addT(trans_size, input_ptr, input_b_ptr, input_c_ptr, stream); aclrtMemcpyAsync(input_ptr, trans_size * sizeof(T), input_c_ptr, trans_size * sizeof(T), ACL_MEMCPY_DEVICE_TO_DEVICE, stream); allgather_demoT(BLOCK_NUM, stream, ffts_addr, (uint8_t*)input_ptr, (uint8_t*)output_ptr, (uint8_t*)ptr, trans_size, magic_ptr); allgather_demoT(BLOCK_NUM, stream, ffts_addr, (uint8_t*)input_b_ptr, (uint8_t*)output_b_ptr, (uint8_t*)ptr, trans_size, magic_ptr2); run_vector_addT(n_pes * trans_size, output_ptr, output_b_ptr, output_c_ptr, stream); aclmdlRICaptureEnd(stream, model); aclmdlRIExecuteAsync(model, stream); } else { aclmdlRIExecuteAsync(model, stream); }捕获完成后model保存了整条算子链后续每次循环只调用aclmdlRIExecuteAsync(model, stream)完成重放。样例默认loop_times 3main.cpp即 1 次捕获 2 次重放每次循环结束后都会执行一次精度校验。三、核心适配改造magic入参从int值改为设备地址为了适配 ACLGraphallGather算子用于同步控制的magic入参需要从「Host 侧传入的int值」改为「设备侧地址__gm__ int *」。这样做的原因非常直接被捕获进图内的 kernel 参数在重放时是固定的只有通过指针引用设备内存才能在每次重放前由 Host 修改该内存中的值从而让图内通信逻辑感知到循环次数的变化。从源码可以清楚看到这一改造的落点内核侧签名中magic已是设备地址__gm__ int* magic并在内核入口通过int magic_data *magic;读取aclgraph_demo_kernel.cppHost 侧在每次循环开始前先计算新的 magic 值并拷贝到设备内存再触发图执行main.cppfor (int zz 0; zz loop_times; zz) { magic; magic_value magic * MAGIC_MULTIPLIER; // MAGIC_MULTIPLIER 1024 magic_value_2 (magic loop_times 10) * MAGIC_MULTIPLIER; *magic_host_ptr magic_value; aclrtMemcpy(magic_ptr, sizeof(int), magic_host_ptr, sizeof(int), ACL_MEMCPY_HOST_TO_DEVICE); *magic_host_ptr magic_value_2; aclrtMemcpy(magic_ptr2, sizeof(int), magic_host_ptr, sizeof(int), ACL_MEMCPY_HOST_TO_DEVICE); ... }magic在图中承担同步角色设备侧内核用magic_data构造信号signal值通过aclshmemx_signal_op/aclshmem_signal_wait_until完成跨 PE 的同步握手详见下文第五节。由于每次循环magic单调递增可以有效区分不同轮次的同步信号避免重放时旧信号被误当作新一轮完成信号。这种「参数指针化 循环前更新」的模式正是 SHMEM 算子被安全嵌入 ACLGraph 图模型的关键技巧。四、编译与运行4.1 编译构建在仓库根目录shmem/下执行构建脚本并携带-examples选项即可编出所有示例含 aclgraph_demoA2 / A3 平台bash scripts/build.sh -examplesAscend950 平台bash scripts/build.sh -soc_type Ascend950 -examples构建完成后可执行文件输出到build/bin/aclgraph_demo动态库输出到build/librun.sh 中的LD_LIBRARY_PATH即指向这两个目录。4.2 运行 demo在examples/aclgraph_demo目录下执行# 完成 PEs 卡下的 ACLGraphadd allGather allGather add并在每次循环中验证精度。 # 支持 PEs : [2, 4, 8] bash run.sh -pes ${PEs}运行脚本会依次完成三件事run.sh生成 golden 数据调用python3 ./scripts/data_gen.py PE_SIZE TEST_TYPE生成每个 PE 的输入与期望输出设置运行环境导出SHMEM_UID_SESSION_ID127.0.0.1:8899并设置LD_LIBRARY_PATH指向构建产物拉起多进程按GNPU_NUM个进程数用msprof以--application方式启动build/bin/aclgraph_demo每个进程对应一个 PE。4.3 run.sh 可选参数run.sh除-pes外还支持以下参数默认值见 run.sh参数含义默认值说明-pesPE进程/rank数量2必须为正整数样例验证 PEs 2 / 4 / 8-ipport通信服务监听地址tcp://127.0.0.1:8766非空字符串-gnpus本机 NPU 数量进程数8必须为正整数若大于PE_SIZE会自动收敛为PE_SIZE-fnpu起始 NPU 编号0非负整数-type测试数据类型int当前仅支持int-tool性能采集工具msprof仅支持msprof脚本内置了严格的参数校验非法参数、空值、非整数、超出支持范围的值都会报错退出当GNPU_NUM PE_SIZE时自动把GNPU_NUM赋值为PE_SIZE并打印提示。4.4 直接运行二进制也可跳过脚本直接运行二进制参数顺序固定见 main.cpp./build/bin/aclgraph_demo n_pes pe_id ipport g_npus f_pe f_npu data_type loop_times位置参数含义样例值n_pesPE 总数2 / 4 / 8pe_id当前进程的 PE 编号0 ~ n_pes-1ipport通信地址tcp://127.0.0.1:8998main 默认值g_npus本机 NPU 数8f_pe/f_npu起始 PE / NPU 偏移0 / 0data_type数据类型仅支持intloop_times循环次数含捕获3进程内会按device_id pe_id % g_npus f_npu绑定 NPUmain.cpp随后依次执行aclInit、aclrtSetDevice、aclshmemx_init_attr初始化模式为ACLSHMEMX_INIT_WITH_DEFAULT并在结束时按aclshmem_finalize→aclrtResetDevice→aclFinalize的顺序清理。五、主机侧实现细节初始化、捕获与逐轮校验5.1 初始化与公共参数main.cpp开头定义了若干编译期常量main.cpp理解它们有助于把握样例的资源边界constexpr int64_t SYNC_FLAG_INTERVAL 16; // 同步 flag 间隔 constexpr int64_t UB_DMA_MAX_SIZE 190 * 1024; // UB 单次 DMA 搬运上限约 190 KB constexpr int64_t GVA_BUFF_MAX_SIZE 100 * 1024 * 1024; // 对称内存缓冲上限100 MB constexpr uint32_t MAGIC_MULTIPLIER 1024; // magic 信号值放大倍数 constexpr uint32_t DATA_SIZE_THRESHOLD 2097152; // 大数据/小数据内核分界2 MB constexpr uint32_t BLOCK_NUM_SMALL_DATA 8; // 小数据内核 AIV 核数 constexpr uint32_t BLOCK_NUM_LARGE_DATA 16; // 大数据内核 AIV 核数测试共分 2 个 case数据长度为262144 * (1 i)i 0, 1即 256K 与 512K 个int元素main.cpp。根据trans_size * sizeof(T)是否小于 2 MB 阈值选择 8 核或 16 核 AIV 的内核启动配置main.cpp。5.2 设备内存布局每个 case 内样例使用aclrtMalloc分配输入/输出缓冲并使用 SHMEM 对称堆接口分配同步与数据缓冲main.cppvoid *ptr aclshmem_malloc(aiv_num * SYNC_FLAG_INTERVAL * sizeof(T) GVA_BUFF_MAX_SIZE / sizeof(T));这段aclshmem_malloc分配的内存同时承载两块用途起始aiv_num * SYNC_FLAG_INTERVAL个int32用作同步 flag 区每核一个 flag 槽位后续最多GVA_BUFF_MAX_SIZE字节用作GVA 数据区即 allGather 各 PE 写入、读取的对称内存中转区。5.3 逐轮精度校验每次循环在aclrtSynchronizeStream(stream)之后将output_c_ptr拷回 Host与 golden 对比main.cppfor (int zzz 0; zzz n_pes * trans_size; zzz) { if (static_castint(output_host[zzz]) ! static_castint(golden_host[zzz]) * (zz 3)) { std::cout ERROR each loop: ... std::endl; status -1; break; } }期望值是golden * (zz 3)其中zz为循环序号0、1、2golden为所有 PE 输入拼接的全收集正确结果。该期望值与数据流严格自洽第一次add把input更新为2 × 原输入随后两个 allGather 分别贡献2 × golden与1 × golden第二次add得3 × golden此后输入随循环不断累加期望值随之递增为4 × golden、5 × golden。可见校验同时覆盖了两点allGather 图内功能正确以及每次重放时经设备地址更新的magic参数确实生效。每个 case 结束时还会调用aclshmemx_get_prof(nullptr, true)触发性能数据落盘并执行内存与模型销毁清理。六、设备侧 allGather 内核实现信号量同步与大小数据分路内核入口ShmemAllGather_int依据数据量分派两条实现路径aclgraph_demo_kernel.cpp小于 2 MB 走all_gather_small_data否则走all_gather_big_data。两条路径共享同一套「两步走」流水Step 1本地写远端每个 AIV 核把本 PE 输入input的一段数据通过aclshmemx_mte_put_nbi写入自己的对称内存区GVA 数据区随后通过aclshmemx_signal_op(..., ACLSHMEM_SIGNAL_SET, my_rank)发出递增信号Step 2远端读本地各核通过aclshmem_signal_wait_until(..., ACLSHMEM_CMP_EQ, magic_data)等待对端信号到达再通过aclshmemx_mte_get_nbi从对端对称内存区把数据拉回output的对应分片位置。小数据路径直接以elements / aiv_num切分数据配合aclshmem_quiet()与AscendC::SyncAll()保证本端写任务完成后再发信号aclgraph_demo_kernel.cpp大数据路径则将数据按GVA_BUFF_MAX_SIZE分块循环处理并把 AIV 核分成「发送半区」与「接收半区」两组接收组使用 ping-pong 双缓冲ping_buff/pong_buff轮转拉取数据同时用flags_ub1/flags_ub2两级 UB 标志位跟踪各数据块的就绪进度避免等待空转aclgraph_demo_kernel.cpp。可以推断这套内核的同步信号语义正是上文「magic 入参指针化」改造的直接受益者magic_data参与信号值的构造与比较Host 侧每轮循环更新设备内存中的magic即可让重放图内每一轮的信号匹配新值从而实现图内通信算子的多轮安全复用。内核启动封装allgather_demoT位于 aclgraph_demo_kernel.h以block_dimAIV 核数、stream、FFTS 配置地址、输入输出及magic设备地址为参数。图内的两个向量add算子由仓库内的 AscendC 算子实现add_custom.hadd_custom内核按块切分数据使用深度为 2 的TPipe流水队列执行 CopyIn →AscendC::Add→ CopyOut并单独处理尾块Host 封装run_vector_add依据platform_ascendc查询到的 UB 大小与 AIV 核数自动计算numBlocks、blockLength与tileSize从而让一个简单的 add 算子也能在全卡规模上并行。七、golden 数据生成与校验基准run.sh启动前会自动调用 scripts/data_gen.py 生成基准数据避免手工准备 bin 文件固定随机种子np.random.seed(42)保证多节点、多次运行的数据可复现每个 PE 的输入input_gm_i.bin用np.random.uniform(0.0, 10.0)生成全收集的正确结果golden.bin为所有 PE 输入的按序拼接目录按golden/allgather_data_len_pe_size/组织主程序运行时按同样的路径拼接规则加载input_gm_pe_id.bin与golden.binmain.cpptype_map预留了int32/float16/bfloat16等类型的映射但当前样例的run.sh与主程序均限定int类型。八、构建与运行的适用前提构建命令中的-soc_type Ascend950仅用于 Ascend950 平台A2 / A3 平台直接使用默认目标即可运行需要ASCEND_HOME_PATH环境变量指向已安装的 CANN 工具链run.sh用它拼接lib64路径多进程由msprof --application拉起每个进程绑定一块 NPU若PE_SIZE超过单机 NPU 数需调整-gnpus与-fnpu以适配实际机器拓扑样例输出目录为examples/aclgraph_demo/output/msprof 落盘位置可在其中检查采集到的性能与日志数据。总结aclgraph_demo 为「ACLGraph 图模式 SHMEM 集合通信算子」的融合使用提供了完整的参考实现通过将allGather的magic同步参数改造为设备侧地址__gm__ int *成功让带通信算子的图模型在捕获后支持多轮重放且每轮的信号同步与精度均正确。其主机侧覆盖了 ACL 初始化、图捕获/重放、golden 校验与资源回收的完整生命周期设备侧则给出了基于signal_op/signal_wait_until与 UB ping-pong 缓冲的大小数据双路 allGather 内核实现。开发者可直接复用本文的编译、运行与参数说明将该模式迁移到其他集合通信算子如 reduce、allreduce 等与 ACLGraph 的集成场景中。【免费下载链接】shmemCANN SHMEM 是面向昇腾平台的多机多卡内存通信库基于OpenSHMEM 标准协议实现跨设备的高效内存访问与数据同步。项目地址: https://gitcode.com/cann/shmem创作声明：本文部分内容由AI辅助生成（AIGC），仅供参考