CANN SHMEM SIMT 远程内存访问（RMA）接口实战：三种 API 形式与 warp 级协作搬运原理

CANN SHMEM SIMT 远程内存访问RMA接口实战三种 API 形式与 warp 级协作搬运原理【免费下载链接】shmemCANN SHMEM 是面向昇腾平台的多机多卡内存通信库基于OpenSHMEM 标准协议实现跨设备的高效内存访问与数据同步。项目地址: https://gitcode.com/cann/shmem本篇技术指南以 CANN SHMEM 仓库中的simt_rma示例examples/simt_rma为核心系统讲解 SIMD 与 SIMT 混合编译模式下 SIMT 远程内存访问RMA接口的三种典型形式、线程组协作粒度thread / block / warp、示例的完整执行流程与结果校验逻辑并结合 设备侧头文件 与 实现文件 剖析其底层 MTE 搬运机制。读完本文你将掌握如何在昇腾设备上编写、编译并运行基于 SIMT 的跨 PE 数据搬运程序并能根据数据长度描述方式在三种 API 形式之间灵活切换。背景SIMD 与 SIMT 混合编译模式下的 RMACANN SHMEM 是基于 OpenSHMEM 标准协议的跨设备内存通信库。在昇腾平台上当需要以 SIMT单指令多线程编程模型组织线程并发访问对称内存时设备侧代码可以调用simt命名空间下的 RMA 接口。该示例展示的核心能力是将一段连续内存区域的数据从本地 PEProcessing Element搬运到远端 PE 的对称内存或从远端 PE 拉取数据到本地。所有接口均声明为__simt_callee__ inline形式意味着它们由 SIMT 矢量函数__simt_vf__内部的线程调用可内联展开。三种接口形式统一抽象为如下签名__simt_callee__ inline void aclshmem_{NAME}_{op}(__gm__ TYPE *dst, __gm__ TYPE *src, size_t elem_size, int32_t pe)__simt_callee__ inline void aclshmem_{op}{BITS}(__gm__ void *dst, __gm__ void *src, size_t nelems, int32_t pe)__simt_callee__ inline void aclshmem_{op}mem(__gm__ void *dst, __gm__ void *src, size_t elem_size, int32_t pe)三种形式的核心功能完全一致——在连续内存区域间搬运数据唯一区别在于数据长度的指定方式形式长度指定方式示例形式一按类型基于每个传输元素的具体数据类型如half、float、int32_t等描述元素个数aclshmem_int16_put形式二按位宽基于每个传输元素的比特位宽如8、16、128描述元素个数aclshmem_put128形式三按字节直接指定需要传输的总字节数aclshmem_putmem接口占位符的取值约束接口名称中的占位符{}的允许取值如下表所示定义于 include/device_simt/gm2gm/shmem_device_simt_rma.h 的ACLSHMEM_TYPE_FUNC宏与位宽宏中占位符允许取值{op}put、get{NAME}half、float、int8、int16、int32、int64、uint8、uint16、uint32、uint64、char、bfloat16{BITS}8、16、32、64、128其中{NAME}与 C 类型的一一对应关系char同时展开为signed char与unsigned char两种在头文件注释中给出了完整映射表例如int16 - int16_t、bfloat16 - bfloat16_t。三种线程组协作粒度thread / block / warp每种形式又各派生 3 个线程组粒度的变体区别在于参与本次搬运的线程范围变体粒度说明aclshmem_*无后缀线程级由调用线程独立完成整块数据的搬运aclshmemx_*_blockblock 级由同一 block 内的线程协作完成搬运aclshmemx_*_warpwarp 级由同一 warp32 个线程协作完成搬运本示例调用的是 warp 级变体。kernel 以dim3(32)启动即一个 warp 恰好 32 个线程由这 32 个线程协作搬运同一份数据而不是每个线程各自遍历整块内存。这种协作模式可让 MTEMemory Transfer Engine以更大的数据块粒度并行搬移减少线程级重复搬运带来的带宽浪费。示例中三种形式对应的实际调用在 main.cpp 中三种形式分别由test_put_get_mem、test_put_get_type、test_put_get_bits三个函数封装实际调用的接口为// 形式一按元素数据类型描述长度 simt::aclshmemx_int16_get_warp(__gm__ int16_t *dst, __gm__ int16_t *src, size_t elem_size, int32_t pe); simt::aclshmemx_int16_put_warp(__gm__ int16_t *dst, __gm__ int16_t *src, size_t elem_size, int32_t pe); // 形式二按元素比特位宽描述长度 simt::aclshmemx_get128_warp(__gm__ void *dst, __gm__ void *src, size_t nelems, int32_t pe); simt::aclshmemx_put128_warp(__gm__ void *dst, __gm__ void *src, size_t nelems, int32_t pe); // 形式三直接指定传输的总字节数 simt::aclshmemx_getmem_warp(__gm__ void *dst, __gm__ void *src, size_t elem_size, int32_t pe); simt::aclshmemx_putmem_warp(__gm__ void *dst, __gm__ void *src, size_t elem_size, int32_t pe);示例默认调用test_put_get_bits形式二如需体验其他形式直接修改 demo_call_simt 中调用的函数即可。示例执行流程详解simt_rma示例通过以下 4 步演示 RMA 接口的工作机制完整逻辑见 main.cpp 的test_aclshmem_rma_mem函数环境初始化每个 PE 通过aclshmemx_init_attr(ACLSHMEMX_INIT_WITH_DEFAULT, attributes)完成库初始化并借助aclshmemx_malloc在对称堆上申请 3 块大小相同的对称内存origin_device、res_prev_device、res_next_device。origin初始化为[my_pe 0, ..., my_pe size - 1]res_prev与res_next初始化为-1。示例中单块数据量为COPY_SIZE 4096个int32_t对称堆大小为 1 GiB。GET 操作演示每个 PE 调用 warp 级get接口将逻辑上属于上一个 PE的origin数据拉取并写入自身的res_prev。上、下邻居的求法为环形取模prev_pe (my_pe - 1 n_pes) % n_pes。PUT 操作演示每个 PE 调用 warp 级put接口将自身origin数据推送到上一个 PE的res_next。之所以推送给上一个 PE是因为对prev_pe而言当前 PE 正是它的“下一个 PE”——这样环形传递一圈后每个 PE 的res_next最终存放的就是它下一个 PE的origin数据。结果校验通信完成后将数据从设备端拷回 Host逐元素比对res_prev[i]必须等于prev_pe i上一个 PE 的 originres_next[i]必须等于next_pe i下一个 PE 的 origin其中next_pe (my_pe 1) % n_pes。全部通过则打印[SUCCESS]否则打印[FAILURE]及首个出错位置。需要注意的是put与get之间以及不同 PE 之间的执行顺序通过aclshmem_barrier_all()全局屏障保证见 main.cpp通信前与通信后各调用一次屏障确保所有 PE 的对称内存均已就绪、且数据搬运全部完成后才进入校验阶段。Kernel 启动方式SIMD 与 SIMT 的桥接示例展示了 SIMD 与 SIMT 混合编译的典型写法main.cpp__simt_vf__ __launch_bounds__(1024) inline void demo_call_simt( __gm__ int32_t* origin, __gm__ int32_t* res_prev, __gm__ int32_t* res_next, __gm__ uint64_t* dbg) { int32_t mype simt::aclshmem_my_pe(); int32_t npes simt::aclshmem_n_pes(); int32_t prev_pe (mype - 1 npes) % npes; test_put_get_bits(origin, res_prev, res_next, prev_pe); } __global__ __vector__ void demo_call( __gm__ int32_t* origin, __gm__ int32_t* res_prev, __gm__ int32_t* res_next, __gm__ uint64_t* dbg) { asc_vf_calldemo_call_simt(dim3(32), origin, res_prev, res_next, dbg); }demo_call是 SIMD 侧的__global__ __vector__kernelasc_vf_calldemo_call_simt(dim3(32), ...)以32 线程一个 warp启动 SIMT 矢量函数demo_call_simt与 warp 级 RMA 变体恰好匹配在 SIMT 函数内部通过simt::aclshmem_my_pe()与simt::aclshmem_n_pes()获取当前 PE 号与 PE 总数供地址寻址使用。三种测试函数的长度换算对于 4096 个int32_t元素合计 16 KiB的搬运三个函数的长度参数换算关系如下main.cppmem 形式直接传字节数COPY_SIZE * sizeof(int32_t)type 形式以int16_t视角传元素个数COPY_SIZE * sizeof(int32_t) / sizeof(int16_t)即 8192 个int16_tbits 形式以 128 bit 为单位传元素个数COPY_SIZE * 32 / 128即 1024 个 128-bit 元素。可以看到三种形式表达的是同一块连续内存只是“度量单位”不同。源码级原理宏展开与 MTE 引擎声明宏自动生成接口族include/device_simt/gm2gm/shmem_device_simt_rma.h通过ACLSHMEM_PUT_TYPENAME_MEM、ACLSHMEM_PUT_SIZE_MEM等宏为每种类型 / 位宽自动生成_put/_put_block/_put_warp三件套声明get 同理。例如#define ACLSHMEM_PUT_TYPENAME_MEM(NAME, TYPE) \ __simt_callee__ inline void aclshmem_##NAME##_put( \ __gm__ TYPE* dst, __gm__ TYPE* src, size_t elem_size, int32_t pe); \ __simt_callee__ inline void aclshmemx_##NAME##_put_block( \ __gm__ TYPE* dst, __gm__ TYPE* src, size_t elem_size, int32_t pe); \ __simt_callee__ inline void aclshmemx_##NAME##_put_warp( \ __gm__ TYPE* dst, __gm__ TYPE* src, size_t elem_size, int32_t pe)同一头文件还声明了配套的低延迟单元素接口aclshmem_{NAME}_p单元素 put与aclshmem_{NAME}_g单元素 get以及全部异步nbinon-blocking变体供有不同同步语义需求的场景使用。实现统一收敛到 MTE 搬运函数所有类型级、位宽级、字节级接口在src/device_simt/gm2gm/shmem_device_simt_rma.hpp中最终都收敛到两个模板函数aclshmemi_mte_put_nbiT, SCOPE与aclshmemi_mte_get_nbiT, SCOPE以线程组枚举ACLSHMEMI_THREADGROUP_THREAD / WARP / BLOCK区分协作粒度。例如__simt_callee__ inline void aclshmemx_int16_put_warp(__gm__ int16_t *dst, __gm__ int16_t *src, size_t elem_size, int32_t pe) { aclshmemi_mte_put_nbiint16_t, ACLSHMEMI_THREADGROUP_WARP(dst, src, elem_size, pe); }几个值得注意的实现细节位宽形式与类型的对应put128 / get128内部以 128 bit 类型int4作为模板参数实现文件8/16/32/64位分别对应int8_t / int16_t / int32_t / int64_t字节形式按char搬运putmem / getmem以char为模板参数逐字节描述实现文件协作机制在 MTE 引擎src/device_simt/gm2gm/engine/shmem_device_simt_mte.hpp内部通过aclshmemi_thread_id_in_threadgroupscope()获取线程在组内的索引、aclshmemi_threadgroup_sizescope()获取组大小从而把总数据量按线程编号切分再由aclshmemi_threadgroup_syncscope()做组内同步实现“组内线程协作搬完同一份数据”的语义。对称内存寻址RMA 的关键前提是“对称地址”——dst/src指针在本地 PE 与远端 PE 的对称堆中拥有相同的偏移语义。单元素接口_p/_g的实现在此体现得最为直观实现文件auto ptr simt::aclshmem_ptr(dst, pe); // 将本地对称地址换算为远端 PE 的地址 __gm__ TYPE *addr_gm reinterpret_cast__gm__ TYPE *(ptr); asc_stwt(addr_gm, value); // 单元素写入simt::aclshmem_ptr负责完成跨 PE 的地址换算这也是所有 RMA 接口能够“用本地指针访问远端内存”的根基。补充UB 缓冲版本除gm2gmGlobal Memory 到 Global Memory外仓库还提供了ub2gm方向的 SIMT RMA 接口include/device_simt/ub2gm/shmem_device_simt_rma.h其源 / 目的指针为__ubuf__UB 统一缓冲区类型适用于数据先驻留 UB、再搬运到全局对称内存的流水场景接口命名与变体规则与gm2gm完全一致。编译与运行环境要求支持设备Ascend 950-soc_type Ascend950编译需要开启 SIMT 支持选项-enable_simt。编译项目在 SHMEM 仓库根目录执行编译脚本bash scripts/build.sh -examples -enable_simt -soc_type Ascend950该命令将编译包括simt_rma在内的示例产物位于build/bin/simt_rma。simt_rma目标由 examples/simt_rma/CMakeLists.txt 中的aclshmem_add_fusion_example(simt_rma main.cpp)定义。运行示例进入示例目录并执行运行脚本cd examples/simt_rma bash run.sh运行细节见 run.shrun.sh默认启动 2 个独立进程每个进程对应一个 PE 并绑定一张卡执行build/bin/simt_rma n_pes my_pe脚本自动设置SHMEM_UID_SESSION_ID127.0.0.1:8899作为进程间握手标识并加载build/lib与$ASCEND_HOME_PATH/lib64下的动态库使用-pes参数指定参与运行的 PE卡数量例如启动 4 卡bash run.sh -pes 4-pes仅接受正整数传入非法值非数字、0、负数时脚本会报错退出。多卡运行时请确保执行环境中实际可用的 NPU 卡数不少于指定的 PE 数量否则进程无法绑定到对应设备。结果验证与预期输出程序运行结束后每个 PE 会打印自身三块缓冲区的数据origin、res_prev、res_next以及校验结论。校验通过时输出形如[SUCCESS] PE 0: Verification passed for RMA transfers. [SUCCESS] PE 1: Verification passed for RMA transfers.main.cpp中的sleep(my_pe 1)用于错峰打印避免多进程输出互相干扰。校验失败的 PE 会打印首个不符元素的位置及期望值 / 实际值便于定位问题。测试覆盖与扩展阅读仓库针对 SIMT RMA 提供了完整的单元测试覆盖gm2gm与ub2gm两个方向设备侧 Kerneltests/unittest/device/simt_rma/simt_rma_gm2gm_kernel.cpp、tests/unittest/device/simt_rma/simt_rma_ub2gm_kernel.cpp主机侧驱动与断言tests/unittest/host/simt_rma/simt_rma_gm2gm_test.cpp、tests/unittest/host/simt_rma/simt_rma_ub2gm_test.cpp。测试通过宏批量展开所有类型half、float、int8至int64、各类uint在 thread / block / warp 三种粒度下的 put / get 及 nbi 变体调用与示例形成“示例演示 单测全量覆盖”的互补。若需要对比 SIMD 模式下的同类接口可参考 simt_rma_scalar 与 simt_rma_ub2gm 示例关于对称堆分配与初始化可参阅 初始化原理文档。【免费下载链接】shmemCANN SHMEM 是面向昇腾平台的多机多卡内存通信库基于OpenSHMEM 标准协议实现跨设备的高效内存访问与数据同步。项目地址: https://gitcode.com/cann/shmem创作声明：本文部分内容由AI辅助生成（AIGC），仅供参考