ascend-transformer-boost SliceOperation C++ Demo 实战：从环境搭建到源码级切片原理

ascend-transformer-boost SliceOperation C Demo 实战从环境搭建到源码级切片原理【免费下载链接】ascend-transformer-boost本项目是CANN提供的是一款高效、可靠的Transformer加速库基于华为Ascend AI处理器提供Transformer定制化场景的高性能融合算子。项目地址: https://gitcode.com/cann/ascend-transformer-boost本篇技术指南围绕 ascend-transformer-boost 仓库中 slice Demo 展开讲解如何在华为 Ascend 环境下用 C 调用加速库的 SliceOperation从输入张量中按指定起始位置与大小提取切片。读完本文你将掌握 ATBAscend Transformer Boost算子的标准调用范式Context / VariantPack / Setup / Execute、SliceParam 参数语义与校验规则并能结合 算子实现 与 runner 实现 理解其底层执行路径。一、SliceOperation 是什么SliceOperation 是加速库中的张量切片算子功能为从输入张量某个起始位置中提取指定大小的切片是 Transformer 推理链路中常用的数据搬运算子例如从长序列的 KV Cache 或隐藏状态中截取片段。输入1 个 ND 格式张量x输出1 个 ND 格式张量output其 shape 由参数offsets与size推导得出。该算子的参数类型atb::infer::SliceParam定义在 include/atb/infer_op_params.h支持的数据类型与格式约束记录在 ops_configs/atb_ops_info.ini 的[SliceOperation]配置节中[SliceOperation] input0.namex input0.dtypefloat16,float,int8,bool,int32,uint32,bf16 input0.formatnd,nd,nd,nd,nd,nd,nd output0.nameoutput output0.dtypefloat16,float,int8,bool,int32,uint32,bf16 output0.formatnd,nd,nd,nd,nd,nd,nd即输入输出均要求 ND 格式支持 float16、float、int8、bool、int32、uint32、bf16 等数据类型。二、环境准备source 两套环境脚本运行 Demo 前需要先加载 CANN 与 nnal 两套安装路径的环境变量。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说明nnal 指以 nnal 形式安装的加速库运行环境源码编译场景下构建产物output/atb/中的 set_env.sh 会同时解决库路径与头文件路径问题。三、编译与运行 Demo在 example/op_demo/slice 目录下执行bash build.sh注意Demo 编译时D_GLIBCXX_USE_CXX11_ABI宏必须与加速库链接时的 ABI 设置保持一致使用cxx_abi0默认时设置g -D_GLIBCXX_USE_CXX11_ABI0 -I ...使用cxx_abi1时改为g -D_GLIBCXX_USE_CXX11_ABI1 -I ...该宏控制 libstdc 新旧两种 ABI 的符号命名涉及std::string等类型若与加速库编译期不一致运行阶段会出现链接或符号解析失败。运行成功后终端输出slice demo success!四、Demo 源码逐段解析完整代码见 slice_demo.cpp其调用流程代表了 ATB 算子的标准使用范式共五个阶段。4.1 准备输入张量atb::Status PrepareInTensor(atb::Context *contextPtr, aclrtStream stream, atb::SVectoratb::Tensor inTensors) { uint32_t dim0 3; uint32_t dim1 6; // 创建tensor0 std::vectorfloat tensorzero{1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15, 16, 17, 18}; atb::Tensor tensorZero; CHECK_STATUS(CreateTensorFromVector(contextPtr, stream, tensorzero, ACL_FLOAT, aclFormat::ACL_FORMAT_ND, {dim0, dim1}, tensorZero)); // 根据顺序将所有输入tensor放入SVector inTensors {tensorZero}; return atb::ErrorType::NO_ERROR; }示例构造了一个 3×6、共 18 个元素的 float 张量数值 118并通过demo_util.h中的工具函数CreateTensorFromVector完成 device 侧内存分配与 H2D 拷贝详见 example/op_demo/demo_util.h。4.2 创建算子并设置参数atb::Status PrepareOperation(atb::Operation **op) { atb::infer::SliceParam sliceParam; sliceParam.offsets {1, 2}; sliceParam.size {2, 4}; return atb::CreateOperation(sliceParam, op); }对 3×6 的输入张量offsets{1,2}表示从第 0 维索引 1、第 1 维索引 2 处开始size{2,4}表示两维分别截取 2 和 4 个元素因此输出 shape 为 2×4元素为输入中第 1、2 行与第 25 列的交叉区域。4.3 初始化运行环境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)); CHECK_STATUS(context-SetExecuteStream(stream));标准初始化顺序为aclInit→ 设置 device →atb::CreateContext创建 ATB 上下文 →aclrtCreateStream创建执行流 → 通过SetExecuteStream将流绑定到 context。4.4 组装 VariantPack 并执行atb::Operation *op nullptr; CHECK_STATUS(PrepareOperation(op)); atb::VariantPack variantPack; CHECK_STATUS(PrepareInTensor(context, stream, variantPack.inTensors)); atb::Tensor tensorOut; CHECK_STATUS(CreateTensor(ACL_FLOAT, aclFormat::ACL_FORMAT_ND, {2, 4}, tensorOut)); variantPack.outTensors.push_back(tensorOut); 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)); } CHECK_STATUS(op-Execute(variantPack, workspacePtr, workspaceSize, context)); CHECK_STATUS(aclrtSynchronizeStream(stream));关键点Setup阶段会校验输入输出张量描述并推导输出 shape本例输出{2, 4}与参数一致同时返回所需 workspace 大小workspace 大于 0 时必须用aclrtMalloc显式申请。Execute阶段将算子提交到绑定流上执行之后必须aclrtSynchronizeStream同步等待 device 侧任务完成才能安全读取/释放结果。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)); CHECK_STATUS(aclFinalize());释放顺序有讲究先释放张量 device 内存与 workspace再释放 Operation对象概念随后销毁 stream最后销毁 context全局资源并aclFinalize。示例注释明确说明“operation 对象概念先释放context 全局资源后释放”。五、SliceParam 参数语义与校验规则SliceParam定义于 include/atb/infer_op_params.h字段语义如下字段含义说明offsets每个维度切片的起始位置元素为负数时表示从该维最高索引处倒推如offsets[i]-1表示起始位置为dimNum-1要求dimNum offsets[i] 0size每个维度切片的大小元素须 -1size[i]-1表示切到该维最后一个位置dimNum-1rsv预留参数8 字节预留置 0约束条件offsets与size的长度必须相等且必须等于输入张量的维数每个维度上offset size不得超过该维大小offset size计算结果不得发生 int64 溢出。这些校验在SliceOperation::ParamCheck中实现见 src/ops/ops_infer/slice/slice_operation.cpp任一维度不满足都会返回ERROR_INVALID_PARAM并打印详细日志。5.1 输出 shape 的推导逻辑InferShapeImplslice_operation.cpp先复制输入 desc再逐维处理if (offsetValue 0) { offsetValue offsetValue xDim; // 负偏移转正 } if (sizeValue -1) { sizeValue xDim - offsetValue; // -1 表示截到末尾 } outTensorDescs.at(0).shape.dims[i] sizeValue;SetupCheckImplslice_operation.cpp则进一步核对用户提供的输出张量 shape 是否与参数推导结果一致不一致时返回ERROR_INVALID_TENSOR_DIM。因此 Demo 中手动创建{2, 4}输出张量时必须与offsets{1,2}, size{2,4}的推导结果吻合。六、底层执行路径两种 Runner 的分发机制SliceOperation::CreateRunnerslice_operation.cpp按平台类型分发执行器if (Mki::PlatformInfo::Instance().GetPlatformType() Mki::PlatformType::ASCEND_950) { return std::make_sharedSliceAclnnRunner(param_); } return std::make_sharedSliceOpsRunner(param_);6.1 SliceOpsRunnerkernelGraph 直通路径非 950 平台如 910B 系列走 slice_ops_runner.cpp在构造阶段把SliceParam转换为AsdOps::OpParam::Slice并构建包含单节点的 kernelGraph随后由 OpsRunner 统一调度底层 kernel。6.2 SliceAclnnRunneraclnn 适配路径ASCEND_950 平台走 slice_aclnn_runner.cpp通过动态加载aclnnSliceV2与aclnnCast实现LoadAclnnFuncsL56-L77从共享库中加载aclnnSliceV2GetWorkspaceSize/aclnnSliceV2/aclnnCastGetWorkspaceSize/aclnnCast四个符号加载失败时CreateOperation会返回ERROR_CANN_ERROR并提示检查 CANN 版本见 slice_operation.cpp。SetAclNNWorkspaceExecutorL182-L294将 offsets/size 转换为 aclnnSliceV2 所需的 starts/ends/axes/steps 数组其中steps恒为 1、axes为逐维索引由于 aclnnSliceV2 不支持 uint32输入为 uint32 时先用 Cast 转 int64输出再转回 uint32对应selfBufferSize_/outBufferSize_的临时 buffer 逻辑。最终 workspace 大小为 slice 与两次 cast 的 workspace 之和叠加临时转换 buffer 的容量L289-L291。由此可见同一算子在不同硬件平台上拥有 kernel 直通与 aclnn 适配两条实现路径对外 C 接口完全一致。七、数据生成与测试用例参考README 明确指出示例中生成的数据不代表实际场景如需更贴近真实推理场景的数据生成参考可查看 tests/apitest/opstest/python/operations/slice 目录下的 Python 用例test_slice.pyoffsets{2,8}, size{10,100}输入 32×128 的 fp16 张量golden 为x[2:12, 8:108]验证了二维大切片场景并包含仅 Ascend910B 支持的 bf16 用例test_slice2.pyoffsets{0,0,0}, size{-1,-1,2048}输入 256×2×2304 的三维张量golden 为x[:, :, 0:2048]验证了size-1截到末尾的语义对三维场景尤为实用。这些用例通过torch_npu构造输入并给出 golden 计算可当作切片语义的“教科书”参考。八、常见问题与排查建议运行报错提示加载 aclnn 函数失败在 ASCEND_950 平台执行CreateOperation时若LoadAclnnFuncs返回失败说明当前 CANN 版本缺少aclnnSliceV2相关符号请升级或更换与加速库配套的 CANN 版本。Setup 阶段报ERROR_INVALID_PARAM优先检查offsets与size长度是否相等且等于输入维数、每个维度offset size dim、offsets[i] -dim。Setup 阶段报ERROR_INVALID_TENSOR_DIM手动创建的输出张量 shape 与参数推导结果不一致应让输出 shape 由InferShapeImpl的规则计算得出。编译或链接失败核对D_GLIBCXX_USE_CXX11_ABI取值是否与加速库一致cxx_abi0对应宏 0cxx_abi1对应宏 1。结果读取为空或异常确认执行后调用了aclrtSynchronizeStream进行流同步否则 device 侧任务可能尚未完成。英文版说明可参考同目录 README_en.md。其它算子的调用范式如 Elewise、Transdata、RMSNorm 等与 Slice 一致均可对照 example/op_demo 下的同类 Demo 快速上手。【免费下载链接】ascend-transformer-boost本项目是CANN提供的是一款高效、可靠的Transformer加速库基于华为Ascend AI处理器提供Transformer定制化场景的高性能融合算子。项目地址: https://gitcode.com/cann/ascend-transformer-boost创作声明：本文部分内容由AI辅助生成（AIGC），仅供参考