自定义通信算子 - AllGather 通信
样例介绍
本样例展示如何基于 HCCL AIV 通信编程接口开发 AllGather 自定义通信算子,主要功能点:
- 基于 AIV (AI Vector) 通信引擎实现 AllGather 集合通信算子。
- 包含 Host 侧算子逻辑与 Device 侧 Kernel 实现。
- 提供完整的编译构建与测试验证流程。
目录结构
├── CMakeLists.txt # 根目录编译/构建配置文件
├── op_host/
│ ├── CMakeLists.txt
│ ├── all_gather.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 # 公共类型定义与宏
│ ├── aiv_all_gather_mesh_1d.h # AIV AllGather 核心算法实现
│ ├── aiv_communication_base_v2.h # AIV 通信基类
│ ├── log.h # 日志工具
│ ├── extra_args.h # 额外参数定义
│ └── sync_interface.h # 同步接口定义
└── testcase/
├── CMakeLists.txt # 测试用例 CMake 配置文件
├── Makefile # 测试用例 Makefile (用于编译运行)
└── main.cc # 测试用例主程序
一、环境准备
1. 环境要求
本样例支持以下产品,组网为单机N卡(N>=2):
- Ascend 950PR / Ascend 950DT
2. 安装 CANN Toolkit 开发套件包
参考 昇腾文档中心-CANN软件安装指南,安装最新版本 CANN Toolkit 开发套件包。
3. 配置环境变量
以 root 用户默认安装路径为例:
source /usr/local/Ascend/cann/set_env.sh
此外,运行测试用例需要 MPI 环境支持,请确保已安装并配置好 MPI。
二、编译与运行
本样例提供了基于 CMake 的构建流程以及基于 Makefile 的测试运行脚本。
1. 编译自定义算子库
在样例根目录下执行以下命令:
# 1. 创建构建目录
mkdir build
# 2. 进入构建目录
cd build
# 3. 执行 CMake 配置
cmake ..
# 4. 编译项目 (生成 libhccl_custom_allgather.so)
make
2. 运行测试用例
编译完成后,进入 testcase 目录执行测试:
# 5. 进入测试用例目录
cd ../testcase
# 6. 编译并运行测试用例
# 该命令会自动编译测试程序,设置 LD_LIBRARY_PATH 并使用 mpirun 运行
make run
3. 预期结果
运行成功后,终端将输出类似以下的日志信息(以 2 卡运行为例):
[INFO] MPI Initialized. World Size: 2
[INFO] Device 0 selected (Total devices: 8)
[INFO] Device 1 selected (Total devices: 8)
[INFO] HCCL Comm Initialized
[INFO] Buffers allocated and initialized
[INFO] Starting HcclAllGatherCustom...
[INFO] HcclAllGatherCustom completed and synchronized
[INFO] Test Passed!