ARTICLE DETAIL

资讯详情

深耕网站建设、视觉设计与SEO优化的一线实战洞察。

CANN SHMEM SDMA 数据搬运接口实战指南:put/get 编程模型、显式 QP 多核并发与同步机制

CANN SHMEM SDMA 数据搬运接口实战指南:put/get 编程模型、显式 QP 多核并发与同步机制 CANN SHMEM SDMA 数据搬运接口实战指南put/get 编程模型、显式 QP 多核并发与同步机制【免费下载链接】shmemCANN SHMEM 是面向昇腾平台的多机多卡内存通信库基于OpenSHMEM 标准协议实现跨设备的高效内存访问与数据同步。项目地址: https://gitcode.com/cann/shmem本文基于 CANN SHMEM 开源仓库中的 examples/sdma 示例系统讲解面向昇腾平台的 SDMA put/get 数据搬运能力。读完本文你将掌握 SDMA 接口的适用平台与版本前提、示例工程的编译运行方法、aclshmemx_sdma_put_nbi与显式 QP 接口aclshmemx_sdma_qp_put/get_nbi的参数语义与典型调用模式以及非阻塞接口的两种可靠完成quiet / notify机制可直接迁移到自己的多核数据搬运算子中。概述SDMA 在 SHMEM 中的角色CANN SHMEM 是面向昇腾平台的多机多卡内存通信库基于 OpenSHMEM 标准协议实现跨设备的内存访问与数据同步。设备侧 RMARemote Memory Access存在多种底层数据搬运引擎data_op_engine_type_t见 include/host_device/shmem_common_types.h其中ACLSHMEM_DATA_OP_MTE0x01MTE 引擎SDMA 之外的另一类搬运通道ACLSHMEM_DATA_OP_SDMA0x02本文主题SDMASynchronous Direct Memory Access引擎ACLSHMEM_DATA_OP_UDMA0x08UDMA 引擎Ascend950 上可用。SDMA put/get 接口用于在对称内存symmetric memory与本地设备内存之间进行 read/write 数据搬运需要 CANN 9.0.0-beta.2 及以上版本支持。与更高层 RMA 接口的区别在于SDMA 接口面向算子内device-side直接下发数据搬运请求需要调用方显式提供 UB workspace 与同步 ID并自行负责任务的完成确认。平台限制A2/A3 平台支持 SDMA put/getAscend950 仅支持 SDMA get不支持 SDMA put。该限制同时体现在设备侧头文件 include/device/gm2gm/engine/shmem_device_sdma.h 的接口注释中“SDMA write is not supported on Ascend950”。环境要求与准备SDMA put/get 接口需要CANN 9.0.0-beta.2 及以上版本支持可用于 read/write 数据搬运。请参考 CANN 版本说明 下载并安装对应版本的 toolkit 包使能 SDMA 时还需要安装与 toolkit 版本和设备类型匹配的ops 包。从源码编译使用示例前还需满足昇腾设备环境Atlas 系列推理/训练产品已配置ASCEND_HOME_PATH等环境变量示例 run.sh 会使用ASCEND_HOME_PATH/lib64加入LD_LIBRARY_PATH。支持设备SDMA put/get 接口在Atlas 200I A2/A3、Atlas 300T A2/A3等 A2/A3 平台可用Ascend950 仅支持 SDMA getSDMA put 不可用。规划多机多卡拓扑时需要先确认目标设备属于哪一平台再决定能否使用 put 方向。example 使用方式examples/sdma示例完成两件事先用单核 QP0 的普通 put 接口验证“向下一 PE 写完整 segment、从上一 PE 校验数据”再运行显式 QP 的多核 allgather可切换 put/get 两种方向。整体流程如下1. 编译软件包并安装在仓库根目录下称shmem/执行bash scripts/build.sh -package ./install/*/SHMEM_1.0.0_linux-*.run --install2. 编译 examples在shmem/目录下编译全部示例包含 sdmabash scripts/build.sh -examplesSDMA 示例的 CMake 接入非常简单通过aclshmem_add_fusion_example(sdma main.cpp)将 main.cpp 注册为名为sdma的 fusion example见 examples/sdma/CMakeLists.txt编译产物为build/bin/sdma。3. 运行 demo在shmem/examples/sdma目录执行bash run.sh -pes ${PES} -type ${TYPES}参数说明参数含义取值范围/说明PES用于运行的设备NPU数量仅支持 2、4、8 卡限定单台机器内TYPES传输数据类型支持int、uint8、int64、fp32除上述两个最常用参数外run.sh 还支持以下可选参数参数含义默认值-ipportbootstrap 通信 IP:端口tcp://127.0.0.1:8766-fpe起始 PE 编号0-gnpus参与运行的 NPU 总数8当大于 PES 时自动收敛为 PES-fnpu起始 NPU 编号0-pe_tablePE 映射表空运行脚本会导出SHMEM_UID_SESSION_ID127.0.0.1:8899随后按GNPU_NUM个进程并行拉起build/bin/sdma每个进程传入参数依次为PES idx IPPORT GNPU_NUM FIRST_PE FIRST_NPU TEST_TYPE PE_TABLE对应 main.cpp 的 argv 解析并等待全部进程结束、汇总返回码。main函数中还会依据data_type将数据类型分发到不同的模板实例int→int、uint8→uint8_t、int64→int64_t、fp32→float其他类型直接报错退出见 main.cpp。SDMA 接口使用说明SDMA 设备侧接口均声明于 include/device/gm2gm/engine/shmem_device_sdma.h完整实现位于src/device/gm2gm/engine/shmem_device_sdma.hpp。所有模板接口的通用约束包括buf为本地 UB workspace地址必须 64 字节对齐大小至少 64 字节sync_id为硬件事件 ID用于流水线同步elem_size * sizeof(T)不得超过UINT32_MAX字节传输的源/目标范围必须完整落在同一个对称内存分配内接口均为非阻塞正常返回只代表请求已提交不代表传输完成。aclshmemx_sdma_put_nbi单核固定 QP0 的普通 put普通 put 接口用于单核提交并固定使用 QP 0原型指针版本为ACLSHMEM_DEVICE 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);示例 kernel 如下template typename T __global__ __aicore__ void sdma_put_single_qp(GM_ADDR gva, int elem_size, int target_pe) { if ASCEND_IS_AIV { if (AscendC::GetBlockIdx() ! 0) { return; } __ubuf__ T* tmp reinterpret_cast__ubuf__ T*(uint64_t(1024)); __gm__ T* src reinterpret_cast__gm__ T*(gva) aclshmem_my_pe() * elem_size; aclshmemx_sdma_put_nbi(src, src, tmp, 64, elem_size, target_pe, EVENT_ID0); aclshmemx_sdma_quiet(tmp, 64, EVENT_ID0); } }要点通过if ASCEND_IS_AIV限定 AIV 核执行并用AscendC::GetBlockIdx() ! 0保证只有一个核提交请求避免多核并发共享 QP 0dst传入的是本地对称地址接口内部会将其换算为target_pe上的对称地址调用后立即跟随aclshmemx_sdma_quiet(tmp, 64, EVENT_ID0)等待 QP 0 上已提交的 SDMA 操作完成然后才允许读dst或复用 workspace本用例先通过该 kernel 让每个 PE 使用 QP 0 向下一 PE 写入一个完整 segment并在 host 侧校验上一 PE 写入的数据见 main.cpp随后再运行显式 QP 的多核 allgather。重要结论普通接口不能由多个 AIV 并发共享 QP 0多核场景应使用aclshmemx_sdma_qp_put_nbi。aclshmemx_sdma_qp_put_nbi显式 QP 的多核 put以指针类型参数接口为例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)接口功能将当前 PE 本地src的数据传输至目标 PEpe的远端对称地址dst传输elem_size个元素。参数名含义dst目标 PEpe上写入数据的远端对称地址src当前 PE 本地的源地址buf缓冲区地址UB workspace64 字节对齐≥64 字节ub_size缓冲区大小elem_size元素个数pe目标 PEqp_idx当前 AIV 使用的 SDMA QP一般取GetBlockIdx()sync_id同步 ID与普通接口相比显式 QP 版本额外增加qp_idx参数将请求提交到指定的 SDMA QP从而支持多个 AIV 各自绑定独立 QP 并发提交。设备侧头文件进一步说明qp_idx必须小于配置的 SDMA channel 数且QP 索引与 block 索引相互独立见 shmem_device_sdma.h。aclshmemx_sdma_qp_get_nbi显式 QP 的多核 get以指针类型参数接口为例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)接口功能从目标 PEpe的远端对称地址src拉取数据写入当前 PE 本地dst传输elem_size个元素。参数名含义dst当前 PE 本地的写入地址src目标 PEpe上读取数据的远端对称地址buf缓冲区地址UB workspace64 字节对齐≥64 字节ub_size缓冲区大小elem_size元素个数pe目标 PEqp_idx当前 AIV 使用的 SDMA QP一般取GetBlockIdx()sync_id同步 ID指针与 Tensor 两套重载上述接口均提供两套重载裸指针版本__gm__ T*/__ubuf__ T*需显式传ub_sizeTensor 版本AscendC::GlobalTensorT/AscendC::LocalTensorT无需ub_size缓冲区容量由 LocalTensor 的dataLen描述。示例中allgather_sdma裸指针版与allgather_sdma_tensorTensor 版展示了完全等价的两套写法后者通过tmp_local.address_.logicPos VECOUT; bufferAddr ub_offset; dataLen ub_size手工构造一个 64 字节的临时 LocalTensor见 main.cpp。Tensor 版接口对数据长度描述为元素个数elem_size因此需要先按元素数构造 GlobalTensorSetGlobalBuffer(addr, count)。多核 allgather 实现剖析示例在验证完单核 QP0 put 后运行显式 QP 的多核 allgather。核心思路以allgather_sdma为例见 main.cpp数据切分comm_block_dim GetBlockNum() * GetSubBlockNum()为 AIV 总数每个 AIV 将本 PE 的elem_size个元素均分base_per_core 余数extra_bytes分配算得自己的data_offset与base_per_core环形收发对每个非本 PE 的iput 模式aclshmemx_sdma_qp_put_nbi(gva data_length * my_pe data_offset, ..., i, cur_block_idx, EVENT_ID0)把本 PE segment 写给目标 PEget 模式aclshmemx_sdma_qp_get_nbi(gva data_length * i data_offset, ..., i, cur_block_idx, EVENT_ID0)从目标 PE 拉取 segment统一收尾循环结束后调用aclshmemx_sdma_qp_quiet(tmp_buff, ub_size, cur_block_idx, EVENT_ID0)等待本 AIV 对应 QP上全部请求完成。这里每个 AIV 使用cur_block_idxAIV 级全局索引取值0 ~ block_num*2-1作为自己的qp_idx从而让 40 个 AIV 分布在 40 条 SDMA QP 上互不阻塞——这正是显式 QP 接口多核并发的价值所在。注意事项与资源约束对称空间与卡数限制当前用例申请128M * sizeof(T)字节对称空间每个 PE 占用16M * sizeof(T)字节见 main.cpp。文档支持矩阵为2、4、8 卡实际可用卡数还需满足对称空间和运行环境容量条件卡数越多总对称空间越大。底层 SDMA 最多支持72 个 AIV/QP。显式 QP 数量的配置aclshmemx_set_qp_num如需限制显式 QP 数量可在aclshmemx_init_attr之前调用aclshmemx_set_qp_num(ACLSHMEM_DATA_OP_SDMA, qp_num);该接口声明于 include/host/init/shmem_host_init.h语义要点qp_num的有效范围是1 到当前设备 vector-core 数量未调用时默认仅创建一个 SDMA stream/QP因此多核并发时必须显式设置所需 QP 数量SDMA 创建的是本地 device-only stream每个 QP 一条 stream且 stream 不绑定特定 peer接口需在任何 ACLSHMEM 实例初始化之前调用配置在实例存活期间冻结最后一个实例 finalize 后所有引擎SDMA/UDMA/ROCE的 QP 数会重置为 1可再次配置所有 PE 上的取值必须一致不一致会产生不兼容的元数据布局该函数进程内线程安全与初始化/反初始化串行化。示例中每个 AIV 使用GetBlockIdx()AIV 级全局索引作为qp_idx因此qp_num应设置为block 数 × 每 block AIV 数。当前为20 × 2 40constexpr uint32_t SDMA_AIVS_PER_BLOCK 2; constexpr uint32_t SDMA_BLOCK_NUM 20; constexpr uint32_t SDMA_QP_NUM SDMA_BLOCK_NUM * SDMA_AIVS_PER_BLOCK;并在初始化前调用见 main.cpp 与 main.cppattributes.option_attr.data_op_engine_type ACLSHMEM_DATA_OP_SDMA; CHECK_RET(aclshmemx_set_qp_num(ACLSHMEM_DATA_OP_SDMA, SDMA_QP_NUM)); CHECK_RET(aclshmemx_init_attr(ACLSHMEMX_INIT_WITH_DEFAULT, attributes));注意示例中test_set_attr里option_attr初始的data_op_engine_type为ACLSHMEM_DATA_OP_MTE随后在main中显式改写为ACLSHMEM_DATA_OP_SDMA并配置 QP 数二者必须在初始化前完成。非阻塞接口的两种完成方式aclshmemx_sdma_qp_put_nbi和aclshmemx_sdma_qp_get_nbi都是非阻塞接口调用后立即返回不等待数据传输完成。用户可通过以下两种方式确保数据传输完成QP 内 quiet 等待算子内同步所有调用aclshmemx_sdma_qp_put/get_nbi的核在 sdma 任务结束后算子内调用相同qp_idx的aclshmemx_sdma_qp_quiet接口等待该 QP 上的 SDMA 操作完成。aclshmemx_sdma_qp_quiet(tmp_buff, ub_size, cur_block_idx, EVENT_ID0);适用场景算子内后续操作依赖 sdma 任务完成例如后续算子需要使用 sdma 传输好的数据。注意aclshmemx_sdma_qp_quiet只等待指定的一个 QP且不会从GetBlockIdx()推导 QP必须显式传与前面请求相同的qp_idx见 shmem_device_sdma.h而普通版aclshmemx_sdma_quiet只排空 QP 0不能用于确认qp_idx 0的请求。notify_record host 侧等待跨 stream 同步所有调用aclshmemx_sdma_qp_put/get_nbi的核在 sdma 任务结束后算子内调用相同qp_idx的aclshmemx_sdma_qp_notify_record接口然后在host 侧调用aclrtWaitAndResetNotify接口等待指定的同步 ID 完成。适用场景其它 stream 上的 kernel 需要等待 sdma 任务完成后才能继续执行。notify record 会排在该 QP 上更早的 SQE 之后保证顺序性。详细用法可查看 NotifyWait 机制使用说明。两类接口配套的完成接口一览均声明于 shmem_device_sdma.h提交接口配套完成接口aclshmemx_sdma_put_nbi/aclshmemx_sdma_get_nbiQP0aclshmemx_sdma_quiet/aclshmemx_sdma_notify_recordaclshmemx_sdma_qp_put_nbi/aclshmemx_sdma_qp_get_nbi显式 QPaclshmemx_sdma_qp_quiet/aclshmemx_sdma_qp_notify_record使用相同qp_idx小结SDMA put/get 为算子内数据搬运提供了不同于高层 RMA 的精细化控制普通接口aclshmemx_sdma_put_nbi适合单核、固定 QP0 的简单场景显式 QP 接口aclshmemx_sdma_qp_put/get_nbi通过qp_idx一般取 AIV 级GetBlockIdx()将多核并发分摊到多条 SDMA QP 上是构建多核 allgather 等集合通信的基础。使用时务必遵循初始化前用aclshmemx_set_qp_num(ACLSHMEM_DATA_OP_SDMA, qp_num)显式配置 QP 数默认仅 1 条、所有 PE 配置一致、QP 数取 block 数 × 每 block AIV 数、UB workspace 64 字节对齐且不小于 64 字节、非阻塞提交后按“算子内 quiet”或“notify_record host 侧等待”两种方式之一确认完成。同时注意平台差异A2/A3 支持 put/getAscend950 仅支持 get。【免费下载链接】shmemCANN SHMEM 是面向昇腾平台的多机多卡内存通信库基于OpenSHMEM 标准协议实现跨设备的高效内存访问与数据同步。项目地址: https://gitcode.com/cann/shmem创作声明:本文部分内容由AI辅助生成(AIGC),仅供参考
返回列表