CANN HCCL 自定义集合通信算子开发：基于 AIV 引擎的 AllGather 算子实战解析-开发者社区

CANN HCCL 自定义集合通信算子开发：基于 AIV 引擎的 AllGather 算子实战解析

【免费下载链接】hccl集合通信库（Huawei Collective Communication Library，简称HCCL）是基于昇腾AI处理器的高性能集合通信库，为计算集群提供高性能、高可靠的通信方案项目地址: https://gitcode.com/cann/hccl

HCCL（Huawei Collective Communication Library）是昇腾 AI 处理器上的高性能集合通信库，其提供的 AIV（AI Vector）通信编程接口允许开发者绕过内置算法，自行实现通信算子的 Host 侧逻辑与 Device 侧 Kernel。本文以仓库中的 aiv 版自定义 AllGather 示例为主线，完整讲解从环境准备、算子库编译安装、MPI 测试运行，到 Host 侧资源管理与 Device 侧 Mesh 1D 通信算法的源码级实现原理。读完本文，你将掌握基于 HCCL AIV 接口从零开发一个可编译、可安装、可验证的自定义集合通信算子的完整路径。

一、示例概述：做什么、支持什么

该示例展示了如何基于 HCCL AIV 通信编程接口开发 AllGather 自定义通信算子，核心特性如下：

基于 AIV（AI Vector）通信引擎实现 AllGather 集合通信算子；
同时包含 Host 侧算子逻辑与 Device 侧 Kernel 实现；
提供完整的编译、构建与测试验证流程。

支持的产品与场景（单机 N 卡配置，N >= 2）：

Ascend 950PR / Ascend 950DT；
Atlas A3 训练/推理产品（仅支持超节点内通信场景）；
Atlas A2 训练/推理产品（仅支持单设备通信场景）。

这里需要特别注意：示例的能力边界与硬件代际强相关，A2 产品上仅支持单设备（单卡内多核）通信，A3 及以上才支持多卡间通信，跨服务器场景不在本示例支持范围内。

二、工程目录结构与职责划分

examples/05_custom_ops_allgather/aiv/ ├── CMakeLists.txt # 示例根目录编译配置 ├── op_host/ │ ├── CMakeLists.txt │ ├── allgather.cc # HcclAllGatherCustom 算子 Host 侧实现 │ ├── launch_kernel.cc # Kernel 提交（加载二进制、启动）逻辑 │ └── launch_kernel.h # Kernel 提交接口声明 ├── op_kernel/ │ ├── CMakeLists.txt │ └── launch_kernel_asc.asc # 算子 Kernel 侧实现（Ascend C） └── inc/ ├── hccl_custom_allgather.h # 自定义算子对外接口头文件 ├── common.h # 公共类型定义与宏（OpParam、SIZE_TABLE 等） ├── aiv_allgather_mesh_1d.h # AIV AllGather 核心算法实现（Mesh 1D） ├── aiv_communication_base_v2.h # AIV 通信基类（同步原语、GM2GM 搬运） ├── log.h # 日志工具 ├── extra_args.h # 扩展参数定义（rank 计数/位移数组） └── sync_interface.h # 同步接口定义

目录划分遵循了典型的"Host 侧算子工程（op_host）+ Device 侧 Kernel 工程（op_kernel）+ 公共头文件（inc）"三段式结构，与仓库中 04_custom_ops_p2p、06_custom_ops_reduce_scatter 等示例保持一致的规范，便于对照学习其他算子的写法。

三、环境准备

3.1 安装 CANN Toolkit 开发套件包

按昇腾文档中心的 CANN 软件安装指南安装最新版本 CANN Toolkit 开发套件包（本示例依赖其中的 HCCL 运行时库、ACL 运行时接口以及 Ascend C 算子编译工具链）。

3.2 配置环境变量

以 root 用户默认安装路径为例：

source /usr/local/Ascend/cann/set_env.sh

该脚本会导出ASCEND_HOME_PATH、ASCEND_OPP_PATH等关键环境变量，后续算子安装路径解析与测试运行都会用到。

3.3 安装 MPI

运行测试用例需要 MPI 环境，请确保系统已安装并配置好 MPI（如 OpenMPI），测试程序通过mpirun拉起多进程、每个进程绑定一个 rank。

四、编译与安装自定义算子库

4.1 编译自定义算子库

bash build.sh --vendor=cust --ops=allgather_aiv --custom_ops_path=./examples/05_custom_ops_allgather/aiv

参数说明：

--vendor：指定自定义算子标识符，示例中为cust，它决定了安装后 OPP 厂商目录名（vendors/cust）；
--ops：指定自定义算子名，示例中为allgather_aiv，用于生成算子库名称与安装包名称；
--custom_ops_path：指定自定义算子工程路径。

构建配置方面，aiv/CMakeLists.txt 会通过find_package(ASC REQUIRED)引入昇腾算子编译工具链，并以project(hccl_custom_${OP_NAME} ...)生成名为hccl_custom_allgather的工程，随后通过add_subdirectory(op_kernel)和add_subdirectory(op_host)分别编译 Kernel 侧.asc源文件与 Host 侧.cc源文件。

4.2 安装自定义算子包

编译完成后，安装包位于./build_out目录，使用--install参数安装：

./build_out/cann-hccl_custom_allgather_aiv_linux-<arch>.run --install --install-path=<ascend_cann_path>

参数说明：

<arch>：当前编译环境的系统架构（如aarch64或x86_64）；
<ascend_cann_path>：可选参数，指定 CANN 软件包安装目录；缺省时取ASCEND_CUSTOM_OPP_PATH或ASCEND_OPP_PATH环境变量指向的 CANN 软件包路径。

安装完成后，算子产物位置如下：

头文件：${ASCEND_HOME_PATH}/opp/vendors/cust/include/hccl_custom_allgather.h
动态库：${ASCEND_HOME_PATH}/opp/vendors/cust/lib64/libhccl_custom_allgather.so

其中${ASCEND_HOME_PATH}即 CANN-Toolkit 安装路径。对照 aiv/CMakeLists.txt 可以看到，头文件正是通过install(FILES ../inc/${PROJECT_NAME}.h DESTINATION ${CUSTOM_OPS_OPP_INC_PATH} COMPONENT hccl)安装到 OPP 厂商 include 目录的。

五、运行测试用例与结果验证

测试源码位于 examples/05_custom_ops_allgather/testcase，在第 4.1 节编译时已一并构建，测试二进制路径为：

./build/examples/05_custom_ops_allgather/testcase/custom_allgather_test

在仓库根目录执行：

export LD_LIBRARY_PATH=${ASCEND_HOME_PATH}/opp/vendors/cust/lib64:${LD_LIBRARY_PATH} cd build/examples/05_custom_ops_allgather/testcase mpirun -n rank_size ./custom_allgather_test data_len

参数说明：

rank_size：使用的卡数（即参与通信的 rank 数）；
data_len：每个 rank 发送的数据长度（元素个数，测试中按float计）。

从 testcase/main.cc 可以看出测试程序的完整执行流程：MPI_Init初始化 MPI → rank 0 通过HcclGetRootInfo生成根节点信息并经MPI_Bcast广播 → 每个 rank 调用HcclCommInitRootInfo初始化通信域 → 通过aclrtMalloc分配收发缓冲区（发送数据填充为各 rank 的 rank 号）→ 调用HcclAllGatherCustom(sendBuf, recvBuf, dataLen, HCCL_DATA_TYPE_FP32, hcclComm, stream)执行算子 →aclrtSynchronizeStream等待完成 →VerifyResult逐元素校验 recvBuf 中第 r 段是否等于 rank r 的原始数据（误差阈值 1e-5）→ 销毁通信域并释放资源。

预期输出

执行成功后，终端输出类似以下日志（以 2 卡为例）：

[1787902520.136766] [Rank 1] MPI Initialized. World Size: 2 [1787902520.136768] [Rank 0] MPI Initialized. World Size: 2 [1787902520.145917] [Rank 0] Device 0 selected (Total devices: 8) [1787902520.145918] [Rank 1] Device 1 selected (Total devices: 8) [1787902520.724696] [Rank 0] Root info generated [1787902520.724744] [Rank 0] HCCL set device[0] [1787902520.727436] [Rank 1] HCCL set device[1] [1787902522.982323] [Rank 0] HCCL Comm Initialized [1787902522.982908] [Rank 0] Buffers allocated and initialized [1787902523.008164] [Rank 1] HCCL Comm Initialized [1787902523.008742] [Rank 1] Buffers allocated and initialized rank1 dataLen=32 time=439 ms [1787902523.447898] [Rank 1] VerifyResult Passed! rank0 dataLen=32 time=465 ms [1787902523.447966] [Rank 0] VerifyResult Passed!

关键判断依据：每个 rank 打印VerifyResult Passed!即代表 AllGather 结果与预期一致（各 rank 数据按 rank 顺序拼接无误）。示例还支持无 MPI 模式——main.cc 中未定义ENABLE_MPI时，会使用多线程（每线程一个设备）模拟多 rank 执行。

六、Host 侧实现原理剖析

6.1 对外接口

对外接口声明于 inc/hccl_custom_allgather.h：

HcclResult HcclAllGatherCustom( void* sendBuf, void* recvBuf, uint64_t sendCount, HcclDataType dataType, HcclComm comm, aclrtStream stream);

签名与 HCCL 内置的HcclAllGather保持一致，便于替换。接口使用extern "C"导出，确保动态库符号可被 C/C++ 共同链接。

6.2 算子入口与参数装配

入口实现在 aiv/op_host/allgather.cc 的HcclAllGatherCustom中，主要完成三件事：

通过HcclGetCommName获取通信域名，拼出tag（格式为<commName>_opbase），用于后续引擎上下文与信道的标识；
调用PrepareResources完成 AIV 缓冲区、信道与对端内存地址的准备工作；
装配OpParam结构体并调用LaunchKernel(param, stream)提交 Kernel。

OpParam定义于 aiv/inc/common.h，它同时是 Host 与 Device 之间传递的"契约"，关键字段包括：

input/output：发送/接收缓冲区地址（GM 地址）；
rank/rankSize：当前 rank 号与通信域大小；
xRankSize/yRankSize/zRankSize：三维拓扑各维度大小（本示例为一维 Mesh，仅xRankSize = rankSize，y/z 置 0）；
len：数据总字节数，由sendCount * SIZE_TABLE[dataType]计算得到，SIZE_TABLE覆盖 int8/int16/int32/fp16/fp32/int64/uint64 等 HCCL 数据类型的大小；
inputSliceStride/outputSliceStride：数据切片步长，AllGather 场景下等于len；
tagId：映射到 Kernel 内的同步 tag；
isOpBase：标记是否为算子基座（op base）执行模式。

6.3 资源准备：引擎上下文、信道与对端地址

资源准备是整个 Host 侧的核心，流程如下：

第一步，初始化 AIV 缓冲区（InitAivBuffer）：调用HcclEngineCtxGet查询引擎上下文，若不存在则调用HcclEngineCtxCreate(comm, aivTag, COMM_ENGINE_AIV, ...)创建，随后aclrtMemset清零通信信息区，并通过HcclCommMemReg将该缓冲区注册为通信内存。该缓冲区在 Kernel 中承载各 rank 的 GM 地址表、flag 同步区与拓扑信息。

第二步，构建信道请求（BuildChannelRequests）：对除自身外的每个远端 rank，通过HcclRankGraphGetLayers获取网络分层、HcclRankGraphGetLinks获取本 rank 与远端 rank 之间的链路，筛选出协议为COMM_PROTOCOL_UB_MEM（UB 内存协议，即通过统一内存池直达通信）的链路，填充HcclChannelDesc（本地/远端端点协议、通信地址、位置、notifyNum = 3等）。

第三步，获取信道与对端缓冲区（AcquireChannelsAndBuffers）：调用HcclChannelAcquire批量获取信道，再通过HcclChannelGetHcclBuffer拿到远端 rank 的 HCCL 通信缓冲区地址（存入buffersIn），通过HcclChannelGetRemoteMems拿到远端 rank 的内存区中最后一个CommMem的地址（存入buffersOut，供 flag 同步使用）。

第四步，下发地址表：PrepareResources将buffersIn与buffersOut两张地址表分别通过aclrtMemcpy写入 AIV 通信信息区（aivCommInfoPtr）及其AIV_TAG_ADDR_OFFSET = 16KB偏移处，Device 侧 Kernel 即可据此访问各 rank 的输入输出内存。

6.4 Kernel 提交

aiv/op_host/launch_kernel.cc 负责 Kernel 的注册与提交：

RegisterKernel：读取二进制文件hccl_custom_allgather_kernels.o，依次调用aclrtCreateBinary、aclrtBinaryLoad、aclrtBinaryGetFunction获取名为HcclAllGatherAivKernel的函数句柄；通过静态标志g_init与互斥锁保证单次注册；
ExecuteKernelLaunch：构造aclrtLaunchKernelCfg，设置三个 launch 属性——ACL_RT_LAUNCH_KERNEL_ATTR_SCHEM_MODE = 1（使用算子的 scheme 调度模式）、ACL_RT_LAUNCH_KERNEL_ATTR_TIMEOUT_US（超时时间由CUSTOM_TIMEOUT = 1836秒换算为微秒）、ACL_RT_LAUNCH_KERNEL_ATTR_ENGINE_TYPE = ACL_RT_ENGINE_TYPE_AIV（明确指定提交到 AIV 引擎），最后通过aclrtLaunchKernelWithHostArgs将OpParam作为 Host 参数随 Kernel 一起提交。

七、Device 侧 Kernel 实现原理剖析

7.1 Kernel 入口与数据类型分发

Kernel 入口为 aiv/op_kernel/launch_kernel_asc.asc 中的HcclAllGatherAivKernel（extern "C" __global__ __aicore__），按param.dataType分发到不同模板实例（int8/int32/fp16/fp32），并通过EXPORT_AIV_META_INFO宏将 Kernel 元信息写入.ascend.meta段，供运行时识别其为 AIV 类型 Kernel（K_TYPE_AIV）。

7.2 算法主体：Mesh 1D 全收集

核心算法在 aiv/inc/aiv_allgather_mesh_1d.h 的AivAllGatherMesh1D类中，继承自AivCommBase：

InitCoreInfo：按GetBlockIdx()（核号）与核数numBlocks_将总长度均分，处理余数，计算本核负责的数据段偏移coreOffset与长度curCount；
Run：首先用CpGM2GM把自己的输入数据搬移到本 rank 在远端地址表中的目标位置，Record(rank_, ...)置起自身就绪 flag；然后遍历所有 rank，WaitFlag(rank, ...)等待该 rank 的数据就绪后，将其数据通过CpGM2GM搬入输出缓冲output_ + rank * stride的对应分片。这正是 AllGather"每个 rank 收集齐全部 rank 的数据分片"的语义；
Process：当核数numBlocks_ >= rankSize_时走单核对齐路径（Run），否则走控制核辅助路径（RunCtrlCore），后者由 block 0 作为协调者，先收集本 rank 各核就绪信号，再统一触发各 rank 的数据搬运，保证核数小于 rank 数时依然正确。

7.3 通信基类：GM 搬运与 flag 同步原语

aiv/inc/aiv_communication_base_v2.h 中的AivCommBase提供了两个关键能力：

GM→GM 数据搬运（CpGM2GM）：数据经 UB（统一缓冲区）中转，单次搬运上限为UB_MAX_DATA_SIZE = 190KB（双缓冲模式取半），通过inOutQue队列以"EnQue/DeQue"流水方式分段搬运，避免一次性占用过多 UB 空间。

基于 flag 的 rank 间同步（Record/WaitFlag）：Record(targetRank, flagOffset, curTag)将当前 tag 写入目标 rank 的 flag 区；WaitFlag循环DataCopyGM2UB读取 flag 并比对 tag，直到相等才继续。BarrierAll与BarrierForFirstOP则在此基础上实现全 rank 汇聚屏障，其中GetTag通过读写AIV_FLAG_CLEAR_OFFSET处的计数器维护单调递增的 tag（达到TAG_RESET_COUNT = 4096后回绕为 1），保证多轮算子调用之间 flag 不会被误判。

八、总结

本示例完整展示了 HCCL AIV 自定义通信算子的标准开发范式：Host 侧通过HcclEngineCtxCreate/HcclCommMemReg/HcclChannelAcquire等 HCCL 资源接口准备引擎上下文、信道与对端地址表，通过aclrtLaunchKernelWithHostArgs将结构化参数提交至 AIV 引擎；Device 侧以 Ascend C 编写__aicore__Kernel，借助 flag 同步原语与 UB 中转搬运实现 Mesh 1D 拓扑下的 AllGather 全收集语义，并以 MPI 多进程测试完成端到端正确性验证。

作为进一步学习的入口，建议对照阅读仓库中的 06_custom_ops_reduce_scatter（ReduceScatter 的 AIV/CCU 变体）、04_custom_ops_p2p（点对点 Send/Recv）以及 experimental/ops/all_reduce（含递归执行器的实验性实现），从不同算子与不同执行路径的对比中加深对 HCCL 通信框架整体架构的理解。

创作声明：本文部分内容由AI辅助生成（AIGC），仅供参考