CANN 自定义算子框架插件Framework Plugin开发实战基于 cann-samples 的 ONNX 与 TensorFlow 算子适配全流程【免费下载链接】cann-samplesCANN高性能实战演进样例与体系化调优知识库项目地址: https://gitcode.com/cann/cann-samples导读本文以 cann-samples 仓库中 custom_op_in_graph 样例为主线系统讲解 CANN 自定义算子工程中「框架插件Framework Plugin」模块的完整开发与部署流程如何编写 ONNX 与 TensorFlow 两种框架的算子适配插件告诉 CANN 编译器识别和映射第三方框架模型中的自定义算子。读完本文你将掌握REGISTER_CUSTOM_OP注册宏的用法、一对一映射与一对多子图映射ParseOpToGraphFn两种适配模式、算子属性的跨框架解析技巧以及插件编译、部署与 ATC 模型转换验证的完整实战链路。背景为什么需要框架插件在昇腾 AI 处理器上运行第三方框架ONNX / TensorFlow模型时模型中的算子必须被 CANN 编译器识别并映射为 CANN 算子才能完成图编译与算子下发。对于框架原生算子CANN 已内置映射但对于自定义算子——无论是用户自己开发的算子还是框架中尚未被 CANN 支持的算子——就需要通过「框架插件」告诉编译器原始算子类型OriginOpType与 CANN 侧算子之间的对应关系属性如何转换例如 ONNX 的alpha属性映射为 CANN 的negative_slope算子的实现类型TVM、AI_CPU 等。本样例 custom_op_in_graph 正是从完整自定义算子工程中独立出来的框架插件模块包含 ONNX 与 TensorFlow 两个框架的自定义算子适配插件代码覆盖了映射、注册、属性解析、Scope 融合等典型场景是最直接可读、可编译、可验证的框架插件参考实现。目录结构与构建目标样例整体结构如下注释为模块职责说明custom_op_in_graph/ ├── CMakeLists.txt // 构建入口 ├── build.sh // 便捷编译脚本 ├── README_CN.md ├── image/ // 模型转换效果对比图 │ ├── addn_original.svg // AddN 原始模型图 │ └── addn_subgraph.svg // AddN 拆分为子图效果 ├── onnx_plugin // ONNX 框架算子适配插件 │ ├── CMakeLists.txt │ ├── add_plugin.cc // Add 算子映射 │ ├── addn_plugin.cc // AddN 算子映射 │ ├── leaky_relu_plugin.cc // LeakyRelu 算子映射 └── tf_plugin // TensorFlow 框架算子适配插件 ├── CMakeLists.txt ├── add_block_cust_plugin.cc // AddBlockCust 算子注册 ├── add_dsl_plugin.cc // AddDsl 算子注册 ├── decode_bbox_v2_scope_fusion_plugin.cc // DecodeBboxV2 融合算子适配 ├── lstm_tik_plugin.cc // LSTMTik 算子注册 ├── reshape_cust_plugin.cc // ReshapeCust 算子注册 ├── scatter_nd_add_plugin.cc // ScatterNdAdd 算子注册 └── unique_cust_plugin.cc // UniqueCust 算子注册顶层 CMakeLists.txt 定义了本模块的两个构建目标cust_onnx_parsers与cust_tf_parsers产物输出目录为build_out/makepkg/packages/vendors/customize/framework/并自动生成部署用的set_env.bash环境变量脚本该脚本将packages/vendors/customize路径追加到ASCEND_CUSTOM_OPP_PATH。两个子目录的 onnx_plugin/CMakeLists.txt 与 tf_plugin/CMakeLists.txt 均以add_library(... SHARED)生成动态库链接 CANN 的libgraph.so与libregister.so其中 ONNX 插件还需引入 nlohmann/json 头文件用于解析算子属性见下文 LeakyRelu 样例。ONNX 框架插件onnx_pluginONNX 插件的作用是将第三方 ONNX 模型中的算子映射为 CANN 算子。三个样例分别演示了两种最核心的映射模式。一对一映射Add 算子add_plugin.ccadd_plugin.cc 是最简模式——将 ONNX 的ai.onnx::11::Add直接映射为 CANN 的Add算子#include register/register.h namespace domi { Status ParseParamsAdd(const ge::Operator op_src, ge::Operator op_dest) { return SUCCESS; } REGISTER_CUSTOM_OP(Add) .FrameworkType(ONNX) .OriginOpType(ai.onnx::11::Add) .ParseParamsByOperatorFn(ParseParamsAdd) .ImplyType(ImplyType::TVM); } // namespace domi关键点拆解REGISTER_CUSTOM_OP(Add)注册一个名为Add的 CANN 算子映射。.FrameworkType(ONNX)声明该映射针对 ONNX 框架。.OriginOpType(ai.onnx::11::Add)指定被映射的原始算子类型。注意此处带了 opset 版本前缀即只匹配 opset 11 的 ONNX Add当需要跨 opset 兼容时应像 LeakyRelu 样例那样传入版本列表见下文。.ParseParamsByOperatorFn(ParseParamsAdd)注册参数解析回调。由于 Add 算子没有需要额外转换的属性回调直接返回SUCCESS参数映射由框架自动完成即文档中所说的由框架自动完成映射。.ImplyType(ImplyType::TVM)声明该算子使用 TVM 实现。一对多映射AddN 算子addn_plugin.ccaddn_plugin.cc 演示了更复杂的一对多映射将 ONNXAddN(x, y, z)拆解为Add(Add(x, y), z)组成的子图注册类型为PartitionedCall。它同时使用了两个回调1.ParseParamsByOperatorFnParseParamsAddn负责参数与输入输出结构的传递。由于 AddN 是变长输入算子这里通过DynamicInputOutputInfo声明动态输入输出并调用AutoMappingByOpFnDynamic让框架按原始节点的输入输出个数in_num3、out_num1自动完成映射ge::Operator op_ori const_castge::Operator(op_src); std::string in_name args; std::string in_value in_num; std::string out_name output; std::string out_value out_num; op_ori.SetAttr(in_value, 3); op_ori.SetAttr(out_value, 1); DynamicInputOutputInfo in_values(kInput, in_name.c_str(), in_name.size(), in_value.c_str(), in_value.size()); DynamicInputOutputInfo out_values(kOutput, out_name.c_str(), out_name.size(), out_value.c_str(), out_value.size()); AutoMappingByOpFnDynamic(op_ori, op_dest, {in_values, out_values});同时设置original_type属性为ai.onnx::11::AddN保留原始算子类型信息。2.ParseOpToGraphFnParseOpToGraphAddn这是映射为子图的核心。回调中构造一个子图先用 3 个Data节点声明输入set_attr_index与原始节点输入顺序对应再用两个 CANN 内置Add算子add0、add1级联最后通过graph.SetInputs(inputs).SetOutputs(output_indices)将add1的第 0 个输出声明为子图输出auto data_0 ge::op::Data().set_attr_index(0); auto data_1 ge::op::Data().set_attr_index(1); auto data_2 ge::op::Data().set_attr_index(2); auto add0 ge::op::Add(add0) .set_input_x1(data_0) .set_input_x2(data_1); auto add1 ge::op::Add(add1) .set_input_x1(data_2) .set_input_x2(add0); std::vectorge::Operator inputs{data_0, data_1, data_2}; std::vectorstd::pairge::Operator, std::vectorsize_t output_indices; output_indices.emplace_back(add1, std::vectorstd::size_t{0}); graph.SetInputs(inputs).SetOutputs(output_indices);注册时同时挂接两个回调REGISTER_CUSTOM_OP(PartitionedCall) .FrameworkType(ONNX) .OriginOpType(ai.onnx::11::AddN) .ParseParamsByOperatorFn(ParseParamsAddn) .ParseOpToGraphFn(ParseOpToGraphAddn) .ImplyType(ImplyType::TVM);这种将算子映射为子图PartitionedCall的能力使一个第三方框架算子可以透明地展开为一组 CANN 原生算子组合是算子移植中非常实用的手段。跨 opset 属性解析LeakyRelu 算子leaky_relu_plugin.ccleaky_relu_plugin.cc 演示了属性解析与多版本兼容.OriginOpType(...)传入一个版本列表{ai.onnx::8::LeakyRelu, ..., ai.onnx::13::LeakyRelu}一次性兼容 opset 8~13 共 6 个版本ParseOnnxParamsLeakyRelu回调从 ONNX 节点的 JSON 属性串中解析alpha默认值0.01f转换为 CANN 算子的negative_slope属性float negative_slope 0.01f; string negative_slope_str; AscendString attrs_string; if (ge::GRAPH_SUCCESS op_src.GetAttr(attribute, attrs_string)) { json attrs json::parse(attrs_string.GetString()); for (json attr : attrs[attribute]) { if (attr[name] alpha attr[type] kTypeFloat) { negative_slope_str attr[f]; // float type in json has accuracy loss, so we use string type to store it negative_slope atof(negative_slope_str.c_str()); } } } op_dest.SetAttr(negative_slope, negative_slope);实现细节值得注意源码特意以字符串形式取出attr[f]再atof转换并注释说明JSON 中 float 类型有精度损失因此用字符串存储这是跨框架属性传递中避免浮点精度丢失的工程经验。这也解释了为什么 onnx_plugin/CMakeLists.txt 需要额外引入 nlohmann/json 依赖——属性解析依赖 JSON 库。TensorFlow 框架插件tf_pluginTF 插件的目标与 ONNX 插件不同将 TensorFlow 自定义算子注册到 CANN 框架。七个样例可归为三类模式。直接注册指定实现类型REGISTER_CUSTOM_OPAutoMappingByOpFn 指定ImplyType参数自动映射。例如add_block_cust_plugin.cc注册AddBlockCust.ImplyType(ImplyType::AI_CPU)即该算子由 AI CPU 实现add_dsl_plugin.cc注册AddDsl.ImplyType(ImplyType::TVM)lstm_tik_plugin.cc注册LSTMTikTVM 实现reshape_cust_plugin.cc注册ReshapeCustAI_CPU 实现scatter_nd_add_plugin.cc注册ScatterNdAddTVM 实现unique_cust_plugin.cc注册UniqueCustAI_CPU 实现。这些样例的共同点是ParseParamsByOperatorFn(AutoMappingByOpFn)——完全交给框架按算子定义自动映射参数插件只需声明原始算子类型与实现类型。当自定义算子同时存在多套实现时可通过多个REGISTER_CUSTOM_OP分别声明不同ImplyType为编译期选择实现提供依据。Scope 融合算子适配DecodeBboxV2decode_bbox_v2_scope_fusion_plugin.ccdecode_bbox_v2_scope_fusion_plugin.cc 是最具代表性的一个它适配的是一个融合算子DecodeBboxV2FusionOp原始模型中被融合进同一个 Scope 的是若干个小算子Unstack、RealDiv等。与前述样例不同它注册的回调是.FusionParseParamsFn(DecodeBboxV2ParseParams)。该回调接收inside_nodesScope 内的小算子列表从中提取融合所需的缩放参数CollectScalesInfo遍历 Scope 内节点识别类型为RealDiv且第一个输入名包含/unstack的节点将其第二个输入常量节点记录为 scale 来源DecodeBboxV2ParseParams若收集到 4 个 scalekScaleSize 4则通过ParseFloatFromConstNode从常量节点的value属性中读取 float 数值写入scales列表否则维持默认{1.0, 1.0, 1.0, 1.0}最终通过op_dest.SetAttr(scales, scales_list)将提取结果设置到融合算子。该样例展示了框架插件的进阶用法从 Scope 内的已有子图中反推融合算子所需参数这在将框架侧已融合的复合算子适配到 CANN 时非常典型。环境要求环境要求与主仓一致详见主仓 README.md 中的「环境部署」章节。本样例额外需要onnx 1.12.0用于生成验证模型编译本模块属于 Host 侧代码不依赖具体 NPU 架构但主仓构建仍要求填写NPU_ARCH参数当前仓库支持 Ascend 950dav-3510与 Ascend 910B/Cdav-2201等平台具体取值见主仓 README 的环境部署章节。编译方式一通过 build.sh 便捷编译推荐在custom_op_in_graph目录下执行chmod x build.sh ./build.shbuild.sh 会自动调用主仓统一构建系统完成配置和编译默认以NPU_ARCHdav-3510在仓库根目录执行cmake -S . -B build然后构建cust_onnx_parsers与cust_tf_parsers两个目标产物输出到当前目录的build_out/makepkg/下。若重新编译先执行./build.sh clean清理产物。方式二通过主仓统一构建从项目根目录启动构建参考项目 README.md# 1. 配置项目NPU_ARCH 为必填参数本模块为 Host 侧代码不依赖具体架构 cmake -S . -B build -DNPU_ARCHdav-3510 # 2. 编译插件 cmake --build build --target cust_onnx_parsers cust_tf_parsers编译后产物直接输出到custom_op_in_graph/build_out/makepkg/目录下。编译产物custom_op_in_graph/build_out/makepkg/ ├── set_env.bash // 环境变量脚本 └── packages/vendors/customize/ └── framework/ ├── onnx/libcust_onnx_parsers.so └── tensorflow/libcust_tf_parsers.so部署编译完成后custom_op_in_graph/build_out/makepkg/目录结构与算子包一致包含所有框架的插件库。部署方式有两种指定目录安装推荐用于验证执行source build_out/makepkg/set_env.bash将编译输出路径追加到ASCEND_CUSTOM_OPP_PATH环境变量。ONNX/TF 的描述转换成 GEOP 时会遍历该路径下的framework/子目录自动发现所有框架的插件库。在当前终端生效后可直接进行模型转换验证。多个厂商的算子包共存时按照ASCEND_CUSTOM_OPP_PATH中从左到右的顺序搜索后 source 的路径优先级更高。默认安装将packages/vendors/customize/目录整体拷贝到 CANN OPP 算子库路径CANN/opp/vendors/下注意必须从customize目录层级拷贝。编辑CANN/opp/vendors/config.ini将customize写入load_priority值头部以逗号分隔其他厂商例如load_prioritycustomize,other_vendor。从 CMakeLists.txt 源码可以看到set_env.bash由构建系统自动生成其内容即vendor_path$(cd $(dirname ${BASH_SOURCE[0]})/packages/vendors/customize; pwd) export ASCEND_CUSTOM_OPP_PATH${vendor_path}:${ASCEND_CUSTOM_OPP_PATH}理解这一点有助于排查部署问题所谓source 即生效本质就是把customize目录写入ASCEND_CUSTOM_OPP_PATH搜索路径。将算子映射为子图一对多映射验证用户可使用 ATC 模型转换工具对算子映射为子图的效果进行验证。下面给出完整验证方法1. 构造包含 AddN 算子的 ONNX 模型假设用户工作路径为work_dir在工作路径下创建 Python 脚本gen_addn.py脚本内容参考import os import numpy as np import onnx def gen_onnx(): X onnx.helper.make_tensor_value_info(X, onnx.TensorProto.FLOAT, [5]) Y onnx.helper.make_tensor_value_info(Y, onnx.TensorProto.FLOAT, [5]) Z onnx.helper.make_tensor_value_info(Z, onnx.TensorProto.FLOAT, [5]) output onnx.helper.make_tensor_value_info(output, onnx.TensorProto.FLOAT, [5]) node0 onnx.helper.make_node(AddN, inputs[X, Y, Z], outputs[output]) inputs [X, Y, Z] outputs [output] graph_def onnx.helper.make_graph( [node0], addn_model, inputs, outputs ) model_def onnx.helper.make_model(graph_def) model_def.opset_import[0].version 11 onnx.save(model_def, addn_model.onnx) print(model_def) if __name__ __main__: gen_onnx()执行脚本生成的 ONNX 模型文件addn_model.onnx位于work_dir目录下python3 gen_addn.py2. 通过 ATC 模型转换功能验证算子映射子图效果1设置环境变量。完成 CANN 软件基础环境变量配置后还需要额外配置如下环境变量export DUMP_GE_GRAPH2 # 控制 dump 图的内容多少 export DUMP_GRAPH_LEVEL2 # 控制 dump 图的个数2进行模型转换atc --model./addn_model.onnx --framework5 --output./addn --input_formatNCHW --soc_version${soc_version}其中soc_version昇腾 AI 处理器的型号请根据实际情况替换。可从 ATC 安装路径下的arch-linux/data/platform_config目录下查看支持的昇腾 AI 处理器的类型对应*.ini文件的名字即为soc_version。模型转换完成后会在执行 atc 命令的当前目录下生成一系列按ge_onnx*.pbtxt命名方式命名的文件。这些文件是基于 ONNX 的开源模型描述结构可以使用 Netron 等可视化软件打开。3结果验证。ge_onnx_00000000_graph_0_PreRunBegin.pbtxt是 GE 获取到的经过 parse 处理的整张下沉图。使用 Netron 等可视化软件打开原始模型和ge_onnx_00000000_graph_0_PreRunBegin.pbtxt可以看到算子映射子图的实际效果转换前原始 AddN 模型转换后AddN 被拆分为 Add Add 子图验证原理与源码呼应转换后的图中出现PartitionedCall节点正是 addn_plugin.cc 中ParseOpToGraphAddn构造的子图结构——Data声明输入、两级Add级联、SetOutputs声明输出。若 dump 图中 AddN 被正确展开为两个级联的 Add即说明插件注册、编译、部署、加载全链路均已打通。小结插件开发的可复用模式从本样例可以归纳出框架插件开发的三种可复用模式场景核心回调代表样例一对一算子映射无属性转换ParseParamsByOperatorFn返回 SUCCESS 或AutoMappingByOpFnadd_plugin.cc、tf_plugin 下各注册样例一对多映射展开为子图ParseParamsByOperatorFnParseOpToGraphFnaddn_plugin.cc属性跨框架转换 / 从 Scope 提取融合参数自定义ParseParamsByOperatorFn/FusionParseParamsFnleaky_relu_plugin.cc、decode_bbox_v2_scope_fusion_plugin.cc在实际自定义算子工程中建议按此模式对照本样例先确定原始框架与 opset 版本、选定实现类型TVM / AI_CPU、判断是直接注册还是需要子图展开再决定采用哪个回调组合。完成插件编写后按照编译 → source set_env.bash → ATC 转换 DUMP_GE_GRAPH 对比的流程即可快速验证映射效果。【免费下载链接】cann-samplesCANN高性能实战演进样例与体系化调优知识库项目地址: https://gitcode.com/cann/cann-samples创作声明：本文部分内容由AI辅助生成（AIGC），仅供参考