CANN SHMEM Device 侧 AtomicAdd 累加模式实战MTE 批量硬件累加、UDMA/RDMA 标量原子与 SDMA Fallback 全解析【免费下载链接】shmemCANN SHMEM 是面向昇腾平台的多机多卡内存通信库基于OpenSHMEM 标准协议实现跨设备的高效内存访问与数据同步。项目地址: https://gitcode.com/cann/shmem本文基于 CANN SHMEM 开源仓库中的 AtomicAdd 累加模式参考手册atomic-add-pattern.md系统讲解 SHMEM Device 侧使用 MTE / UDMA / RDMA / SDMA 引擎实现跨 PE reduce 累加的正确模式与效率决策并结合 shmem_device_amo.h、shmem_device_amo.hpp 等真实源码印证统一 API 的拓扑分发逻辑与同步模型。读完后你将能够区分标量 atomic与批量 atomicSetAtomicAddT()两类机制、为不同链路拓扑HCCS / PCIE / SIO / RoCE选择正确的累加路径、写出带完整 barrier 与 MTE3 drain 的安全调用序列并规避丢失累加的经典反模式。1. 核心概念两类必须区分的 Atomic Add 机制在 SHMEM Device 侧存在两类atomic add 机制混用它们是最常见的错误来源类型机制操作对象数据规模标量 atomic通过引擎 APIaclshmemx_mte_atomic_add等对远端对称地址做标量加法远端 PE 的__gm__地址单个元素批量 atomicSetAtomicAddT()切换当前 core 的 GM 写入模式后续 MTE3 写入如mte_get_nbi的目的端自动变为D src本地GM 地址tile 级别关键语义SetAtomicAddT()不是执行一次加法而是把当前 core 后续的所有 GM 写入模式切换为 atomic add真正发生加法的是切换之后 MTE3 的写入操作。概念说明SetAtomicAddT()切换当前 core 的 GM 写模式为 atomic addT 为元素类型half/float/int32/bfloat16SetAtomicNone()恢复 GM 写为普通覆盖模式标量 atomic 触发者aclshmemx_mte_atomic_add/aclshmemx_udma_atomic_add/aclshmemx_roce_atomic_add批量 atomic 触发者aclshmemx_mte_get_nbi的目的端 GM 写入数学语义普通写D srcatomic 模式D D src2. 三引擎 AtomicAdd 接口全景与源码分发逻辑2.1 统一 APIaclshmem_NAME_atomic_add推荐首选入口统一 API 声明在 shmem_device_amo.h通过类型宏批量生成各 dtype 的重载ACLSHMEM_TYPE_FUNC_ATOMIC_ADDint8/int16/bfloat16/halfACLSHMEM_TYPE_FUNC_ATOMIC_ADD_910int32/floatAscend 910 系列ACLSHMEM_TYPE_FUNC_ATOMIC_ADD_EXTuint32/uint64/int64Ascend950同时头文件提供 OpenSHMEM 风格别名映射见 shmem_device_amo.h#L212-L217#define shmem_int8_atomic_add aclshmem_int8_atomic_add #define shmem_int16_atomic_add aclshmem_int16_atomic_add #define shmem_int32_atomic_add aclshmem_int32_atomic_add #define shmem_half_atomic_add aclshmem_half_atomic_add #define shmem_bfloat16_atomic_add aclshmem_bfloat16_atomic_add #define shmem_float_atomic_add aclshmem_float_atomic_add函数签名为void aclshmem_NAME_atomic_add(__gm__ TYPE* dst, TYPE value, int32_t pe)内部根据topo_list[pe]自动分发。这一点在实现文件 shmem_device_amo.hpp 中可以清楚看到#define ACLSHMEM_ATOMIC_ADD_910_TYPENAME(NAME, TYPE) ACLSHMEM_DEVICE void aclshmem_##NAME##_atomic_add(__gm__ TYPE* dst, TYPE value, int32_t pe) { __gm__ aclshmem_device_host_state_t* device_state aclshmemi_get_state(); if (device_state-topo_list[pe] ACLSHMEM_TRANSPORT_MTE) { /* MTE path - supports A2/A3/Ascend_950 */ ... } else if (device_state-topo_list[pe] ACLSHMEM_TRANSPORT_ROCE) { /* ROCE path - supports Ascend_950 only */ ... aclshmemx_roce_atomic_add(dst, value, pe); } else if (device_state-topo_list[pe] ACLSHMEM_TRANSPORT_UDMA) { /* UDMA path - supports Ascend_950 only */ aclshmemx_udma_atomic_add(dst, value, pe); } else { ACLSHMEM_DEBUG_FUNC(aclshmemi_kernel_abort, atomic_add is only supported on MTE, ROCE and UDMA path.); } }即分发规则为topo_list[pe] ACLSHMEM_TRANSPORT_MTE → MTE 路径 topo_list[pe] ACLSHMEM_TRANSPORT_ROCE → RDMARoCE路径 topo_list[pe] ACLSHMEM_TRANSPORT_UDMA → UDMA 路径注意一个从源码可以确认的重要事实int8/int16/bfloat16/half四个直接支持类型走纯 MTE 路径非 MTE 拓扑下会直接kernel_abort见 shmem_device_amo.hpp#L25-L45 中atomic_add for int8/int16/bfloat16/half is MTE-only的分支——小位宽浮点原子加只存在于 HCCS 机内链路场景。MTE 路径的内部实现shmem_device_amo.hpp#L29-L41分六步aclshmem_ptr(dst, pe)把对称地址转译为远端可访问地址Scalar 把value写入 UB bufferub_tensor.SetValue(0, value)SetFlagS_MTE3/WaitFlagS_MTE3隔离 Scalar→MTE3SetAtomicAddTYPE()开启 atomic 模式aclshmemi_copy_ub2gm触发 MTE3 写入——这是内部接口用户代码 MUST 调用aclshmem_*_atomic_add/aclshmemx_mte_atomic_add禁止直接调用aclshmemi_copy_ub2gmSetAtomicNone()关闭 atomic 模式。UDMA 路径则直接转调aclshmemx_udma_atomic_add(dst, value, pe)且后续必须调用aclshmemx_udma_quiet(pe)确保远端完成。统一 API 的约束操作远端对称地址dst 是远程 PE 的地址通过aclshmem_ptr转译每次调用只加一个标量值不支持批量 tile 累加异步语义void返回的 add 接口返回时远端操作可能尚未完成。头文件中 Atomic Interface Synchronization Model 明确区分了两类接口同步接口fetch_add/fetch_inc/swap/compare_swap等带返回值的完成前不返回内部已做 quiet异步接口void返回的add/inc/set等必须自行同步——RDMA 用aclshmemx_roce_quietUDMA 用aclshmemx_udma_quietMTE 按数据依赖插入SetFlag/WaitFlag读 GM 用MTE3_MTE2事件复用 UB 用MTE3_S事件。另外头文件注释多次强调MTE 传输不支持跨 PCIeinter-node通信跨 PCIe 场景必须走 RDMA 或 UDMA 路径且 RDMA 路径不支持对同一 PE 的并发 RMA/AMO。2.2 MTE 引擎标量与批量两条路径MTE 是机内 HCCS full-mesh 场景的主力引擎提供标量和批量两条路径。2.2.1 标量aclshmemx_mte_atomic_addaclshmemx_mte_atomic_add((__gm__ T*)dst, value, pe);属性说明操作方向本 PE 写入远端 PE的 dst 地址执行remote_dst value数据规模单个标量值执行单元910B/910C: MTE2Ascend950: UB MTE3uint32/int64/uint64 为 Scalar AtomicAdd见下方平台差异表同步需SetFlagS_MTE3/WaitFlagS_MTE3隔离 Scalar→MTE3读结果需MTE3_MTE2隔离dtypeint8_t/int16_t/int32_t/float/half/bfloat16_tAscend950 额外支持uint32_t/int64_t/uint64_t平台差异与头文件 shmem_device_mte.h#L254-L281 的文档注释完全一致dtype910B/910CAscend950int8_t/int16_t/half/bfloat16_tMTE2 atomicUB MTE3 atomicint32_t/floatMTE2 atomicUB MTE3 atomicuint32_t/int64_t/uint64_t不支持Scalar AtomicAdd2.2.2 批量SetAtomicAddT()mte_get_nbiAscendC::SetAtomicAddElementD(); // 所有 mte_get_nbi 的目的端 GM 写入自动变为 D src reduceScatter(gmBlockSrc, ..., gmBlockDst, ..., remoteRankIdx); // drain MTE3 AscendC::SetAtomicNone();属性说明操作方向本地core 对本地D做 atomic add 写入数据规模tile 级别如32 × 256元素可多次 get 并发 in-flight执行单元MTE3目的端写回阶段同步PipeBarrierPIPE_ALL()隔离模式切换SetFlagMTE3_S/WaitFlagMTE3_Sdraindtypehalf/float/int32/bfloat16SetAtomicAdd支持的 AscendC 类型2.2.3 两条路径对比标量mte_atomic_add批量SetAtomicAdd get操作对象远端 PE 地址本地 GM 地址目的端数据粒度单个元素tile多元素批量并发模型逐元素调用各自独立完成多个 get 可全部 in-flightMTE3 管道并行累加典型场景signal 计数器、flags、单元素累加reduce-scatter / allreduce 的 tile 级批量累加吞吐低标量高批量流水2.3 UDMA 引擎aclshmemx_udma_atomic_add((__gm__ T*)dst, value, pe); aclshmemx_udma_quiet(pe); // 确保操作在远端完成属性说明操作方向写入远端 PE 地址数据规模单个标量值批量能力无UDMA 不提供类似SetAtomicAdd的批量模式同步必须在读取结果前调用aclshmemx_udma_quiet(pe)dtypeint32_t/uint32_t/int64_t/uint64_t/float注意事项并发 RMA/AMO 到同一 PE 不支持接口声明见 shmem_device_udma.h#L1039。仓库中还有可直接运行的 UDMA 原子加示例工程 examples/udma_atomic_add含 main.cpp、udma_atomic_add_kernel.cpp 与 run.sh适合对照本节的调用与同步要求做实操验证。2.4 RDMA 引擎aclshmemx_roce_atomic_add((__gm__ T*)dst, value, pe);属性说明平台Ascend950操作方向跨 server 写入远端 PE 地址数据规模单个标量值批量能力无dtype仅int32_t/uint32_t/int64_t/uint64_t不支持浮点/half统一 API 中对float会显式 abort见 shmem_device_amo.hpp#L71-L75注意事项并发 RMA/AMO 到同一 PE 不支持读结果前需aclshmemx_roce_quiet2.5 SDMA 引擎现状SDMA没有 device 侧 atomic add API——代码库中不存在aclshmemx_sdma_atomic_add或等价物shmem_device_sdma.h 只提供aclshmemx_sdma_get_nbi等搬运类接口目的端为普通覆盖写。因此在 SHMEM AIV kernel 中实现跨 PE reduce 时如果 MTE 不可用跨 PCIE 或跨 server必须使用第 4 节的 SDMA fallback 策略。2.6 引擎能力总表引擎标量 atomic API批量 reduce支持 dtype链路及平台MTEaclshmemx_mte_atomic_addSetAtomicAddT()mte_get_nbiint8/16/32, float, half, bf16 uint32/64、int64 on 950HCCS full-mesh / 910B-950UDMAaclshmemx_udma_atomic_add—int32/uint32/int64/uint64, floatPCIE/SIO / 910B-950RDMAaclshmemx_roce_atomic_add—int32/uint32/int64/uint64RoCE 跨机 / Ascend950SDMA———HCCS/PCIE/SIO仅 Host 侧 inline reduce/ 910B-910C3. MTE 标量 vs 批量如何选择场景推荐路径理由单个 signal/flag 累加aclshmemx_mte_atomic_add标量一条调用即完成无需切换 pipeline 模式tile 级批量 reducereduce-scatter / allreduceSetAtomicAddT()mte_get_nbi批量所有 peer 的 get in-flightMTE3 管道并行少量元素的远端累加 64 元素标量mte_atomic_addSetAtomicAdd 模式切换开销 收益需要读取累加前旧值fetch_addaclshmem_NAME_atomic_fetch_add统一 APIMTE/UDMA/ROCE 自动分发Ascend950批量模式无 fetch 语义fetch 系列为同步接口内部已含 quiet4. SDMA 如何实现并发 ReduceFallback 策略当 MTE 不可用跨 PCIE/SIO 链路SDMA 仍然可以做跨 PE 数据搬运aclshmemx_sdma_get_nbi但缺少 inline reduce 能力必须退化为本地收数据 → UB 累加 → 写回的模式。4.1 约束无 device 侧 inline reduceSDMA get 的目的端写入是普通覆盖写不是 atomic add如果多个 peer 同时 get 到同一 destination 且用普通 DataCopy 写回结果是覆盖而非累加UB 累加路径导致 peer 维度必须串行每轮 get → UB →AscendC::Add→ 释放 UB → 下一轮 get。4.2 策略chunk 分拆 核间并行 核内串行步骤 1将本地输出按 offset 切分为互不重叠的 chunk chunk 数 参与的 AIV core 数 步骤 2每个 AIV core 分配一个 chunk 不同 core 的 destination 地址不重叠 → 核间无写冲突 → 核间可并行 步骤 3每个 core 内部对 R-1 个 peer 串行处理 for peer in [0, R) 且 peer ! self: sdma_get_nbi(peer Sym[chunk]) → UB AscendC::Add(UB, local_D[chunk]) → UB DataCopy(UB → local_D[chunk])核间并行示意core 0: D[chunk0] ← peer0 → peer1 → ... → peer(R-1) ┐ core 1: D[chunk1] ← peer0 → peer1 → ... → peer(R-1) ─┼─ 并行 core 2: D[chunk2] ← peer0 → peer1 → ... → peer(R-1) ┘4.3 对比 MTE bulk atomicAddMTE bulk atomicAddSDMA fallback并发来源peer 维度所有 peer 的 get 全部 in-flightchunk 维度不同 core 写不同 chunkpeer 维度并行MTE3 atomic add串行UB 累加get 引擎MTE3SDMA写回方式MTE3 atomic add 到DDataCopy 普通写回Dchunk 数要求无即使单 chunk 也高效chunk 数 ≥ core 数时才能发挥并行度适用链路HCCSHCCS / PCIE / SIO4.4 何时使用 SDMA 路径SHOULD在 MTE 不可用时使用如跨 PCIE/SIO 的 PE 拓扑chunk 数 min(AIV core 数, ceil(output_rows / chunk_rows))SHOULD尽量接近或超过 peer 数量当 chunk 数 ≥ peer 数时每个 core 只处理有限 peer 数接近 MTE 路径效率当 chunk 数 peer 数时串行瓶颈显著性能差于 MTE 路径。5. 什么场景必须使用 AtomicAdd满足以下任一条件时MUST使用 atomic add标量或批量视数据规模选择场景为什么reduce 语义多个 PE 的贡献需要累加到一个 destination普通 write 会覆盖D丢失已有值多 source PE 写同一目的地址多个 remote rank 的 contribution 落到D的同一 tile顺序写也只是顺序覆盖不是累加异步 copy 多 UB stageUB_STAGES2时多个 get 可能 in-flight 且乱序完成普通写可能丢 update反例不需要atomic add单 PE 写独享 destination如 AIC 直接写 localD写 symmetric workspaceAIC 写 remote contribution 到自己的 workspace 是单 writer采用 SDMA fallback 策略§4chunk 分拆保证不同 core 写不同地址核内串行保证同一地址单 writer。5.1 即使无写冲突多 source 聚合也 SHOULD 优先 MTE 批量 AtomicAdd有一个容易误判的场景allreduce 的 reduce-scatter 阶段每个 PE 最终拥有全量输出直觉上每个 PE 写自己的D无冲突可以走 UB 累加。这个判断在正确性上成立但在效率上不成立。UB 累加路径含 SDMA fallback每轮必须先get远端数据到 UB、在 UB 内AscendC::Add、再释放 UBR-1个 remote source 的 get 被强制串行化UB 累加串行 get PE0 → UB → Add │ get PE1 → UB → Add │ ... │ get PE(R-1) → UB → Add → 写回 D ↑ 必须等上一轮 UB 累加完成 ↑ MTE 批量 AtomicAdd并行 SetAtomicAddT() mte_get_nbi(PE0) ╮ mte_get_nbi(PE1) ─┼── 全部 in-flightMTE3 硬件累加 mte_get_nbi(...) ╯ drain → SetAtomicNone()SHOULD优先 MTE 批量 atomicAdd 的场景多个 remote source 的数据最终需要聚合到同一个 destination无论该 destination 是否有跨 PE 写冲突。仅在 MTE 不可用时才退回到 SDMA fallback§4。6. 安全调用顺序6.1 MTE 批量SetAtomicAddT()mte_get_nbi1. aclshmem_barrier_all() // 确保所有 PE 的 workspace 已就绪 2. SetAtomicAddT() // 开启 atomic 模式 3. PipeBarrierPIPE_ALL() // 隔离 atomic 模式切换 4. reduceScatter.InitBlockLoop() // 初始化 epilogue block loop 5. for each remote rank r ! self: 6. aclshmemx_mte_get_nbi( // 远端数据 → 本地 Datomic add 模式 7. dstTensor local D tile, 8. srcTensor remote symmetric tile, 9. ... 10. ) 11. reduceScatter.FinalizeBlockLoop() // 等待 UB_STAGES 个 copy event 完成 12. SetFlagMTE3_S(EVENT_ID0) // 插 MTE3 drain event 13. WaitFlagMTE3_S(EVENT_ID0) // 等待 MTE3 写完成 14. SetAtomicNone() // 关闭 atomic 模式 15. PipeBarrierPIPE_ALL() // 隔离 atomic 模式退出 16. aclshmem_barrier_all() // 通知所有 PE 本 PE 已完成读取6.2 MTE 标量aclshmemx_mte_atomic_add1. value_to_ub: Scalar 写 value 到 UB buffer 2. SetFlagS_MTE3(sync_id) // 确保 Scalar→UB 完成后 MTE3 才能读 UB 3. WaitFlagS_MTE3(sync_id) 4. aclshmemx_mte_atomic_add(dst, value, pe) 5. SetFlagMTE3_MTE2(sync_id) // 如需读结果确保 MTE3→GM 完成后才能读 6. WaitFlagMTE3_MTE2(sync_id)6.3 UDMA 标量aclshmemx_udma_atomic_add1. aclshmemx_udma_atomic_add(dst, value, pe) 2. aclshmemx_udma_quiet(pe) // **MUST** 在读取结果前调用6.4SetAtomicNone()前必须 drain MTE3步骤 12-13SetFlagMTE3_S→WaitFlagMTE3_S是必选项。如果还有 in-flight MTE3 写时提前执行SetAtomicNone()后续完成的写入会退化为覆盖写丢失累加。NEVER省略 MTE3 drain。7. Atomic Add 的作用边界作用于不作用于标量 API远端 PE 的 GM 地址批量模式源端 GM 读取批量模式本地目的端 GM 写入localDAIC 写 self contribution 到D的路径aclshmemx_mte_get_nbi的目的端AIC 写 remote contribution 到gmSymmetric的路径标量 API通过 MTE2/MTE3 或 UDMA 写入远端barrier、cross-core flag、signal记住批量 atomic mode 只影响本地写入动作标量 API 只影响远端指定的单个地址。8. Self Contribution 与 Atomic Add 的配合标准 ReduceScatter 的实现采用分叉路径AIC 计算的 contribution / \ / \ targetRank self targetRank ! self | | v v D 直接写入 写入 symmetric workspace D self_contrib Sym[stage][targetself] | | --- AIV atomic add ------ 从 remote PE 读取 Sym 并 D src关键约束AIC 对 self contributionMUST直接写D不是 atomic addAIC 对 remote contributionMUST写 symmetric workspace不是DAIVMUSTcontinue跳过remoteRankIdx params.rankIdx的 task跳过 self 不是漏算——self 已被 AIC 直接写入。NEVER把 self contribution 也走 atomic add 路径。9. 并发粒度与 Tile 划分参考手册中 reduce-scatter 场景的 tile 划分参数源自实际算子 matmul_reduce_scatter 的实现参数AIC output tile: 128 × 256 L1TileShape communication block: 64 × 256 COMM_BLOCK_ROWS × COMM_BLOCK_COLUMNS scatter tile: 32 × 256 SCATTER_TILE_ROWS × SCATTER_TILE_COLUMNS一个 AIC output tile 被拆为 2 个 communication block每个 communication block 再拆为 2 个 scatter tile128 × 256 AIC output tile ----------------------------- || comm block 0: 64 × 256 | || ------------------------ | || | tile 0: 32 × 256 GET | | || ------------------------ | || | tile 1: 32 × 256 GET | | || ------------------------ | ----------------------------- || comm block 1: 64 × 256 | || ------------------------ | || | tile 2: 32 × 256 GET | | || ------------------------ | || | tile 3: 32 × 256 GET | | || ------------------------ | -----------------------------每个32 × 256tile 对应一次aclshmemx_mte_get_nbi调用所有 tile 的 get 都在同一个 atomic 区间内发出copyParams指定每行的repeat行数和length列数及 stride。9.1 pre barrier 的作用AIV on PE p 等待本 PE AIC flag 完成但这只保证本 PE。 其他 PE r 的 workspace 也必须已写完PE p 才能安全读取 Sym_r[stage][targetp]。 pre aclshmem_barrier_all(): 确保所有 PE 都已完成当前 stage 的 AIC 写入。 只有此后才能开始 atomic add GET。9.2 post barrier 的作用PE p 完成读取后其他 PE 可能还在读 PE p 的 workspace。 post aclshmem_barrier_all(): 确保所有 PE 都已读取完彼此的 stage。 只有此后 AIC 才能复用该 stage buffer。因此flagAivFinishCompute[stage]在 post barrier之后设置——AIC 复用前等待此 flag间接保证 workspace 不被过早覆盖。10. 风险点与反模式以下改动会直接破坏正确性改动风险去掉SetAtomicAddT()remote contribution覆盖D不是累加过早SetAtomicNone()MTE3 drain 前in-flight MTE3 写可能退化为覆盖丢加去掉 preaclshmem_barrier_all()可能读到 remote PE 尚未写完的 workspace去掉 postaclshmem_barrier_all()AIC 可能复用 stage覆盖仍被 remote PE 读取的数据不跳过 self rankself contribution 不在 symmetric workspace 中读不到或重复AIC 不直接写 selfD但 AIV 仍跳过 selfD初值缺失结果少 self contributionatomic 区间内混入普通 DataCopy 写 GM该次写入也会变成 atomic add破坏预期值改变blockPerComm但不保证可被rankSize整除target rank 分桶和 workspace offset 错乱UDMA/RDMA atomic add 后未 quiet/drain 就读结果远端操作可能未完成读到旧值SDMA fallback 时多 core 写同一 chunk覆盖而非累加10.1 必须的配套机制配套机制用途aclshmem_barrier_all()× 2pre: workspace 就绪post: workspace 可复用CrossCoreFlagAIC ↔ AIVAIC 通知 AIV 当前 stage 已写完AIV 通知 AIC 当前 stage 已读完tail clipping尾块actualCommBlockShape裁剪到真实有效范围不对 padding 区域做 atomic addaclshmemx_udma_quiet(pe)UDMA 路径确保远端操作完成chunk 不重叠SDMA fallback 路径确保核间无写冲突11. 完整调用示例11.1 MTE 批量 AtomicAddreduce-scatter / allreduce RS 阶段// commIdx 循环内部 uint32_t stageIdx commIdx % WORKSPACE_STAGES; // 1. 等待 AIC 写完当前 stage Catlass::Arch::CrossCoreWaitFlag(flagAicFinishStore[stageIdx]); // 2. 全局 barrier确保所有 PE 的 workspace 可读 aclshmem_barrier_all(); // 3. 进入 atomic add 模式 AscendC::SetAtomicAddElementD(); AscendC::PipeBarrierPIPE_ALL(); // 4. 初始化 epilogue reduceScatter.InitBlockLoop(); // 5. 遍历 remote PE读取并 atomic add 到本地 D for (uint32_t commLoopIdx aicoreIdx; commLoopIdx commCoreLoops; commLoopIdx commAicoreNum) { DistMatrixCoord commBlockCoord commScheduler.GetBlockCoord(commLoopIdx); uint32_t remoteRankIdx commBlockCoord.rank(); if (remoteRankIdx params.rankIdx) { continue; // self contribution 已由 AIC 直接写入 D } auto gmBlockSrc gmSymmetric[layoutSymmetric.GetOffset(offsetSrc)]; auto gmBlockDst gmD[params.layoutD.GetOffset(offsetDst)]; reduceScatter( gmBlockSrc, layoutBlockSrc, gmBlockDst, layoutBlockDst, actualCommBlockShape, remoteRankIdx ); } // 6. 等待所有 copy event 完成 reduceScatter.FinalizeBlockLoop(); // 7. Drain MTE3 AscendC::SetFlagAscendC::HardEvent::MTE3_S(EVENT_ID0); AscendC::WaitFlagAscendC::HardEvent::MTE3_S(EVENT_ID0); // 8. 退出 atomic 模式 AscendC::SetAtomicNone(); AscendC::PipeBarrierPIPE_ALL(); // 9. 全局 barrier aclshmem_barrier_all(); // 10. 通知 AIC 当前 stage 可复用 Catlass::Arch::CrossCoreSetFlag0x2, PIPE_MTE3(flagAivFinishCompute[stageIdx]);11.2 MTE 标量 AtomicAdd单元素 flag 累加// 推荐使用统一 API自动根据 topo_list 分发 MTE/UDMA 路径 // 内部已封装 S_MTE3 sync、SetAtomicAdd、MTE3 copy 等细节 for (int64_t peer 0; peer rankSize; peer) { if (peer myRank) continue; aclshmem_int32_atomic_add((__gm__ int32_t*)(dst_addr), value, peer); } aclshmem_barrier_all(); // 或者直接使用 engine-specific API // aclshmemx_mte_atomic_add((__gm__ T*)dst, value, pe); // aclshmemx_udma_atomic_add((__gm__ T*)dst, value, pe); // aclshmemx_udma_quiet(pe);11.3 SDMA Fallback 并发 Reducechunk 分拆 核内串行// 前置所有 PE 已将 contribution 写入 symmetric workspace aclshmem_barrier_all(); uint32_t myCore AscendC::GetBlockIdx(); uint32_t totalCores AscendC::GetBlockNum(); uint32_t chunkRows CeilDiv(totalRows, totalCores); uint32_t rowStart myCore * chunkRows; uint32_t rowEnd Min(rowStart chunkRows, totalRows); uint32_t myChunkRows rowEnd - rowStart; // 每个 core 负责 [rowStart, rowEnd) 的 chunk不同 core 写不同地址 for (uint32_t peer 0; peer rankSize; peer) { if (peer myRank) continue; // self contribution 已直接写入 D // 从 remote PE 的 workspace 读取本 PE 对应的 chunk 数据 aclshmemx_sdma_get_nbi( localUb, // UB 接收 buffer /* remote symmetric address for this peers contribution to my chunk */, chunkBytes, peer); // 等待 SDMA get 完成 AscendC::WaitFlagAscendC::HardEvent::SDMA_S(eventId); // UB 内向量累加 AscendC::Add(localUb, localUb, gmD[rowStart], myChunkRows * cols); // 写回本地 D AscendC::DataCopy(gmD[rowStart], localUb, myChunkRows * cols); } aclshmem_barrier_all();12. 快速参考按决策优先级降序优先选择编号小的方案优先级场景推荐方法并发方式1首选多 remote source → 同 destinationreduce-scatter / allreduce RS 阶段SetAtomicAddT()MTE 批量mte_get_nbi 双 barrierpeer 维度全部 in-flightMTE3 硬件累加2单元素远端累加signal/flag/counteraclshmem_*_atomic_add统一 API自动分发逐元素调用各自独立3跨 PCIE/SIO 单元素远端累加aclshmemx_udma_atomic_addaclshmemx_udma_quietUDMA 标量逐元素需 quiet 同步4MTE 不可用时多 PE reduce跨 PCIE/SIO 或跨 serverSDMA fallbackchunk 分拆 → 核间并行 / 核内串行 getUB 累加chunk 维度并行peer 维度串行5单 remote source 或无 remote 传输UB get →AscendC::Add→ DataCopy 写回无并发需求13. 相关源码与延伸阅读统一 AMO API 声明与同步模型注释include/device/gm2gm/shmem_device_amo.h拓扑分发实现MTE/ROCE/UDMA 三路径src/device/gm2gm/shmem_device_amo.hppMTE 引擎接口与平台 dtype 差异include/device/gm2gm/engine/shmem_device_mte.hUDMA / SDMA / RDMA 引擎接口shmem_device_udma.h、shmem_device_sdma.h、shmem_device_rdma.h可运行的 UDMA 原子加示例工程examples/udma_atomic_add同步/异步原子 API 行为对比文档docs/api/atomic_api_sync_async_comparison.md本模式参考手册原文.agents/skills/shmem-ops-code-gen/references/atomic-add-pattern.md【免费下载链接】shmemCANN SHMEM 是面向昇腾平台的多机多卡内存通信库基于OpenSHMEM 标准协议实现跨设备的高效内存访问与数据同步。项目地址: https://gitcode.com/cann/shmem创作声明：本文部分内容由AI辅助生成（AIGC），仅供参考