ATVC ReduceSum 自定义算子样例：基于 ATVC Reduce 模板开发、编译与 AclNN 调用的完整流程

ATVC ReduceSum 自定义算子样例基于 ATVC Reduce 模板开发、编译与 AclNN 调用的完整流程【免费下载链接】atvcATVCAscend C Templates for Vector Compute是为基于Ascend C开发的典型Vector算子封装的一系列模板头文件的集合可帮助用户快速开发典型Vector算子。项目地址: https://gitcode.com/cann/atvc本文基于 ATVC 仓库中的examples/ops_aclnn/reduce_sum样例讲解如何基于 ATVC 的 Reduce 算子模板开发一个名为ReduceSumCustom的自定义算子工程从算子原型 JSON 定义、msOpGen 工程生成、Host 侧 Tiling 计算、Kernel 侧按 policy 分派模板到最终通过单算子 APIAclNN安装部署与调用验证。读完后你将能够独立完成一个 Reduce 类算子的工程化开发、编译部署与功能自测全流程并理解 ATVC 框架在 host/kernel 两侧各自承担的职责。1. 样例概述与目录结构本样例基于ReduceSumCustom算子工程介绍了基于 ATVC 的自定义算子工程调用先用 msOpGen 工具根据 JSON 原型文件生成算子工程骨架再填入 ATVC 框架的 host 与 kernel 实现编译出.run格式的自定义算子包安装到 opp 算子库后通过 AclNN 两段式 API 调用执行并校验结果。样例目录结构如下来源examples/ops_aclnn/reduce_sum/README.md├── reduce_sum // 使用框架调用的方式调用ReduceSum算子 │ ├── AclNNInvocation // 通过aclnn调用的方式调用ReduceSumCustom算子 │ ├── ReduceSumCustom // ReduceSumCustom算子工程 │ ├── ReduceSumCustom.json // ReduceSumCustom算子的原型定义json文件 │ └── install.sh // 脚本调用msOpGen生成自定义算子工程并编译其中AclNNInvocationNaive是一个简化了工程配置的调用工程包含 CMakeLists.txt、单算子调用入口 main.cpp 和编译运行脚本 run.sh。样例整体支持 Ascend 910B 系列处理器ops_aclnn 目录说明。2. ReduceSum 算子描述与规格ReduceSum 是对输入 tensor 的指定轴进行规约累加并输出结果的 Reduce 类算子。样例中的算子规格如下项目取值算子类型OpTypeReduceSum输入 xshape8 * 2048数据类型floatformatND输出 yshape1 * 2048数据类型floatformatND核函数名reduce_sum_custom与规格对应的算子原型定义见 ReduceSumCustom.json它是 msOpGen 生成工程的输入关键信息包括输入x、输出y均为必选参数format 为ND支持的数据类型为float32与int32属性dim为必选类型为list_int即需要规约的轴列表样例中调用时传入{0}对第 0 轴做规约输出 shape 变为1 * 2048。{ op: ReduceSumCustom, language: cpp, input_desc: [ { name: x, param_type: required, format: [ND, ND], type: [float32, int32] } ], output_desc: [ { name: y, param_type: required, format: [ND, ND], type: [float32, int32] } ], attr: [ { name: dim, param_type: required, type: list_int } ] }3. ATVC Reduce 模块Host 与 Kernel 的分工从 ATVC Reduce 数据流图 可以看出 ATVC 框架的典型分工Host 侧内置Tiling 计算 API根据输入 shape、dim 等计算场景参数推导出ReducePolicy分派策略运行态参数并产出描述数据搬配行为的ReduceParam分派策略 API内置 Adapter / 自定义 Adapter从源码结构看负责在不同 policy 间选择实现。Kernel 侧内置Reduce 算子模板消费ReduceParam运行态参数按ReducePolicy编译态参数选择最佳 Kernel 实现计算逻辑由ReduceCompute内置或自定义的仿函数结构体通过ATVC::OpTraits描述算子原型提供。Tiling 的计算过程如下图所示tiling 流程输入shape、ReduceDim、PlatformInfo后先做数据规整删除 shape 为 1 的维度、连续 AR 轴合轴再计算基本块大小basicBlock完成 c 轴tile 层切轴与从 C 轴往外找 A 轴、R 轴block 层切轴最终填充TilingData、Policy、各维度的 stride 与 shape并计算 workspace 大小。4. 算子工程实现剖析算子工程目录ReduceSumCustom包含算子的实现文件├── ReduceSumCustom // ReduceSumCustom自定义算子工程 │ ├── op_host // host侧实现文件 │ ├── op_kernel // kernel侧实现文件 │ ├── build.sh // 算子构建入口 │ └── CMakeLists.txt // 算子的cmake文件CANN 软件包中的工程创建工具 msOpGen 可以基于 ReduceSumCustom.json 自动创建算子工程。创建完成后开发者重点需要完成host 与 kernel 两个文件的功能开发本样例已在ReduceSumCustom目录下准备好了这两份实现下面逐一看它们如何用 ATVC 模板完成开发。4.1 Host 侧TilingFunc 与算子原型注册Host 侧实现位于 reduce_sum_custom.cpp它引入 reduce_host.h 后完成三件事定义算子描述原型L14-L15分别为 float 与 int32 定义ATVC::OpTraits作为模板编译态参数描述算子原型using ReduceOpTraitsFloat ATVC::OpTraitsATVC::OpInputsfloat, ATVC::OpOutputsfloat; using ReduceOpTraitsInt ATVC::OpTraitsATVC::OpInputsint32_t, ATVC::OpOutputsint32_t;实现 TilingFuncL18-L46从TilingContext中读取输入 shape 与dim属性通过GetListInt(0)解析list_int类型的dim按输入数据类型分派到对应 traits 调用ATVC::Host::CalcReduceTiling计算 Tiling 数据并将 policy 的 id 写入tiling-policyId作为 kernel 侧分支判断的依据随后设置 blockDim等于tiling-tilingData.coreNum并声明 workspace 为 0ATVC::ReducePolicy policy {0, 0, 0}; // ... 读取输入shape与dim属性 ... ATVC::ReduceParam *tiling context-GetTilingDataATVC::ReduceParam(); tiling-nBufferNum 2; // 2 : default value if (context-GetInputDesc(0)-GetDataType() ge::DataType::DT_FLOAT) { (void)ATVC::Host::CalcReduceTilingReduceOpTraitsFloat(shapeIn, dim, policy, tiling); } else if (context-GetInputDesc(0)-GetDataType() ge::DataType::DT_INT32) { (void)ATVC::Host::CalcReduceTilingReduceOpTraitsInt(shapeIn, dim, policy, tiling); } tiling-policyId policy.getID(); context-SetBlockDim(tiling-tilingData.coreNum);注册算子原型L66-L85ReduceSumCustom继承OpDef声明输入x、输出yfloat/int32、ND与必选list_int属性dim绑定InferShape/InferDataType推导函数输出 shape 与输入一致、输出类型跟随输入并通过this-AICore().SetTiling(optiling::TilingFunc).AddConfig(ascend910b)将 TilingFunc 注册到 AI Core、适配 ascend910b 平台最后以OP_ADD(ReduceSumCustom)完成算子注册。4.2 Kernel 侧按 policyId 分派 Reduce 算子模板Kernel 侧实现位于 reduce_sum_custom.cpp仅约 90 行。核函数reduce_sum_custom引入 reduce_kernel.h 后声明任务类型KERNEL_TASK_TYPE_DEFAULT(KERNEL_TYPE_MIX_AIV_1_0)AI 与 AIV 混合使能REGISTER_TILING_DEFAULT(ATVC::ReduceParam)与GET_TILING_DATA(param, tiling)获取 host 侧算出的运行态 Tiling 数据按param.policyId分派到ATVC::REDUCE_POLICY0至ATVC::REDUCE_POLICY22的 23 个分支每个分支实例化ATVC::Kernel::ReduceOpTemplateATVC::ReduceSumComputeReduceOpTraits, 对应Policy()并调用op.Run(x, y, param)执行计算extern C __global__ __aicore__ void reduce_sum_custom(GM_ADDR x, GM_ADDR y, GM_ADDR workspace, GM_ADDR tiling) { KERNEL_TASK_TYPE_DEFAULT(KERNEL_TYPE_MIX_AIV_1_0); REGISTER_TILING_DEFAULT(ATVC::ReduceParam); GET_TILING_DATA(param, tiling); if (param.policyId ATVC::REDUCE_POLICY0.ID) { auto op ATVC::Kernel::ReduceOpTemplateATVC::ReduceSumComputeReduceOpTraits, ATVC::REDUCE_POLICY0(); op.Run(x, y, param); } else if (param.policyId ATVC::REDUCE_POLICY1.ID) { // ... 依次分派至 REDUCE_POLICY22 ... } }从源码结构看这里刻意使用param.policyId而不是TILING_KEY_IS做分支判断原因是 Reduce 的 policy 分支过多若用 tilingKey 判断会遇到爆栈问题这一点在 ops_aclnn 开发说明 中也有明确提示。计算逻辑ATVC::ReduceSumCompute由 ATVC 框架内置提供开发者无需手写规约计算本身。4.3 构建配置ATVC 头文件路径Kernel 侧 CMakeLists.txt 体现了样例中ATVC_PATH环境变量的作用优先读取环境变量ATVC_PATH未设置时回退到仓库内的include目录../../../../../include并通过npu_op_kernel_options以-I方式将其加入 kernel 编译包含路径if(DEFINED ENV{ATVC_PATH}) set(_ATVC_PATH $ENV{ATVC_PATH}) else() set(_ATVC_PATH ${CMAKE_CURRENT_SOURCE_DIR}/../../../../../include) endif() npu_op_kernel_options(ascendc_kernels ALL OPTIONS -I ${_ATVC_PATH})此外ATVC 使用自定义的 Tiling 定义ATVC::ReduceParam等因此样例工程在生成骨架后会删除 msOpGen 自动生成的*_tiling.h这一点由install.sh完成见下一节。5. 编译运行完整流程针对自定义算子工程编译运行包含五个步骤调用 msOpGen 工具生成自定义算子工程基于 ATVC 框架完成算子 host 和 kernel 实现本样例已提供编译自定义算子工程生成自定义算子包安装自定义算子包到自定义算子库中调用执行自定义算子。5.1 获取源码包及环境配置编译运行样例前请先参考 ops_aclnn 样例准备说明按 快速上手文档 完成 CANN 软件包安装、环境变量配置以及 ATVC 源码获取。5.2 生成自定义算子工程、复制 host 和 kernel 实现并编译算子导入 ATVC 环境变量如果导入会覆盖 CMake 中的默认回退路径./atvc/include# 如果不导入默认使用./atvc/include路径 export ATVC_PATH${atvc}/include切换到 msOpGen 脚本 install.sh 所在目录# 若开发者以git命令行方式clone了master分支代码并切换目录 cd ./examples/ops_aclnn/reduce_sum运行脚本生成工程并编译算子bash install.sh -v [SOC_VERSION]参数说明以 install.sh 源码为准参数说明-v/--soc-versionSOC_VERSION昇腾 AI 处理器型号。若无法确定可在安装昇腾 AI 处理器的服务器上执行npu-smi info查询在查询到的Name前增加Ascend信息例如Name为xxxyy则配置为Ascendxxxyy。从脚本的版本校验逻辑看取值为Ascend910B1、Ascend910B2、Ascend910B3、Ascend910B4见 install.sh L40-L44-i/--install-pathASCEND_INSTALL_PATHCANN 软件包安装路径可选未指定时依次回退到环境变量ASCEND_HOME_PATH、$HOME/Ascend/ascend-toolkit/latest、/usr/local/Ascend/ascend-toolkit/latest脚本内部依次执行install.sh L60-L70删除旧的CustomOp目录 → 执行msopgen gen -i ReduceSumCustom.json -c ai_core-${SOC_VERSION} -lan cpp -out CustomOp生成工程骨架 → 将样例的op_host、op_kernel实现复制到CustomOp对应目录 → 删除 msOpGen 生成的CustomOp/op_kernel/*_tiling.hATVC 使用自己的 tiling 定义→ 进入CustomOp执行bash build.sh完成编译。脚本运行成功后会在当前目录创建CustomOp目录编译完成后在CustomOp/build_out中生成自定义算子安装包custom_opp_target os_target architecture.run例如custom_opp_ubuntu_x86_64.run。备注CustomOp目录为生成目录每次执行 install.sh 脚本都会删除该目录并重新生成切勿在该目录下编码算子会存在丢失风险。5.3 部署自定义算子包部署前请确保存在自定义算子包默认部署路径环境变量ASCEND_OPP_PATHecho $ASCEND_OPP_PATH # 输出示例 /usr/local/Ascend/latest/opp # 若没有则需导出CANN环境变量 source [ASCEND_INSTALL_PATH]/bin/setenv.bash # 例如 source /usr/local/Ascend/latest/bin/setenv.bash其中ASCEND_INSTALL_PATH为 CANN 软件包安装路径一般和上一步中指定的路径保持一致。然后在安装包所在路径下执行安装cd CustomOp/build_out ./custom_opp_target os_target architecture.run命令执行成功后自定义算子包中的相关文件将部署至 opp 算子库环境变量ASCEND_OPP_PATH指向的vendors/customize目录中同时该目录下的op_api子目录会为调用方生成单算子 API 头文件与库AclNNInvocationNaive 的 CMakeLists.txt 中即链接了${INC_PATH}/opp/vendors/customize/op_api。5.4 通过 AclNN 单算子 API 调用执行算子调用工程为 AclNNInvocationNaive。自定义算子编译部署后会自动生成单算子 API其形式为“两段式接口”// 获取算子使用的workspace空间大小 aclnnStatus aclnnReduceSumCustomGetWorkspaceSize(const aclTensor *x, const aclIntArray *dim, const aclTensor *out, uint64_t *workspaceSize, aclOpExecutor **executor); // 执行算子 aclnnStatus aclnnReduceSumCustom(void *workspace, int64_t workspaceSize, aclOpExecutor *executor, aclrtStream stream);第一段接口aclnnReduceSumCustomGetWorkspaceSize用于计算本次调用所需的 workspace 大小获取后按workspaceSize申请 Device 侧内存再调用第二段接口aclnnReduceSumCustom执行计算。main.cpp 按 6 个阶段组织L125-L195初始化aclInit→aclrtSetDevice(0)→aclrtCreateStream构造张量输入inputXShape {8, 2048}、输出outputYShape {1, 2048}aclrtMalloc申请 device 内存、aclrtMemcpy拷入主机数据aclCreateTensor创建 ND 格式 float 张量规约轴属性dim {0}通过aclCreateIntArray传入样例将输入填充为全 1.0golden 值按“8 行逐列求和”取 8.0L94-L105与算子规格一致两段式调用先aclnnReduceSumCustomGetWorkspaceSize获取 workspaceSize大于 0 时分配 workspace 内存再aclnnReduceSumCustom执行算子随后aclrtSynchronizeStream同步回拷校验aclrtMemcpy将输出从 device 拷回 host与 golden 逐元素比较打印前 10 个结果并输出test pass/test failedL107-L122资源释放销毁张量、释放 device 内存与 workspace、销毁 stream、复位设备并aclFinalize。运行步骤先完成 5.2 节的算子工程编译锚点参见 编译算子工程进入样例目录cd atvc/examples/ops_aclnn/reduce_sum/AclNNInvocationNaive修改 CMakeLists.txt将其中默认路径/usr/local/Ascend/ascend-toolkit/latest替换为 CANN 软件包安装后的实际路径如/home/HwHiAiUser/Ascend/ascend-toolkit/latest。从源码看该文件通过环境变量DDK_PATH确定头文件路径、NPU_HOST_LIB确定库路径均未设置时回退到上述默认值L22-L42执行产物execute_reduce_sum_op需链接ascendcl、cust_opapi、acl_op_compiler、nnopbase、stdcL64-L70参考 run.sh 编译与运行bash run.shrun.sh会自动导出DDK_PATH、NPU_HOST_LIB执行cmakecmake --build并在运行时将$_ASCEND_INSTALL_PATH/opp/vendors/customize/op_api/lib加入LD_LIBRARY_PATH后执行./execute_reduce_sum_op最终打印算子输出与校验结果。6. 小结该样例展示了 ATVC 框架开发 Reduce 类算子的标准姿势Host 侧只需实现 TilingFunc 调用ATVC::Host::CalcReduceTiling并注册 policyIdKernel 侧只需按policyId分派到ReduceOpTemplateReduceSumCompute, Policy规约的数据搬配、切轴与计算逻辑全部由框架模板完成。配合 msOpGen 工程生成、.run算子包部署与 AclNN 两段式 API 调用即可在 Ascend 910B 系列处理器上完成一个自定义 ReduceSum 算子从开发到功能自测的闭环。【免费下载链接】atvcATVCAscend C Templates for Vector Compute是为基于Ascend C开发的典型Vector算子封装的一系列模板头文件的集合可帮助用户快速开发典型Vector算子。项目地址: https://gitcode.com/cann/atvc创作声明：本文部分内容由AI辅助生成（AIGC），仅供参考