已合并
refactor: 迁移 TF 插件注册至新框架 API(part 3) #9758
refactor: 迁移 TF 插件注册至新框架 API(part 3) #9758
已合并
Nice try创建于 8 天前
Nice try
Nice try成员
8 天前

描述

将 TF 框架侧自定义算子注册(REGISTER_CUSTOM_OP)从旧工具链迁移至 ops-nn 新仓(TF 插件迁移 part 3),共新增/调整 22 个 *_tf_plugin.cpp 注册文件。

改动原因

  1. 旧注册 API(ParseParamsFn/AutoMappingFn)与标量 OriginOpType({...}) 写法在新工具链已不适用,统一迁移为 ParseParamsByOperatorFn(AutoMappingByOpFn) 与 OriginOpType(std::vectorge::AscendString{...}) 形式;
  2. 依赖 TF protobuf NodeDef(node_def.pb.h)的旧解析逻辑在新工具链中已不存在,相关解析改为基于 ge::Operator 的子图分解(ParseOpToGraphFn/FusionParseParamsFn)实现;
  3. 承接本仓已完成的 TF 插件迁移 part 1/2,补齐剩余算子的 TF 注册。

改动方法

按模块汇总:

  • common:新增 AscendWeightQuant(TF 字符串 dst_type 到 ge DataType 的转换解析)、GatherPoint(分解为 GatherV2+Const 子图)、GatherPointGrad(分解为 ScatterUpdate 子图)注册;
  • index:gather_v2 新增 Gather/ResourceGather 到自定义 Gather 的注册;scatter_add/scatter_sub 新增 TensorScatterAdd/TensorScatterSub 注册;新增 BroadcastGradientArgs、SparseSegmentMean、SparseSegmentMeanGrad、SparseSlice、SparseToDense 注册;
  • norm:BNInfer/BNTrainingReduceGrad/BNTrainingUpdateGrad/LayerNorm/LayerNormGrad 由 ParseParamsFn 迁移到 ParseParamsByOperatorFn;LayerNorm/LayerNormGrad 增补 FusedLayerNorm/FusedLayerNormGrad 双 origin 映射;bn_training_update_v2 新增 BNTrainingUpdate 注册;
  • optim:ApplyAdagrad 迁移新 API(保留 ApplyAdagrad/ResourceApplyAdagrad 双 origin);
  • matmul:新增 Einsum 注册;
  • rnn:新增 BlockLSTM(分解为 DynamicRNN 子图,含 seq_length 的 Cast 与输出重排)、DynamicRNN 融合解析(FusionParseParamsFn)与 DynamicRnn 注册、DynamicRNNV2/DynamicRnnV2 系列注册(含 is_misplaced 属性处理)。

关联的Issue

关联 Issue #5505:https://gitcode.com/cann/ops-nn/issues/5505

测试

本 PR 仅涉及注册/图解析代码迁移,无测试文件变更;依赖仓内 CI(编译、静态检查、UT/ST 工作流)验证。

文档更新

无。

类型标签

AI/Agent生成声明

likedislike
Pull Request已成功合入, 合并人@CANN-robot
(感谢 Nice try 的贡献)
Nice tryNice try成员
8 天前 创建了 pull request,commit 5bf4ada7
atomgit-bot
atomgit-bot
8 天前 评论:

变更摘要

本 PR 为 TensorFlow 算子插件迁移工作("migrate TF plugin part 3")的第三部分,主体思路是:将一批已有算子注册从旧接口 ParseParamsFn(AutoMappingFn) 迁移到新的 ParseParamsByOperatorFn(AutoMappingByOpFn),同时新增多个以自动映射方式注册的 TF 自定义算子,并对部分需要图改写或特殊属性处理的算子提供自定义解析函数(如 ParseOpToGraphFn 图转换、融合解析与类型转换等)。

主要改动

  • 注册接口迁移bn_infer_tf_plugin.cppbn_training_reduce_grad_tf_plugin.cppbn_training_update_grad_tf_plugin.cppbn_training_update_v2_tf_plugin.cppapply_adagrad_tf_plugin.cpp 等文件将 .ParseParamsFn(AutoMappingFn) 替换为 .ParseParamsByOperatorFn(AutoMappingByOpFn),并同步将 OriginOpType 改为 std::vector<ge::AscendString> 形式。
  • 新增自动映射算子注册:新增 RandomStandardNormalRandomUniformIntBroadcastGradientArgsSparseSegmentMeanSparseSegmentMeanGradSparseSliceSparseToDenseEinsum 等文件的注册,统一使用 AutoMappingByOpFn 映射并声明 ImplyType::TVM
  • 扩展变体算子覆盖:在既有文件中追加注册以支持同名变体算子,包括 GatherV2 文件新增 Gather/ResourceGatherScatterAdd 新增 TensorScatterAddScatterSub 新增 TensorScatterSubBNTrainingUpdateV2 新增 BNTrainingUpdateLayerNorm/LayerNormGrad 增加 FusedLayerNorm/FusedLayerNormGrad 对应关系。
  • 基于图改写的复合算子映射:新增 GatherPoint(改写为 GatherV2Const 轴输入)、GatherPointGrad(改写为 ScatterUpdate)及 BlockLSTM(改写为 DynamicRNN,含 Cast 序列长度输入与 forget_bias/cell_clip/use_peephole 属性透传)三个注册,均通过 ParseOpToGraphFn 构造子图并设置输入输出。
  • DynamicRNN 系列注册与融合解析dynamic_rnn_tf_plugin.cpp 新增 DynamicRNN 注册并附带 FusionParseParamsFn(设置 time_major/forget_bias 属性及输入 HWCN 格式);dynamic_rnnv2_tf_plugin.cppDynamicRnnV2/DynamicRnnv2WithoutSeqlength 设置 is_misplaced 属性;ascend_weightquant_tf_plugin.cpp 将字符串属性 dst_type(如 "INT4")转换为对应整型数据类型后写回算子。
likedislike
不准确?
atomgit-bot
atomgit-bot
8 天前 评论:

代码审查

✅ 未发现问题

likedislike
不准确?
CANN-robotCANN-robot成员
8 天前 添加了label:cann-cla/yes
CANN-robot
CANN-robot成员
8 天前 评论:

Thanks for your pull-request.
The full list of commands accepted by me can be found at here.
You can get sig-info at here.
You can self-configure the PR merge rules for this repository. For more details, please refer to Here.
For more, you also can visit HICANN.


PR Approval Progress

Congratulations! All modules have met the lgtm and approve requirements.

Module Approval Details

module lgtm status approve status
common 杨阳, 汤平川 (2/2) 杨阳 (1/1)
index 汤平川, 钱泽洪, 杨阳 (3/2) 汤平川, 钱泽洪 (2/1)
matmul 杨阳, 汤平川 (2/2) 杨阳 (1/1)
norm 钱泽洪, 汤平川, 杨阳 (3/2) 钱泽洪, 汤平川 (2/1)
optim 杨阳, 汤平川, 钱泽洪 (3/2) 汤平川, 钱泽洪 (2/1)
rnn 汤平川, 钱泽洪, 杨阳 (3/2) 汤平川, 钱泽洪 (2/1)

💡 Tip:

  • Committer can comment /approve or /lgtm
  • Commenting /approve implies both code review (lgtm) and intent to merge (approve)

CLA Signature Pass

Nice_try, thanks for your pull request. All authors of the commits have signed the CLA. 👍

likedislike
此处折叠了44条消息 查看更多
yangyang016成员
7 天前 评论:

/lgtm
/approve

likedislike
CANN-robotCANN-robot成员
7 天前 添加了label:lgtmapproved
CANN-robotCANN-robot成员
7 天前 关闭了关联的issue
CANN-robotCANN-robot成员
7 天前 合入了pull request
CANN-robot
CANN-robot成员
7 天前 评论:

Pull Request 已合并或已关闭。

If you want to solve this problem, you can click here to do it in the FAQs.

likedislike