昇腾 Ascend C 算子接入 TensorFlow：内置算子映射与自定义算子两种实现方式详解

• 人工智能
• 深度学习
• 算子库
• CANN
• Ascend

导读

本文基于 asc-devkit 开源仓库中的 TensorFlow 框架接入样例，系统讲解如何将 Ascend C 自定义算子接入 TensorFlow 并在昇腾 NPU 上执行。文章覆盖两条典型路径：把 Add 计算映射到 TensorFlow 内置算子 AddV2（tensorflow_builtin），以及映射到 TensorFlow 自定义算子 AddCustom（tensorflow_custom），并深入剖析框架适配插件、TensorFlow 调用脚本、算子注册实现与编译运行全流程。读完本文，你将掌握 Ascend C 算子接入 TensorFlow 的完整工程方法、关键配置项与调试验证手段。

概述

在昇腾 AI 处理器上开发自定义算子后，需要将其接入主流 AI 框架才能被业务模型调用。asc-devkit 仓库在 examples/01_simd_cpp_api/02_features/00_framework 目录下提供了 PyTorch、TensorFlow、ONNX、GE 原生构图、ACLGraph 等多种接入场景的样例，其中01_tensorflow目录专门介绍 TensorFlow 框架的自定义算子实现方法。

本样例以最基础的 Add 计算（z = x + y）为载体，展示了两种接入思路：

• 映射到 TensorFlow 内置算子：Ascend C 自定义算子通过框架适配插件注册为 TensorFlow 内置算子（如AddV2）的替代实现，调用侧无需任何算子名改动；
• 映射到 TensorFlow 自定义算子：Ascend C 算子注册为 TensorFlow 的自定义算子（AddCustom），通过tf.load_op_library加载后直接调用。

样例列表

两个样例均要求 CANN 版本不低于：Ascend 950PR/Ascend 950DT 为 CANN 9.1.0，Atlas A2/A3 训练与推理系列为 CANN 9.0.0。

前置条件：先完成自定义算子工程编译、打包与部署

两条接入路径都依赖一个已部署到 CANN 环境中的 Ascend C 自定义算子工程。在运行 TensorFlow 样例之前，需先进入自定义样例工程 custom_op 目录完成编译、打包和部署，操作如下：

source ${install_path}/cann/set_env.sh mkdir -p build && cd build cmake .. && make -j binary package ./custom_opp_*.run

执行结果显示SUCCESS即部署成功。其中${install_path}为 CANN 包安装目录，未指定安装目录时默认安装至/usr/local/Ascend下。该工程在framework/tf_plugin目录下提供 TensorFlow 框架适配插件源文件tensorflow_add_custom_plugin.cc，这是两个 TensorFlow 样例能够识别 Ascend C 算子的关键枢纽。

提示：custom_op 工程还支持 ccache 缓存编译加速，可通过cmake -DENABLE_CCACHE=ON .. && make -j binary package开启，详见 custom_op README。

方式一：映射到 TensorFlow 内置算子（tensorflow_builtin）

该样例展示基于 Add 计算，将 Ascend C 自定义算子映射到 TensorFlow 内置算子AddV2，从而在 TensorFlow 侧以标准tf.math.add接口完成 NPU 上的张量加法。

目录结构

├── tensorflow_builtin │ ├── run_add_custom.py // TensorFlow 调用脚本（映射到内置 AddV2 算子） │ └── README.md // 样例说明文档

完成自定义算子工程创建与部署后，样例工程目录下会生成framework/tf_plugin目录，用于存放 TensorFlow 框架适配插件实现文件。

框架适配插件：将 Ascend C 算子注册为内置算子

本样例的核心适配点在 tensorflow_add_custom_plugin.cc。由于内置 AddV2 在 TensorFlow 侧已经存在，需要把插件代码中 TensorFlow 调用算子的名称OriginOpType修改为"AddV2"：

#include "register/register.h" namespace domi { // register op info to GE REGISTER_CUSTOM_OP("AddCustom") .FrameworkType(TENSORFLOW) // type: TENSORFLOW .OriginOpType("AddV2") // name in tf module .ParseParamsByOperatorFn(AutoMappingByOpFn); } // namespace domi

各字段含义如下：

• REGISTER_CUSTOM_OP("AddCustom")：注册到 GE（图引擎）的 Ascend C 算子名，与算子工程中的 OpType 保持一致；
• .FrameworkType(TENSORFLOW)：声明该算子接入的框架类型为 TensorFlow；
• .OriginOpType("AddV2")：指定 TensorFlow 侧对应的算子名。填内置算子名时，框架图编译阶段会把计算图中出现的AddV2自动映射到 Ascend C 的AddCustom实现；
• .ParseParamsByOperatorFn(AutoMappingByOpFn)：使用自动映射函数解析算子参数，将 TensorFlow 算子的输入、输出、属性按约定映射到 Ascend C 算子。

正是通过这一注册机制，TensorFlow 内置算子 AddV2 的调用最终落到了 Ascend C 的add_custom核函数上，而调用方完全感知不到算子被替换。

TensorFlow 调用脚本核心流程

run_add_custom.py 完整走通了"构造输入 → 构建计算图 → CPU/NPU 双会话对比 → 校验结果"的验证链路：

1. 使用np.random.uniform在(-2, 2)区间生成形状为(8, 2048)、数据类型为np.float16的随机输入数据，并通过tf.compat.v1.placeholder定义输入节点；
2. 构建计算图：使用tf.math.add实现张量加法运算；
3. 分别创建 CPU 和 NPU 会话，通过session.run执行计算图并传入输入数据，其中 NPU 会话通过ConfigProto注入NpuOptimizer配置；
4. 使用np.allclose(result_ai_core, result_cpu, ATOL, RTOL)对比 NPU 与 CPU 计算结果，容差参数ATOL = 0.001、RTOL = 0.001，一致则打印test pass。

NPU 会话的关键配置如下（脚本中config("ai_core")分支）：

session_config = tf.compat.v1.ConfigProto(allow_soft_placement=True, log_device_placement=False) custom_op = session_config.graph_options.rewrite_options.custom_optimizers.add() custom_op.name = "NpuOptimizer" custom_op.parameter_map["enable_data_pre_proc"].b = True custom_op.parameter_map["mix_compile_mode"].b = True custom_op.parameter_map["use_off_line"].b = True custom_op.parameter_map["min_group_size"].b = 1

这些参数控制着图编译与下沉行为：enable_data_pre_proc使能数据预处理、mix_compile_mode开启混合编译模式、use_off_line指定离线模型执行、min_group_size设置算子融合的最小分组。CPU 会话则只需基础的ConfigProto，作为数值基准。

编译运行步骤

在 tensorflow_builtin 样例根目录下依次执行：

1. 安装 TensorFlow 插件包：参考 CANN 文档《TensorFlow 2.6.5 模型迁移》中的"安装框架插件包"章节，将框架适配插件安装到 TensorFlow 环境；
2. 配置环境变量：

source ${install_path}/cann/set_env.sh

3. 执行样例：

python3 run_add_custom.py

执行结果如下，说明执行成功：

test pass

方式二：映射到 TensorFlow 自定义算子（tensorflow_custom）

该样例展示将 Ascend C 自定义算子 Add 映射为 TensorFlow 侧的自定义算子AddCustom，通过tf.load_op_library加载算子库后直接调用。

目录结构

├── tensorflow_custom │ ├── CMakeLists.txt // 编译工程文件 │ ├── custom_assign_add_custom.cc // TensorFlow 自定义算子实现 │ ├── run_add_custom_tf.py // TensorFlow 调用脚本（映射到自定义 AddCustom 算子） │ └── README.md // 样例说明文档

与内置算子方式不同，该目录是一个独立的 CMake 工程，需要先编译出 TensorFlow 自定义算子动态库libcustom_ops.so，再执行调用脚本。

框架适配插件：注册为自定义算子

同样修改 tensorflow_add_custom_plugin.cc，此时OriginOpType填自定义算子名：

#include "register/register.h" namespace domi { // register op info to GE REGISTER_CUSTOM_OP("AddCustom") .FrameworkType(TENSORFLOW) // type: TENSORFLOW .OriginOpType("AddCustom") // name in tf module .ParseParamsByOperatorFn(AutoMappingByOpFn); } // namespace domi

与方式一的差异仅在于OriginOpType：填"AddV2"是接管内置算子，填"AddCustom"则是建立 TensorFlow 自定义算子与 Ascend C 算子之间的映射关系。

TensorFlow 侧自定义算子注册实现

custom_assign_add_custom.cc 使用 TensorFlow 原生 Op 注册机制定义了算子接口与 CPU 占位实现：

#include "tensorflow/core/framework/op.h" #include "tensorflow/core/framework/op_kernel.h" #include "tensorflow/core/framework/shape_inference.h" #include "tensorflow/core/framework/common_shape_fns.h" using namespace tensorflow; // 注册 TensorFlow 自定义算子 REGISTER_OP("AddCustom") // TensorFlow 自定义算子名称 .Input("x: T") // 输入 tensor x .Input("y: T") // 输入 tensor y .Output("z: T") // 输出 tensor z .Attr("T: {half}") // 属性 T，支持 half 数据类型 // 设置 shape 函数，BroadcastBinaryOpShapeFn 处理输入、输出 tensor shape 相同的情况 .SetShapeFn(shape_inference::BroadcastBinaryOpShapeFn); // TensorFlow 自定义算子的 CPU 实现 class AddCustomOp : public OpKernel { public: explicit AddCustomOp(OpKernelConstruction* context) : OpKernel(context) {} // 当前算子不支持 CPU 设备，实现该函数以抛出异常 void Compute(OpKernelContext* context) override { OP_REQUIRES(context, false, errors::Unimplemented("AddCustomOp is not supported on CPU")); } }; // 注册 TensorFlow 自定义算子的 CPU 实现 REGISTER_KERNEL_BUILDER(Name("AddCustom").Device(DEVICE_CPU), AddCustomOp);

这里的关键设计是：算子接口定义（REGISTER_OP）用于图构建与 shape 推断，AddCustomOp的 CPUCompute实现则故意抛出Unimplemented异常——因为真实计算由 NPU 上的 Ascend C 核函数完成，注册一个"不支持 CPU"的占位实现是为了让 TensorFlow 图能够正常构建，同时保证算子不会错误地落到 CPU 设备执行。

CMake 编译工程

tensorflow_custom/CMakeLists.txt 负责将.cc源文件编译为共享库libcustom_ops.so，核心逻辑包括：

• 通过 Python3 的tf.sysconfig.get_compile_flags()与tf.sysconfig.get_link_flags()动态获取当前 TensorFlow 的编译与链接参数；
• file(GLOB_RECURSE SOURCE_FILES CONFIGURE_DEPENDS "*.cc")收集源码，add_library(custom_ops SHARED ...)构建动态库；
• 通过set_target_properties将输出目录设为outputs、输出名设为custom_ops，最终生成outputs/libcustom_ops.so。

TensorFlow 调用脚本核心流程

run_add_custom_tf.py 的验证思路是"标准算子 vs 自定义算子"双路对比：

1. 通过tf.load_op_library加载outputs/libcustom_ops.so，获取自定义算子接口add_custom；
2. 构造输入数据，使用tf.compat.v1.placeholder定义输入张量，分别计算tf.math.add（标准 TensorFlow 加法）与add_custom（Ascend C 自定义算子）的结果；
3. 配置ConfigProto，启用NpuOptimizer，并关闭重映射与内存优化，确保自定义算子按预期执行：

config = tf.compat.v1.ConfigProto() custom_op = config.graph_options.rewrite_options.custom_optimizers.add() custom_op.name = "NpuOptimizer" config.graph_options.rewrite_options.remapping = RewriterConfig.OFF config.graph_options.rewrite_options.memory_optimization = RewriterConfig.OFF

关闭remapping与memory_optimization是为了避免 TensorFlow 图优化阶段改写自定义算子节点，从而保证add_custom节点原样下沉到 NPU 执行；

4. 在两个独立会话中分别运行tf_z与ac_z，调用np.allclose(tf_golden, ac_golden, ATOL, RTOL)对比结果，一致则打印test pass。

编译运行步骤

在 tensorflow_custom 样例根目录下依次执行：

source ${install_path}/cann/set_env.sh mkdir -p build; cd build cmake .. && make -j python3 ../run_add_custom_tf.py

执行结果如下，说明执行成功：

test pass

两种接入方式的对比与选型建议

从源码结构可以推断，两种方式的差异本质上只体现在算子名映射与调用方式上：内置算子方式复用 TensorFlow 既有算子语义、对用户完全透明，适合算子语义与框架内置算子一致的场景；自定义算子方式显式暴露算子名、可自由扩展框架不存在的算子语义，适合新增算子类型。两者共用同一份 Ascend C 算子工程与tensorflow_add_custom_plugin.cc适配插件，改动成本集中在OriginOpType一行与调用脚本的配置项上。

总结

本文以 Add 计算为例，完整梳理了 asc-devkit 仓库中 Ascend C 算子接入 TensorFlow 的两条路径。核心要点可归纳为：

• 插件是枢纽：REGISTER_CUSTOM_OP+FrameworkType(TENSORFLOW)+OriginOpType决定了 Ascend C 算子以什么名字接入 TensorFlow；
• 算子工程是底座：两个样例均依赖 custom_op 工程的编译、打包与部署；
• 验证方法是闭环：内置算子方式对比"NPU vs CPU"结果，自定义算子方式对比"标准算子 vs 自定义算子"结果，均以np.allclose+test pass收尾。

按上述流程实践，即可将任意 Ascend C 算子接入 TensorFlow，并在此基础上进一步扩展 tiling、多输入输出、属性解析等更复杂的算子形态。

• 人工智能
• 深度学习
• 算子库
• CANN
• Ascend