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

描述

将 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成员
22 天前 创建了 pull request,commit 5bf4ada7
atomgit-bot
atomgit-bot
22 天前 评论:

变更摘要

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

主要改动

  • 注册接口迁移:bn_infer_tf_plugin.cpp、bn_training_reduce_grad_tf_plugin.cpp、bn_training_update_grad_tf_plugin.cpp、bn_training_update_v2_tf_plugin.cpp、apply_adagrad_tf_plugin.cpp 等文件将 .ParseParamsFn(AutoMappingFn) 替换为 .ParseParamsByOperatorFn(AutoMappingByOpFn),并同步将 OriginOpType 改为 std::vector<ge::AscendString> 形式。
  • 新增自动映射算子注册:新增 RandomStandardNormal、RandomUniformInt、BroadcastGradientArgs、SparseSegmentMean、SparseSegmentMeanGrad、SparseSlice、SparseToDense、Einsum 等文件的注册,统一使用 AutoMappingByOpFn 映射并声明 ImplyType::TVM。
  • 扩展变体算子覆盖:在既有文件中追加注册以支持同名变体算子,包括 GatherV2 文件新增 Gather/ResourceGather、ScatterAdd 新增 TensorScatterAdd、ScatterSub 新增 TensorScatterSub、BNTrainingUpdateV2 新增 BNTrainingUpdate、LayerNorm/LayerNormGrad 增加 FusedLayerNorm/FusedLayerNormGrad 对应关系。
  • 基于图改写的复合算子映射:新增 GatherPoint(改写为 GatherV2 加 Const 轴输入)、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.cpp 为 DynamicRnnV2/DynamicRnnv2WithoutSeqlength 设置 is_misplaced 属性;ascend_weightquant_tf_plugin.cpp 将字符串属性 dst_type(如 "INT4")转换为对应整型数据类型后写回算子。
likedislike
不准确?
atomgit-bot
atomgit-bot
22 天前 评论:

代码审查

✅ 未发现问题

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

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成员
21 天前 评论:

/lgtm
/approve

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

Pull Request 已合并或已关闭。

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

likedislike