CANN ascend-transformer-boost GatherOperation C++ Demo 实战指南：从环境搭建到算子调用全流程

CANN ascend-transformer-boost GatherOperation C++ Demo 实战指南：从环境搭建到算子调用全流程

【免费下载链接】ascend-transformer-boost本项目是CANN提供的是一款高效、可靠的Transformer加速库，基于华为Ascend AI处理器，提供Transformer定制化场景的高性能融合算子。项目地址: https://gitcode.com/cann/ascend-transformer-boost

本指南以 example/op_demo/gather/README.md 为骨架，系统讲解在 CANN ascend-transformer-boost 加速库中如何编译、运行 GatherOperation（按索引收集切片）C++ 调用示例。文章不仅完整覆盖官方 README 的环境配置与编译步骤，还深入 gather_demo.cpp 的每一行核心逻辑，并结合算子实现、Kernel 源码与 Python 测试用例，帮助读者真正掌握 ATB 算子调用的标准五步流程、GatherParam 参数语义以及双 Runner（Ascend 950 aclnn 与通用 Ops）实现机制。读完本文，你将具备独立编写和运行任意 ATB 推理算子 Demo 的能力。

一、Demo 概览：GatherOperation 是什么

Gather 算子（对应 infer_op_params.h 中atb::infer::GatherParam）的功能是从输入张量中根据索引（indices）收集切片，并将这些切片组合成一个新的张量，是 Transformer 推理中 embedding 查表、token 重排、batch 内索引选择等场景的常用算子。

本目录（example/op_demo/gather/）下即为该算子的 C++ 调用示例，包含三个文件：

• README.md：本说明文档；
• README_en.md：英文版说明；
• gather_demo.cpp：完整的 C++ Demo 源码。

示例的数学模型非常简单：输入x为 shape[3, 3]的张量，内容为[[1,2,3],[4,5,6],[7,8,9]]（存储为扁平向量{1..9}）；索引indices为[2, 1]；在axis=0上收集，输出 shape 为[2, 3]，即[[7,8,9],[4,5,6]]。运行成功的标志是终端打印Gather demo success!。

二、环境准备：source CANN 与 nnal 安装路径

运行 Demo 前，必须保证 CANN 工具包与加速库（nnal/atb）的运行时环境已就绪。README 给出的标准流程是依次 source 两个环境脚本：

1. source CANN 安装路径下的set_env.sh，默认路径为：

   source /usr/local/Ascend/ascend-toolkit/set_env.sh

2. source nnal（加速库）安装路径下的set_env.sh，默认路径为：

   source /usr/local/Ascend/nnal/atb/set_env.sh

   特别说明：如果你使用的是加速库源码编译产物，则改为 source 源码编译输出目录下的脚本，例如：

   source ./ascend-transformer-boost/output/atb/set_env.sh

两个脚本缺一不可：CANN 的set_env.sh提供 ACL（Ascend Computing Language）运行时头文件与库；加速库的set_env.sh提供 ATB 头文件（atb/atb_infer.h、atb/operation.h、atb/types.h等）以及编译好的libatb库。关于整体编译方式可参考仓库根目录文档 编译与构建。

三、编译与运行 Demo

3.1 构建命令

在示例目录下直接执行：

bash build.sh

即可完成编译与运行。需要注意的是：仓库中该目录并未提交build.sh文件，需要读者参照编译要求自行编写构建脚本，核心是使用g++编译 gather_demo.cpp，并链接 CANN 的 ACL 库与 ATB 库。仓库其他示例（如 example/op_demo/linear/）已提供现成 CMake 构建方式可供参考。

3.2 关于D_GLIBCXX_USE_CXX11_ABI的编译注意点

README 明确强调：ATB 加速库在编译时对 C++ 标准库 ABI 有要求，编译 Demo 时必须与加速库本体保持一致，否则会出现链接或运行时符号不匹配问题。具体分为两种情况：

• 使用 cxx_abi=0（默认）：设置D_GLIBCXX_USE_CXX11_ABI=0，即编译命令形如：

  g++ -D_GLIBCXX_USE_CXX11_ABI=0 -I ... gather_demo.cpp ... -o gather_demo

• 使用 cxx_abi=1：更改宏为 1，即：

  g++ -D_GLIBCXX_USE_CXX11_ABI=1 -I ... gather_demo.cpp ... -o gather_demo

D_GLIBCXX_USE_CXX11_ABI是 GCC 控制新旧两种std::string/std::listABI 的宏：值为 0 时使用旧版（C++03 风格）ABI，值为 1 时使用新版（C++11 风格）ABI。ATB 库在构建安装时选择了哪个 ABI，你的 Demo 就必须用同样的宏编译，这是新手最容易踩的坑。

3.3 头文件与链接库

Demo 源码通过 demo_util.h 引入：

• acl/acl.h：ACL 运行时 API（设备管理、Stream、内存分配）；
• atb/atb_infer.h、atb/operation.h、atb/types.h：ATB 算子框架、atb::Operation、atb::Tensor、atb::VariantPack等核心类型。

链接时需要同时链接 ACL 库（如-lascendcl）与 ATB 库（如-latb），具体库名以 CANN/nnal 安装环境为准。

四、Demo 源码逐段精读：ATB 算子调用的标准五步流程

gather_demo.cpp 完整展示了 ATB 算子调用范式。其主流程可分为五个阶段，这也是所有 ATB 算子 Demo 的统一骨架。

4.1 第一阶段：初始化 ACL、创建 Context 与 Stream

CHECK_STATUS(aclInit(nullptr)); int32_t deviceId = 0; CHECK_STATUS(aclrtSetDevice(deviceId)); atb::Context *context = nullptr; CHECK_STATUS(atb::CreateContext(&context)); void *stream = nullptr; CHECK_STATUS(aclrtCreateStream(&stream)); context->SetExecuteStream(stream);

要点解析：

• aclInit(nullptr)初始化 ACL 运行时；aclrtSetDevice(deviceId)指定使用的昇腾设备（默认 0 号卡）；
• atb::CreateContext(&context)创建 ATB 上下文，context是算子执行期间持有资源（如 tiling buffer 池、runner pool）的全局对象；
• aclrtCreateStream创建 ACL 流，并通过context->SetExecuteStream(stream)绑定到 Context，后续算子执行将异步提交到该流；
• 所有调用均用CHECK_STATUS宏包裹，该宏定义在 demo_util.h：失败时打印出错文件行号与错误码，并根据错误码区间（100000~999999判定为 ACL 错误，否则为 ATB 错误）给出对应的官方错误码文档指引后直接返回。

4.2 第二阶段：构造算子并准备输入/输出张量

先创建 Gather 算子并设置参数：

atb::infer::GatherParam gatherParam; gatherParam.axis = 0; gatherParam.batchDims = 0; return atb::CreateOperation(gatherParam, op);

GatherParam仅有两个业务字段（见 infer_op_params.h）：

随后通过PrepareInTensor准备输入：

uint32_t dim0 = 3; uint32_t dim1 = 3; std::vector<float> tensorzero{1, 2, 3, 4, 5, 6, 7, 8, 9}; atb::Tensor tensorZero; CHECK_STATUS(CreateTensorFromVector(contextPtr, stream, tensorzero, ACL_FLOAT16, aclFormat::ACL_FORMAT_ND, {dim0, dim1}, tensorZero)); std::vector<int64_t> tensorone{2, 1}; atb::Tensor tensorOne; CHECK_STATUS(CreateTensorFromVector(contextPtr, stream, tensorone, ACL_INT64, aclFormat::ACL_FORMAT_ND, {2}, tensorOne)); inTensors = {tensorZero, tensorOne};

这里有两个值得注意的细节：

1. 输入 0（x）声明为ACL_FLOAT16：虽然主机侧数据是float（{1..9}），但通过 CreateTensorFromVector 封装会先用aclrtMemcpy把主机数据拷到设备，再调用 Elewise 的ELEWISE_CAST算子将ACL_FLOAT中间类型转为ACL_FLOAT16，因此最终喂给 Gather 的是 FP16 数据——这正是 Transformer 推理中权重/激活普遍采用 FP16 的典型用法；
2. 输入 1（indices）为ACL_INT64：索引张量使用 64 位整型，与算子实现的索引类型要求一致。

输出张量直接按推断出的 shape 创建：

atb::Tensor tensorOut; CHECK_STATUS(CreateTensor(ACL_FLOAT16, aclFormat::ACL_FORMAT_ND, {2, 3}, tensorOut)); variantPack.outTensors.push_back(tensorOut);

variantPack.inTensors与variantPack.outTensors共同组成atb::VariantPack，输入张量的顺序必须与算子定义一致：第 0 个是x，第 1 个是indices（见下文源码佐证）。CreateTensor会根据 shape 与数据类型调用atb::Utils::GetTensorSize计算字节数并分配设备内存。

4.3 第三阶段：Setup——推导 shape 与计算 workspace

uint64_t workspaceSize = 0; CHECK_STATUS(op->Setup(variantPack, workspaceSize, context)); uint8_t *workspacePtr = nullptr; if (workspaceSize > 0) { CHECK_STATUS(aclrtMalloc((void **)(&workspacePtr), workspaceSize, ACL_MEM_MALLOC_HUGE_FIRST)); }

Setup阶段完成两件事：校验输入参数并推导输出 shape（InferShape），同时返回算子执行所需的 workspace 大小。若workspaceSize > 0，必须用aclrtMalloc显式分配工作空间，并在 Execute 阶段传入。

4.4 第四阶段：Execute——异步执行算子

op->Execute(variantPack, workspacePtr, workspaceSize, context); CHECK_STATUS(aclrtSynchronizeStream(stream));

Execute将算子任务异步提交到 Context 绑定的 Stream 上，因此必须调用aclrtSynchronizeStream(stream)等待设备侧任务执行完成，再读取结果或释放资源。

4.5 第五阶段：资源释放（顺序敏感）

for (atb::Tensor &inTensor : variantPack.inTensors) { CHECK_STATUS(aclrtFree(inTensor.deviceData)); } if (workspaceSize > 0) { CHECK_STATUS(aclrtFree(workspacePtr)); } CHECK_STATUS(atb::DestroyOperation(op)); CHECK_STATUS(aclrtDestroyStream(stream)); CHECK_STATUS(atb::DestroyContext(context)); std::cout << "Gather demo success!" << std::endl;

释放顺序有明确讲究，注释也给出了提示：

• 先释放设备内存（输入张量、workspace）；
• 再释放Operation（算子对象，先释放）；
• 最后释放Stream 与 Context（Context 是全局资源，后释放）。

若顺序颠倒（如先销毁 Context 再销毁 Operation），可能导致悬垂引用或未定义行为。

五、源码纵深：GatherOperation 的校验、形状推导与双 Runner 实现

5.1 算子对象：输入/输出数量与参数校验

GatherOperation实现在 src/ops/ops_infer/gather/gather_operation.cpp，其GetInputNum()/GetOutputNum()分别返回 2 和 1（输入为x、indices，输出为y）。

atb::CreateOperation(gatherParam, op)模板特化（见 gather_operation.cpp）在创建对象前会做三重校验，与 README 中的参数约束一一对应：

• axis < 0 || batchDims < 0：直接报ERROR_INVALID_PARAM（对应"必须 >= 0"）；
• batchDims > axis：报ERROR_INVALID_PARAM（对应"必须 <= axis"）；
• 此外还调用OP_PARAM_RSV_CHECK校验预留字段rsv必须为 0，并为 Ascend 950 平台预加载 aclnn 函数（详见 5.3 节）。

ParamCheck（gather_operation.cpp）还进一步约束：

• axis必须小于x的维度数；
• batchDims必须不大于indices的维度数；
• x.dimNum + indices.dimNum - 1 - batchDims <= 8，即输出维度数不能超过 ATB 支持的最大维度。

InferShapeCheckImpl额外校验indices的前batchDims个维度必须与x对应维度相等——这正体现了batchDims的语义：indices的前batchDims个维度与x保持一致，作为 batch 维度逐元素收集。

5.2 形状推导公式：输出 shape 如何算出

InferShapeImpl（gather_operation.cpp）给出了输出维度的构造规则，输出由三段拼接而成：

1. x在axis之前的维度（dims[0..axis-1]）；
2. indices从batchDims开始的维度（dims[batchDims..]）；
3. x在axis+1之后的维度（dims[axis+1..]）。

用公式表达即：

out.dimNum = axis + (indices.dimNum - batchDims) + (x.dimNum - axis - 1) out.dims = x.dims[0:axis] + indices.dims[batchDims:] + x.dims[axis+1:]

代入本 Demo：x=[3,3]、indices=[2]、axis=0、batchDims=0，则out.dims = [] + [2] + [3] = [2,3]，与源码中手动创建的输出{2, 3}完全一致。SetupCheckImpl则会在 Setup 阶段反向逐段核对用户提供的输出 shape 是否满足上述拼接规则。

5.3 双 Runner 机制：Ascend 950 走 aclnn，其余平台走自研 Kernel

CreateRunner（gather_operation.cpp）根据平台类型选择执行后端：

• Ascend 950 平台：创建GatherAclnnRunner（gather_aclnn_runner.cpp）。该 Runner 通过动态加载共享库中的aclnnGatherV3GetWorkspaceSize与aclnnGatherV3两个函数（见 gather_aclnn_runner.cpp），把 ATB 的VariantPack转换成 aclnn 的aclTensor后下发执行。注意其中mode = 1表示"索引聚集场景性能优化"（gather_aclnn_runner.cpp），若 CANN 版本不支持该函数，CreateOperation阶段会报错提示检查 CANN 版本（对应 README 中"source 对应 CANN 版本"的要求）。
• 其他平台：创建GatherOpsRunner（gather_ops_runner.cpp），构造AsdOps::OpParam::Gather = {batchDims, {axis}}参数，组织一张单节点 KernelGraph 交给 ops runner 执行。

5.4 Kernel 实现：模板实例化与索引类型支持

底层 Kernel 位于 src/kernels/kernels/gather/gather_kernel/gather_kernel.cpp，采用GatherKernel<ELEM_SIZE, INDICE_TYPE>模板，按"数据元素字节数 × 索引类型"组合实例化出 6 个 Kernel：

• Gather16I64Kernel / Gather16I32Kernel：数据元素 2 字节（FP16/BF16），索引支持 INT64/INT32；
• Gather32I64Kernel / Gather32I32Kernel：数据元素 4 字节，索引支持 INT64/INT32；
• Gather64I64Kernel / Gather64I32Kernel：数据元素 8 字节，索引支持 INT64/INT32。

CanSupport在启动前校验输入/输出张量数量（2 入 1 出）与数据类型宽度；InitImpl调用GatherTiling（src/kernels/kernels/gather/tiling/gather_tiling.cpp）完成切分计算，并把axis以 int64 常量张量形式写入 tiling buffer。这解释了 Demo 中x用 FP16、indices用 INT64 的组合正是 Kernel 完全支持的标准配置。

六、数据生成参考：Python 测试用例

README 末尾指出：示例中生成的数据（{1..9}等）不代表实际场景，如需贴近真实业务的数据生成方式，请参考根目录下的 Python 用例目录：

• tests/apitest/opstest/python/operations/gather/

其中 test_gather.py 展示了 Gather 算子的 Python 侧用例组织方式，包括随机 shape、随机数据类型组合下的数据构造与算子执行流程，可作为生成更复杂测试数据、验证边界条件的参考。此外，tests/apitest/kernelstest/gather/ 下还提供了 Kernel 级别的单测脚本，用于直接验证 Kernel 行为。

七、常见问题排查

八、小结

通过本指南，读者应能完整掌握：

1. CANN 与 nnal 环境脚本的 source 顺序与源码编译产物的路径差异；
2. D_GLIBCXX_USE_CXX11_ABI编译宏对齐的坑位与正确姿势；
3. ATB 算子调用的五步流程：初始化/建 Context → 建算子与张量 → Setup → Execute → 按序释放资源；
4. GatherParam中axis/batchDims的语义约束与输出 shape 拼接公式，并能在 gather_operation.cpp 中逐行找到对应实现；
5. Gather 算子在 Ascend 950（aclnnGatherV3）与其他平台（自研 KernelGraph）之间的双 Runner 分发机制。

这套 Demo 骨架在仓库中高度复用——example/op_demo/ 下的所有算子示例（elewise、linear、rms_norm、rope 等）均遵循同一模式。掌握了 Gather 示例，你就能快速迁移到任何 ATB 推理算子的开发与验证中。

【免费下载链接】ascend-transformer-boost本项目是CANN提供的是一款高效、可靠的Transformer加速库，基于华为Ascend AI处理器，提供Transformer定制化场景的高性能融合算子。项目地址: https://gitcode.com/cann/ascend-transformer-boost