CANN SHMEM 设备侧DeviceAPI 全览从 AMO/RMA 原子操作到同步模型的源码级解读【免费下载链接】shmemCANN SHMEM 是面向昇腾平台的多机多卡内存通信库基于OpenSHMEM 标准协议实现跨设备的高效内存访问与数据同步。项目地址: https://gitcode.com/cann/shmem导读本文以 docs/api/device_api.rst 这份 Doxygen API 索引文档为骨架系统梳理 CANN SHMEM昇腾平台多机多卡内存通信库在NPU 设备侧Kernel 内提供的全部头文件与接口族原子操作AMO、远程内存访问RMA、点对点信号同步P2P Sync、集体同步Barrier/Sync、内存序控制Quiet/Fence、信号操作Signal Operation与团队Team查询并逐一对照仓库头文件给出签名、参数语义、支持的数据类型与硬件平台、MTE/RDMA/UDMA 三种传输路径下的同步要求。读者读完可掌握在昇腾 KernelAscendC内正确选用与调用各设备侧 API并理解何时需要quiet、何时需要SetFlag/WaitFlag等同步手段避免出现数据竞争或 UB 覆盖问题。1. 文档定位设备侧 API 的 Doxygen 索引入口docs/api/device_api.rst本身是一份面向 Doxygen 生成:project: ACLSHMEM_CPP_API的 API 索引页按头文件组织设备侧接口。它同时承担一个重要作用该文档是仓库中设备侧Kernel 内可调用API 的清单与之相对的 docs/api/host_api.rst 则索引 Host 侧 API。设备侧接口的物理位置集中在 include/device 下按通信域分为gm2gmGM to GM跨 PE 全局内存通信ub2gmUB to GMUB 缓冲到全局内存的通信变体如 MTE 引擎team团队管理相关engine底层传输引擎MTE、RDMA、SDMA、UDMA。include/device/ ├── gm2gm/ │ ├── shmem_device_amo.h # 原子操作 AMO │ ├── shmem_device_cc.h # 集体同步 Collective Communication │ ├── shmem_device_mo.h # 内存序 Memory Orderingquiet/fence │ ├── shmem_device_p2p_sync.h # 点对点同步 signal/wait/test │ ├── shmem_device_rma.h # 远程内存访问 RMA put/get │ ├── shmem_device_so.h # 信号操作 Signal Operationput_signal │ └── engine/ # mte/rdma/sdma/udma 引擎层 ├── team/shmem_device_team.h # 团队查询 └── ub2gm/ # UB→GM 路径的 rma 与 mte 接口以下按该文档列出的头文件顺序逐族展开。2. 原子操作shmem_device_amo.h对应文档条目shmem_device_amo.h源码位于 include/device/gm2gm/shmem_device_amo.h。该头文件覆盖 OpenSHMEM 规范中的全部 AMO 语义且大量函数通过宏批量生成以支持多种数据类型。2.1 接口族与宏生成机制通过ACLSHMEM_TYPE_FUNC_*系列宏自动生成不同数据类型的原子函数宏名生成的接口支持类型ACLSHMEM_TYPE_FUNC_ATOMIC_ADDatomic_addint8、int16、bfloat16、halfACLSHMEM_TYPE_FUNC_ATOMIC_ADD_910atomic_addint32、floatAscend910 系列ACLSHMEM_TYPE_FUNC_ATOMIC_ADD_EXTatomic_adduint32、uint64、int64Ascend950ACLSHMEM_TYPE_FUNC_ATOMIC_ADD_950atomic_fetch_add/atomic_inc/atomic_fetch_inc/atomic_fetchuint32、uint64、int32、int64ACLSHMEM_TYPE_FUNC_ATOMIC_SWAPatomic_set/atomic_swap/atomic_compare_swapuint32、uint64ACLSHMEM_TYPE_FUNC_ATOMIC_SWAP_CAST同上CAST 版本int32、int64、floatACLSHMEM_TYPE_FUNC_ATOMIC_CAS_CASTatomic_compare_swapCASTint32、int64ACLSHMEM_TYPE_FUNC_ATOMIC_LOGICatomic_and/atomic_or/atomic_xor及 fetch 版本uint32、uint64、int32、int64完整接口清单包括atomic_add、atomic_fetch_add、atomic_inc、atomic_fetch_inc、atomic_and、atomic_or、atomic_xor、atomic_fetch_and、atomic_fetch_or、atomic_fetch_xor、atomic_fetch、atomic_set、atomic_swap、atomic_compare_swap。这些接口同时在头文件末尾通过宏导出 OpenSHMEM 风格别名例如#define shmem_int32_atomic_add aclshmem_int32_atomic_add #define shmem_float_atomic_add aclshmem_float_atomic_add2.2 典型签名所有原子接口都符合统一形态以atomic_add为例ACLSHMEM_DEVICE void aclshmem_type_atomic_add(__gm__ TYPE* dst, TYPE value, int32_t pe);dst对称内存中目标数据的指针__gm__全局内存限定value要叠加/运算的值pe远端 PE 编号。带返回值的 fetch 类接口atomic_fetch_add、atomic_swap、atomic_compare_swap等返回操作前dst处的旧值。2.3 三条传输路径与硬件平台支持头文件注释明确给出了路径 × 类型 × 平台的矩阵路径atomic_add 支持类型硬件平台MTEint8、int16、bf16、half、int32、floatAscend910B/Ascend910C/Ascend950MTEuint32、uint64、int64Ascend950RDMAint32、uint32、int64、uint64Ascend950UDMAint32、uint32、int64、uint64、floatAscend950实现根据传输拓扑在 MTE、RDMA、UDMA 之间分发实现见 src/device/gm2gm/shmem_device_amo.hpp。2.4 原子接口同步模型关键这是使用 AMO 最需要理解的规则在头文件amo_sync_model章节有系统说明带返回值同步接口fetch_add、fetch_inc、fetch_and、fetch_or、fetch_xor、fetch、swap、compare_swap在返回前已通过内部 quiet 保证完成并可见调用方无需额外同步void 返回异步接口add、inc、and、or、xor、set返回时可能尚未完成调用方必须自行同步RDMA 路径调用aclshmemx_roce_quietUDMA 路径调用aclshmemx_udma_quietMTE 路径按数据依赖插入SetFlag/WaitFlag。MTE 路径的数据流为Scalar → UB → MTE2/MTE3 → GM具体同步点调用前若其他单元如 MTE2也写同一段 UB需先对 MTE2 写 UB 的事件做SetFlag/WaitFlag同步防止 UB 数据被覆盖调用后若对结果有数据依赖读 GM 前需对 MTE3 事件做SetFlag/WaitFlag如MTE3_MTE2事件若将向同一 UB 区域写入新值需在覆盖 UB 前同步 MTE3如MTE3_S事件。其中 Ascend950 上的uint32/uint64/int64类型走Scalar AtomicAdd直写 GM 的路径不经 UB 与 MTE3同步模型与 UBMTE3 路径不同需按实际数据流在 Scalar 计算单元与 MTE2/MTE3 传输单元之间插入SetFlag/WaitFlag。MTE 路径的 UB 缓冲偏移默认 0可用aclshmemx_set_mte_config(offset, ub_size, sync_id)调整。⚠️ 注意事项MTE 传输路径不支持跨 PCIe节点间通信跨节点场景必须走 RDMA 或 UDMA 路径。3. 集体同步shmem_device_cc.h对应文档条目shmem_device_cc.h源码位于 include/device/gm2gm/shmem_device_cc.h实现见 src/device/gm2gm/shmemi_device_cc_kernel.cpp。3.1 接口清单接口语义参与者范围aclshmem_barrier(team)团队内集体同步等待团队所有 PE 到达后返回确保此前所有 store 与远程更新含 AMO/RMA完成全部核心aclshmem_barrier_all()ACLSHMEM_TEAM_WORLD上的 barrier全部核心aclshmem_sync(team)类似 barrier但只保证此前 store 的完成与可见不保证 ACLSHMEM 通信接口发出的远程更新完成全部核心aclshmem_sync_all()ACLSHMEM_TEAM_WORLD上的 sync全部核心aclshmemx_sync_vec(team)仅 VEC 核参与同步不保证 ACLSHMEM 远程更新的完成CUBE 核调用无效VEC 核aclshmemx_sync_vec_all()ACLSHMEM_TEAM_WORLD上的sync_vecVEC 核aclshmemx_barrier_vec(team)/aclshmemx_barrier_all_vec()aclshmemx_sync_vec的兼容封装已标记[[deprecated]]VEC 核3.2 使用限制头文件顶部 WARNING全核同步 API 只能在 MIX Kernel 中使用如果 Kernel 缺少有效的 CUBE 指令如Mmad或向量指令如DataCopy编译器可能将其优化为 VEC 或 CUBE 类型因此需要插入有效的Mmad/DataCopy调用ACLSHMEM 的 PE 间同步与SyncAll冲突避免在同一 Kernel 内混用aclshmemx_sync_vec系列只同步 VEC 核之间对内存的 store 可见性不会完成 ACLSHMEM 通信 API 发出的远程更新CUBE 核的 Scalar 单元不被这些 API 同步不要依赖它在集体同步之间传递数据。3.3 可见性语义仅存在 scale-up 网络HCCS时barrier 后更新全局可见同时存在 HCCSscale-up与 RDMAscale-out时只保证更新对目标 PE 可见在 CPU 上发起的 barrier 只完成 CPU 发起的通信操作若需从 CPU 侧保证 NPU 侧操作完成应使用aclrtSynchronizeStream/aclrtDeviceSynchronize或基于 Stream 的接口。4. 内存序控制shmem_device_mo.h对应文档条目shmem_device_mo.h源码位于 include/device/gm2gm/shmem_device_mo.h。4.1 aclshmem_quiet确保调用 PE 此前对对称数据对象发出的所有操作完成。同样遵循仅有 HCCS 时全局可见HCCSRDMA 时只保证对目标 PE 可见的语义CPU/NPU 各自只完成本侧发起的操作。4.2 aclshmem_fenceOpenSHMEM 规范中 fence 只保证 Put/AMO/store 的到达顺序而不保证完成。由于硬件能力限制本实现中aclshmem_fence与aclshmem_quiet行为一致同时保证顺序与完成使用起来比规范语义更强。ACLSHMEM_DEVICE void aclshmem_quiet(void); ACLSHMEM_DEVICE void aclshmem_fence(void);5. 点对点同步shmem_device_p2p_sync.h对应文档条目shmem_device_p2p_sync.h源码位于 include/device/gm2gm/shmem_device_p2p_sync.h是设备侧做细粒度 Producer-Consumer 同步的核心。5.1 接口族接口族语义返回aclshmemx_signal_op对远端sig_addr用sig_opACLSHMEM_SIGNAL_SET/ACLSHMEM_SIGNAL_ADD更新signal值voidaclshmem_signal_wait_until阻塞直到本地sig_addr满足比较条件满足条件的信号值type_wait_until阻塞直到ivar满足条件voidtype_wait阻塞直到ivar不等于cmp_valuevoidtype_wait_until_all等待集合内所有元素满足条件voidtype_wait_until_any等待集合内任一元素满足条件满足条件的下标空集返回SIZE_MAXtype_wait_until_some等待集合内至少一个元素满足条件返回满足条件的下标个数type_wait_until_*_vector逐元素比较值版本cmp_values为数组同前type_test非阻塞测试单个元素1/0type_test_any/type_test_some/type_test_*_vector非阻塞测试集合下标/个数/1-05.2 比较操作符所有 wait/test 接口的cmp参数支持六种比较符ACLSHMEM_CMP_EQ、ACLSHMEM_CMP_NE、ACLSHMEM_CMP_GT、ACLSHMEM_CMP_GE、ACLSHMEM_CMP_LT、ACLSHMEM_CMP_LE。5.3 status 掩码语义对于 wait_set/test_set 类接口status数组长度nelems决定哪些元素参与等待status[i] 0ivars[i]参与等待status[i] ! 0ivars[i]被排除status NULL所有元素都参与。典型组合用法Producer 端调用aclshmemx_signal_op(sig_addr, 1, ACLSHMEM_SIGNAL_ADD, pe)Ascend950 上以 RDMA 为通信引擎保证signal_set的原子性Consumer 端调用aclshmem_signal_wait_until(sig_addr, ACLSHMEM_CMP_EQ, 1)实现高效点对点同步。完整接口生成见 src/device/gm2gm/shmem_device_p2p_sync.hpp。6. 远程内存访问shmem_device_rma.h对应文档条目gm2gm/shmem_device_rma.h及ub2gm/shmem_device_rma.h源码位于 include/device/gm2gm/shmem_device_rma.h。这是设备侧 put/get 数据搬运的主体底层支持MTE、RDMA、SDMA、UDMA四种引擎。6.1 单元素低延迟接口p / gaclshmem_type_p(dst, value, pe)向远端对称地址写入单个元素aclshmem_type_g(src, pe)从远端对称地址读回单个元素并返回。两者支持 half、float、double、int8/16/32/64、uint8/16/32/64、char、bfloat16 共 13 种类型适合标志位、标量通信等低延迟场景。6.2 批量接口put / get 系列接口语义同步属性aclshmem_type_put(dst, src, elem_size, pe)本地 → 远端连续拷贝同步aclshmem_type_get(dst, src, elem_size, pe)远端 → 本地连续拷贝同步aclshmem_putmem/aclshmem_getmemvoid* 版本连续拷贝同步aclshmem_type_iput/aclshmem_iget及putBITS/getBITS带步长strided拷贝dst为 dest 步长、sst为 source 步长同步aclshmem_type_put_nbi/aclshmem_type_get_nbi及putmem_nbi/getmem_nbi非阻塞non-blocking版本异步aclshmem_type_put_nbi/get_nbi的non_contiguous_copy_param重载高性能非连续数据拷贝异步GlobalTensorT模板重载AscendCGlobalTensor接口形态同前elem_size参数是元素个数实际字节数 elem_size * sizeof(T)。同步版本内部会保证完成NBI 异步版本返回后需要用 quiet 类接口RDMAaclshmemx_roce_quietUDMAaclshmemx_udma_quiet或信号协议来确认完成。6.3 地址与传输约束put 操作中dst必须指向对称内存会按 PE 换算远端地址src可以是本地任意有效 GM 地址get 操作中src必须指向对称内存dst可以是本地任意有效 GM 地址若初始化时通过ACLSHMEM_DATA_OP_ROCE启用了 RDMA则两操作数的完整传输区间都必须落在各自的对称内存分配范围内运行时按目标 PE 自动分发调用方无需关心具体引擎选择⚠️ 使用 RDMA 作为底层传输时不支持对同一 PE 的并发 RMA/AMO 操作使用 RDMA 或 SDMA 时通过device_state的rdma_config/sdma_config中的sync_id做流水线同步。6.4 MTE 配置接口ACLSHMEM_DEVICE void aclshmemx_set_mte_config(uint64_t offset, uint32_t ub_size, uint32_t sync_id); ACLSHMEM_DEVICE void aclshmemx_set_udma_config(uint64_t offset, uint32_t ub_size, uint32_t sync_id);aclshmemx_set_mte_config设置 put/get 使用的临时 UB 缓冲起始偏移、大小与同步 IDaclshmemx_set_udma_config为 UDMA MTE 中转操作设置 UB 暂存区ub_size至少为ACLSHMEM_UDMA_MTE_STAGING_UB_SIZE字节sync_id用于 MTE3→S 同步。UB→GM 方向的接口位于 include/device/ub2gm/shmem_device_rma.h实现见 src/device/ub2gm/shmem_device_rma.hpp。7. 信号操作shmem_device_so.h对应文档条目shmem_device_so.h源码位于 include/device/gm2gm/shmem_device_so.h。信号操作把数据拷贝与信号置位绑定在同一个请求内实现数据到位即通知的流水线化同步。7.1 接口族接口语义aclshmem_putmem_signal(dst, src, elem_size, sig_addr, signal, sig_op, pe)同步 put 后在远端更新信号字aclshmem_putmem_signal_nbi(...)异步版本aclshmem_type_put_signal(dst, src, elem_size, sig_addr, signal, sig_op, pe)13 种类型化版本aclshmem_type_put_signal_nbi(...)异步类型化版本aclshmem_type_put_signal的non_contiguous_copy_param/GlobalTensorT重载非连续与张量形态aclshmem_putBITS_signal/aclshmem_putBITS_signal_nbi按位宽8/16/...的 void* 版本统一签名形态类型化版本ACLSHMEM_DEVICE void aclshmem_type_put_signal( __gm__ TYPE* dst, __gm__ TYPE* src, size_t elem_size, __gm__ int32_t* sig_addr, int32_t signal, int sig_op, int pe);sig_addr要更新的信号字所在对称地址signal用于更新信号字的值sig_op更新操作仅支持ACLSHMEM_SIGNAL_SET与ACLSHMEM_SIGNAL_ADD。⚠️ 这些接口仅支持单核调用多核或多个 writer 并发调用是不支持的。UDMA 路径下超过 256 MB 的传输会在更新sig_addr之前内部切分。8. 团队查询shmem_device_team.h对应文档条目shmem_device_team.h源码位于 include/device/team/shmem_device_team.h。接口语义aclshmem_my_pe()返回本地 PE 编号0 ~ npes-1aclshmem_n_pes()返回程序中的 PE 总数aclshmem_team_my_pe(team)返回调用 PE 在指定团队内的编号ACLSHMEM_TEAM_INVALID时返回 -1aclshmem_team_n_pes(team)返回团队内 PE 数无效句柄返回 -1aclshmem_team_translate_pe(src_team, src_pe, dest_team)将一个团队中的 PE 编号映射到另一团队aclshmem_team_pe_mapping(team, pe)将团队内 PE 编号映射到全局团队编号团队概念的 Host 侧创建与拆分见 src/host/team/shmem_team.cpp 与 docs/principles/team.md。9. 传输引擎层与 UDMA 256 MB 上限文档重点device_api.rst在gm2gm/engine系列头文件条目下特别强调了 UDMA 的容量约束这是设备侧 RMA 编程必须遵守的硬性限制UDMA put 接口每次异步 put 请求最多传输 256 MB256 × 1024 × 1024 字节。elem_size是元素个数因此实际传输字节数为elem_size * sizeof(T)。对于更大的数据需要将操作拆分为多个单次不超过 256 MB 的 UDMA put 请求并在读取目标数据或复用内容必须保持稳定的缓冲区之前使用配套的完成/同步接口如aclshmemx_udma_quiet(pe)或文档规定的信号协议。引擎层头文件清单与对应实现头文件docs 索引源码位置gm2gm/engine/shmem_device_mte.hinclude/device/gm2gm/engine/shmem_device_mte.h实现 src/device/gm2gm/engine/shmem_device_mte.hppgm2gm/engine/shmem_device_rdma.hinclude/device/gm2gm/engine/shmem_device_rdma.h实现 src/device/gm2gm/engine/shmem_device_rdma.hppgm2gm/engine/shmem_device_sdma.hinclude/device/gm2gm/engine/shmem_device_sdma.h实现 src/device/gm2gm/engine/shmem_device_sdma.hppgm2gm/engine/shmem_device_udma.hinclude/device/gm2gm/engine/shmem_device_udma.h实现 src/device/gm2gm/engine/shmem_device_udma.hppub2gm/engine/shmem_device_mte.hinclude/device/ub2gm/engine/shmem_device_mte.hub2gm/shmem_device_rma.hinclude/device/ub2gm/shmem_device_rma.h9.1 UDMA 引擎操作配置UDMA 引擎提供了精细的 WQE 配置结构include/device/gm2gm/engine/shmem_device_udma.htypedef struct { uint32_t cqe; /// 完成队列事件使能合法值 0/1 uint32_t se; /// 提示事件使能保留位支持前必须为 0 uint32_t fence; /// Fence 使能合法值 0/1 uint32_t odr; /// 排序域限制合法值 0..7 } aclshmemx_udma_op_config_t; inline constexpr uint32_t ACLSHMEMX_UDMA_ODR_NO 0b000U; // 无排序 inline constexpr uint32_t ACLSHMEMX_UDMA_ODR_RO 0b101U; // 松弛排序 inline constexpr uint32_t ACLSHMEMX_UDMA_ODR_SO 0b110U; // 强排序预置配置包括ACLSHMEMX_UDMA_OP_CONFIG_DEFAULT、ACLSHMEMX_UDMA_OP_CONFIG_NO_CQE与ACLSHMEMX_UDMA_OP_CONFIG_SO。ACLSHMEMX_UDMA_OP_CONFIG_SO发布强排序 WQE常用于数据后置标志的排序点保证此前同一 PE/QP 的 Relax Order 或 Strong Order WQE 先完成No Order WQE 不受此保证覆盖。典型数据-标志用法// 先发数据强排序再发原子加操作置标志强排序确保数据先于标志到达 aclshmemx_udma_put_nbiT, PIPE_MTE3, ACLSHMEMX_UDMA_OP_CONFIG_SO(dst, src, buf, elem_size, pe, sync_id); aclshmemx_udma_atomic_adduint32_t, ACLSHMEMX_UDMA_OP_CONFIG_SO(flag, 1U, pe);UDMA 还支持基于aclshmemx_defer_t/aclshmemx_submit_t的批量batch提交defer 操作要求cqe0submit 调用要求cqe1每批调用必须使用相同的操作类型、状态、PE 参数、buf 基址与sync_idbuf容量至少为64 * n字节提交成功后需对目标 PE 调用aclshmemx_udma_quiet后再消费 Get 目标或复用 Put 源。10. 设备侧 API 同步模型速查与最佳实践综合上述各头文件设备侧编程可遵循以下决策表场景接口需要的外部同步单元素写远端aclshmem_type_p无内部完成单元素读远端aclshmem_type_g无内部完成批量 put/get同步aclshmem_type_put/get无批量 put/get异步aclshmem_type_put_nbi/get_nbiRDMAaclshmemx_roce_quietUDMAaclshmemx_udma_quietMTESetFlag/WaitFlagvoid 返回的 AMOatomic_add/atomic_inc/atomic_set等同异步规则fetch 类 AMOatomic_fetch_add/atomic_swap等无需额外同步内置 quiet团队级同步aclshmem_barrier/aclshmem_sync全核同步需 MIX Kernel不要与SyncAll混用VEC 核同步aclshmemx_sync_vec不完成远程通信更新数据-标志流水线aclshmem_type_put_signal单核调用UDMA 单请求 ≤ 256 MB几点关键实践建议跨节点必须避开 MTEMTE 传输路径不支持跨 PCIe节点间 AMO/RMA 需保证走 RDMA/UDMA 路径如设置ACLSHMEM_DATA_OP_ROCE启用 RDMAUB 复用前先同步 MTE3MTE 路径下向 UB 写新值前、或读 GM 结果前务必按事件MTE2_S、MTE3_MTE2、MTE3_S插入SetFlag/WaitFlag防止数据错乱RDMA 下不要并发访问同一 PE并发 RMA/AMO 到同一 PE 不受支持需用device_state.rdma_config中的sync_id做流水线同步大块数据拆分为 ≤256 MB 的 UDMA 请求并在读取目标或复用缓冲区前调用aclshmemx_udma_quiet异步接口一律显式确认完成同步接口依赖内置 quiet不要在异步接口上假设返回即完成。11. 延伸阅读设备侧 AMO/RMA 的宏生成与分发实现src/device/gm2gm 下的shmem_device_amo.hpp、shmem_device_rma.hpp、shmem_device_p2p_sync.hpp引擎层 UDMA 批量提交与排序域契约include/device/gm2gm/engine/shmem_device_udma.hHost 侧对应 API 索引docs/api/host_api.rst通用类型与枚举定义include/device/shmem_def.h 与 include/host_device/shmem_common_types.h团队机制的完整说明docs/principles/team.md完整 API 演示示例examples 下的simt_rma、rdma_*、udma_*、notifywait等目录展示了设备侧 API 的实际调用方式。【免费下载链接】shmemCANN SHMEM 是面向昇腾平台的多机多卡内存通信库基于OpenSHMEM 标准协议实现跨设备的高效内存访问与数据同步。项目地址: https://gitcode.com/cann/shmem创作声明：本文部分内容由AI辅助生成（AIGC），仅供参考