CANN SHMEM NotifyWait 机制实战基于 SDMA QP 的异步搬运与 Host 通知同步【免费下载链接】shmemCANN SHMEM 是面向昇腾平台的多机多卡内存通信库基于OpenSHMEM 标准协议实现跨设备的高效内存访问与数据同步。项目地址: https://gitcode.com/cann/shmem导读本文基于 CANN SHMEM 开源仓库中的 notifywait 示例examples/notifywait/README.md及其配套源码系统讲解 SHMEM 在昇腾平台上基于 SDMAStream Direct Memory Access引擎的NotifyWait 异步通信机制如何在指定 QPQueue Pair上以 record 类型 SQE 记录完成事件如何通过 Host 侧aclrtWaitAndResetNotify等待通知以及带 QP 与不带 QP 两类 SDMA 接口的差异与选型。读完本文你将掌握 notifywait 示例的编译运行方法、显式 QP 多核 allgather 与单核 no-QP allgather 两种编程范式并能依据容量与 AIV 约束正确评估自己的用例是否可行。平台限制提示当前仓库暂不支持在 Ascend950 平台配套编译运行 NotifyWait 相关 SDMA 写操作SDMA write 不支持 Ascend950。一、环境要求与准备NotifyWait 依赖 SDMA 功能需要CANN 9.0.0-beta.2 及以上版本支持。请参考 CANN 版本说明 下载并安装对应版本的 toolkit 包使能 SDMA 时还需要安装与 toolkit 版本和设备类型匹配的 ops 包。从 quickstart.md 版本矩阵 可以看到 SDMA 的版本依赖关系驱动固件CANN 版本SDMA 可用能力附加依赖Ascend HDK 25.5.1 及以上9.0.0 及以上MTE / RDMA / SDMAD2H、H2D、D2rH/rH2D 的 SDMA 仅支持 A3需下载与 toolkit 版本和设备类型匹配的 ops 包Ascend HDK 25.5.6 及以上9.1.0 及以上SDMAAscend950 仅支持 CMOAscend950 下使能 SDMA 需额外下载对应的 toolkit 和 Ascend HDK 包并安装与 toolkit 版本和设备匹配的 ops 包以 CANN 9.0.0、x86_64 为例各 SoC 对应的 ops 包文件名形如Ascend-cann-910b-ops_9.0.0_linux-x86_64.run、Ascend-cann-A3-ops_9.0.0_linux-x86_64.run、Ascend-cann-950-ops_9.0.0_linux-x86_64.run具体可查阅 quickstart.md。二、example 编译与执行说明2.1 编译安装软件包在shmem/根目录编译软件包并安装bash scripts/build.sh -package ./install/*/SHMEM_1.0.0_linux-*.run --install2.2 编译 examplesbash scripts/build.sh -examples编译由 examples/CMakeLists.txt 统一组织notifywait 示例通过aclshmem_add_fusion_example(notifywait main.cpp)接入构建体系见 examples/notifywait/CMakeLists.txt产物为build/bin/notifywait。2.3 运行 demo在shmem/examples/notifywait目录执行bash run.sh -pes ${PES} -type ${TYPES}参数说明-pes指定用于运行的设备NPU数量仅支持 2、4、8 卡限定单台机器内。-type指定传输数据类型当前支持int、uint8、int64、fp32。run.sh 还提供了更完整的可调参数便于在不同环境上复现参数默认值含义-ipporttcp://127.0.0.1:8766引导阶段 TCP 地址与端口-pes2参与的 PE进程数量-fpe0起始 PE 编号-gnpus8大于-pes时自动收敛为-pes的值启动的 NPU 进程数-fnpu0起始 NPU 编号-typeint传输数据类型-pe_table空可选 PE 映射表脚本会导出SHMEM_UID_SESSION_ID127.0.0.1:8899作为 UID 引导的会话标识并设置LD_LIBRARY_PATH指向build/lib与ASCEND_HOME_PATH/lib64随后按-gnpus数量并行拉起build/bin/notifywait进程逐个wait回收并汇总退出码。2.4 容量与 AIV 限制用例申请128M * sizeof(T)字节对称空间其中输入和结果各需要PES * 8M * sizeof(T)字节经过容量评估文档支持矩阵为 2、4、8 卡。实际可用卡数还需满足对称空间和运行环境的容量条件。SDMA 共享 workspace 为 28 KiB。按 A5 最大 72 个 AIV/QP 计算notify ID 和三组 flag 区需要14 KiB 72 * 4 B 3 * 72 * 64 B 28,448 B剩余 224 B空间足够。当前 kernel 启动 20 个 block按 2 个 subblock 计算实际使用 40 个 AIV/QP底层基础设施和 notify 数组已按最多 72 个 AIV/QP 预留。超过 72 个 vector core 的设备当前返回不支持。以上约束与 main.cpp 中的常量一一对应每个 block 含 2 个 AIVSDMA_AIVS_PER_BLOCK 2启动 20 个 blockSDMA_BLOCK_NUM 20因此需要的 QP 数为SDMA_QP_NUM 20 * 2 40并在初始化阶段通过aclshmemx_set_qp_num(ACLSHMEM_DATA_OP_SDMA, SDMA_QP_NUM)显式配置见 main.cpp。该 Host 侧接口声明于 include/host/init/shmem_host_init.h。三、NotifyWait 用法说明3.1 用法示例整体流程如下对应图中三个步骤// 步骤1 // stream1上的kernel1调用显式QP的sdma接口搬运数据 aclshmemx_sdma_qp_notify_record // 步骤2 // host: aclrtWaitAndResetNotify(notify_id, stream2, 0) // 步骤3 // stream2上的kernel2使用sdma搬运好的数据核心思路是在数据搬运所在的 stream 上以“搬运 record”的方式把完成事件记录下来Host 侧在目标 stream 上等待该 notify 被记录收到通知后后续 kernel 才在该 stream 上消费数据从而以“事件通知”替代“轮询等待”实现流间依赖的同时释放设备侧计算资源。3.2 用法说明aclshmemx_sdma_qp_notify_record会在指定 QP 上下发 record 类型的 SQE后续在 Host 侧等待 notify 记录完成再继续执行后续 kernel。相比aclshmemx_sdma_qp_quiet使用 AIV 轮询 flag 的方式可及时释放 AIV 资源——AIV 只需提交 record SQE 即可继续干别的完成状态的感知完全交给 Host 侧事件通知机制。从实现看notify record 的下发最终落到 STARS 通道的 SQE 提交逻辑src/device/gm2gm/engine/shmem_device_sdma.hpp根据qp_idx从通道基址取到对应的stars_channel_info_t并断言qp_idx 通道数从共享 notify 地址区按notify_addr[qp_idx]取出该 QP 专属的 notify ID在 SQ 环形队列尾部填写 record 类型的stars_notify_sqe_tnotify_id写入 SQE见 shmem_device_sdma.hpp通过set_value回写 SQ tail 寄存器完成“摇铃”提交。每个 QP 对应独立的 notify ID因此 Host 侧可以按 QP 逐一等待这正是“多 QP 各自记录、各自等待”并行模型的底层基础。3.3 Host 侧等待与校验在 main.cpp 中显式 QP 版本的 Host 等待方式为对 40 个 QPtotal_block_num * sub_block_num逐一执行aclrtWaitAndResetNotify(g_state_host.notify_arr[i], g_state_host.default_stream, 0);收到全部 notify 后调用aclshmem_barrier_all()同步再启动copy_demo将结果整块拷贝到校验区ptr_A最后用aclrtSynchronizeStream收尾并从设备回读校验。kernel 内部则对每个远端 PE 依次调用aclshmemx_sdma_qp_put_nbi或aclshmemx_sdma_qp_get_nbi并按 AIV 编号切分数据最后提交aclshmemx_sdma_qp_notify_record见 main.cpp。四、显式 QP 接口签名与参数详解以下接口声明与参数约束来自 include/device/gm2gm/engine/shmem_device_sdma.h两者均提供指针与GlobalTensor/LocalTensor两种重载。4.1 显式 QP 的异步 PUT / GETtemplate typename T ACLSHMEM_DEVICE void aclshmemx_sdma_qp_put_nbi( __gm__ T* dst, __gm__ T* src, __ubuf__ T* buf, uint32_t ub_size, uint32_t elem_size, int pe, uint32_t qp_idx, uint32_t sync_id); template typename T ACLSHMEM_DEVICE void aclshmemx_sdma_qp_get_nbi( __gm__ T* dst, __gm__ T* src, __ubuf__ T* buf, uint32_t ub_size, uint32_t elem_size, int pe, uint32_t qp_idx, uint32_t sync_id);参数说明参数含义与约束dst/src对称内存地址PUT 的dst、GET 的src会按pe做地址翻译与本地 GM 地址buf本地 UB workspace地址必须 64 字节对齐ub_sizeUB workspace 大小至少 64 字节elem_size搬运的元素个数要求elem_size * sizeof(T)不超过UINT32_MAX字节pe目标/源 PE必须在已初始化的 PE 范围内qp_idxSDMA QP 索引必须小于已配置的 SDMA 通道数QP 索引与 block 索引相互独立sync_id用于流水线同步的硬件事件 ID完成语义函数正常返回只代表 SQE 已提交、并非搬运完成且aclshmemx_sdma_quiet只排空 QP 0对qp_idx 0的请求无效。显式 QP 请求必须使用同一qp_idx的aclshmemx_sdma_qp_quiet或aclshmemx_sdma_qp_notify_record来完成。4.2 显式 QP 的 notify recordtemplate typename T ACLSHMEM_DEVICE void aclshmemx_sdma_qp_notify_record( AscendC::LocalTensorT buf, uint32_t qp_idx, uint32_t sync_id); // 指针重载 template typename T ACLSHMEM_DEVICE void aclshmemx_sdma_qp_notify_record( __ubuf__ T* buf, uint32_t ub_size, uint32_t qp_idx, uint32_t sync_id);notify record 在同一 QP上按 FIFO 顺序排在上一个 SQE即对应的put_nbi/get_nbi之后因此 Host 收到 notify 即可确认该 QP 上此前的搬运全部完成。源码中aclshmemx_sdma_qp_notify_record通过aclshmemi_stars_submit_notify_record(ub_tensor, sync_id, qp_idx)携带qp_idx下发见 shmem_device_sdma.hpp每个 QP 使用独立的notify_addr[qp_idx]。五、不带 QP 的 SDMA 接口单核接口除显式 QP 接口外notifywait 示例还演示了不带 QP 的 SDMA 接口见 main.cpp 中allgather_sdma_noqp内核。二者接口形态接近区别是不带 QP 的接口固定使用 QP 0无需传入qp_idx属于单核单 AIV接口// 异步搬运固定使用 QP 0 template typename T void aclshmemx_sdma_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); // 在 QP 0 上追加 notify recordHost侧等待 notify_arr[0] 即可 template typename T void aclshmemx_sdma_notify_record(__ubuf__ T* buf, uint32_t ub_size, uint32_t sync_id);对应实现为main.cpp中的allgather_sdma_noqp内核仅由 0 号 AIV 执行GetBlockIdx() ! 0直接返回对本 PE 数据整块搬运无需按 AIV 切分并在 QP 0 上记录 notify// kernel内仅0号AIV执行 aclshmemx_sdma_put_nbi(dst, src, tmp_buff, ub_size, size, pe, EVENT_ID0); aclshmemx_sdma_notify_record(tmp_buff, ub_size, EVENT_ID0); // host侧只等待1个notifyQP 0对应notify_arr[0] aclrtWaitAndResetNotify(g_state_host.notify_arr[0], stream, 0);Host 侧只启动 1 个 block、等待notify_arr[0]一个 notify 即可见 main.cpp。从源码看不带 QP 的notify_record内部等价于qp_idx 0的提交aclshmemi_stars_submit_notify_record(ub_tensor, sync_id, 0)见 shmem_device_sdma.hpp。六、带 QP 与不带 QP 接口的对比与选型对比项不带 QP 接口显式 QP 接口使用的 QP固定 QP 0通过qp_idx指定可用满已创建的全部 QP执行方式单 AIV 执行多 AIV 并发每个 AIV 使用独立 QP数据切分无需切分整块搬运需按 AIV 切分数据Host 等待仅notify_arr[0]每个 QP 各等待一次 notify适用场景单核简单收发、快速验证多核并发、带宽敏感场景此外还有一处易被忽略的差异不带 QP 接口的完成原语是aclshmemx_sdma_quiet只排空 QP 0而显式 QP 接口必须配对使用同一qp_idx的aclshmemx_sdma_qp_quiet或aclshmemx_sdma_qp_notify_record二者不可混用。接口头文件include/device/gm2gm/engine/shmem_device_sdma.h对此均有明确注释。选型建议需要快速验证、数据量小、单核搬运即可满足带宽需求时优先用不带 QP 接口代码简单、无需切分追求多核并发吞吐、带宽敏感、或需要与 AIV 数量对齐做流水线切分时使用显式 QP 接口并按“每个 AIV 一个独立 QP、Host 按 QP 逐次等待”的模型组织代码。七、运行结果验证运行run.sh时demo 先执行显式 QP 的多核 allgather 并校验再执行不带 QP 的单核 allgather 并校验控制台输出两段校验结果Pe {id} AllGather result in ptr_A after notify_wait显式 QP 多核 allgather 的结果统计Pe {id} AllGather result in ptr_A after sdma_put_nbi (no QP)不带 QP 单核 allgather 的结果统计。校验逻辑位于 main.cpp 的check_resultlambda将设备侧结果ptr_A回读到 Host与期望值num10 i即10 PE 编号逐一比对并输出unexpected values个数。只有当每个 PE 的输入数据都通过 SDMA 正确广播到所有 PE、且 notify 等待时机正确时unexpected 数量才为 0。每个 PE 运行结束还会打印[SUCCESS] demo run success in pe {my_pe}run.sh汇总所有进程退出码任一进程失败即返回非 0。总结NotifyWait 机制为 CANN SHMEM 的 SDMA 数据面提供了一条“设备侧提交 record SQE → Host 侧事件等待 → 后续 kernel 按依赖消费”的高效同步路径。通过 notifywait 示例你可以同时掌握两种编程范式多 AIV 显式 QP 的高并发 allgather数据按 AIV 切分、每个 QP 独立 notify以及单 AIV QP 0 的极简收发整块搬运、仅等待一个 notify。使用时请务必核对 CANN 版本与 ops 包、对称空间容量、AIV/QP 数量≤72以及当前不支持 Ascend950 的平台约束即可在自己的昇腾环境上快速复现并扩展该模式。参考文件示例文档examples/notifywait/README.md示例主程序examples/notifywait/main.cpp运行脚本examples/notifywait/run.sh构建接入examples/notifywait/CMakeLists.txtSDMA 设备侧接口声明include/device/gm2gm/engine/shmem_device_sdma.hSDMA 设备侧实现notify record SQE 提交src/device/gm2gm/engine/shmem_device_sdma.hppQP 数量配置接口include/host/init/shmem_host_init.hCANN 版本说明docs/quickstart.md【免费下载链接】shmemCANN SHMEM 是面向昇腾平台的多机多卡内存通信库基于OpenSHMEM 标准协议实现跨设备的高效内存访问与数据同步。项目地址: https://gitcode.com/cann/shmem创作声明：本文部分内容由AI辅助生成（AIGC），仅供参考