CANN SHMEM 算子性能优化模式从瓶颈定位到通信、流水线与分核策略的实战指南【免费下载链接】shmemCANN SHMEM 是面向昇腾平台的多机多卡内存通信库基于OpenSHMEM 标准协议实现跨设备的高效内存访问与数据同步。项目地址: https://gitcode.com/cann/shmem本文是 .agents/skills/shmem-ops-performance-optim/references/optimization-patterns.md 的完整展开版。文中examples/、include/等路径均以 CANN SHMEM 仓库根目录为起点阅读前建议先通过 shmem-repo-resolution.md 定位仓库结构。CANN SHMEM 是基于 OpenSHMEM 标准的多机多卡内存通信库其算子的性能上限往往取决于 kernel 内数据搬运、同步与分核方式的设计。本文面向在昇腾 NPU 上编写或优化 SHMEM 通信算子的开发者系统梳理了六类典型瓶颈及其对应的优化模式通信重叠、流水线与 Double Buffer、同步替换、内存路径精简、分核策略重构以及优先级排序。读完本文你将能够根据 prof 数据快速判断瓶颈类型并在 examples/allgather、examples/kv_shuffle 等仓内示例中找到可复用的工程化写法。优化总览先定位瓶颈类型再选择手段优化的第一原则是每轮先定位瓶颈类型再选择对应手段而不是盲目套用技巧。根据瓶颈所处的位置可把优化方向分为六类类别典型瓶颈对应章节通信优化数据必须先落 symmetric heap 再由远端读取拷贝成为串行开销大消息双拷贝导致带宽封顶§1流水线优化单 buffer 时 NBI 搬运与使用必须串行搬运气泡明显§2同步优化全局 barrier 让所有 PE 等待最慢者整块传完才 signal 导致消费者空等§3内存优化远端 get 落 GM、再经 DataCopy 进 UB 的多余读写§4分核优化单 AIV 内for peer in 0..n_pes串行搬运无法利用多条链路并行§5综合排序多个优化点并存时收益与风险不同§6下面逐类展开每一类都给出 ❌ 错误做法、✅ 正确做法、关键点以及仓内源码佐证。1. 通信优化1.1 Copy-to-Symmetric-Heap Overlap瓶颈数据MUST先从本地 GM 拷贝到 symmetric heap再由远端 PE 读取在没有重叠设计时拷贝时间成为串行开销。优化将 core 分为 sender 组和 receiver 组。sender 按 chunk 把数据拷贝到 symmetric heap并逐 chunk 发 signalreceiver 轮询 signal一旦某 chunk 就绪立即拉取。该模式源码见 examples/allgather/allgather_kernel.cpp 中的all_gather_originkernel。❌ 错误做法——所有数据拷贝完 → 全局 barrier → 再开始通信copy_all_to_symmetric(input, gva_data, total_size); aclshmem_barrier_all(); fetch_all_from_peers(output, gva_data, total_size);✅ 正确做法——sender 逐 chunk 拷贝 逐 chunk signalreceiver 轮询就绪即拉取// Sender core逐 chunk 拷贝 逐 chunk signal if (aivIndex core_group_num) { while (remaining chunk_size) { aclshmemx_mte_put_nbi(gva_data offset, input offset, tmp_buff, ub_size, chunk_elems, my_rank, EVENT_ID0); SetFlagHardEvent::MTE3_S(EVENT_ID0); WaitFlagHardEvent::MTE3_S(EVENT_ID0); flag times magic; aclshmemx_signal_op(gva_sync flag_offset, flag, SIGNAL_SET, my_rank); } } // Receiver core轮询 signal就绪即拉取 while (!all_done) { for (int g 0; g core_group_num; g) { aclshmem_int32_get_nbi(flag_ub, gva_sync g * INTERVAL, 1, peer); PipeBarrierPIPE_ALL(); int64_t ready *flag_ub - magic; if (ready fetched[g]) { // 立即拉取已就绪 chunk使用 ping-pong buffer fetch_chunks(output, gva_data, fetched[g], ready, ping_pong_buf); } } }在 allgather_kernel.cpp 中可以看到该模式的关键实现参数int core_group_num aivNum / 2;sender 和 receiver 的 core 数量各占一半aclshmemx_mte_put_nbi(..., my_rank, EVENT_ID0)搭配SetFlag/WaitFlagHardEvent::MTE3_S(EVENT_ID0)保证逐 chunk 拷贝完成flag times magic; aclshmemx_signal_op(gva_sync_gm flag_offset, flag, ACLSHMEM_SIGNAL_SET, my_rank);signal 使用magic chunk_count编码receiver 通过ready_flag - magic解码得到已就绪 chunk 数receiver 侧用aclshmem_int32_get_nbi轮询远端 flag并用 ping/pong bufferping_buff/pong_buff各约 95KB交替接收。关键点sender 和 receiver 的 core 数量各占一半signal 使用magic chunk_count编码receiver 通过解码判断就绪程度同时用(ready_flag shift) (magic shift)校验轮次有效性不需要全局 barrier数据逐 chunk 流式可用。1.2 AllToAllV 大消息 — 禁止双拷贝 staging瓶颈sendbuf → 本地 symm → put 远端 symm这条路径使有效带宽封顶约 50% 目标值——同样的数据被搬运了两次。优化远端直接mte_put(sendbuf displ, remote_symm put_off, pe dst)只有dst my_rank时才 staging 到本地symm_off避免对每个 dst 都先拷贝到symm_off再拷贝到远端put_off的双倍流量。同步send 半核quiet后执行全体aclshmem_barrier_all()再让 recv 半核发起 get。❌ 错误对每个 dst 先 Copy 到symm_off再 Copysymm_off到远端put_off双倍流量。 ✅ 正确远端一次 put半核 send/recv 之间aclshmem_barrier_all()。引擎选择数据量、拓扑与推荐 Engine 的对应关系如下数据量拓扑推荐 Engine原因 2MB任意MTEaclshmemx_mte_put_nbi/aclshmemx_mte_get_nbi延迟低setup 开销小≥ 2MB节点内 P2P 可达SDMAaclshmemx_sdma_put_nbi/aclshmemx_sdma_get_nbi带宽高不占 MTE 通道≥ 2MB跨节点 / P2P 不可达RDMA/RoCEaclshmemx_roce_put_nbi/aclshmemx_roce_get_nbi唯一可达路径混合场景节点内跨节点MTE 节点内 RDMA 跨节点按拓扑分层选择Engine 类型在 include/host_device/shmem_common_types.h 中定义enum data_op_engine_type_t { ACLSHMEM_DATA_OP_MTE 0x01, ACLSHMEM_DATA_OP_SDMA 0x02, ACLSHMEM_DATA_OP_ROCE 0x04, ACLSHMEM_DATA_OP_UDMA 0x08, ACLSHMEM_DATA_OP_MAX 0x08, };MTE 接口的签名见 include/device/gm2gm/engine/shmem_device_mte.h为template typename T ACLSHMEM_DEVICE void aclshmemx_mte_put_nbi( __gm__ T* dst, __gm__ T* src, __ubuf__ T* buf, uint32_t ub_size, uint32_t elem_size, int pe, uint32_t sync_id);其中buf为本地 UB 临时缓冲、ub_size为 UB 缓冲字节数、elem_size为搬运元素个数、sync_id为流水线同步 ID同文件还提供了non_contiguous_copy_param形式的非连续搬运重载。注意SDMA 需要预留至少 64B 的 UB bufferRDMA 不支持同 PE 上的 RMA/AMO 并发乱序写option_attr.data_op_engine_type可设置默认引擎SDMA 路径 MUST 先保证 MTE 路径大档 PASS再启用 SDMA验证建议用 uniformbase_count 8388608的正确性用例 perf-workflow.md 的带宽表。1.3 Chunk Size 调优规则说明MTE 最优 chunk190KB190 * 1024bytesallgather example 经验值最小搬运≥ 16KB 才能接近峰值带宽最大单次受 UB 可用空间限制double buffer 时为 UB/2调优方向chunk 太小 → 增大减少 setup 开销chunk 太大 → 缩小增加流水级数仓库实现中该值被定义为编译期常量见 examples/allgather/allgather_kernel.cppconstexpr int64_t UB_DMA_MAX_SIZE 190 * 1024;receiver 侧 double buffer 时单 buffer 使用UB_DMA_MAX_SIZE / 2约 95KB保证搬运期间另一块 buffer 可被使用。1.4 非连续搬运优化场景优化固定 stride 访问使用iput/iget或non_contiguous_copy_paramrepeat length src_ld dst_ld多段小块合并为连续 buffer 后一次搬运减少 put/get 调用次数分散 gather在本 PE 做 local pack → symmetric → 一次 put 到远端non_contiguous_copy_param重载同样定义在 shmem_device_mte.h 中核心思想是用repeat重复次数、length单次长度与src_ld/dst_ld源/目的 stride描述非连续排布把多次小搬运合并为一次 DMA 描述。2. 流水线与 Double Buffer2.1 Ping-Pong Buffer 模式瓶颈单 buffer 时NBI 搬运和使用MUST串行硬件在等待期间空转。优化两组 UB buffer 两组 event id 交替使用——搬运 buffer A 时使用 buffer B。标准写法可参考 examples/kv_shuffle/kv_shuffle_kernel.cpp其缓冲区布局为 K 缓存 ping/pongUB 0 / 32KB与 V 缓存 ping/pongUB 64KB / 96KB共 4 个 event idconstexpr uint64_t kPingUb 0; constexpr uint64_t kPongUb 32 * 1024; constexpr uint32_t kUbSize 32 * 1024; int ping_pong_flag 0; for (int block 0; block block_nums; block) { uint64_t ub_offset ping_pong_flag 0 ? kPingUb : kPongUb; TEventID event ping_pong_flag 0 ? EVENT_ID0 : EVENT_ID1; WaitFlagHardEvent::MTE3_MTE2(event); aclshmemx_mte_put_nbi(dst, src block * chunk, ub_offset, kUbSize, chunk_elems, target_pe, event); SetFlagHardEvent::MTE3_MTE2(event); ping_pong_flag 1 - ping_pong_flag; }关键约束两个 event idMUST成对SetFlag/WaitFlag不跨未完成 NBI 复用kv_shuffle 中 K/V 各用一对即 EVENT_ID0/1 与 EVENT_ID2/3互不干扰首次进入循环前需预设 flag或首轮跳过 Waitbuffer 大小相等避免 tail 时一侧溢出。2.2 通信-计算 Overlap瓶颈通信和计算串行执行硬件利用率低。优化让 chunk k 的通信和 chunk k-1 的计算重叠执行。前提≥ 2 个 chunk 迭代且通信和计算时间均不可忽略。时间线示意comm: [chunk0] [chunk1] [chunk2] [chunk3] compute: [chunk0] [chunk1] [chunk2] [chunk3]实现要求需要两组 data buffercompute 用当前、comm 填下一组使用 event 或 signal/wait 在 chunk 完成时通知计算端尾部需 drain最后一个 chunk 通信完成后单独计算。2.3 Multi-Stage PipelineSTAGES 2场景存储层次多级如 L1 → L0A → L0B → Compute时两级 ping-pong 不够用需要更深流水。模式使用 circular index 管理多级 buffer。constexpr uint32_t STAGES 2; // 可扩展为 3、4 uint32_t bufId 0; for (uint32_t iter 0; iter total_iters; iter) { auto tile tensorList[bufId]; WaitFlagHardEvent::M_MTE1(eventList[bufId]); // 使用当前 buffer 计算 compute(tile); SetFlagHardEvent::M_MTE1(eventList[bufId]); bufId (bufId 1 STAGES) ? (bufId 1) : 0; // circular }每个 stage 有独立的 buffer 与 eventbufId循环递增使搬运stage k1 的填充与计算stage k 的使用深度重叠。3. 同步优化3.1 Barrier → Signal/Wait瓶颈全局 barrier 让所有 PE 等待最慢者且与数据依赖无关的 PE 也被拖住。优化替换为 producer-consumer 点对点 signal/wait。❌ 错误做法aclshmem_barrier_all(); // 每个 phase 全员等待✅ 正确做法// Producer 完成后只通知消费者 aclshmemx_signal_op(remote_sig, magic, ACLSHMEM_SIGNAL_SET, target_pe); // Consumer 只等自己需要的数据就绪 aclshmem_signal_wait_until(local_sig, ACLSHMEM_CMP_EQ, magic);这两个接口在 include/host/data_plane/shmem_host_so.hsignal_op与 include/host/data_plane/shmem_host_p2p_sync.hsignal_wait_until中有对应实现ACLSHMEM_CMP_EQ为比较操作符枚举值。allgather 的 small-data 路径也展示了aclshmem_quiet()SyncAll()signal_wait_until的组合用法见 allgather_kernel.cpp。3.2 Per-Chunk Signaling瓶颈整块数据传输完才发 signal消费者长时间空等。优化每个 chunk 传输完立即 signal消费者可逐 chunk 消费。signal 编码建议flag chunk_count magicmagic 区分轮次chunk_count 表示已就绪数量receiver 通过(flag shift) (magic shift)验证轮次有效性每个 core 使用独立 flag offsetflag_offset aivIndex * SYNC_FLAG_INTERVAL。仓内 allgather 定义了SYNC_FLAG_INTERVAL 16allgather_kernel.cppflag_offset aivIndex * SYNC_FLAG_INTERVAL每轮 sender 发flag times magicreceiver 解码ready_num ready_flag - magic后与本地*flags_ub1[group_idx]比较即可知道本次该拉取哪些新 chunkallgather_kernel.cpp。3.3 Phase 合并场景合并方式reduce-scatter allgather合并为一个 fused kernelphase 间在 Device 侧同步copy-to-symmetric barrier fetch合并为 sender/receiver overlap 模式见 §1.1compute comm epilogue计算 kernel 直接写 symmetric state通信 epilogue 在同一 kernel 完成Phase 合并的核心收益是消除 kernel 间往返与多次全局同步allgather 的all_gather_origin就是 copy-to-symmetric、per-chunk signal、remote fetch 三个阶段在一个 kernel 内合并的实例。4. 内存优化4.1 Avoid GM Scratch瓶颈远端 get → GM tmp → DataCopy to UB → compute → DataCopy to GM产生了多余的 GM 读写。优化使用get_nbi直接把数据搬到 UB在 UB 上做 streaming reduce。❌ 错误做法aclshmem_getmem(gm_tmp, remote_src, size, peer); // 落 GM aclshmemx_mte_quiet(); DataCopy(ub_buf, gm_tmp, size); // 再搬 UB compute(ub_buf);✅ 正确做法// 直接搬到 UBUB2GM get aclshmem_float_get_nbi((__ubuf__ float *)ub_buf, remote_src, count, peer); SetFlagHardEvent::MTE2_MTE3(EVENT_ID0); WaitFlagHardEvent::MTE2_MTE3(EVENT_ID0); compute(ub_buf);注意这里使用MTE2_MTE3事件get 数据到达 UB 后用该事件把 MTE2搬运与 MTE3计算串起来确保 compute 读到的是完整数据。4.2 Buffer 对齐对齐要求带宽影响512B 对齐接近峰值带宽Atlas A2 可比 32B 对齐提升 30%32B 对齐最低要求带宽受损// 分配时指定对齐 void *sym_buf aclshmem_align(512, total_bytes);aclshmem_align在 include/host/mem/shmem_host_heap.h 中声明并提供了shmem_align别名分配的是对称内存对齐参数alignment需为 2 的幂且不小于默认对齐。4.3 最小搬运量规则说明单次搬运 ≥ 16KB低于此值 DMA setup 开销占比过大合并小块多个小搬运合并为一次大搬运padding 对齐必要时 padding 到 16KB 倍数搬运后忽略 padding 部分4.4 UB Buffer Fusion项目说明瓶颈多步 Vector 计算之间把中间结果写回 GM 再读入优化中间结果保留在 UB连续执行多步 Vector 操作UB Buffer Fusion 与 §4.1 一脉相承凡是中间结果只被后续计算使用、不跨 PE 消费的数据都应留在 UB 内完成整个计算链避免任何 GM 往返。5. 分核策略优化分核策略优化的出发点是检查当前 kernel 是否存在单 AIV 内for peer in 0..n_pes的串行循环同一个 chunk 或同一段任务需要该 AIV 逐个、串行访问不同 PE 的数据。如果 block_dim 和 chunk size 已调到局部最优但性能仍未达标瓶颈大概率就在这条 AIV 内跨 PE 串行链路上。以下模式给出把 AIV 内串行链路转化为多核并行的方案。5.1 Sender/Receiver Core Split项目说明模式一半 core 做 local→symmetric 拷贝sender另一半做 remote fetchreceiver来源allgatherall_gather_origincore_group_num aivNum / 2适用数据量大、需要 overlap local copy 和 remote fetch 的场景不适用小数据量所有 core 做同一操作更快5.2 AIV 内跨 PE 串行链路消除瓶颈单个 AIV 内存在for peer in 0..n_pes的串行循环每次迭代发起一次跨 PE 搬运后必须等待该次搬运完成才能进入下一 peer无法利用多条 HCCS 链路同时传输。检测在 kernel 的 prof 打点范围内查找逐 peer 的 get/put 循环且循环体内部有WaitFlag或PipeBarrier等同步点阻断并发。影响8PE reduce_scatter 实测 bus bandwidth 仅 65%~71% HCCL根源即 AIV 内逐 peer 串行 getadd。优先级block_dim 和 chunk size 调至局部最优后若未达标MUST优先尝试本节方案。当前典型串行模式// AIV 内逐 peer 串行get PE0 → wait → add → get PE1 → wait → add ... for (int peer 0; peer n_pes; peer) { get 当前 chunk 从 peer → tmp_ub; WaitFlag(完成); acc_ub tmp_ub; }以下两种子模式将这条串行链路拆分到多核或多阶段。子模式 1按源 PE 分组并行拉取将 AIV 按源 PE 分组每组专责从一个 PE 拉取数据各组并行拉取完成后做一次本地跨 PE 汇总。// n_pes 组 AIV 并行组 0 拉 PE0、组 1 拉 PE1、... int peer aiv_idx / core_per_peer; // 该 AIV 属于哪组负责哪个 PE int lane aiv_idx % core_per_peer; // 组内编号 int my_len chunk_len / core_per_peer (lane extra); // 该 AIV 处理的数据量 int my_off 该 lane 的数据起始偏移; get symm[peer][shard_off my_off] → local_scratch[lane]; // 全部组完成拉取后做一次本地跨 PE 汇总 for (int p 0; p n_pes; p) { acc local_scratch[p * core_per_peer lane]; }要求AIV 总数 ≥ n_pes 且能被 n_pes 整除。原本 1 个 AIV 串行 8 次 getwaitadd现在变为 8 组 AIV 并行 1 次 getwait汇总阶段改为本地读。子模式 2Sender put Receiver 本地汇总将 AIV 对半分为 sender 和 receiver。sender 不等待远端读请求——它主动把自己的 contributionput到 owner PE 的对称内存完毕后 signal 对应的 receiver。receiver 在本 PE 的对称内存中直接读取各 sender 推来的数据并汇总无需发起跨 PE 读取、无需逐 peer 等待。// Sender AIV推数据到 owner PE if (aiv_idx sender_num) { put 我的 contribution → owner_pe.sym_buf[my_slot]; signal(receiver_on_owner_pe, my_slot_ready); } // Receiver AIV在本地对称内存中汇总 else { for (int src 0; src n_pes; src) { wait(signal from src); acc owner_pe.sym_buf[src * slot_size my_offset]; } }注意每个 sender 需要独立的对称内存槽位slot由 receiver 侧预先分配好避免多 sender 写同一地址冲突。子模式选择条件选择AIV 总数 ≥ n_pes 且可整除子模式 1按源 PE 分组AIV 总数 n_pes 或需要避免多路并行 get 的链路拥塞子模式 2Sender put Receiver 本地汇总n_pes ≥ 8 且对称内存有限子模式 1若 AIV 不足则配合 §5.3 分层规约5.3 分层规约瓶颈全量 peer 参与单级规约AIV 不足时要么串行要么分配不均。优化先做 PE pair 局部规约得到中间结果再对中间结果做 reduce_scatter。适用n_pes ≥ 4AIV 不足以按 §5.2 子模式 1 均分到每组至少 1 核时。约束每层增加一次 barrier层级过多反而抵消收益。先在最小 PE 数如 2PE 或 4PE验证再用到大的 PE 配置。8PE 两层示例 Layer 1PE pairPE0PE1 局部 sum → PE0PE2PE3 → PE2... Layer 2RS4 个 pair owner 做 4-PE reduce_scatter → 各自 shard5.4 Dynamic Block Dim数据量Block Dim 策略 2MB固定 8 核小数据 setup 开销大于并行收益≥ 2MB2 * n_pes或按 chunk 数动态计算纯通信AIV 按 PE、数据维度均分仓内 allgather 的做法可作为参考kernel 入口按数据量分流——elements * sizeof(type) 2097152即 2MB走all_gather_small_data否则走all_gather_big_data后者把大 buffer 切成 100MiB 窗口GVA_BUFF_MAX_SIZE分批执行见 allgather_kernel.cpp 与 allgather_kernel.cpp。5.5 负载均衡❌ 错误做法int per_core total / core_num; // tail 核可能负载翻倍✅ 正确做法int base total / core_num; int extra total % core_num; int my_count base (core_idx extra ? 1 : 0); int my_offset core_idx * base min(core_idx, extra);核心思想是把余数extra平摊到前extra个 core避免最后一个 core 独自承担全部余量从而让各核负载差不超过 1 个单元。5.6 Tail 合并策略适用场景合并到最后一个正常 chunktail chunk 极小独立一轮循环开销不值得特殊处理tail chunk 大小差异大需要单独搬运参数填充到对齐tail 很小但对齐后可复用标准路径allgather 的 sender 侧对尾部做了典型处理主循环按UB_DMA_MAX_SIZE满 chunk 搬运剩余部分用copy_total_size / sizeof(T)作为elem_size单独发最后一次 put见 allgather_kernel.cpp保证尾块不浪费也不越界。6. 优化优先级对于典型 SHMEM 通信算子建议按以下顺序尝试优化优先级方向预期收益风险1串行 Peer 迭代消除§5.2拆 AIV 内逐 peer 串行链路为多核并行reduce 类预期 20%分核策略变更大验证正确性2Copy-to-symmetric overlap§1.1掩盖对称内存拷贝延迟分核策略变更大3Double buffer / ping-pong§2.1掩盖搬运气泡需要 2× UB 空间4Barrier → signal/wait§3.1减少全局等待正确性风险较高5Engine 选择 chunk 调优§1.2/§1.3提升有效带宽不同拓扑结果不同6动态 block_dim / 负载均衡§5.4/§5.5平衡各核负载需要多组实验7内存对齐 减少 GM 中转§4提升搬运效率收益相对固定排优先级的原则是先解决结构性瓶颈再处理参数类调优串行链路消除和通信重叠解决的是“链路结构”问题收益最大但改动面也大需要充分的正确性验证例如先跑通 uniform 大数据用例再对比 perf-workflow.md 的带宽基准表Double buffer、Engine 选择、block_dim、内存对齐等则是改动局部、可快速实验的参数类优化。每一轮优化完成后都应回到起点重新 prof 定位瓶颈再迭代下一轮。【免费下载链接】shmemCANN SHMEM 是面向昇腾平台的多机多卡内存通信库基于OpenSHMEM 标准协议实现跨设备的高效内存访问与数据同步。项目地址: https://gitcode.com/cann/shmem创作声明：本文部分内容由AI辅助生成（AIGC），仅供参考 
长河君 · 2026/9/18 12:50:32 · 阅读 1.2k CANN SHMEM 算子性能优化模式：从瓶颈定位到通信、流水线与分核策略的实战指南
CANN SHMEM 算子性能优化模式从瓶颈定位到通信、流水线与分核策略的实战指南【免费下载链接】shmemCANN SHMEM 是面向昇腾平台的多机多卡内存通信库基于OpenSHMEM 标准协议实现跨设备的高效内存访问与数据同步。项目地址: https://gitcode.com/cann/shmem本文是 .agents/skills/shmem-ops-performance-optim/references/optimization-patterns.md 的完整展开版。文中examples/、include/等路径均以 CANN SHMEM 仓库根目录为起点阅读前建议先通过 shmem-repo-resolution.md 定位仓库结构。CANN SHMEM 是基于 OpenSHMEM 标准的多机多卡内存通信库其算子的性能上限往往取决于 kernel 内数据搬运、同步与分核方式的设计。本文面向在昇腾 NPU 上编写或优化 SHMEM 通信算子的开发者系统梳理了六类典型瓶颈及其对应的优化模式通信重叠、流水线与 Double Buffer、同步替换、内存路径精简、分核策略重构以及优先级排序。读完本文你将能够根据 prof 数据快速判断瓶颈类型并在 examples/allgather、examples/kv_shuffle 等仓内示例中找到可复用的工程化写法。优化总览先定位瓶颈类型再选择手段优化的第一原则是每轮先定位瓶颈类型再选择对应手段而不是盲目套用技巧。根据瓶颈所处的位置可把优化方向分为六类类别典型瓶颈对应章节通信优化数据必须先落 symmetric heap 再由远端读取拷贝成为串行开销大消息双拷贝导致带宽封顶§1流水线优化单 buffer 时 NBI 搬运与使用必须串行搬运气泡明显§2同步优化全局 barrier 让所有 PE 等待最慢者整块传完才 signal 导致消费者空等§3内存优化远端 get 落 GM、再经 DataCopy 进 UB 的多余读写§4分核优化单 AIV 内for peer in 0..n_pes串行搬运无法利用多条链路并行§5综合排序多个优化点并存时收益与风险不同§6下面逐类展开每一类都给出 ❌ 错误做法、✅ 正确做法、关键点以及仓内源码佐证。1. 通信优化1.1 Copy-to-Symmetric-Heap Overlap瓶颈数据MUST先从本地 GM 拷贝到 symmetric heap再由远端 PE 读取在没有重叠设计时拷贝时间成为串行开销。优化将 core 分为 sender 组和 receiver 组。sender 按 chunk 把数据拷贝到 symmetric heap并逐 chunk 发 signalreceiver 轮询 signal一旦某 chunk 就绪立即拉取。该模式源码见 examples/allgather/allgather_kernel.cpp 中的all_gather_originkernel。❌ 错误做法——所有数据拷贝完 → 全局 barrier → 再开始通信copy_all_to_symmetric(input, gva_data, total_size); aclshmem_barrier_all(); fetch_all_from_peers(output, gva_data, total_size);✅ 正确做法——sender 逐 chunk 拷贝 逐 chunk signalreceiver 轮询就绪即拉取// Sender core逐 chunk 拷贝 逐 chunk signal if (aivIndex core_group_num) { while (remaining chunk_size) { aclshmemx_mte_put_nbi(gva_data offset, input offset, tmp_buff, ub_size, chunk_elems, my_rank, EVENT_ID0); SetFlagHardEvent::MTE3_S(EVENT_ID0); WaitFlagHardEvent::MTE3_S(EVENT_ID0); flag times magic; aclshmemx_signal_op(gva_sync flag_offset, flag, SIGNAL_SET, my_rank); } } // Receiver core轮询 signal就绪即拉取 while (!all_done) { for (int g 0; g core_group_num; g) { aclshmem_int32_get_nbi(flag_ub, gva_sync g * INTERVAL, 1, peer); PipeBarrierPIPE_ALL(); int64_t ready *flag_ub - magic; if (ready fetched[g]) { // 立即拉取已就绪 chunk使用 ping-pong buffer fetch_chunks(output, gva_data, fetched[g], ready, ping_pong_buf); } } }在 allgather_kernel.cpp 中可以看到该模式的关键实现参数int core_group_num aivNum / 2;sender 和 receiver 的 core 数量各占一半aclshmemx_mte_put_nbi(..., my_rank, EVENT_ID0)搭配SetFlag/WaitFlagHardEvent::MTE3_S(EVENT_ID0)保证逐 chunk 拷贝完成flag times magic; aclshmemx_signal_op(gva_sync_gm flag_offset, flag, ACLSHMEM_SIGNAL_SET, my_rank);signal 使用magic chunk_count编码receiver 通过ready_flag - magic解码得到已就绪 chunk 数receiver 侧用aclshmem_int32_get_nbi轮询远端 flag并用 ping/pong bufferping_buff/pong_buff各约 95KB交替接收。关键点sender 和 receiver 的 core 数量各占一半signal 使用magic chunk_count编码receiver 通过解码判断就绪程度同时用(ready_flag shift) (magic shift)校验轮次有效性不需要全局 barrier数据逐 chunk 流式可用。1.2 AllToAllV 大消息 — 禁止双拷贝 staging瓶颈sendbuf → 本地 symm → put 远端 symm这条路径使有效带宽封顶约 50% 目标值——同样的数据被搬运了两次。优化远端直接mte_put(sendbuf displ, remote_symm put_off, pe dst)只有dst my_rank时才 staging 到本地symm_off避免对每个 dst 都先拷贝到symm_off再拷贝到远端put_off的双倍流量。同步send 半核quiet后执行全体aclshmem_barrier_all()再让 recv 半核发起 get。❌ 错误对每个 dst 先 Copy 到symm_off再 Copysymm_off到远端put_off双倍流量。 ✅ 正确远端一次 put半核 send/recv 之间aclshmem_barrier_all()。引擎选择数据量、拓扑与推荐 Engine 的对应关系如下数据量拓扑推荐 Engine原因 2MB任意MTEaclshmemx_mte_put_nbi/aclshmemx_mte_get_nbi延迟低setup 开销小≥ 2MB节点内 P2P 可达SDMAaclshmemx_sdma_put_nbi/aclshmemx_sdma_get_nbi带宽高不占 MTE 通道≥ 2MB跨节点 / P2P 不可达RDMA/RoCEaclshmemx_roce_put_nbi/aclshmemx_roce_get_nbi唯一可达路径混合场景节点内跨节点MTE 节点内 RDMA 跨节点按拓扑分层选择Engine 类型在 include/host_device/shmem_common_types.h 中定义enum data_op_engine_type_t { ACLSHMEM_DATA_OP_MTE 0x01, ACLSHMEM_DATA_OP_SDMA 0x02, ACLSHMEM_DATA_OP_ROCE 0x04, ACLSHMEM_DATA_OP_UDMA 0x08, ACLSHMEM_DATA_OP_MAX 0x08, };MTE 接口的签名见 include/device/gm2gm/engine/shmem_device_mte.h为template typename T ACLSHMEM_DEVICE void aclshmemx_mte_put_nbi( __gm__ T* dst, __gm__ T* src, __ubuf__ T* buf, uint32_t ub_size, uint32_t elem_size, int pe, uint32_t sync_id);其中buf为本地 UB 临时缓冲、ub_size为 UB 缓冲字节数、elem_size为搬运元素个数、sync_id为流水线同步 ID同文件还提供了non_contiguous_copy_param形式的非连续搬运重载。注意SDMA 需要预留至少 64B 的 UB bufferRDMA 不支持同 PE 上的 RMA/AMO 并发乱序写option_attr.data_op_engine_type可设置默认引擎SDMA 路径 MUST 先保证 MTE 路径大档 PASS再启用 SDMA验证建议用 uniformbase_count 8388608的正确性用例 perf-workflow.md 的带宽表。1.3 Chunk Size 调优规则说明MTE 最优 chunk190KB190 * 1024bytesallgather example 经验值最小搬运≥ 16KB 才能接近峰值带宽最大单次受 UB 可用空间限制double buffer 时为 UB/2调优方向chunk 太小 → 增大减少 setup 开销chunk 太大 → 缩小增加流水级数仓库实现中该值被定义为编译期常量见 examples/allgather/allgather_kernel.cppconstexpr int64_t UB_DMA_MAX_SIZE 190 * 1024;receiver 侧 double buffer 时单 buffer 使用UB_DMA_MAX_SIZE / 2约 95KB保证搬运期间另一块 buffer 可被使用。1.4 非连续搬运优化场景优化固定 stride 访问使用iput/iget或non_contiguous_copy_paramrepeat length src_ld dst_ld多段小块合并为连续 buffer 后一次搬运减少 put/get 调用次数分散 gather在本 PE 做 local pack → symmetric → 一次 put 到远端non_contiguous_copy_param重载同样定义在 shmem_device_mte.h 中核心思想是用repeat重复次数、length单次长度与src_ld/dst_ld源/目的 stride描述非连续排布把多次小搬运合并为一次 DMA 描述。2. 流水线与 Double Buffer2.1 Ping-Pong Buffer 模式瓶颈单 buffer 时NBI 搬运和使用MUST串行硬件在等待期间空转。优化两组 UB buffer 两组 event id 交替使用——搬运 buffer A 时使用 buffer B。标准写法可参考 examples/kv_shuffle/kv_shuffle_kernel.cpp其缓冲区布局为 K 缓存 ping/pongUB 0 / 32KB与 V 缓存 ping/pongUB 64KB / 96KB共 4 个 event idconstexpr uint64_t kPingUb 0; constexpr uint64_t kPongUb 32 * 1024; constexpr uint32_t kUbSize 32 * 1024; int ping_pong_flag 0; for (int block 0; block block_nums; block) { uint64_t ub_offset ping_pong_flag 0 ? kPingUb : kPongUb; TEventID event ping_pong_flag 0 ? EVENT_ID0 : EVENT_ID1; WaitFlagHardEvent::MTE3_MTE2(event); aclshmemx_mte_put_nbi(dst, src block * chunk, ub_offset, kUbSize, chunk_elems, target_pe, event); SetFlagHardEvent::MTE3_MTE2(event); ping_pong_flag 1 - ping_pong_flag; }关键约束两个 event idMUST成对SetFlag/WaitFlag不跨未完成 NBI 复用kv_shuffle 中 K/V 各用一对即 EVENT_ID0/1 与 EVENT_ID2/3互不干扰首次进入循环前需预设 flag或首轮跳过 Waitbuffer 大小相等避免 tail 时一侧溢出。2.2 通信-计算 Overlap瓶颈通信和计算串行执行硬件利用率低。优化让 chunk k 的通信和 chunk k-1 的计算重叠执行。前提≥ 2 个 chunk 迭代且通信和计算时间均不可忽略。时间线示意comm: [chunk0] [chunk1] [chunk2] [chunk3] compute: [chunk0] [chunk1] [chunk2] [chunk3]实现要求需要两组 data buffercompute 用当前、comm 填下一组使用 event 或 signal/wait 在 chunk 完成时通知计算端尾部需 drain最后一个 chunk 通信完成后单独计算。2.3 Multi-Stage PipelineSTAGES 2场景存储层次多级如 L1 → L0A → L0B → Compute时两级 ping-pong 不够用需要更深流水。模式使用 circular index 管理多级 buffer。constexpr uint32_t STAGES 2; // 可扩展为 3、4 uint32_t bufId 0; for (uint32_t iter 0; iter total_iters; iter) { auto tile tensorList[bufId]; WaitFlagHardEvent::M_MTE1(eventList[bufId]); // 使用当前 buffer 计算 compute(tile); SetFlagHardEvent::M_MTE1(eventList[bufId]); bufId (bufId 1 STAGES) ? (bufId 1) : 0; // circular }每个 stage 有独立的 buffer 与 eventbufId循环递增使搬运stage k1 的填充与计算stage k 的使用深度重叠。3. 同步优化3.1 Barrier → Signal/Wait瓶颈全局 barrier 让所有 PE 等待最慢者且与数据依赖无关的 PE 也被拖住。优化替换为 producer-consumer 点对点 signal/wait。❌ 错误做法aclshmem_barrier_all(); // 每个 phase 全员等待✅ 正确做法// Producer 完成后只通知消费者 aclshmemx_signal_op(remote_sig, magic, ACLSHMEM_SIGNAL_SET, target_pe); // Consumer 只等自己需要的数据就绪 aclshmem_signal_wait_until(local_sig, ACLSHMEM_CMP_EQ, magic);这两个接口在 include/host/data_plane/shmem_host_so.hsignal_op与 include/host/data_plane/shmem_host_p2p_sync.hsignal_wait_until中有对应实现ACLSHMEM_CMP_EQ为比较操作符枚举值。allgather 的 small-data 路径也展示了aclshmem_quiet()SyncAll()signal_wait_until的组合用法见 allgather_kernel.cpp。3.2 Per-Chunk Signaling瓶颈整块数据传输完才发 signal消费者长时间空等。优化每个 chunk 传输完立即 signal消费者可逐 chunk 消费。signal 编码建议flag chunk_count magicmagic 区分轮次chunk_count 表示已就绪数量receiver 通过(flag shift) (magic shift)验证轮次有效性每个 core 使用独立 flag offsetflag_offset aivIndex * SYNC_FLAG_INTERVAL。仓内 allgather 定义了SYNC_FLAG_INTERVAL 16allgather_kernel.cppflag_offset aivIndex * SYNC_FLAG_INTERVAL每轮 sender 发flag times magicreceiver 解码ready_num ready_flag - magic后与本地*flags_ub1[group_idx]比较即可知道本次该拉取哪些新 chunkallgather_kernel.cpp。3.3 Phase 合并场景合并方式reduce-scatter allgather合并为一个 fused kernelphase 间在 Device 侧同步copy-to-symmetric barrier fetch合并为 sender/receiver overlap 模式见 §1.1compute comm epilogue计算 kernel 直接写 symmetric state通信 epilogue 在同一 kernel 完成Phase 合并的核心收益是消除 kernel 间往返与多次全局同步allgather 的all_gather_origin就是 copy-to-symmetric、per-chunk signal、remote fetch 三个阶段在一个 kernel 内合并的实例。4. 内存优化4.1 Avoid GM Scratch瓶颈远端 get → GM tmp → DataCopy to UB → compute → DataCopy to GM产生了多余的 GM 读写。优化使用get_nbi直接把数据搬到 UB在 UB 上做 streaming reduce。❌ 错误做法aclshmem_getmem(gm_tmp, remote_src, size, peer); // 落 GM aclshmemx_mte_quiet(); DataCopy(ub_buf, gm_tmp, size); // 再搬 UB compute(ub_buf);✅ 正确做法// 直接搬到 UBUB2GM get aclshmem_float_get_nbi((__ubuf__ float *)ub_buf, remote_src, count, peer); SetFlagHardEvent::MTE2_MTE3(EVENT_ID0); WaitFlagHardEvent::MTE2_MTE3(EVENT_ID0); compute(ub_buf);注意这里使用MTE2_MTE3事件get 数据到达 UB 后用该事件把 MTE2搬运与 MTE3计算串起来确保 compute 读到的是完整数据。4.2 Buffer 对齐对齐要求带宽影响512B 对齐接近峰值带宽Atlas A2 可比 32B 对齐提升 30%32B 对齐最低要求带宽受损// 分配时指定对齐 void *sym_buf aclshmem_align(512, total_bytes);aclshmem_align在 include/host/mem/shmem_host_heap.h 中声明并提供了shmem_align别名分配的是对称内存对齐参数alignment需为 2 的幂且不小于默认对齐。4.3 最小搬运量规则说明单次搬运 ≥ 16KB低于此值 DMA setup 开销占比过大合并小块多个小搬运合并为一次大搬运padding 对齐必要时 padding 到 16KB 倍数搬运后忽略 padding 部分4.4 UB Buffer Fusion项目说明瓶颈多步 Vector 计算之间把中间结果写回 GM 再读入优化中间结果保留在 UB连续执行多步 Vector 操作UB Buffer Fusion 与 §4.1 一脉相承凡是中间结果只被后续计算使用、不跨 PE 消费的数据都应留在 UB 内完成整个计算链避免任何 GM 往返。5. 分核策略优化分核策略优化的出发点是检查当前 kernel 是否存在单 AIV 内for peer in 0..n_pes的串行循环同一个 chunk 或同一段任务需要该 AIV 逐个、串行访问不同 PE 的数据。如果 block_dim 和 chunk size 已调到局部最优但性能仍未达标瓶颈大概率就在这条 AIV 内跨 PE 串行链路上。以下模式给出把 AIV 内串行链路转化为多核并行的方案。5.1 Sender/Receiver Core Split项目说明模式一半 core 做 local→symmetric 拷贝sender另一半做 remote fetchreceiver来源allgatherall_gather_origincore_group_num aivNum / 2适用数据量大、需要 overlap local copy 和 remote fetch 的场景不适用小数据量所有 core 做同一操作更快5.2 AIV 内跨 PE 串行链路消除瓶颈单个 AIV 内存在for peer in 0..n_pes的串行循环每次迭代发起一次跨 PE 搬运后必须等待该次搬运完成才能进入下一 peer无法利用多条 HCCS 链路同时传输。检测在 kernel 的 prof 打点范围内查找逐 peer 的 get/put 循环且循环体内部有WaitFlag或PipeBarrier等同步点阻断并发。影响8PE reduce_scatter 实测 bus bandwidth 仅 65%~71% HCCL根源即 AIV 内逐 peer 串行 getadd。优先级block_dim 和 chunk size 调至局部最优后若未达标MUST优先尝试本节方案。当前典型串行模式// AIV 内逐 peer 串行get PE0 → wait → add → get PE1 → wait → add ... for (int peer 0; peer n_pes; peer) { get 当前 chunk 从 peer → tmp_ub; WaitFlag(完成); acc_ub tmp_ub; }以下两种子模式将这条串行链路拆分到多核或多阶段。子模式 1按源 PE 分组并行拉取将 AIV 按源 PE 分组每组专责从一个 PE 拉取数据各组并行拉取完成后做一次本地跨 PE 汇总。// n_pes 组 AIV 并行组 0 拉 PE0、组 1 拉 PE1、... int peer aiv_idx / core_per_peer; // 该 AIV 属于哪组负责哪个 PE int lane aiv_idx % core_per_peer; // 组内编号 int my_len chunk_len / core_per_peer (lane extra); // 该 AIV 处理的数据量 int my_off 该 lane 的数据起始偏移; get symm[peer][shard_off my_off] → local_scratch[lane]; // 全部组完成拉取后做一次本地跨 PE 汇总 for (int p 0; p n_pes; p) { acc local_scratch[p * core_per_peer lane]; }要求AIV 总数 ≥ n_pes 且能被 n_pes 整除。原本 1 个 AIV 串行 8 次 getwaitadd现在变为 8 组 AIV 并行 1 次 getwait汇总阶段改为本地读。子模式 2Sender put Receiver 本地汇总将 AIV 对半分为 sender 和 receiver。sender 不等待远端读请求——它主动把自己的 contributionput到 owner PE 的对称内存完毕后 signal 对应的 receiver。receiver 在本 PE 的对称内存中直接读取各 sender 推来的数据并汇总无需发起跨 PE 读取、无需逐 peer 等待。// Sender AIV推数据到 owner PE if (aiv_idx sender_num) { put 我的 contribution → owner_pe.sym_buf[my_slot]; signal(receiver_on_owner_pe, my_slot_ready); } // Receiver AIV在本地对称内存中汇总 else { for (int src 0; src n_pes; src) { wait(signal from src); acc owner_pe.sym_buf[src * slot_size my_offset]; } }注意每个 sender 需要独立的对称内存槽位slot由 receiver 侧预先分配好避免多 sender 写同一地址冲突。子模式选择条件选择AIV 总数 ≥ n_pes 且可整除子模式 1按源 PE 分组AIV 总数 n_pes 或需要避免多路并行 get 的链路拥塞子模式 2Sender put Receiver 本地汇总n_pes ≥ 8 且对称内存有限子模式 1若 AIV 不足则配合 §5.3 分层规约5.3 分层规约瓶颈全量 peer 参与单级规约AIV 不足时要么串行要么分配不均。优化先做 PE pair 局部规约得到中间结果再对中间结果做 reduce_scatter。适用n_pes ≥ 4AIV 不足以按 §5.2 子模式 1 均分到每组至少 1 核时。约束每层增加一次 barrier层级过多反而抵消收益。先在最小 PE 数如 2PE 或 4PE验证再用到大的 PE 配置。8PE 两层示例 Layer 1PE pairPE0PE1 局部 sum → PE0PE2PE3 → PE2... Layer 2RS4 个 pair owner 做 4-PE reduce_scatter → 各自 shard5.4 Dynamic Block Dim数据量Block Dim 策略 2MB固定 8 核小数据 setup 开销大于并行收益≥ 2MB2 * n_pes或按 chunk 数动态计算纯通信AIV 按 PE、数据维度均分仓内 allgather 的做法可作为参考kernel 入口按数据量分流——elements * sizeof(type) 2097152即 2MB走all_gather_small_data否则走all_gather_big_data后者把大 buffer 切成 100MiB 窗口GVA_BUFF_MAX_SIZE分批执行见 allgather_kernel.cpp 与 allgather_kernel.cpp。5.5 负载均衡❌ 错误做法int per_core total / core_num; // tail 核可能负载翻倍✅ 正确做法int base total / core_num; int extra total % core_num; int my_count base (core_idx extra ? 1 : 0); int my_offset core_idx * base min(core_idx, extra);核心思想是把余数extra平摊到前extra个 core避免最后一个 core 独自承担全部余量从而让各核负载差不超过 1 个单元。5.6 Tail 合并策略适用场景合并到最后一个正常 chunktail chunk 极小独立一轮循环开销不值得特殊处理tail chunk 大小差异大需要单独搬运参数填充到对齐tail 很小但对齐后可复用标准路径allgather 的 sender 侧对尾部做了典型处理主循环按UB_DMA_MAX_SIZE满 chunk 搬运剩余部分用copy_total_size / sizeof(T)作为elem_size单独发最后一次 put见 allgather_kernel.cpp保证尾块不浪费也不越界。6. 优化优先级对于典型 SHMEM 通信算子建议按以下顺序尝试优化优先级方向预期收益风险1串行 Peer 迭代消除§5.2拆 AIV 内逐 peer 串行链路为多核并行reduce 类预期 20%分核策略变更大验证正确性2Copy-to-symmetric overlap§1.1掩盖对称内存拷贝延迟分核策略变更大3Double buffer / ping-pong§2.1掩盖搬运气泡需要 2× UB 空间4Barrier → signal/wait§3.1减少全局等待正确性风险较高5Engine 选择 chunk 调优§1.2/§1.3提升有效带宽不同拓扑结果不同6动态 block_dim / 负载均衡§5.4/§5.5平衡各核负载需要多组实验7内存对齐 减少 GM 中转§4提升搬运效率收益相对固定排优先级的原则是先解决结构性瓶颈再处理参数类调优串行链路消除和通信重叠解决的是“链路结构”问题收益最大但改动面也大需要充分的正确性验证例如先跑通 uniform 大数据用例再对比 perf-workflow.md 的带宽基准表Double buffer、Engine 选择、block_dim、内存对齐等则是改动局部、可快速实验的参数类优化。每一轮优化完成后都应回到起点重新 prof 定位瓶颈再迭代下一轮。【免费下载链接】shmemCANN SHMEM 是面向昇腾平台的多机多卡内存通信库基于OpenSHMEM 标准协议实现跨设备的高效内存访问与数据同步。项目地址: https://gitcode.com/cann/shmem创作声明：本文部分内容由AI辅助生成（AIGC），仅供参考