CANN通信层SHMEM参考

通信层SHMEM / UDMA参考【免费下载链接】cannbot-skillsCANNBot 是面向 CANN 开发的用于提升开发效率的系列智能体本仓库为其提供可复用的 Skills 模块。项目地址: https://gitcode.com/cann/cannbot-skills本文档承载 MC2 skill 的通信子能力。涵盖SHMEM API 目录、host/device 侧用法、禁止 HCCL 清单的执行细则、扩展其他通信原语AllReduce / ReduceScatter / AllGather的思路。官方文档https://shmem-doc.pages.dev/。本文件不重抄官方 API只记录参考工程用到的子集 MC2 场景下的使用要点。1. SHMEM 初始化host 侧参考工程的 host 侧初始化在src/all_to_all_matmul.cpp::runAllToAllMatmul按以下顺序// 1. ACL 标准 init ACL_CHECK(aclInit(nullptr)); ACL_CHECK(aclrtSetDevice(deviceId)); ACL_CHECK(aclrtCreateStream(stream)); // 2. SHMEM init关键 aclshmemx_init_attr_t attributes; aclshmemx_uniqueid_t defaultFlagUid; test_set_attr(rankId, rankNum, localMemSize, ipport, defaultFlagUid, attributes); attributes.option_attr.data_op_engine_type ACLSHMEM_DATA_OP_UDMA; // 选 UDMA 引擎 ACL_CHECK(aclshmemx_init_attr(ACLSHMEMX_INIT_WITH_DEFAULT, attributes)); // 3. 分配 SHMEM 空间所有 rank 对齐 void *shmemSpace aclshmem_align(PACKAGE_SIZE, SHMEM_SPACE_SIZE); // 512B 对齐1 GB多 rank 启动参考工程用fork启动多进程src/all_to_all_matmul.cpp::main每个进程对应一个 rankfor (int rankId 0; rankId rankNum; rankId) { pid_t pid fork(); if (pid 0) { // 子进程setDevice(rankId) SHMEM init 跑 kernel runAllToAllMatmul(rankNum, rankId, ...); exit(ret); } }tcp://127.0.0.1:8998是 rank 间发现的默认地址可在test_set_attr里改。关键参数参数含义参考工程取值SHMEM_SPACE_SIZE单卡 SHMEM 空间大小1024 * 1024 * 1024(1 GB)PACKAGE_SIZESHMEM 对齐粒度512bytesUDMA 要求ACLSHMEM_DATA_OP_UDMA数据搬运引擎选 UDMA不是 RDMA必须选 UDMA资源释放aclshmem_free(shmemSpace); ACL_CHECK(aclshmem_finalize()); ACL_CHECK(aclrtDestroyStream(stream)); ACL_CHECK(aclrtResetDevice(deviceId)); ACL_CHECK(aclFinalize());顺序不能错先 free SHMEM 空间再 finalize SHMEM再销毁 stream/device。2. SHMEM 设备侧 API 目录参考工程在include/kernel/all_to_all_comm_udma.h用到的设备侧 APIAPI作用调用位置aclshmem_n_pes()返回总 rank 数InitParamsaclshmem_my_pe()返回当前 rank IDInitParamsaclshmem_ptr(shmemSpace, rankId)拿到本 rank 视角下指定 rank 的 SHMEM 基地址GM 地址GetDataAddrGm/GetScaleAddrGmaclshmemx_udma_put_nbi(dst, src, ubuf, size, remoteRank)非阻塞 Put本地 src → 远程 dstPutSegmentToRankaclshmemx_udma_quiet(remoteRank)等本次 Put 下发完不等对端收PutSegmentToRankaclshmemx_barrier_all_vec()全卡 Barrier等所有 Put 被对端收到BarrierAllPut 的 dst 地址语义UDMA Put 的dst是对端 rank 的 GM 地址但本 rank 拿不到对端 GM 物理地址。SHMEM 的设计所有 rank 通过aclshmem_align分配等大空间本 rank 通过aclshmem_ptr(shmemSpace, remoteRank)拿到本 rank 视角下 remoteRank 那块 SHMEM 的 GM 地址写到这个地址就等于写到对端。// all_to_all_comm_udma.h::GetDataAddrGm GM_ADDR baseAddr (GM_ADDR)(aclshmem_ptr(shmemContextGM_, rankId_)); return baseAddr bufferId * bufferBlockSize_;注意aclshmem_ptr第二个参数是rankId_本卡自己的 rank不是 remoteRank。本卡视角下本卡的 SHMEM 基地址就是其他卡 Put 过来的数据落点。ubuf 参数aclshmemx_udma_put_nbi的第三个参数(__ubuf__ uint8_t*)nullptr在参考工程中始终为nullptr。它的本意是提供 UB 缓冲用于小消息聚合MC2 的大块 Put 直接走 GM→GM不需要 UB 中转。3. AllToAll 通信实现剖析参考工程的AllToAllComm类include/kernel/all_to_all_comm_udma.h实现了All rank 互发 M 段的 AllToAll 语义。核心逻辑3.1 Put 数据// PutToAllRanks: 当前 rank 把自己的 M 段发给所有其他 rank templatetypename XType __aicore__ inline void AllToAllCommXType::PutToAllRanks( uint64_t srcMOffset, uint64_t mSize, uint32_t bufferId) { // 每个 Block 负责发送到对应的 remoteRank避免卡间竞争 if ((AscendC::GetBlockIdx() rankSize_) (AscendC::GetBlockIdx() ! rankId_)) { PutSegmentToRank(AscendC::GetBlockIdx(), srcMOffset, mSize, bufferId); } BarrierAll(); // 等所有卡完成 Put }关键设计不是单个 Block 串行 Put 给所有 remoteRank而是多 Block 并发——Block i 负责发给 remoteRank i。这样 UDMA 引擎能并行下发多条指令。3.2 Put 单段templatetypename XType __aicore__ inline void AllToAllCommXType::PutSegmentToRank( uint32_t remoteRank, uint64_t srcMOffset, uint64_t mSize, uint32_t bufferId) { GM_ADDR remoteWinAddr GetDataAddrGm(bufferId); // 本 rank 视角下自己的 SHMEM // 注意Put 的 dst 是本 rank 视角下本 rank 的 SHMEM 地址 // UDMA 引擎会把这个地址映射到 remoteRank 的 SHMEM // 但 SHMEM 的地址翻译是按 rank 偏移做的所以要加上 rankId 的偏移 uint64_t dstDataOffset rankId_ * commMSize_ * bytesPerMRow_; // 对端按 srcRank 区分 uint64_t srcDataOffset remoteRank * M_ * bytesPerMRow_ srcMOffset * bytesPerMRow_; uint64_t dataSize mSize * bytesPerMRow_; aclshmemx_udma_put_nbi( remoteWinAddr dstDataOffset, inputDataAddr_ srcDataOffset, (__ubuf__ uint8_t*)nullptr, dataSize, remoteRank); aclshmemx_udma_quiet(remoteRank); // 保序 }地址语义陷阱dstDataOffset rankId_ * commMSize_ * bytesPerMRow_对端收数据时按我从哪个 rank 来放到对应位置避免多 rank 数据互相覆盖srcDataOffset remoteRank * M_ * bytesPerMRow_ srcMOffset * bytesPerMRow_本 rank 的 A 矩阵逻辑上按remoteRank分段第 i 段是要发给 remoteRanki 的数据。3.3 Scale 一次性 PutScale量化系数不参与 M 轴流水参考工程在 AIV 启动时一次性 Put 全部 Scale// all_to_all_matmul_impl.h::AllToAllProcess 开头 allToAllComm_.PutScaleToAllRanks(0, axisM_); // 一次 Put 全 M 行的 Scale因为 Scale 数据量小M * CeilDiv(K, 64) * 2bytes相比 A 的M * K小 32 倍不值得做流水。4. 扩展其他通信原语MC2 算子不限于 AllToAll。下面给出扩展思路参考工程未直接实现但架构允许4.1 AllReduceAllReduce AllToAll Reduce。两步AllToAll 把每卡的段发到所有其他卡复用参考工程的PutToAllRanks接收端做本地 reduceAIV 用 Vector 指令做加法。// 伪代码AllReduce 通信层 templatetypename XType __aicore__ inline void AllReduceComm::ReduceAllRanks(uint32_t bufferId) { // bufferId 处的数据 layout: [rankSize][reduceLen] // 用 AIV Vector 把 rankSize 个段加到一起 // ...用 AscendC::Add 做 UB 内 reduce }4.2 ReduceScatterReduceScatter AllReduce 每卡只保留一段。在 AllReduce 基础上每张卡只 Reduce 自己负责的那段输出落 GM。4.3 AllGatherAllGather 每卡 Put 自己段给所有其他卡但不 reduce。本质就是 AllToAll 不带 reduce 步骤——参考工程的PutToAllRanks已经覆盖。4.4 自定义原语注意点必须保留aclshmemx_udma_quiet(remoteRank)调用否则 Put 不保序Barrier 类操作必须用aclshmemx_barrier_all_vec不能用 HCCL 的同步原语跨核同步AIV↔AIC必须用CrossCoreSetFlag/CrossCoreWaitFlag不能用 SHMEM 的 barrier 替代。5. 禁止 HCCL API 速查以下 API 出现在operators/{op}/任意文件即视为违反约束 1Step 4 直接 FAIL# Reviewer 标准检查命令 grep -rn Hccl:: operators/{op}/完整禁止清单13 类类别API初始化/终结Hccl::Init()/Hccl::InitV2()/Hccl::Finalize()任务调度Hccl::Commit()/Hccl::Wait()/Hccl::Query()/Hccl::Iterate()集合通信原语Hccl::AllReduce()/Hccl::AllGather()/Hccl::ReduceScatter()/Hccl::AlltoAll()/Hccl::AlltoAllV()写操作Hccl::BatchWrite()/Hccl::AlltoAllvWrite()TilingHccl::SetCcTiling()/Hccl::SetCcTilingV2()跨组同步Hccl::InterHcclGroupSync()ContextGetHcclContext()历史背景这些 API 来自asc-devkit/adv_api/hccl/在 CANN 早期版本是集合通信的主要入口。它们的特点是服务端调度API 内部把任务发给通信 CPUKernel 不能干预时序。MC2 场景下通信必须与计算在同一 Kernel 内深度耦合HCCL 无法满足。6. SHMEM 编译与链接参考工程的cmake/shmem.cmake自动处理 SHMEM 的拉取和构建。关键步骤third_party/shmem是 symlink 到外部 shmem 仓库CMake 首次配置时触发add_custom_target(cann_samples_shmem_dependencies)如果 submodule 未初始化运行git submodule update --init --recursive用cmake -S ... -B .../build -DSOC_TYPEAscend950构建 shmem 静态库cmake --build ... --target install安装到third_party/shmem/install/shmem/生成两个 IMPORTED 共享库cann_samples_shmem_bootstrap_uid基于 UID 的 bootstrapcann_samples_shmem_bootstrap_config_store配置存储cann_samples::shmemalias 提供头文件路径 链接库。Agent 开发注意不要手动修改third_party/shmem/内部它会被外部 symlink 覆盖新算子的 CMakeLists 应直接include(${CMAKE_CURRENT_SOURCE_DIR}/cmake/shmem.cmake REQUIRED)与参考工程一致若 shmem 构建失败常见原因ASCEND_HOME_PATH未设、SOC_TYPE不是Ascend950。7. 排错速查现象可能原因排查方向aclshmem_align返回 nullptrSHMEM 空间已被前次分配占满检查aclshmem_finalize是否被调用减小SHMEM_SPACE_SIZEaclshmemx_udma_put_nbi后立即读 buffer 拿不到数据Put 是非阻塞的没等 quiet/barrier在读前调aclshmemx_barrier_all_vec()多卡运行 deadlockAIV 的CrossCoreSetFlagidx 与 AIC 的CrossCoreWaitFlagidx 不匹配检查mLoopIdx在两侧是否一致Put 数据错位dstDataOffset算错rankId 没乘对 commMSize打印 bufferId、srcMOffset、dstDataOffset 在 host 端验证编译报aclshmemx_udma_put_nbi未定义没链接cann_samples::shmem检查target_link_libraries是否含cann_samples::shmem运行报aclshmemx_init_attr failedtcp://地址被占用或 rankNum 不匹配改ipport确认所有 rank 进程同时启动8. 后续阅读想了解读Blaze Matmul 接入matmul_blaze.md通算融合整体架构mc2_architecture.md性能采集L2 flush msprofprofiling_mc2.md参考工程改造食谱codebase_map.md【免费下载链接】cannbot-skillsCANNBot 是面向 CANN 开发的用于提升开发效率的系列智能体本仓库为其提供可复用的 Skills 模块。项目地址: https://gitcode.com/cann/cannbot-skills创作声明：本文部分内容由AI辅助生成（AIGC），仅供参考