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 参数语义以及双 RunnerAscend 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]在axis0上收集输出 shape 为[2, 3]即[[7,8,9],[4,5,6]]。运行成功的标志是终端打印Gather demo success!。二、环境准备source CANN 与 nnal 安装路径运行 Demo 前必须保证 CANN 工具包与加速库nnal/atb的运行时环境已就绪。README 给出的标准流程是依次 source 两个环境脚本source CANN 安装路径下的set_env.sh默认路径为source /usr/local/Ascend/ascend-toolkit/set_env.shsource 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提供 ACLAscend Computing Language运行时头文件与库加速库的set_env.sh提供 ATB 头文件atb/atb_infer.h、atb/operation.h、atb/types.h等以及编译好的libatb库。关于整体编译方式可参考仓库根目录文档 编译与构建。三、编译与运行 Demo3.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_abi0默认设置D_GLIBCXX_USE_CXX11_ABI0即编译命令形如g -D_GLIBCXX_USE_CXX11_ABI0 -I ... gather_demo.cpp ... -o gather_demo使用 cxx_abi1更改宏为 1即g -D_GLIBCXX_USE_CXX11_ABI1 -I ... gather_demo.cpp ... -o gather_demoD_GLIBCXX_USE_CXX11_ABI是 GCC 控制新旧两种std::string/std::listABI 的宏值为 0 时使用旧版C03 风格ABI值为 1 时使用新版C11 风格ABI。ATB 库在构建安装时选择了哪个 ABI你的 Demo 就必须用同样的宏编译这是新手最容易踩的坑。3.3 头文件与链接库Demo 源码通过 demo_util.h 引入acl/acl.hACL 运行时 API设备管理、Stream、内存分配atb/atb_infer.h、atb/operation.h、atb/types.hATB 算子框架、atb::Operation、atb::Tensor、atb::VariantPack等核心类型。链接时需要同时链接 ACL 库如-lascendcl与 ATB 库如-latb具体库名以 CANN/nnal 安装环境为准。四、Demo 源码逐段精读ATB 算子调用的标准五步流程gather_demo.cpp 完整展示了 ATB 算子调用范式。其主流程可分为五个阶段这也是所有 ATB 算子 Demo 的统一骨架。4.1 第一阶段初始化 ACL、创建 Context 与 StreamCHECK_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字段默认值含义约束axis0指定要收集切片的轴必须 0且小于输入x的维度数batchDims0允许从 batch 的每个元素中收集不同的项目必须 0且 axis随后通过PrepareInTensor准备输入uint32_t dim0 3; uint32_t dim1 3; std::vectorfloat 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::vectorint64_t tensorone{2, 1}; atb::Tensor tensorOne; CHECK_STATUS(CreateTensorFromVector(contextPtr, stream, tensorone, ACL_INT64, aclFormat::ACL_FORMAT_ND, {2}, tensorOne)); inTensors {tensorZero, tensorOne};这里有两个值得注意的细节输入 0x声明为ACL_FLOAT16虽然主机侧数据是float{1..9}但通过 CreateTensorFromVector 封装会先用aclrtMemcpy把主机数据拷到设备再调用 Elewise 的ELEWISE_CAST算子将ACL_FLOAT中间类型转为ACL_FLOAT16因此最终喂给 Gather 的是 FP16 数据——这正是 Transformer 推理中权重/激活普遍采用 FP16 的典型用法输入 1indices为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 与计算 workspaceuint64_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阶段完成两件事校验输入参数并推导输出 shapeInferShape同时返回算子执行所需的 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 与 ContextContext 是全局资源后释放。若顺序颠倒如先销毁 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对应必须 0batchDims axis报ERROR_INVALID_PARAM对应必须 axis此外还调用OP_PARAM_RSV_CHECK校验预留字段rsv必须为 0并为 Ascend 950 平台预加载 aclnn 函数详见 5.3 节。ParamCheckgather_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 如何算出InferShapeImplgather_operation.cpp给出了输出维度的构造规则输出由三段拼接而成x在axis之前的维度dims[0..axis-1]indices从batchDims开始的维度dims[batchDims..]x在axis1之后的维度dims[axis1..]。用公式表达即out.dimNum axis (indices.dimNum - batchDims) (x.dimNum - axis - 1) out.dims x.dims[0:axis] indices.dims[batchDims:] x.dims[axis1:]代入本 Demox[3,3]、indices[2]、axis0、batchDims0则out.dims [] [2] [3] [2,3]与源码中手动创建的输出{2, 3}完全一致。SetupCheckImpl则会在 Setup 阶段反向逐段核对用户提供的输出 shape 是否满足上述拼接规则。5.3 双 Runner 机制Ascend 950 走 aclnn其余平台走自研 KernelCreateRunnergather_operation.cpp根据平台类型选择执行后端Ascend 950 平台创建GatherAclnnRunnergather_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 版本的要求。其他平台创建GatherOpsRunnergather_ops_runner.cpp构造AsdOps::OpParam::Gather {batchDims, {axis}}参数组织一张单节点 KernelGraph 交给 ops runner 执行。5.4 Kernel 实现模板实例化与索引类型支持底层 Kernel 位于 src/kernels/kernels/gather/gather_kernel/gather_kernel.cpp采用GatherKernelELEM_SIZE, INDICE_TYPE模板按数据元素字节数 × 索引类型组合实例化出 6 个 KernelGather16I64Kernel / Gather16I32Kernel数据元素 2 字节FP16/BF16索引支持 INT64/INT32Gather32I64Kernel / Gather32I32Kernel数据元素 4 字节索引支持 INT64/INT32Gather64I64Kernel / Gather64I32Kernel数据元素 8 字节索引支持 INT64/INT32。CanSupport在启动前校验输入/输出张量数量2 入 1 出与数据类型宽度InitImpl调用GatherTilingsrc/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 行为。七、常见问题排查现象可能原因与处理编译报std::string相关链接错误或符号不匹配D_GLIBCXX_USE_CXX11_ABI与加速库编译选项不一致按 3.2 节对齐 ABI运行报 ACL/ATB 错误码用CHECK_STATUS打印的错误码到 CANN/ATB 官方文档检索确认已按第二节 source 两个环境脚本CreateOperation报Load aclnn function failedAscend 950 平台需对应 CANN 版本提供aclnnGatherV3升级/更换匹配的 CANN 版本输出 shape 校验失败检查axis、batchDims约束0 batchDims axis x.dimNum且输出 shape 必须严格满足 5.2 节的拼接公式结果不对但未报错检查aclrtSynchronizeStream是否调用异步执行需同步后再读结果检查输入张量顺序是否为{x, indices}八、小结通过本指南读者应能完整掌握CANN 与 nnal 环境脚本的 source 顺序与源码编译产物的路径差异D_GLIBCXX_USE_CXX11_ABI编译宏对齐的坑位与正确姿势ATB 算子调用的五步流程初始化/建 Context → 建算子与张量 → Setup → Execute → 按序释放资源GatherParam中axis/batchDims的语义约束与输出 shape 拼接公式并能在 gather_operation.cpp 中逐行找到对应实现Gather 算子在 Ascend 950aclnnGatherV3与其他平台自研 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创作声明：本文部分内容由AI辅助生成（AIGC），仅供参考