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），每个进程对应一个 rank：

for (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里改。

关键参数

资源释放

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用到的设备侧 API：

Put 的 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 template<typename XType> __aicore__ inline void AllToAllComm<XType>::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 单段

template<typename XType> __aicore__ inline void AllToAllComm<XType>::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 段是要发给 remoteRank=i 的数据。

3.3 Scale 一次性 Put

Scale（量化系数）不参与 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 AllReduce

AllReduce = AllToAll + Reduce。两步：

1. AllToAll 把每卡的段发到所有其他卡（复用参考工程的PutToAllRanks）；
2. 接收端做本地 reduce（AIV 用 Vector 指令做加法）。

// 伪代码：AllReduce 通信层 template<typename XType> __aicore__ inline void AllReduceComm::ReduceAllRanks(uint32_t bufferId) { // bufferId 处的数据 layout: [rankSize][reduceLen] // 用 AIV Vector 把 rankSize 个段加到一起 // ...（用 AscendC::Add 做 UB 内 reduce） }

4.2 ReduceScatter

ReduceScatter = AllReduce + 每卡只保留一段。在 AllReduce 基础上，每张卡只 Reduce 自己负责的那段，输出落 GM。

4.3 AllGather

AllGather = 每卡 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}/任意文件即视为违反约束 1，Step 4 直接 FAIL：

# Reviewer 标准检查命令 grep -rn "Hccl::" operators/{op}/

完整禁止清单（13 类）：

历史背景

这些 API 来自asc-devkit/adv_api/hccl/，在 CANN 早期版本是集合通信的主要入口。它们的特点是服务端调度：API 内部把任务发给通信 CPU，Kernel 不能干预时序。MC2 场景下通信必须与计算在同一 Kernel 内深度耦合，HCCL 无法满足。

6. SHMEM 编译与链接

参考工程的cmake/shmem.cmake自动处理 SHMEM 的拉取和构建。关键步骤：

1. third_party/shmem是 symlink 到外部 shmem 仓库；
2. CMake 首次配置时触发add_custom_target(cann_samples_shmem_dependencies)：
   • 如果 submodule 未初始化，运行git submodule update --init --recursive；
   • 用cmake -S ... -B .../build -DSOC_TYPE=Ascend950构建 shmem 静态库；
   • cmake --build ... --target install安装到third_party/shmem/install/shmem/；
3. 生成两个 IMPORTED 共享库：
   • cann_samples_shmem_bootstrap_uid：基于 UID 的 bootstrap；
   • cann_samples_shmem_bootstrap_config_store：配置存储；
4. 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. 排错速查

8. 后续阅读

【免费下载链接】cannbot-skillsCANNBot 是面向 CANN 开发的用于提升开发效率的系列智能体，本仓库为其提供可复用的 Skills 模块。项目地址: https://gitcode.com/cann/cannbot-skills