CANN SHMEM UDMA 原子加（Atomic Add）示例：Ascend950 平台跨设备原子操作的编译、运行与实现解析

CANN SHMEM UDMA 原子加Atomic Add示例Ascend950 平台跨设备原子操作的编译、运行与实现解析【免费下载链接】shmemCANN SHMEM 是面向昇腾平台的多机多卡内存通信库基于OpenSHMEM 标准协议实现跨设备的高效内存访问与数据同步。项目地址: https://gitcode.com/cann/shmemCANN SHMEM 是基于 OpenSHMEM 标准协议、面向昇腾平台的多机多卡内存通信库。udma_atomic_add是仓库中展示如何使用UDMA统一 DMA传输引擎在远端设备对称内存上执行原子加操作的完整样例一个 PE 通过aclshmemx_udma_atomic_add将固定值原子累加到相邻 PE 的对称地址上并用aclshmemx_udma_quiet保证操作完成。读完本文你将掌握该样例在 Ascend950 平台上的编译与双进程运行方法、全部命令行参数的含义以及从 Host 侧初始化、Device 侧 Kernel 下发到 UDMA 引擎底层实现的完整调用链。样例定位与核心能力udma_atomic_add是 CANN SHMEM 中面向 UDMA 引擎的原子内存操作Atomic Memory Operation示例位于 examples/udma_atomic_add/。它的核心价值在于验证两件事UDMA 引擎下的原子加能力通过aclshmemx_udma_atomic_add在远端 PE 的对称内存地址上执行读-改-写原子累加UDMA 的同步与完成语义原子加属于非阻塞提交之后必须调用aclshmemx_udma_quiet等待该 PE 上的 UDMA 操作完成才能安全地消费结果。从实现上看该样例选择将attributes.option_attr.data_op_engine_type显式设置为ACLSHMEM_DATA_OP_UDMA见 main.cpp从而强制整个会话的数据面走 UDMA 引擎而不是默认的 MTEDMA 路径。这意味着 UDMA 原子加不仅依赖 CANN SHMEM 自身的初始化还依赖 CANN 9.1.0 提供的 HCOMM 底层资源接口。版本与平台支持在使用该样例前请先确认运行环境满足以下约束仅支持 Ascend950 平台其他平台不支持运行本样例CANN 版本必须为 9.1.0 或更高。CANN 9.1.0 起提供 UDMA 所需的 HCOMM 资源接口包括HcommEndpointCreate、HcommMemReg、HcommChannelCreate、HcommChannelGetStatus等。低于 9.1.0 的 CANN 版本不在支持范围内请从 CANN 9.1.0 官方资源页面下载并安装 Ascend950 对应的 toolkit 包和 ops 包初始化时会加载并检查 HCOMM 符号。若 HCOMM 运行时库或符号不完整例如仅安装了不匹配版本的 CANNaclshmemx_init_attr初始化会直接失败。此外编译前需要先加载 CANN 环境变量source /usr/local/Ascend/ascend-toolkit/set_env.sh若 toolkit 为自定义安装路径请使用对应的set_env.sh。编译构建示例与运行环境准备样例已集成进仓库根目录的顶层构建脚本无需单独编译。在仓库根目录shmem/下执行bash scripts/build.sh -examples -soc_type Ascend950关键点说明-examples指定只构建 examples 下的样例其中 CMakeLists.txt 通过aclshmem_add_collective_example(udma_atomic_add)一条语句注册该示例-soc_type Ascend950与仅支持 Ascend950的约束一一对应编译出的二进制才能正确链接 UDMA 所需的设备侧算子与 HCOMM 资源编译产物输出到build/bin/udma_atomic_add运行时需要的库位于build/lib。运行双进程启动与命令行参数双进程运行示例在仓库根目录执行以下命令启动 2 个全局 PEPE 0 与 PE 1它们将通过tcp://127.0.0.1:8899完成 SHMEM 初始化握手export PROJECT_ROOTshmem-root-directory export LD_LIBRARY_PATH${PROJECT_ROOT}/build/lib:$LD_LIBRARY_PATH export SHMEM_UID_SESSION_ID127.0.0.1:8899 ./build/bin/udma_atomic_add 2 0 tcp://127.0.0.1:8899 2 0 0 # PE 0 ./build/bin/udma_atomic_add 2 1 tcp://127.0.0.1:8899 2 0 0 # PE 1各环境变量的作用PROJECT_ROOT仓库根目录用于定位build/lib下的 SHMEM 动态库LD_LIBRARY_PATH将 SHMEM 构建产物目录加入动态库搜索路径否则运行时无法加载libshmem等库SHMEM_UID_SESSION_ID以IP:port形式指定 UID唯一标识会话端点多个进程使用相同的会话 ID 才能互相发现并完成初始化。仓库还提供了自动化运行脚本 run.sh它替你完成PROJECT_ROOT、LD_LIBRARY_PATH、SHMEM_UID_SESSION_ID的设置后台启动 PE 0 与 PE 1 并逐个wait等待退出码方便在 CI 或脚本中直接复用。命令行参数说明./udma_atomic_add n_pes pe_id ipport g_npus f_pe f_npu参数含义示例n_pes全局 PE进程总数2pe_id当前进程的 PE 号取值范围[0, n_pes)0或1ipportSHMEM 初始化所需的 IP 与端口格式为tcp://IP地址:端口号tcp://127.0.0.1:8899g_npus当前服务器上启动的 NPU 卡数量2f_pe当前服务器上使用的第一个 PE 号0f_npu当前服务器上执行本样例使用的第一张 NPU 卡卡号0参数与运行时行为的对应关系可以从 main.cpp 的解析逻辑中看出device_id pe_id % g_npus f_npu决定每个 PE 绑定到哪张 NPU 卡当单机启动的 PE 数不超过 NPU 卡数时g_npus与f_npu的配合保证各 PE 独占一张卡。默认的对称内存大小为 1 GiB1024UL * 1024UL * 1024见 main.cpp。运行流程解析Host 初始化到结果校验样例的执行流程清晰展示了ACL 初始化 → SHMEM 初始化 → 数据准备 → Kernel 下发 → 同步 → 结果校验的完整链路对应 main.cppACL 初始化aclInit、aclrtSetDevice(device_id)、aclrtCreateStream完成设备侧运行环境准备SHMEM 初始化通过test_set_attr填充aclshmemx_init_attr_t属性并显式设置data_op_engine_type ACLSHMEM_DATA_OP_UDMA随后调用aclshmemx_init_attr(ACLSHMEMX_INIT_WITH_DEFAULT, attributes)完成初始化。正是这一步会加载并检查前述 HCOMM 符号对称内存分配aclshmem_malloc(1024)在对称堆上分配目标缓冲区ptr输入数据准备Host 侧构造input[i] pe_id 10的 int32 数据通过aclrtMemcpyH2D拷贝到设备侧ptrKernel 下发调用launch_udma_atomic_add(1, stream, ptr, dump, trans_size * sizeof(int32_t))启动设备核函数随后aclrtSynchronizeStream等待执行完成再aclshmem_barrier_all做全局同步结果校验将设备侧结果 D2H 拷贝回 Host校验y_host[0] pe_id 10 * 2通过则打印check transport result success, pepe_id最终进程输出[SUCCESS] demo run success in pe pe_id资源释放依次执行aclshmem_free、aclshmem_finalize、aclrtDestroyStream、aclrtResetDevice、aclFinalize。当启用了ENABLE_ASCENDC_DUMP宏编译时还会通过Adx::AdumpPrintWorkSpace将调试 workspace 落盘便于排查设备侧执行问题。设备侧 Kernel原子加如何被提交设备侧核函数位于 udma_atomic_add_kernel.cpp是理解原子加语义的关键int32_t my_rank aclshmem_my_pe(); int32_t pe_size aclshmem_n_pes(); int32_t peer_id (my_rank 1) % pe_size; AscendC::PipeBarrierPIPE_ALL(); int32_t value 10; aclshmemx_udma_atomic_add((__gm__ int32_t *)gva, value, peer_id); aclshmemx_udma_quiet(peer_id);它完成三件事通过aclshmem_my_pe()/aclshmem_n_pes()获取当前 PE 号与 PE 总数计算出目标 PE 为peer_id (my_rank 1) % pe_size形成 PE 0 → PE 1 → PE 0 的环形目标调用aclshmemx_udma_atomic_add把常量10原子累加到peer_id的对称地址gva上。以 PE 0 为例它把自己的0 10 10加到 PE 1 的缓冲区上而 PE 1 最初写入的是1 10 11因此最终校验值为10 11 21 pe_id(1) 20与 Host 侧pe_id 10 * 2的预期完全一致随后调用aclshmemx_udma_quiet(peer_id)等待该 PE 方向上的 UDMA 操作完成保证原子加在后续读取结果前已经生效。底层实现UDMA 原子加的引擎语义接口契约aclshmemx_udma_atomic_add是设备侧模板接口声明于 include/device/gm2gm/engine/shmem_device_udma.htemplate typename T, const aclshmemx_udma_op_config_t CONFIG ACLSHMEMX_UDMA_OP_CONFIG_DEFAULT ACLSHMEM_DEVICE void aclshmemx_udma_atomic_add(__gm__ T* dst, T value, int32_t pe);dst指向远端 PE 的对称地址本地传gva内部会自动完成对称地址翻译value为累加操作数支持的数据类型为int32、uint32、int64、uint64模板参数CONFIG控制操作配置其结构体aclshmemx_udma_op_config_t包含四个字段cqe完成队列事件使能取值 0/1、se保留位必须为 0、fence栅栏使能取值 0/1、odr排序域取值 0~7并提供ACLSHMEMX_UDMA_OP_CONFIG_DEFAULT、ACLSHMEMX_UDMA_OP_CONFIG_NO_CQE、ACLSHMEMX_UDMA_OP_CONFIG_SO等预置配置。头文件中还给出了数据先行、标志位随后的典型用法先用aclshmemx_udma_put_nbi..., ACLSHMEMX_UDMA_OP_CONFIG_SO传数据再用aclshmemx_udma_atomic_adduint32_t, ACLSHMEMX_UDMA_OP_CONFIG_SO更新标志位实现跨设备的 release 语义。引擎内部实现从 src/device/gm2gm/engine/shmem_device_udma.hpp 的实现看aclshmemx_udma_atomic_add的底层行为包括先执行aclshmemi_udma_require_cqeCONFIG()即原子操作要求配置中cqe 1默认配置满足该要求通过aclshmemi_udma_check_atomic_len按UDMA_OPCODE_FAAFetch-Add-Add 语义校验类型长度是否被 UDMA 引擎支持不支持时调用aclshmem_kernel_abort中止内核调用aclshmem_ptr(dst, pe)完成本地地址到远端 PE 对称地址的翻译最终以UDMA_OP_WRITE_WITH_REDUCE操作码构造原子参数value与cond0提交到 UDMA 引擎由硬件在远端内存上完成读-改-写原子累加。需要注意的是UDMA 传输引擎下的并发 RMA/AMO 操作存在约束对同一 PE 的并发 RMA/AMO 操作不被支持见 shmem_device_udma.h因此示例中每个 PE 只对单一目标 PE 发起原子加。此外UDMA 原子加是典型的非阻塞提交必须配合aclshmemx_udma_quiet才能获得完成保证——这正是核函数在原子加之后立即调用quiet的原因。结果验证与故障排查要点运行成功后两个 PE 的终端会分别输出check transport result success, pe0 [SUCCESS] demo run success in pe 0 check transport result success, pe1 [SUCCESS] demo run success in pe 1若校验失败main.cpp会打印pepe_id: 实际值 ! 期望值并最终输出[FAILED] demo run failed in pe pe_id。常见排查方向初始化失败检查 CANN 是否为 9.1.0 及以上、Ascend950 对应的 toolkit/ops 包是否完整安装、set_env.sh是否已 source、HCOMM 符号是否完整编译报错或产物缺失确认编译时使用了-soc_type Ascend950且LD_LIBRARY_PATH指向build/lib多进程无法互相发现确认所有 PE 使用了相同的SHMEM_UID_SESSION_ID与ipport端口未被占用设备绑定错误检查g_npus、f_pe、f_npu与实际 NPU 卡资源是否匹配。延伸阅读原子加的同族接口还包括aclshmemx_udma_atomic_inc等价于value 1的原子加、aclshmemx_udma_atomic_fetch_add等全部声明于 include/device/gm2gm/engine/shmem_device_udma.h若要观察 UDMA 引擎更完整的操作集合Put/Get/NBI/聚合提交等可阅读 src/device/gm2gm/engine/shmem_device_udma.hpp 及同目录下的shmemi_device_udma.h其他基于 UDMA 的样例可参考 examples/udma_demo/、examples/udma_qp_demo/ 与 examples/tp_allreduce_udma/仓库整体编译与运行方式见 compilation_build_guide.md示例统一构建入口为 scripts/build.sh。【免费下载链接】shmemCANN SHMEM 是面向昇腾平台的多机多卡内存通信库基于OpenSHMEM 标准协议实现跨设备的高效内存访问与数据同步。项目地址: https://gitcode.com/cann/shmem创作声明：本文部分内容由AI辅助生成（AIGC），仅供参考