基于 HCCL AIV 通信引擎开发 ReduceScatter 自定义通信算子：完整实战指南-开发者社区

基于 HCCL AIV 通信引擎开发 ReduceScatter 自定义通信算子：完整实战指南

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

导读

本文基于 CANN HCCL 开源仓库中的examples/06_custom_ops_reduce_scatter/aiv样例，系统讲解如何基于 AIV（AI Vector）通信引擎从零开发一个 ReduceScatter 自定义集合通信算子，涵盖工程目录设计、Host 侧算子逻辑、Device 侧 Kernel 实现、资源申请与通道建立、编译安装包生成，以及基于 MPI 的多卡测试验证全流程。读完本文，你将掌握 HCCL 自定义通信算子的完整开发范式，能够照此方法扩展实现 AllReduce、AllGather 等其它集合通信算子。

一、样例总体介绍

本样例的核心目标是：绕过 HCCL 内置算子逻辑，直接基于 AIV 通信编程接口实现 ReduceScatter 集合通信算子。ReduceScatter 的语义是：通信域内每个 rank 将自己的完整输入数据按 rank 数均分，对每个分片在所有 rank 上做规约（如 SUM），最终每个 rank 得到对应自己分片的规约结果。

样例的主要功能点：

基于 AIV（AI Vector）通信引擎实现 ReduceScatter 集合通信算子；
包含 Host 侧算子逻辑与 Device 侧 Kernel 实现，两者通过OpParam结构体传递参数；
提供完整的 CMake 编译构建流程与基于 MPI 的多卡测试验证流程。

从仓库结构看，整个样例位于 examples/06_custom_ops_reduce_scatter，其下又分为aiv/（AIV 引擎实现，本文主题）、ccu/（CCU 实现）与testcase/（测试用例）三个子工程，体现了 HCCL 自定义算子支持多种通信引擎的设计理念。

二、工程目录结构解析

样例的目录结构如下（来自 aiv/README.md 并补充了源码文件说明）：

├── CMakeLists.txt # 根目录编译/构建配置文件 ├── op_host/ │ ├── CMakeLists.txt │ ├── reduce_scatter.cc # HcclReduceScatterCustom 算子Host侧实现 │ ├── launch_kernel.cc # Kernel 下发逻辑实现 │ └── launch_kernel.h # Kernel 下发接口定义 ├── op_kernel/ │ ├── CMakeLists.txt │ └── launch_kernel_asc.asc # 算子 Kernel 侧实现 (Ascend C) └── inc/ ├── hccl_custom_reduce_scatter.h # 自定义算子对外接口头文件 ├── common.h # 公共类型定义与宏 ├── aiv_reduce_scatter_mesh_1d.h # AIV ReduceScatter 核心算法实现 ├── aiv_communication_base_v2.h # AIV 通信基类 ├── log.h # 日志工具 ├── extra_args.h # 额外参数定义 └── sync_interface.h # 同步接口定义

各模块职责划分：

op_host/：运行在 Host 侧的算子逻辑。reduce_scatter.cc负责组装参数、申请通信资源、建立通道；launch_kernel.cc负责加载 kernel 二进制并通过 ACL 接口下发到 Device。
op_kernel/：Device 侧 Ascend C kernel。launch_kernel_asc.asc定义HcclReduceScatterAivKernel入口，按数据类型分发到具体的模板实现。
inc/：公共头文件。common.h中定义了 Host 与 Device 共享的OpParam参数结构体和数据类型字节大小表；aiv_reduce_scatter_mesh_1d.h是 1D Mesh 拓扑下的核心通信算法。

三、环境准备

3.1 环境要求

本样例支持以下产品，组网为单机 N 卡（N>=2）：

Ascend 950PR / Ascend 950DT

需要说明的是，AIV 通信引擎的可用性与硬件平台强相关，以上产品列表是当前样例明确声明支持的范围，换用其它型号硬件时需以实际支持的通信引擎为准。

3.2 安装 CANN Toolkit 开发套件包

参考昇腾文档中心的 CANN 软件安装指南，安装最新版本 CANN Toolkit 开发套件包。CANN Toolkit 提供了编译自定义算子所需的hccl_comm.h、hccl_res.h、hcomm_primitives.h等头文件以及libhccl.so、libascendcl.so等运行库。

3.3 配置环境变量

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

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

此外，运行测试用例需要 MPI 环境支持，请确保已安装并配置好 MPI。MPI 配置请参考配套版本的昇腾文档中心 HCCL 性能测试工具使用指南中的"MPI 安装与配置"章节。MPI 在本样例中承担两个职责：一是通过MPI_Init/MPI_Comm_rank/MPI_Comm_size获取全局 rank 与总进程数；二是通过MPI_Bcast将 0 号 rank 生成的HcclRootInfo根节点信息广播给所有进程，用于建立通信域。

四、编译与运行

4.1 编译自定义算子库

样例提供了基于 CMake 的构建流程，在仓库根目录下执行以下命令：

bash build.sh --vendor=cust --ops=reduce_scatter_aiv --custom_ops_path=./examples/06_custom_ops_reduce_scatter/aiv

参数说明：

--vendor：自定义算子标识，这里为cust，对应 CANN 的opp/vendors/cust自定义算子厂商目录；
--ops：自定义算子名称，reduce_scatter_aiv对应aiv子工程的实现；
--custom_ops_path：自定义算子工程路径，即仓库根目录下的./examples/06_custom_ops_reduce_scatter/aiv。

从 aiv/CMakeLists.txt 可以看到，构建过程会通过find_package(ASC REQUIRED)引入 Ascend C 编译工具链，project(hccl_custom_${OP_NAME} LANGUAGES CXX ASC)同时编译 C++ 与 ASC（Ascend C kernel 源文件），并将对外头文件hccl_custom_reduce_scatter.h安装到自定义算子头文件目录。构建中还会将 op_kernel/CMakeLists.txt 编译生成的 kernel 二进制与 op_host/CMakeLists.txt 编译生成的 Host 侧动态库一起打包进安装包。

4.2 安装算子包

自定义算子安装包生成在./build_out目录下，通过--install参数进行安装：

./build_out/cann-hccl_custom_reduce_scatter_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_reduce_scatter.h
动态库：${ASCEND_HOME_PATH}/opp/vendors/cust/lib64/libhccl_custom_reduce_scatter.so

其中${ASCEND_HOME_PATH}为 CANN-Toolkit 安装路径。该动态库即测试程序需要链接的自定义算子库（见 testcase/CMakeLists.txt 中的target_link_libraries(... hccl_custom_reduce_scatter)）。

4.3 运行测试用例

测试代码位于 testcase/main.cc，已在前节"编译自定义算子库"中一并编译。测试样例二进制文件路径为：

./build/examples/06_custom_ops_reduce_scatter/testcase/custom_reduce_scatter_test

从仓库根目录开始执行以下命令：

export LD_LIBRARY_PATH=${ASCEND_HOME_PATH}/opp/vendors/cust/lib64:${LD_LIBRARY_PATH} cd build/examples/06_custom_ops_reduce_scatter/testcase/ mpirun -n rank_size ./custom_reduce_scatter_test data_len

参数说明：

rank_size：使用的卡数；
data_len：数据长度（每个 rank 接收分片包含的元素个数）。

测试用例支持两种运行模式（由是否定义ENABLE_MPI宏决定）：

MPI 模式（ENABLE_MPI）：各 MPI 进程分别绑定设备，0 号进程生成HcclRootInfo后经MPI_Bcast广播，所有进程通过HcclCommInitRootInfo建立通信域；
线程模式：单进程内创建rank_size个线程，每个线程绑定一块设备，共享同一份HcclRootInfo，同样通过HcclCommInitRootInfo建立通信域。

4.4 预期结果

运行成功后，终端将输出类似以下的日志信息（以 2 卡运行为例，来自 aiv/README.md）：

[1786071476.120968] [Rank 0] MPI Initialized. World Size: 2 [1786071476.120968] [Rank 1] MPI Initialized. World Size: 2 [1786071476.127411] [Rank 0] Device 0 selected (Total devices: 8) [1786071476.127411] [Rank 1] Device 1 selected (Total devices: 8) [1786071478.023709] [Rank 0] Root info generated [1786071478.023786] [Rank 0] HCCL set device[0] [1786071478.023778] [Rank 1] HCCL set device[1] [1786071483.214938] [Rank 0] HCCL Comm Initialized [1786071483.221873] [Rank 0] Buffers allocated and initialized [1786071483.254098] [Rank 1] HCCL Comm Initialized [1786071483.259378] [Rank 1] Buffers allocated and initialized rank1 dataLen=1024 time=835 ms [1786071484.095144] [Rank 1] VerifyResult Passed! rank0 dataLen=1024 time=873 ms [1786071484.095200] [Rank 0] VerifyResult Passed!

日志中的关键节点对应测试用例的执行阶段：MPI 初始化 → 设备选择 → RootInfo 生成与广播 → 设置设备 → 通信域初始化 → 缓冲区分配与初始化 → 执行自定义算子 → 结果校验通过。

五、源码级解析：自定义算子如何工作

5.1 对外接口定义

自定义算子的对外 API 定义在 hccl_custom_reduce_scatter.h，签名与 HCCL 内置HcclReduceScatter保持一致的调用习惯：

HcclResult HcclReduceScatterCustom( void* sendBuf, void* recvBuf, uint64_t recvCount, HcclDataType dataType, HcclReduceOp op, HcclComm comm, aclrtStream stream);

各参数含义：

sendBuf：发送缓冲区（Device 侧内存），长度为recvCount * rankSize个元素；
recvBuf：接收缓冲区（Device 侧内存），长度为recvCount个元素；
recvCount：每个 rank 接收分片的元素个数；
dataType：数据类型，如HCCL_DATA_TYPE_FP32；
op：规约操作，如HCCL_REDUCE_SUM；
comm：已初始化的 HCCL 通信域句柄；
stream：ACL 任务流。

5.2 Host 侧实现（reduce_scatter.cc）

Host 侧核心逻辑位于 op_host/reduce_scatter.cc，整个执行流程分为三步：

第一步：申请 AIV 通信上下文缓冲区（InitAivBuffer）

调用HcclEngineCtxGet/HcclEngineCtxCreate以CommEngine::COMM_ENGINE_AIV引擎类型申请一块 2MB 的标签缓冲区（AIV_TAG_BUFF_LEN = 2 * 1024 * 1024，定义于 common.h），清零后通过HcclCommMemReg注册到通信域并获得HcclMemHandle。缓冲区按算子 tag 缓存在g_memHandleCache中（带互斥锁保护），避免重复申请。注意AIV_TAG_ADDR_OFFSET = 16 * 1024的偏移量设计：缓冲区前 16KB 存放远端 rank 的发送缓冲区地址表，偏移之后存放接收缓冲区地址表。

第二步：建立通道（BuildChannelRequests/AcquireChannelsAndBuffers）

通过HcclRankGraphGetLayers获取网络分层、HcclRankGraphGetLinks获取本地 rank 与每个远端 rank 之间的链路信息，筛选出COMM_PROTOCOL_UB_MEM（UB 内存协议）的链路，为每个远端 rank 构建一个HcclChannelDesc通道描述（notifyNum = 3，即每个通道需要 3 个通知信号）。随后调用HcclChannelAcquire批量申请通道，并通过HcclChannelGetHcclBuffer与HcclChannelGetRemoteMems获取本地与远端的通信缓冲区地址，填入buffersIn/buffersOut数组。

第三步：组装参数并下发 Kernel（HcclReduceScatterCustom）

通过HcclGetCommName获取通信域名，拼接生成算子 tag（格式<commName>_opbase）；
通过HcclGetRankId/HcclGetRankSize获取当前 rank 与总 rank 数；
计算param.len = recvCount * SIZE_TABLE[dataType]（字节数），SIZE_TABLE定义了各数据类型的字节大小；
设置切片步长inputSliceStride = param.len、outputSliceStride = param.len、repeatNum = 1等 Mesh 1D 算法参数；
将本 rank 的发送缓冲区地址写入buffersIn[rank]、接收缓冲区地址写入buffersOut[rank]，连同远端地址表一并拷贝到 AIV 缓冲区中，使 Device 侧 Kernel 可以直接索引；
调用LaunchKernel(param, stream)完成下发。

5.3 Kernel 下发机制（launch_kernel.cc）

Kernel 的加载与下发实现在 op_host/launch_kernel.cc，采用"加载一次、多次复用"的注册机制：

RegisterKernel：从hccl_custom_reduce_scatter_kernels.o二进制文件读取编译产物，依次调用aclrtCreateBinary、aclrtBinaryLoad、aclrtBinaryGetFunction获取HcclReduceScatterAivKernel函数句柄。整个过程由互斥锁和原子变量g_init保护，保证单例注册；
ExecuteKernelLaunch：构造aclrtLaunchKernelCfg配置，设置三个关键属性：
- ACL_RT_LAUNCH_KERNEL_ATTR_SCHEM_MODE = 1：启用 scheme 模式调度；
- ACL_RT_LAUNCH_KERNEL_ATTR_TIMEOUT_US：超时时间，取CUSTOM_TIMEOUT * 1000000微秒（CUSTOM_TIMEOUT = 1836，见 common.h）；
- ACL_RT_LAUNCH_KERNEL_ATTR_ENGINE_TYPE = ACL_RT_ENGINE_TYPE_AIV：指定在 AIV 引擎上执行；
最终通过aclrtLaunchKernelWithHostArgs(g_funcHandle, 1, stream, &cfg, &param, sizeof(OpParam), nullptr, 0)将整个OpParam结构体以 Host 参数形式原样传递给 Device Kernel。

5.4 Device 侧 Kernel 实现（launch_kernel_asc.asc）

Device 侧入口定义在 op_kernel/launch_kernel_asc.asc，是一个extern "C" __global__ __aicore__内核函数，按param.dataType分发到具体类型模板：

extern "C" __global__ __aicore__ void HcclReduceScatterAivKernel(OpParam param) { switch(param.dataType) { case AscendC::HCCL_DATA_TYPE_INT8: CALL_AIV_KERNEL(int8_t); break; case AscendC::HCCL_DATA_TYPE_INT32: CALL_AIV_KERNEL(int32_t); break; case AscendC::HCCL_DATA_TYPE_FP16: CALL_AIV_KERNEL(half); break; case AscendC::HCCL_DATA_TYPE_FP32: CALL_AIV_KERNEL(float); break; default: break; } }

CALL_AIV_KERNEL宏将OpParam的各字段展开为AivReduceScatterMesh1D<T>的形参，包括输入/输出地址、rank/rankSize、三维拓扑尺寸（xRankSize/yRankSize/zRankSize，本样例仅使用 1D Mesh，故yRankSize/zRankSize置 0）、数据长度、规约操作、切片步长与 repeat 参数等。

5.5 核心算法：AivReduceScatterMesh1D

核心通信算法位于 aiv_reduce_scatter_mesh_1d.h，类AivReduceScatterMesh1DOp继承自AivCommBase（定义于 aiv_communication_base_v2.h），算法分两个阶段：

阶段一：本 rank 数据搬入通信区

InitCoreInfo按coreId对总数据量做均分（含余数分配），计算出每个 AIV Core 负责的字节偏移coreOffset与元素数curCount。Run中首先执行CpGM2GM将本 rank 输入数据中属于自己分片的部分（rank_ * stride + coreOffset）拷贝到本 rank 的 AIV 通信缓冲区，随后调用Record(rank_, GetBlockIdx() * FLAG_SIZE + flagOffset, curTag)记录完成标志，PipeBarrier<PIPE_ALL>()保证流水线同步。

阶段二：跨 rank 规约

遍历所有 rank：WaitFlag等待对应远端 rank 的数据就绪标志，然后将本 rank AIV 缓冲区中远端 rank 对应分片的数据累加到自己的输出缓冲区output_ + coreOffset（CpGM2GM(output, gmOthers, curCount, reduceOp_)，携带规约操作符）。最终每个 rank 的输出缓冲区即所有 rank 对应分片的规约结果。

从源码实现可以推断：该算法采用"先本地数据上板、再逐 rank 规约累加"的 Mesh 1D 通信模式，通信量在 rank 间直接对等交换，配合Record/WaitFlag标志机制实现跨卡同步，属于典型的 AIV 通信原语组合。

六、测试用例设计要点

测试程序 testcase/main.cc 的验证逻辑设计值得借鉴：

数据构造：每个 rank 的发送缓冲区填充float(rank + TEST_DATA_VALUE)，其中TEST_DATA_VALUE = 2026，即 rank0 全为 2026.0、rank1 全为 2027.0……接收缓冲区清零；
期望值计算：float sum = (TEST_DATA_VALUE + TEST_DATA_VALUE + rankSize - 1) * rankSize / 2，即等差数列求和。对 FP32 SUM 规约而言，每个 rank 分片的期望值都等于所有 rank 数据的累加和；
结果校验：VerifyResult将接收缓冲区拷回 Host，逐元素与期望值比较，容差1e-5，全部通过则打印VerifyResult Passed!；
性能观测：使用high_resolution_clock统计HcclReduceScatterCustom调用到aclrtSynchronizeStream返回的耗时并打印rankN dataLen=... time=... ms。

测试用例通过 testcase/CMakeLists.txt 链接ascendcl、hccl、hccl_custom_reduce_scatter三个库（MPI 模式下追加 MPI 库），并显式设置_GLIBCXX_USE_CXX11_ABI=0与-std=c++17，与 CANN 运行时保持 ABI 一致。

七、小结与扩展建议

本文完整走通了"环境准备 → 编译 → 安装 → 运行 → 验证"的 HCCL AIV 自定义 ReduceScatter 算子开发链路，并从源码层面剖析了 Host 侧资源申请/通道建立/参数组装、Kernel 加载下发、Device 侧 Mesh 1D 规约算法与测试验证机制。

如果想基于此样例做扩展，可参考以下方向：

更换规约操作：修改HcclReduceScatterCustom传入的op参数（如HCCL_REDUCE_MAX），并确认 Device 侧CpGM2GM携带的reduceOp_支持对应操作；
扩展数据类型：在launch_kernel_asc.asc的 switch 分支中增加如HCCL_DATA_TYPE_INT16、HCCL_DATA_TYPE_FP64等类型分支；
移植到其它集合通信算子：本样例的工程骨架（op_host+op_kernel+inc三段式结构、OpParam参数传递、RegisterKernel单例注册）同样适用于 AllGather、AllReduce 等算子的自定义实现，仓库中的 05_custom_ops_allgather 即提供了 AIV 版 AllGather 的对照实现。