CANN SHMEM 设备侧 SDMA NotifyWait 机制使用指南：显式 QP 多核并发与无 QP 单核搬运实战

CANN SHMEM 设备侧 SDMA NotifyWait 机制使用指南显式 QP 多核并发与无 QP 单核搬运实战【免费下载链接】shmemCANN SHMEM 是面向昇腾平台的多机多卡内存通信库基于OpenSHMEM 标准协议实现跨设备的高效内存访问与数据同步。项目地址: https://gitcode.com/cann/shmem本文围绕 CANN SHMEM 开源仓库中的notifywait示例README_en.md展开系统讲解设备侧 SDMASystem DMA异步数据搬运与 NotifyWait 完成通知机制如何在指定 SDMA QP 上下发 record 类型的 SQE、在 Host 侧通过aclrtWaitAndResetNotify等待通知以及显式 QP 接口与无 QP 接口的差异与选型。读完本文你将掌握notifywait示例的编译运行方法、容量与 AIV 限制评估以及基于 SDMA 的 AllGather 多核/单核两种实现路径并能在自己的算子中直接复用这套搬运 通知 等待的流水线同步模式。一、机制概览为什么需要 NotifyWait在昇腾设备侧SDMA 引擎承担 GM 与 GM 之间、GM 与 UB 之间的高速数据搬运。aclshmemx_sdma_qp_put_nbi/aclshmemx_sdma_qp_get_nbi这类接口是**非阻塞nbi non-blocking**的——函数返回只代表 SQE 已提交到 QP并不代表搬运完成。因此调用方必须通过某种手段确认数据就绪后才能安全地消费数据或复用工作区。CANN SHMEM 提供两种完成确认方式aclshmemx_sdma_qp_quiet在 AIV 内轮询 flag直到指定 QP 上的 SQE 全部完成。缺点是该 AIV 被阻塞在轮询循环中无法及时释放。NotifyWait本示例的主题在数据搬运后追加一条 record 类型的 SQE由 Host 侧aclrtWaitAndResetNotify等待通知等待完成后 AIV 已提前释放后续 kernel 可直接使用搬运结果。其核心思路可用三步概括Kernel 1stream 1搬运数据并记录通知 → Host 等待并复位 notify → Kernel 2stream 2消费数据。相比quiet的 AIV 轮询NotifyWait 让 Host 承担等待职责从而及时释放 AIV 资源见 README_en.md。二、环境要求与软件准备SDMA 功能是较新的能力有明确的软件版本门槛CANN 版本SDMA 功能需要 CANN 9.0.0 及以上版本trial 版支持中文版 README.md 中标注为 CANN 9.0.0-beta.2 及以上。需要安装Toolkit 包与ops-legacy 包两类软件包ops 包需根据硬件平台A2/A3、x86_64/aarch64选择与 toolkit 版本匹配的版本。平台限制当前暂不支持 Ascend950平台配套编译运行SDMA 写操作在 Ascend950 上不受支持见 shmem_device_sdma.h 中相关接口注释。运行环境PES 仅支持 2、4、8 卡且限定在单台机器内数据通过 TCP 环回地址默认tcp://127.0.0.1:8766进行初始化通信。三、编译与运行示例3.1 三步编译运行流程按照 README_en.md 的操作顺序在仓库根目录shmem/下编译软件包并安装bash scripts/build.sh -package ./install/*/SHMEM_1.0.0_linux-*.run --install在仓库根目录shmem/下编译所有 examplesbash scripts/build.sh -examples进入shmem/examples/notifywait目录运行 demobash run.sh -pes ${PES} -type ${TYPES}参数说明来自文档PES用于运行 demo 的设备NPU数量仅支持 2、4、8 卡限定单台机器内。TYPES传输的数据类型当前支持int、uint8、int64、fp32。3.2 运行脚本支持的全部参数run.sh 实际还支持更多命令行选项可通过-键 值的方式传入参数含义默认值-pes进程/PE 数量2若-gnpus大于该值会自动收敛为-pes值-type数据类型int / uint8 / int64 / fp32int-ipport初始化通信 IP 端口tcp://127.0.0.1:8766-fpe起始 PE 编号0-gnpus使用的 NPU 数量对应每卡一个进程8-fnpu起始 NPU 编号0-pe_tablePE 映射表空脚本内部会导出SHMEM_UID_SESSION_ID127.0.0.1:8899作为 UID 会话标识并将${PROJECT_ROOT}/build/lib与${ASCEND_HOME_PATH}/lib64加入LD_LIBRARY_PATH随后为每个 NPU 拉起一个后台进程等待全部进程结束后统一返回退出码。四、容量与 AIV 限制运行前必读文档明确给出了本示例的资源预算运行前需要据此评估硬件是否满足条件对称内存容量示例申请128M * sizeof(T)字节的对称空间其中输入区与结果区各需PES * 8M * sizeof(T)字节。以 8 卡、int4 字节为例即每个 PE 需约128M * 4B 512MB对称空间。文档支持矩阵为 2、4、8 卡实际可用卡数还需满足对称空间与运行环境的容量条件。在 main.cpp 中可看到symmetric_elements 128 * 1024 * 1024、trans_size 8 * 1024 * 1024的定义。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空间足够。AIV/QP 数量当前 kernel 启动 20 个 block每个 block 含 2 个 subblockAIV实际使用 40 个 AIV/QP底层基础设施和 notify 数组已按最多72 个 AIV/QP预留。上报 vector core 数超过 72 的设备当前返回不支持。对应的常量定义位于 main.cppSDMA_AIVS_PER_BLOCK 2、SDMA_BLOCK_NUM 20、SDMA_QP_NUM SDMA_BLOCK_NUM * SDMA_AIVS_PER_BLOCK 40并在初始化时通过aclshmemx_set_qp_num(ACLSHMEM_DATA_OP_SDMA, SDMA_QP_NUM)配置 SDMA 通道数main.cpp同时将attributes.option_attr.data_op_engine_type置为ACLSHMEM_DATA_OP_SDMA以启用 SDMA 数据通路。五、NotifyWait 三步用法详解5.1 用法示例notifywait 机制三步流程示意整个机制分为三个步骤对应 README_en.md 中的伪代码// 步骤 1 // stream 1 上的 kernel 1调用显式 QP 的 SDMA 接口搬运数据并追加 aclshmemx_sdma_qp_notify_record // 步骤 2 // HostaclrtWaitAndResetNotify(notify_id, stream2, 0) // 步骤 3 // stream 2 上的 kernel 2使用 SDMA 搬运好的数据5.2 原理record 类型 SQE 与 Host 侧等待在aclshmemx_sdma_qp_notify_record中会向选定的 STARS QP下发一条 record 类型的 SQE。由于 SQE 在 QP 内保序这条 record 通知会排在之前提交的所有搬运 SQE 之后因此 Host 等待到该通知时即可确认此前该 QP 上的搬运全部完成。随后 Host 再调度后续 kernel天然形成了搬运完成 → 数据可用的依赖关系。相比aclshmemx_sdma_qp_quiet依赖AIV 轮询 flag的方式NotifyWait 将等待从 AIV 转移到 Host从而及时释放 AIV 资源让 AIV 可以立即投入其他计算任务。设备侧实现可在 shmem_device_sdma.hpp 中看到aclshmemi_stars_submit_notify_record通过notify_addr[qp_idx]定位到对应 QP 的通知地址并填充 record SQE无 QP 变体则固定向 QP 0 追加通知aclshmemi_stars_submit_notify_record(ub_tensor, sync_id, 0)。Host 侧的 notify 对象由 SDMA 传输管理模块在初始化阶段创建通过aclrtCreateNotify为每个 QP 创建 notify、aclrtGetNotifyId获取 notify ID并在结束时用aclrtDestroyNotify销毁见 device_sdma_transport_manager.cpp。示例中 Host 通过g_state_host.notify_arr[i]数组按 QP 索引一一对应等待main.cpp。六、显式 QP 的 SDMA 接口多核 AllGather 实现6.1 接口形态显式 QP 接口相比无 QP 接口多一个qp_idx参数且提供__ubuf__指针与AscendC::GlobalTensor/LocalTensor两套重载。核心接口签名完整声明见 shmem_device_sdma.htemplate typename T 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 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); template typename T void aclshmemx_sdma_qp_notify_record(__ubuf__ T* buf, uint32_t ub_size, uint32_t qp_idx, uint32_t sync_id);dst/src对称地址会在指定 PE 上做地址翻译或本设备 GM 地址dst/src需落在同一个对称分配块内。buf/ub_sizeUB 工作区地址必须64 字节对齐大小至少 64 字节。elem_size搬运的元素个数elem_size * sizeof(T)不得超过UINT32_MAX字节。pe对端 PE必须处于已初始化的 PE 范围内。qp_idxSDMA QP 索引必须小于已配置的 SDMA 通道数QP 索引与 block 索引相互独立。sync_id流水线同步使用的硬件事件 ID。6.2 多核 allgather_sdma 内核main.cpp 中的allgather_sdma内核演示了显式 QP 的标准用法每个 AIV 根据GetBlockIdx()计算出自己负责的连续数据区间base_per_core/extra_bytes按元素均摊切分保证各 AIV 负载均衡循环向除自身外的每个 PE 调用aclshmemx_sdma_qp_put_nbi或get_nbi每个 AIV 使用与自身编号相同的 QP收发数据全部搬运提交后调用aclshmemx_sdma_qp_notify_record在本 QP 上追加通知。Host 侧对 40 个 QP 逐个等待for (int i 0; i total_block_num * sub_block_num; i) { CHECK_RET(aclrtWaitAndResetNotify(g_state_host.notify_arr[i], g_state_host.default_stream, 0)); }需要说明的是接口的完成语义是函数正常返回仅代表请求已提交不代表搬运完成。在显式 QP 场景下不能用aclshmemx_sdma_quiet它只排空 QP 0来等待qp_idx 0的请求而应使用同 QP 的aclshmemx_sdma_qp_quiet或如本例一样追加aclshmemx_sdma_qp_notify_record并在 Host 等待shmem_device_sdma.h 的接口注释对此有明确说明。示例还提供了allgather_sdma_tensor内核main.cpp展示GlobalTensor/LocalTensor重载的等价用法。七、不带 QP 的 SDMA 接口单核 AllGather 实现除显式 QP 接口外示例还演示了不带 QP 的 SDMA 接口。二者接口形态接近区别是不带 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 执行其余 AIV 直接返回对本 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);与显式 QP 接口的对比对比项不带 QP 接口显式 QP 接口使用的 QP固定 QP 0通过qp_idx指定可用满已创建的全部 QP执行方式单 AIV 执行多 AIV 并发每个 AIV 使用独立 QP数据切分无需切分整块搬运需按 AIV 切分数据Host 等待仅notify_arr[0]每个 QP 各等待一次 notify适用场景单核简单收发、快速验证多核并发、带宽敏感场景八、运行验证与结果解读运行run.sh时demo 的执行顺序是固定的见 main.cpp显式 QP 多核 AllGatherallgather_sdma内核搬运 → Host 等待 40 个 notify →aclshmem_barrier_all()同步 → 用 MTE 将结果拷入结果区 → Host 校验控制台打印after notify_wait段的结果无 QP 单核 AllGatherallgather_sdma_noqp内核整块搬运 → Host 等待notify_arr[0]→aclshmem_barrier_all()同步 → 拷贝并校验控制台打印after sdma_put_nbi (no QP)段的结果。两段校验逻辑相同逐元素比对结果区中每个 PE 贡献的数据是否等于num10 inum10 10并统计异常值个数若异常值均为 0则说明对应阶段的搬运与通知机制工作正常。最终每个 PE 打印[SUCCESS] demo run success in pe N表示整体通过。九、总结NotifyWait 机制为 CANN SHMEM 设备侧 SDMA 异步搬运提供了一条搬运即通知、Host 等待、流间接力的同步链路aclshmemx_sdma_qp_notify_record在指定 QP 追加 record SQEHost 以aclrtWaitAndResetNotify等待避免了quiet方案中 AIV 轮询 flag 的资源占用。notifywait示例同时给出了两种工程范式——显式 QP 的多核并发 AllGather每 AIV 一 QP、按 AIV 切分数据与无 QP 的单核快速验证 AllGather固定 QP 0、整块搬运配合本文给出的容量评估方法可直接迁移到其他基于 SDMA 的集合通信或流水线算子设计中。进一步阅读完整的中文说明见 README.md接口头文件见 shmem_device_sdma.h设备侧实现见 shmem_device_sdma.hppHost 侧 notify 生命周期管理见 device_sdma_transport_manager.cpp。【免费下载链接】shmemCANN SHMEM 是面向昇腾平台的多机多卡内存通信库基于OpenSHMEM 标准协议实现跨设备的高效内存访问与数据同步。项目地址: https://gitcode.com/cann/shmem创作声明：本文部分内容由AI辅助生成（AIGC），仅供参考