CANN SHMEM 纯通信算子 Kernel 开发代码模式与规范从 Host 编排到 Device 通信骨架【免费下载链接】shmemCANN SHMEM 是面向昇腾平台的多机多卡内存通信库基于OpenSHMEM 标准协议实现跨设备的高效内存访问与数据同步。项目地址: https://gitcode.com/cann/shmem本文是面向 CANN SHMEM 开源仓库昇腾平台多机多卡内存通信库开发者的代码模式实战指南。它从仓库examples目录中提炼出可复用的 Host/Device 代码组织方式、put signal wait与get local compute两类核心通信骨架、AllGather/ReduceScatter/AllReduce 通用实现套路以及 symmetric buffer 生命周期、多 PE rank 分支、tail/chunk 循环等高频写法规范。读完本文你将能按照仓库既有标准从零搭建一个可编译、可运行、可校验的 SHMEM 纯通信算子工程并理解每一步背后的源码级依据。本文以 code-patterns.md 为主体辅以仓库中 allgather、kv_shuffle、sdma、one_multi_path 等真实示例的源码佐证。1. Host 侧main.cpp的固定阶段结构一个完整的 SHMEM example 的 Host 侧通常按 11 个阶段组织examples/allgather/main.cpp 是这套结构的典型实现阶段职责allgather 中的实际代码1. 解析 rank 参数读取n_pes、pe_id、ipport、设备数、起始设备、数据类型、shape、重复次数main.cpp 按INDEX1..INDEX8顺序解析 8 个参数2. ACL 初始化aclInit、aclrtSetDevice异步执行时创建aclrtStreammain.cppdevice_id pe_id % g_npus f_npu完成 PE 到物理设备的映射3. SHMEM 初始化参数构造设置my_pe、n_pes、ip_port、local_mem_size、option_attr.data_op_engine_type与超时utils.h 中test_set_attr统一完成4. SHMEM 初始化aclshmemx_init_attr按场景选 default / MPI / uniqueid bootstrapmain.cpp 使用ACLSHMEMX_INIT_WITH_DEFAULT5. 准备 kernel 辅助参数需要 Device barrier 的 kernel 获取util_get_ffts_config()需要 dump 时分配 dump workspacemain.cpp6. 分配输入/输出/symmetric buffer普通 GM 用aclrtMalloc跨 PE 共享状态、通信 staging、signal flags 用aclshmem_mallocmain.cpp7. 输入读取官方 demo 有时内置rank 常量输入生成生成算子优先读取gen_data.py产出的 input binmain.cpp 通过ReadFile读入 golden/input bin8. Kernel launch模板分发或 dtype 分支选择 kernel wrapper传入 stream、FFTS、GM 地址、symmetric 地址和 shapemain.cpp9. 完成等待aclrtSynchronizeStream、aclshmemx_handle_wait、aclshmem_barrier_all或 control barrier取决于通信完成位置main.cpp10. 结果写出拷回 Host 写入按 rank 隔离的输出文件精度校验交给 Python checkermain.cpp11. 释放资源Host pinned memory、普通 device memory、symmetric buffer、stream最后aclshmem_finalize、aclrtResetDevice、aclFinalize[main.cpp](https://link.gitcode.com/i/d184d4262efb3328b5a884a1b9b9f1e1#L150-L156, L199-L201)核心原则顺序不可颠倒SHMEM 初始化早于 symmetric allocation → 所有 kernel 完成后才能释放 symmetric buffer →aclshmem_finalize最后执行。工程组织上还有两条硬约束约束说明单进程单 PEmain.cpp一次启动只绑一个 device/my_pe多 PE 由scripts/run.sh/ launcher 启动多个独立进程Host 逻辑边界复杂 tiling/route/packing 拆到独立.cpp/.hmain.cpp只做BuildPlan → PrepareInputs → LaunchOp → WriteOutputs编排多 PE 启动必须发生在脚本层。allgather 的 run.sh 展示了标准写法用SCRIPT_DIR/PROJECT_ROOT定位工程导出SHMEM_UID_SESSION_ID与LD_LIBRARY_PATH循环后台拉起GNPU_NUM个进程后统一wait并收集返回值最后调用scripts/data_statistic.py做统计。禁止在main.cpp中fork/spawn多 PE 子进程。2. Device kernel 文件拆分习惯与职责边界examples 中常见两类拆分方式类型结构示例纯通信 demomain.cpp承担 Host 逻辑*_kernel.cpp/.h承担 kernel launch wrapperallgatherkernel 在 allgather_kernel.cpp单文件 demokernel 模板、launch wrapper、Host 测试逻辑放在一个main.cppsdmakernel 与allgather_kernelHost wrapper 同文件推荐的职责边界位置职责main.cpp编排初始化、数据准备、launch、同步、结果写出独立.cpp/.h复杂运行时计划route/tiling/packing 等*_kernel.cpp/.hdtype/shape dispatch、launch wrapper、Device kernelparam.h/utils.h公共参数、索引常量、错误检查宏Device kernelPE 查询、地址计算、搬运、计算、kernel 内同步allgather 的拆分非常典型allgather_kernel.cpp内既包含 Device 侧all_gather_origin/all_gather_small_data/all_gather_big_data模板也包含__global__ __aicore__入口宏ALLGATHER_FUNC_DEF/TYPE_FUNC和 Host 侧allgather_demolaunch wrapper而main.cpp只负责生命周期与数据编排并通过 dtype 分支调用allgather_demoint32_t等实例化模板allgather_kernel.cpp。utils.h/param.h这类公共工具如ACL_CHECK、ReadFile、test_set_attr被各 example 复用避免在main.cpp中散写。3. 三种地址表达方式raw pointer / GlobalTensor / LocalTensor3.1 raw pointerraw pointer 写法适合直接控制地址单位、偏移和 UB 预留位置GM 地址使用__gm__ T *。UB 地址使用__ubuf__ T *常通过固定 offset 构造临时 buffer。例如 allgather_kernel.cpp 用reinterpret_cast__ubuf__ T*(uint64_t(1024 32))构造临时 UB buffersdma/main.cpp 在 UB offset 1024 处放 64B 临时 buffer。byte-level engine 接口常用GM_ADDR或uint8_t *做统一地址计算。注意统一偏移单位指针加法通常按元素计uint8_t *或GM_ADDR加法按字节计混用时必须显式乘以sizeof(T)。raw pointer 不能绕过 SHMEM 数据面。即使通过aclshmem_ptr或 engine-specific pointer 得到远端 symmetric 地址跨 PE 搬运也必须调用aclshmem_*typed/putmem/getmem/put_signal 或公开aclshmemx_*MTE/SDMA/RDMA 接口。DataCopy只用于本 PE 内本地 GM/UB 搬运禁止直接读写远端 PE 地址。各引擎对非对称 GM 地址的支持不同直接影响 kernel 是否需要 kernel 外搬运引擎put_nbi 的 src 能是用户 GMget_nbi 的 dst 能是用户 GM含义MTE能能kernel 内直接用put_nbi(symm, input_gm, ...)无需 kernel 外搬运SDMA能能同上UDMA能能同上RDMA不能不能双方数据都必须位于对称内存需要 Host 侧aclrtMemcpy先搬运3.2AscendC::GlobalTensorGlobalTensor写法适合模板化接口和算子框架使用SetGlobalBuffer(reinterpret_cast__gm__ T *(addr), elem_count)绑定 GM。与aclshmemx_mte_*、aclshmemx_sdma_*、typed RMA 的 tensor 重载配合。见 sdma/main.cppsrc_tensor.SetGlobalBuffer(data_addr, base_per_core)后直接调用aclshmemx_sdma_qp_put_nbi(dst_tensor, src_tensor, tmp_local, ...)。3.3AscendC::LocalTensorLocalTensor适合表达 UB buffer通过TPipe/TBuf分配或手动设置logicPos、bufferAddr、dataLen。SDMA/RDMA 通常需要至少 64B 临时 UB buffer。SDMA 示例中正是这样构造tmp_local.address_.logicPos VECOUT; bufferAddr 1024; dataLen 64sdma/main.cpp。ping-pong 搬运时建议显式管理两个LocalTensor或两个 UB offset并配套EVENT_ID0/1。4. put signal wait 生产者-消费者模式该模式适合生产者把数据推到远端或本地 symmetric buffer消费者按 signal 判断数据是否可读。基本流程生产者计算本 core 负责的数据范围。通过aclshmemx_mte_put_nbi、aclshmemx_sdma_put_nbi、aclshmemx_roce_put_nbi或 typedput把数据写入目标 PE 的 symmetric buffer。生产者用 event、quiet 或 engine-specific quiet 保证该段数据写入完成到所需语义。生产者调用aclshmemx_signal_op或put_signal更新目标 PE 的 signal word。消费者调用aclshmem_signal_wait_until或 typed wait 等待 signal 到达预期值。消费者读取 symmetric buffer进入后续搬运或计算。signal 值建议带上 epoch/magic避免不同轮次复用同一个 signal slot 时误读旧值。多 core 场景下每个 core 应使用独立 flag offset 或通过 lane id 做隔离。源码佐证一小数据 AllGatherallgather_kernel.cpp 的all_gather_small_data中生产者用aclshmemx_mte_put_nbi写本地分片到 symmetric 区后aclshmem_quiet()SyncAll()再aclshmemx_signal_op(gva_sync_gm flag_offset, magic, ACLSHMEM_SIGNAL_SET, my_rank)通知消费者用aclshmem_signal_wait_until(aclshmem_ptr(gva_sync_gm, x) flag_offset, ACLSHMEM_CMP_EQ, magic)等待随后get_nbi拉取。flag 区用SYNC_FLAG_INTERVAL 16做 core 间隔离allgather_kernel.cpp。源码佐证二大数据流水allgather_kernel.cpp 的all_gather_origin展示了大数据的生产/消费分离前一半 core 负责把本地 GM 写入 symmetric staging每写完一段就signal_op递增 flagflag 值 times magic后一半 core 轮询远端 signal通过aclshmem_int32_get_nbi拉 flag 到 UB 再判断ready_num确认新 chunk 到达后按aclshmem_ptr(gva_sync_gm, x)语义等待并get_nbi拉数据。mstx 事件EXAMPLE_MSTX_CROSS_CORE_SET_FLAG_REPORT用于跨 core 调试追踪。5. get local compute 模式该模式适合 reduce、scatter、KV shuffle 等“本地掌握调度、从远端拉数据后计算”的场景。基本流程根据 PE、phase、chunk 或 expert id 计算远端地址。通过aclshmem_ptr或 engine-specific pointer 获取远端 symmetric 地址。调用get_nbi或对应aclshmemx_*_get_nbi把远端数据拉到本地 GM 或 UB不得用DataCopy直接读取远端地址。使用 event/quiet 等待当前 tile 到达。在本地做 elementwise add、atomic add、cast、matmul epilogue 或数据重排。将结果写回本地 output 或本地 symmetric state。该模式的优点是写冲突容易控制所有累加都在本 PE 本地完成不要求远端并发 atomic。对于 RDMA/SDMA 这类搬运本身不做 reduce 的通路通常先 get 到tmp_recv再本地累加回state_symm。源码佐证KV shufflekv_shuffle_kernel.cpp 的ShmemKVShuffle是典型的“本地掌握调度、从远端拉/推数据”实现kernel 入口先读global_shuffle_table得到pair_rank与operation0 表示发送依据local_rank/pair_rank与 AIV 编号计算独立的 sync flag offsetaiv_idx 8的 core 负责 K cache、aiv_idx 8负责 V cache各用 32KB UB slot 与EVENT_ID0..3做 ping-pong 搬运通过aclshmemx_mte_put_nbi直接写远端对称地址完成后用aclshmemx_signal_op回发 signal。Host 侧KVShuffleOps构造时用aclshmem_malloc分配n_pes * block_dims * SYNC_FLAG_INTERVAL大小的 sync 区并清零kv_shuffle_kernel.cpp每次 compute 递增count_作为 epoch。6. AllGather / ReduceScatter / AllReduce 通用骨架6.1 AllGather常见骨架每个 PE 把本地 input 写入本 PE 的 symmetric staging 区。通过 quiet/barrier 或 signal 通知其他 PE。每个 PE 按 rank 顺序从各远端 PE 的 staging 区读取对应分片写入本地 output 的[rank * local_count, (rank 1) * local_count)。小数据可用固定 core 数、barrier 和一次搬运完成大数据常按 chunk 循环用一半 core 做本地 GM 到 symmetric staging另一半 core 轮询 signal 并从远端拉取。收益点是读写分离生产阶段只写本 PE staging消费阶段按 PE 拉取避免多个 PE 同时写同一 output 分片。allgather 的大数据路径正是这种结构core_group_num aivNum / 2core_per_rank core_group_num / pe_size生产者按aivIndex * len_per_core切分本 PE 数据消费者按x get_core_idx / core_per_rank决定从哪个远端 PE 拉取[allgather_kernel.cpp](https://link.gitcode.com/i/3369b47e0443aced65c9efeb27104fc2#L61-L64, L129-L131)。6.2 ReduceScatter通用骨架将全量输入按 PE 或 chunk 切成目标分片。每个 phase 根据调度决定本 PE 需要向谁发送、从谁接收哪些 chunk。远端数据先进入本地tmp_recv或 UB tile。本地对目标 chunk 做累加最终只保留属于当前 PE 的 scatter 分片。phase 间需要 barrier 或 schedule-defined signal确保前一轮 partial result 已完成后再覆盖 staging buffer。对于非 ring 或跨机非均匀调度建议 Host 预先生成 per-phase/per-peer range listDevice 按 range list 执行避免 kernel 内复杂解析。这一边界与 code-patterns.md 中“复杂 route 计划放 Host、kernel 只按计划执行”的原则一致。6.3 AllReduceAllReduce 通常由 ReduceScatter AllGather 组成ReduceScatter 阶段把每个 chunk 的归约结果收敛到负责该 chunk 的 PE。AllGather 阶段把各 PE 拥有的归约后 chunk 广播给所有 PE。对 schedule-driven allreduceHost 只 launch 一个 fused kernelDevice 内部执行 init、reduce-scatter、all-gather、finalizephase 间在 Device 侧同步。AllReduce 的关键规范是明确“谁拥有某个 chunk 的写权”。同一目标 slice 不应被多个 AIV 无序写入除非使用受控 atomic 或先聚合再写回。allgather 大路径中 “Split data groups among the get cores assigned to the same PE to avoid overlapping writes” 的注释正是这一写权规则的体现allgather_kernel.cpp。6.4 统一实现 vs 大小分支默认 single path一套 kernel 逻辑覆盖 S/L 档UB 内while分块是内部细节不是small/large两条路径。禁止未证明收益就维护*_small_data/*_big_data、*_small/*_large并行实现。允许仅因GVA_BUFF_MAX_SIZE/ symm 容量触发的外层 tile 循环内存约束不是 perf 分支。size 分支须附 profiling≥5%bus_bandwidth_GBps 或时延收益才保留否则合并并删除死代码。需要说明的是现有 allgather 的ShmemAllGather_*入口按elements * sizeof(type) 20971522MB 阈值在all_gather_small_data与all_gather_big_data间分发其中 big 路径又因GVA_BUFF_MAX_SIZE 100MB容量限制做了外层 tile 循环times (elements max_gva_num - 1) / max_gva_num见 allgather_kernel.cpp并以aclshmemx_sync_vec_all做轮次间 Device 同步。这份既有示例可作为“何时该拆大小路径”的对照样本——新生成代码若要做 size 分支必须先以 profiling 数据证明收益。7. symmetric buffer 生命周期symmetric buffer 一般分为三类数据 staging存放本 PE 待其他 PE 读取的数据。状态/结果保存 partial result、reduce state 或 allreduce 最终结果。同步区signal flags、barrier counters、per-core ready flags。生命周期建议在aclshmemx_init_attr成功后分配。所有 PE 使用相同分配顺序和大小即使某 PE 当前 case 不使用某段 buffer也应保持布局一致。通过固定 layout 管理同一 buffer 内的数据区和 flag 区例如先放aiv_num * flag_interval再放数据区。allgather 中data_offset aivNum * SYNC_FLAG_INTERVAL正是这种布局allgather_kernel.cpp分配侧则是一次aclshmem_malloc(aiv_num * SYNC_FLAG_INTERVAL * sizeof(T) GVA_BUFF_MAX_SIZE / sizeof(T))allgather/main.cpp。每轮复用前重置 signal/state或使用 magic/epoch 区分轮次。allgather 每轮递增magic并以magic * MAGIC_MULTIPLIER作为 flag 值kv_shuffle 每次 compute 递增count_。kernel 和 Host 校验全部完成后再aclshmem_free。8. 多 PE rank 分支写法多 PE kernel 中常见分支my_pe aclshmem_my_pe()、n_pes aclshmem_n_pes()作为所有 rank 逻辑入口。if (peer my_pe) continue;避免对本 PE 走远端路径。SDMA allgather 中即为if (i my_pe) { continue; }sdma/main.cpp。pair_rank、left/right peer、x block / core_per_rank用于 ring、shuffle 或 allgather peer 分配。kv_shuffle 中pair_rank来自global_shuffle_tablekv_shuffle_kernel.cppallgather 中x get_core_idx / core_per_rank决定消费端从哪个 PE 拉数据。core_per_rank comm_core_num / n_pes将通信 core 分配给不同 peer。rank_start/rank_count用于测试或跨机分组只验证/处理局部连续 rank。建议把 PE 映射和 core 映射写成显式变量不把复杂表达式散落在搬运调用参数中。这样更容易检查越界、tail 和跨 PE 地址。9. tail / chunk 循环写法数据搬运通常按 chunk 切分用ceil_div(total, chunk)计算轮数。每个 core 先算base_per_core total / core_num和extra total % core_num前extra个 core 多处理 1 个元素或字节。最后一个 core 或最后一个 chunk 处理 remainder。大块搬运用while (remaining ub_size)处理整块再单独处理尾块。ping-pong buffer 用两个 event id 交替等待和下发降低搬运气泡。SDMA 示例给出了精确的按字节尾块划分实现base_per_core与extra_bytes确定每个 AIV 的data_offset前extra_bytes个 core 多搬 1 字节sdma/main.cppMTE allgather 的大路径则先while (copy_total_size copy_ub_size)处理整块再用if (copy_total_size 0) return;后的尾块put_nbi收尾allgather_kernel.cpp并以EVENT_ID0/EVENT_ID1两个 event 做 ping-pongallgather_kernel.cpp。建议在变量名中体现单位*_bytes表示字节数。*_elems或elem_count表示元素数。offset_bytes和offset_elems不混用。10. 代码规范性要求模式层详细规范以 code-style.md 为准本文只列模式层面的关键要求。生成代码时先对齐官方 CANN samples 和 SHMEM examples 的 license、utils、CMake、run script、输入/golden/output 目录布局。领域关键要求详见 code-style.mdHost 错误检查所有 ACL/SHMEM 返回值检查cleanup 阶段用cleanup_ok累计§1资源管理多资源函数统一 cleanup 或 RAII逆序释放§3main.cpp边界单进程单 PE复杂逻辑拆独立.cpp/.hgolden/checker 放 Python§5.2Device 传输跨 PE 必须aclshmem_*/aclshmemx_*禁止DataCopy远端地址§6.3NBI 完成event/quiet/barrier 显式完成路径§6.2UB/event集中命名常量ping-pong 用不同 event id§6.4signal 隔离按 rank/core/phase 隔离magic/epoch 单调递增—CMake/脚本target-scopedSCRIPT_DIR/PROJECT_ROOT定位输出按 rank 隔离§10可测试性固定 seedrank 特征输入结果文件含 PE id精度阈值与 dtype 相关§8错误检查方面仓库 utils.h 提供了现成的ACL_CHECK/ACL_CHECK_WITH_RET宏配套ERROR_LOG等日志宏SDMA 示例则示范了CHECK_RET宏与“已持有资源时不用早返回宏绕过 cleanup”的写法。同步规范上NBI 操作后必须有明确完成路径单引擎 MTE 用aclshmemx_mte_quiet()、SDMA 用aclshmemx_sdma_quiet()、多引擎混用或收尾用aclshmem_quiet()phase 间同步使用aclshmem_barrier_all()或aclshmem_barrier(team)。附从模式到可运行工程的落地路径参考 allgather 的完整工程布局main.cppallgather_kernel.cpp/.hscripts/data_gen.py、data_statistic.pyrun.shCMakeLists.txt这是纯通信算子最标准的模板。需要 SDMA 引擎时在test_set_attr之后设置attributes.option_attr.data_op_engine_type ACLSHMEM_DATA_OP_SDMA并用aclshmemx_set_qp_num(ACLSHMEM_DATA_OP_SDMA, SDMA_QP_NUM)配置 QP 数QP 数需等于 block 数 × 每 block AIV 数见 [sdma/main.cpp](https://link.gitcode.com/i/37dd7f932041a8578ef722e2e95855c0#L60-L64, L363-L364)。公共 helper 优先复用 examples/utils/utils.h 与 examples/utils/param.h 中的test_set_attr、ReadFile/WriteFile、collect_prof_data_to_csv等工具保持与官方示例一致。生成代码后对照 code-style.md 第 12 节的审查清单逐项自查NBI 同步、跨 PE 接口边界、逆序释放、输出文件含 PE id 等即可达到仓库的交付质量门槛。【免费下载链接】shmemCANN SHMEM 是面向昇腾平台的多机多卡内存通信库基于OpenSHMEM 标准协议实现跨设备的高效内存访问与数据同步。项目地址: https://gitcode.com/cann/shmem创作声明：本文部分内容由AI辅助生成（AIGC），仅供参考