企业官网建设流程全解析

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"与"批量 atomic（SetAtomicAdd<T>()）"两类机制、为不同链路拓扑（HCCS / PCIE / SIO / RoCE）选择正确的累加路径、写出带完整 barrier 与 MTE3 drain 的安全调用序列，并规避丢失累加的经典反模式。

1. 核心概念：两类必须区分的 Atomic Add 机制

在 SHMEM Device 侧存在两类atomic add 机制，混用它们是最常见的错误来源：

关键语义：SetAtomicAdd<T>()不是"执行一次加法"，而是把当前 core 后续的所有 GM 写入模式切换为 atomic add；真正发生加法的是切换之后 MTE3 的写入操作。

2. 三引擎 AtomicAdd 接口全景与源码分发逻辑

2.1 统一 API：aclshmem_NAME_atomic_add（推荐首选入口）

统一 API 声明在 shmem_device_amo.h，通过类型宏批量生成各 dtype 的重载：

• ACLSHMEM_TYPE_FUNC_ATOMIC_ADD：int8/int16/bfloat16/half
• ACLSHMEM_TYPE_FUNC_ATOMIC_ADD_910：int32/float（Ascend 910 系列）
• ACLSHMEM_TYPE_FUNC_ATOMIC_ADD_EXT：uint32/uint64/int64（Ascend950）

同时头文件提供 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 → RDMA（RoCE）路径 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）分六步：

1. aclshmem_ptr(dst, pe)把对称地址转译为远端可访问地址；
2. Scalar 把value写入 UB buffer（ub_tensor.SetValue(0, value)）；
3. SetFlag<S_MTE3>/WaitFlag<S_MTE3>隔离 Scalar→MTE3；
4. SetAtomicAdd<TYPE>()开启 atomic 模式；
5. aclshmemi_copy_ub2gm触发 MTE3 写入——这是内部接口，用户代码 MUST 调用aclshmem_*_atomic_add/aclshmemx_mte_atomic_add，禁止直接调用aclshmemi_copy_ub2gm；
6. SetAtomicNone()关闭 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_quiet，UDMA 用aclshmemx_udma_quiet，MTE 按数据依赖插入SetFlag/WaitFlag（读 GM 用MTE3_MTE2事件，复用 UB 用MTE3_S事件）。

另外，头文件注释多次强调：MTE 传输不支持跨 PCIe（inter-node）通信，跨 PCIe 场景必须走 RDMA 或 UDMA 路径；且 RDMA 路径不支持对同一 PE 的并发 RMA/AMO。

2.2 MTE 引擎：标量与批量两条路径

MTE 是机内 HCCS full-mesh 场景的主力引擎，提供标量和批量两条路径。

2.2.1 标量：aclshmemx_mte_atomic_add

aclshmemx_mte_atomic_add((__gm__ T*)dst, value, pe);

平台差异与头文件 shmem_device_mte.h#L254-L281 的文档注释完全一致：

2.2.2 批量：SetAtomicAdd<T>()+mte_get_nbi

AscendC::SetAtomicAdd<ElementD>(); // 所有 mte_get_nbi 的目的端 GM 写入自动变为 D += src reduceScatter(gmBlockSrc, ..., gmBlockDst, ..., remoteRankIdx); // drain MTE3 AscendC::SetAtomicNone();

2.2.3 两条路径对比

2.3 UDMA 引擎

aclshmemx_udma_atomic_add((__gm__ T*)dst, value, pe); aclshmemx_udma_quiet(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);

2.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 引擎能力总表

3. MTE 标量 vs 批量：如何选择

4. SDMA 如何实现并发 Reduce（Fallback 策略）

当 MTE 不可用（跨 PCIE/SIO 链路），SDMA 仍然可以做跨 PE 数据搬运（aclshmemx_sdma_get_nbi），但缺少 inline reduce 能力，必须退化为"本地收数据 → UB 累加 → 写回"的模式。

4.1 约束

• 无 device 侧 inline reduce：SDMA 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 atomicAdd

4.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（标量或批量，视数据规模选择）：

反例（不需要atomic add）：

• 单 PE 写独享 destination（如 AIC 直接写 localD）；
• 写 symmetric workspace（AIC 写 remote contribution 到自己的 workspace 是单 writer）；
• 采用 SDMA fallback 策略（§4）：chunk 分拆保证不同 core 写不同地址，核内串行保证同一地址单 writer。

5.1 即使无写冲突，多 source 聚合也 SHOULD 优先 MTE 批量 AtomicAdd

有一个容易误判的场景：allreduce 的 reduce-scatter 阶段，每个 PE 最终拥有全量输出，直觉上"每个 PE 写自己的D无冲突，可以走 UB 累加"。这个判断在正确性上成立，但在效率上不成立。UB 累加路径（含 SDMA fallback）每轮必须先get远端数据到 UB、在 UB 内AscendC::Add、再释放 UB，R-1个 remote source 的 get 被强制串行化：

UB 累加（串行）： get PE0 → UB → Add │ get PE1 → UB → Add │ ... │ get PE(R-1) → UB → Add → 写回 D ↑ 必须等上一轮 UB 累加完成 ↑ MTE 批量 AtomicAdd（并行）： SetAtomicAdd<T>() mte_get_nbi(PE0) ╮ mte_get_nbi(PE1) ─┼── 全部 in-flight，MTE3 硬件累加 mte_get_nbi(...) ╯ drain → SetAtomicNone()

SHOULD优先 MTE 批量 atomicAdd 的场景：多个 remote source 的数据最终需要聚合到同一个 destination，无论该 destination 是否有跨 PE 写冲突。仅在 MTE 不可用时才退回到 SDMA fallback（§4）。

6. 安全调用顺序

6.1 MTE 批量（SetAtomicAdd<T>()+mte_get_nbi）

1. aclshmem_barrier_all() // 确保所有 PE 的 workspace 已就绪 2. SetAtomicAdd<T>() // 开启 atomic 模式 3. PipeBarrier<PIPE_ALL>() // 隔离 atomic 模式切换 4. reduceScatter.InitBlockLoop() // 初始化 epilogue block loop 5. for each remote rank r != self: 6. aclshmemx_mte_get_nbi( // 远端数据 → 本地 D（atomic add 模式） 7. dstTensor = local D tile, 8. srcTensor = remote symmetric tile, 9. ... 10. ) 11. reduceScatter.FinalizeBlockLoop() // 等待 UB_STAGES 个 copy event 完成 12. SetFlag<MTE3_S>(EVENT_ID0) // 插 MTE3 drain event 13. WaitFlag<MTE3_S>(EVENT_ID0) // 等待 MTE3 写完成 14. SetAtomicNone() // 关闭 atomic 模式 15. PipeBarrier<PIPE_ALL>() // 隔离 atomic 模式退出 16. aclshmem_barrier_all() // 通知所有 PE 本 PE 已完成读取

6.2 MTE 标量（aclshmemx_mte_atomic_add）

1. value_to_ub: Scalar 写 value 到 UB buffer 2. SetFlag<S_MTE3>(sync_id) // 确保 Scalar→UB 完成后 MTE3 才能读 UB 3. WaitFlag<S_MTE3>(sync_id) 4. aclshmemx_mte_atomic_add(dst, value, pe) 5. SetFlag<MTE3_MTE2>(sync_id) // 如需读结果，确保 MTE3→GM 完成后才能读 6. WaitFlag<MTE3_MTE2>(sync_id)

6.3 UDMA 标量（aclshmemx_udma_atomic_add）

1. aclshmemx_udma_atomic_add(dst, value, pe) 2. aclshmemx_udma_quiet(pe) // **MUST** 在读取结果前调用

6.4SetAtomicNone()前必须 drain MTE3

步骤 12-13（SetFlag<MTE3_S>→WaitFlag<MTE3_S>）是必选项。如果还有 in-flight MTE3 写时提前执行SetAtomicNone()，后续完成的写入会退化为覆盖写，丢失累加。NEVER省略 MTE3 drain。

7. Atomic Add 的作用边界

记住：批量 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][target=self]） | | +<--- AIV atomic add ------+ （从 remote PE 读取 Sym 并 D += src）

关键约束：

• AIC 对 self contributionMUST直接写D（不是 atomic add）；
• AIC 对 remote contributionMUST写 symmetric workspace（不是D）；
• AIVMUSTcontinue跳过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 tile：

128 × 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][target=p]。 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. 风险点与反模式

以下改动会直接破坏正确性：

10.1 必须的配套机制

11. 完整调用示例

11.1 MTE 批量 AtomicAdd（reduce-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::SetAtomicAdd<ElementD>(); AscendC::PipeBarrier<PIPE_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::SetFlag<AscendC::HardEvent::MTE3_S>(EVENT_ID0); AscendC::WaitFlag<AscendC::HardEvent::MTE3_S>(EVENT_ID0); // 8. 退出 atomic 模式 AscendC::SetAtomicNone(); AscendC::PipeBarrier<PIPE_ALL>(); // 9. 全局 barrier aclshmem_barrier_all(); // 10. 通知 AIC 当前 stage 可复用 Catlass::Arch::CrossCoreSetFlag<0x2, 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 并发 Reduce（chunk 分拆 + 核内串行）

// 前置：所有 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 peer's contribution to my chunk */, chunkBytes, peer); // 等待 SDMA get 完成 AscendC::WaitFlag<AscendC::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. 快速参考

按决策优先级降序，优先选择编号小的方案：

13. 相关源码与延伸阅读

• 统一 AMO API 声明与同步模型注释：include/device/gm2gm/shmem_device_amo.h
• 拓扑分发实现（MTE/ROCE/UDMA 三路径）：src/device/gm2gm/shmem_device_amo.hpp
• MTE 引擎接口与平台 dtype 差异：include/device/gm2gm/engine/shmem_device_mte.h
• UDMA / 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