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

CANN HCCL 自定义集合通信算子开发基于 AIV 引擎的 AllGather 算子实战解析【免费下载链接】hccl集合通信库Huawei Collective Communication Library简称HCCL是基于昇腾AI处理器的高性能集合通信库为计算集群提供高性能、高可靠的通信方案项目地址: https://gitcode.com/cann/hcclHCCLHuawei Collective Communication Library是昇腾 AI 处理器上的高性能集合通信库其提供的 AIVAI Vector通信编程接口允许开发者绕过内置算法自行实现通信算子的 Host 侧逻辑与 Device 侧 Kernel。本文以仓库中的 aiv 版自定义 AllGather 示例 为主线完整讲解从环境准备、算子库编译安装、MPI 测试运行到 Host 侧资源管理与 Device 侧 Mesh 1D 通信算法的源码级实现原理。读完本文你将掌握基于 HCCL AIV 接口从零开发一个可编译、可安装、可验证的自定义集合通信算子的完整路径。一、示例概述做什么、支持什么该示例展示了如何基于 HCCL AIV 通信编程接口开发 AllGather 自定义通信算子核心特性如下基于 AIVAI Vector通信引擎实现 AllGather 集合通信算子同时包含 Host 侧算子逻辑与 Device 侧 Kernel 实现提供完整的编译、构建与测试验证流程。支持的产品与场景单机 N 卡配置N 2Ascend 950PR / Ascend 950DTAtlas 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 --vendorcust --opsallgather_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-pathascend_cann_path参数说明arch当前编译环境的系统架构如aarch64或x86_64ascend_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 dataLen32 time439 ms [1787902523.447898] [Rank 1] VerifyResult Passed! rank0 dataLen32 time465 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.hHcclResult 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 rankSizey/z 置 0len数据总字节数由sendCount * SIZE_TABLE[dataType]计算得到SIZE_TABLE覆盖 int8/int16/int32/fp16/fp32/int64/uint64 等 HCCL 数据类型的大小inputSliceStride/outputSliceStride数据切片步长AllGather 场景下等于lentagId映射到 Kernel 内的同步 tagisOpBase标记是否为算子基座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_MEMUB 内存协议即通过统一内存池直达通信的链路填充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 中的HcclAllGatherAivKernelextern C __global__ __aicore__按param.dataType分发到不同模板实例int8/int32/fp16/fp32并通过EXPORT_AIV_META_INFO宏将 Kernel 元信息写入.ascend.meta段供运行时识别其为 AIV 类型 KernelK_TYPE_AIV。7.2 算法主体Mesh 1D 全收集核心算法在 aiv/inc/aiv_allgather_mesh_1d.h 的AivAllGatherMesh1D类中继承自AivCommBaseInitCoreInfo按GetBlockIdx()核号与核数numBlocks_将总长度均分处理余数计算本核负责的数据段偏移coreOffset与长度curCountRun首先用CpGM2GM把自己的输入数据搬移到本 rank 在远端地址表中的目标位置Record(rank_, ...)置起自身就绪 flag然后遍历所有 rankWaitFlag(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/WaitFlagRecord(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_scatterReduceScatter 的 AIV/CCU 变体、04_custom_ops_p2p点对点 Send/Recv以及 experimental/ops/all_reduce含递归执行器的实验性实现从不同算子与不同执行路径的对比中加深对 HCCL 通信框架整体架构的理解。【免费下载链接】hccl集合通信库Huawei Collective Communication Library简称HCCL是基于昇腾AI处理器的高性能集合通信库为计算集群提供高性能、高可靠的通信方案项目地址: https://gitcode.com/cann/hccl创作声明：本文部分内容由AI辅助生成（AIGC），仅供参考