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）展开，系统讲解设备侧 SDMA（System 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 1（stream 1）搬运数据并记录通知 → Host 等待并复位 notify → Kernel 2（stream 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/下编译所有 examples：

bash scripts/build.sh -examples

进入shmem/examples/notifywait目录运行 demo：

bash 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 / fp32	`int`
`-ipport`	初始化通信 IP 端口	`tcp://127.0.0.1:8766`
`-fpe`	起始 PE 编号	`0`
`-gnpus`	使用的 NPU 数量（对应每卡一个进程）	`8`
`-fnpu`	起始 NPU 编号	`0`
`-pe_table`	PE 映射表	空

脚本内部会导出SHMEM_UID_SESSION_ID=127.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 卡、int（4 字节）为例，即每个 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 个 subblock（AIV），实际使用 40 个 AIV/QP；底层基础设施和 notify 数组已按最多72 个 AIV/QP预留。上报 vector core 数超过 72 的设备当前返回"不支持"。

对应的常量定义位于 main.cpp：SDMA_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： // Host：aclrtWaitAndResetNotify(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.h）：

template <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_size：UB 工作区，地址必须64 字节对齐，大小至少 64 字节。
elem_size：搬运的元素个数，elem_size * sizeof(T)不得超过UINT32_MAX字节。
pe：对端 PE，必须处于已初始化的 PE 范围内。
qp_idx：SDMA 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 record，Host 侧等待 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 个 notify（QP 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 多核 AllGather：allgather_sdma内核搬运 → Host 等待 40 个 notify →aclshmem_barrier_all()同步 → 用 MTE 将结果拷入结果区 → Host 校验，控制台打印after notify_wait段的结果；
无 QP 单核 AllGather：allgather_sdma_noqp内核整块搬运 → Host 等待notify_arr[0]→aclshmem_barrier_all()同步 → 拷贝并校验，控制台打印after sdma_put_nbi (no QP)段的结果。

两段校验逻辑相同：逐元素比对结果区中每个 PE 贡献的数据是否等于num10 + i（num10 = 10），并统计异常值个数；若异常值均为 0，则说明对应阶段的搬运与通知机制工作正常。最终每个 PE 打印[SUCCESS] demo run success in pe N表示整体通过。

九、总结

NotifyWait 机制为 CANN SHMEM 设备侧 SDMA 异步搬运提供了一条"搬运即通知、Host 等待、流间接力"的同步链路：aclshmemx_sdma_qp_notify_record在指定 QP 追加 record SQE，Host 以aclrtWaitAndResetNotify等待，避免了quiet方案中 AIV 轮询 flag 的资源占用。notifywait示例同时给出了两种工程范式——显式 QP 的多核并发 AllGather（每 AIV 一 QP、按 AIV 切分数据）与无 QP 的单核快速验证 AllGather（固定 QP 0、整块搬运），配合本文给出的容量评估方法，可直接迁移到其他基于 SDMA 的集合通信或流水线算子设计中。

进一步阅读：完整的中文说明见 README.md，接口头文件见 shmem_device_sdma.h，设备侧实现见 shmem_device_sdma.hpp，Host 侧 notify 生命周期管理见 device_sdma_transport_manager.cpp。

创作声明：本文部分内容由AI辅助生成（AIGC），仅供参考