人工智能深度学习算子库CANNAscend【免费下载链接】asc-devkit本项目是CANN 推出的昇腾AI处理器专用的算子程序开发语言原生支持C和C标准规范主要由类库和语言扩展层构成提供多层级API满足多维场景算子开发诉求。项目地址https://gitcode.com/cann/asc-devkit点击查看免费下载导读本文基于 asc-devkit 开源仓库中的 TensorFlow 框架接入样例系统讲解如何将 Ascend C 自定义算子接入 TensorFlow 并在昇腾 NPU 上执行。文章覆盖两条典型路径把 Add 计算映射到 TensorFlow 内置算子 AddV2tensorflow_builtin以及映射到 TensorFlow 自定义算子 AddCustomtensorflow_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加载后直接调用。样例列表目录名称功能描述支持的产品tensorflow_builtin样例展示将 Ascend C 自定义算子映射到 TensorFlow 内置算子并通过 TensorFlow 调用Ascend 950PR/Ascend 950DTAtlas A3 训练系列产品/Atlas A3 推理系列产品Atlas A2 训练系列产品/Atlas A2 推理系列产品tensorflow_custom样例展示将 Ascend C 自定义算子映射到 TensorFlow 自定义算子并通过 TensorFlow 调用Ascend 950PR/Ascend 950DTAtlas A3 训练系列产品/Atlas A3 推理系列产品Atlas A2 训练系列产品/Atlas A2 推理系列产品两个样例均要求 CANN 版本不低于Ascend 950PR/Ascend 950DT 为 CANN 9.1.0Atlas 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_CCACHEON .. 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 双会话对比 → 校验结果的验证链路使用np.random.uniform在(-2, 2)区间生成形状为(8, 2048)、数据类型为np.float16的随机输入数据并通过tf.compat.v1.placeholder定义输入节点构建计算图使用tf.math.add实现张量加法运算分别创建 CPU 和 NPU 会话通过session.run执行计算图并传入输入数据其中 NPU 会话通过ConfigProto注入NpuOptimizer配置使用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_placementTrue, log_device_placementFalse) 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 样例根目录下依次执行安装 TensorFlow 插件包参考 CANN 文档《TensorFlow 2.6.5 模型迁移》中的安装框架插件包章节将框架适配插件安装到 TensorFlow 环境配置环境变量source ${install_path}/cann/set_env.sh执行样例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 自定义算子双路对比通过tf.load_op_library加载outputs/libcustom_ops.so获取自定义算子接口add_custom构造输入数据使用tf.compat.v1.placeholder定义输入张量分别计算tf.math.add标准 TensorFlow 加法与add_customAscend C 自定义算子的结果配置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 执行在两个独立会话中分别运行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_builtintensorflow_custom插件OriginOpTypeAddV2内置算子名AddCustom自定义算子名调用侧接口tf.math.add无需感知算子被替换add_custom需通过tf.load_op_library显式加载额外编译产物无依赖已部署的算子工程libcustom_ops.so动态库需 CMake 构建图优化配置配置NpuOptimizer并开启数据预处理等开关需额外关闭remapping与memory_optimization防止自定义节点被改写适用场景希望以标准 TensorFlow 算子接口透明使用 Ascend C 实现希望以独立算子名在 TensorFlow 图中显式调用 Ascend C 算子从源码结构可以推断两种方式的差异本质上只体现在算子名映射与调用方式上内置算子方式复用 TensorFlow 既有算子语义、对用户完全透明适合算子语义与框架内置算子一致的场景自定义算子方式显式暴露算子名、可自由扩展框架不存在的算子语义适合新增算子类型。两者共用同一份 Ascend C 算子工程与tensorflow_add_custom_plugin.cc适配插件改动成本集中在OriginOpType一行与调用脚本的配置项上。总结本文以 Add 计算为例完整梳理了 asc-devkit 仓库中 Ascend C 算子接入 TensorFlow 的两条路径。核心要点可归纳为插件是枢纽REGISTER_CUSTOM_OPFrameworkType(TENSORFLOW)OriginOpType决定了 Ascend C 算子以什么名字接入 TensorFlow算子工程是底座两个样例均依赖 custom_op 工程的编译、打包与部署验证方法是闭环内置算子方式对比NPU vs CPU结果自定义算子方式对比标准算子 vs 自定义算子结果均以np.allclosetest pass收尾。按上述流程实践即可将任意 Ascend C 算子接入 TensorFlow并在此基础上进一步扩展 tiling、多输入输出、属性解析等更复杂的算子形态。赞分享人工智能深度学习算子库CANNAscend【免费下载链接】asc-devkit本项目是CANN 推出的昇腾AI处理器专用的算子程序开发语言原生支持C和C标准规范主要由类库和语言扩展层构成提供多层级API满足多维场景算子开发诉求。项目地址https://gitcode.com/cann/asc-devkit点击查看免费下载相关推荐CANN asc-devkit 中通过 TensorFlow 调用 Ascend C 自定义算子内置算子映射与自定义算子映射两种适配路径实战CANN asc devkit 中通过 TensorFlow 调用 Ascend C 自定义算子内置算子映射与自定义算子映射两种适配路径实战 导读 在昇腾 A人工智能深度学习算子库CANNAscendTiny11Builder 完整指南用 PowerShell 制作精简版 Windows 11 镜像Tiny11Builder 完整指南用 PowerShell 制作精简版 Windows 11 镜像 Tiny11Builder 是一个制作精简版 Windo示例工程AscendAscend C 向量编程范式全解以 add_custom 算子为例剖析昇腾自定义向量算子开发全流程Ascend C 向量编程范式全解以 add_custom 算子为例剖析昇腾自定义向量算子开发全流程 本节内容源自 05Paradigm.md https:/文档教程人工智能上一篇Microcks革命Kubernetes原生API测试与Mock全攻略下一篇终极LLM-Engineers-Handbook性能优化指南提升推理速度和降低成本的10个实用技巧创作声明：本文部分内容由AI辅助生成（AIGC），仅供参考