Pull Request已成功合入, 合并人@CANN-robot
(感谢 sunchun 的贡献)变更摘要
本次 PR 主要为 TensorFlow 框架注册一批自定义算子插件,通过 REGISTER_CUSTOM_OP 宏将这些 TF 原生算子映射到 Ascend 后端。所有新增插件均采用 AutoMappingByOpFn 自动映射方式,集中在 index 目录下的各个算子子目录中。此外,对已有的 ScatterMin 插件进行了扩展,新增对 ResourceScatterMin 变体的支持。
主要改动
-
新增 25 个 TF 算子插件文件: 在
index/下为DataFormatDimMap、GatherNd、GatherV2、InplaceUpdate、ReverseSequence、ReverseV2、ScatterAdd、ScatterElements、ScatterNd、ScatterNdAdd、ScatterNdMax、ScatterNdMin、ScatterNdSub、ScatterNdUpdate、ScatterNonAliasingAdd、ScatterSub、ScatterUpdate、SegmentSum、SparseSegmentSumGrad、TensorScatterUpdate、UnsortedSegmentMax/Min/Prod/Sum、Where等算子新增_tf_plugin.cpp文件,每个文件注册一个对应的 TF 自定义算子,ImplyType均为CUSTOM(ScatterElements和LookupTableFind除外,为AI_CPU)。 -
新增 Embedding 系列 8 个 TF 算子插件: 在
index/embedding/framework/embedding_tf_plugin.cpp中批量注册了LookupTableFind(AI_CPU)、FakeRemoteLookupUniqued、InitEmbeddingHashmapV2、DeinitEmbeddingHashmapV2、TableToResourceV2、EmbeddingHashmapSize、EmbeddingHashmapFileSize、EmbeddingHashTableApplyAdamW共 8 个算子,均采用AutoMappingByOpFn自动映射。 -
ScatterMin插件扩展OriginOpType: 在已有的index/scatter_min/framework/scatter_min_tf_plugin.cpp中,将OriginOpType从单个字符串"ScatterMin"改为初始值列表{"ScatterMin", "ResourceScatterMin"},使其同时匹配ResourceScatterMin变体算子。 -
ScatterNonAliasingAdd使用非同名映射: 新增的scatter_non_aliasing_add_tf_plugin.cpp中,注册名为ScatterNonAliasingAdd的算子,但其OriginOpType设置为"ScatterNdNonAliasingAdd",与其他直接同名映射的算子不同。


代码审查
All 27 files have been read and reviewed. Here is my closing summary.
审查总结
发现统计
| 优先级 | 数量 |
|---|---|
| P0 | 0 |
| P1 | 0 |
| P2 | 0 |
| P3 | 2 |
总体风险评估:低风险
本次变更主要是为 26 个 TensorFlow 算子新增 TF plugin 注册文件,以及修改 1 个现有文件(scatter_min_tf_plugin.cpp)以增加 ResourceScatterMin 作为额外的 OriginOpType。所有新增文件均遵循相同的模式(REGISTER_CUSTOM_OP + AutoMappingByOpFn),代码形式一致、无编译期或运行时明显缺陷。发现 2 个 P3 级别的注释不准确问题。
各文件审查结论
| 文件 | 状态 |
|---|---|
| index/data_format_dim_map/framework/data_format_dim_map_tf_plugin.cpp | 无问题 |
| index/embedding/framework/embedding_tf_plugin.cpp | P3:LookupTableFind 注释与实际映射不一致 |
| index/gather_nd/framework/gather_nd_tf_plugin.cpp | 无问题 |
| index/gather_v2/framework/gather_v2_tf_plugin.cpp | 无问题 |
| index/inplace_update/framework/inplace_update_tf_plugin.cpp | 无问题 |
| index/reverse_sequence/framework/reverse_sequence_tf_plugin.cpp | 无问题 |
| index/reverse_v2/framework/reverse_v2_tf_plugin.cpp | 无问题 |
| index/scatter_add/framework/scatter_add_tf_plugin.cpp | 无问题 |
| index/scatter_elements/framework/scatter_elements_tf_plugin.cpp | 无问题 |
| index/scatter_min/framework/scatter_min_tf_plugin.cpp | 无问题 |
| index/scatter_nd/framework/scatter_nd_tf_plugin.cpp | 无问题 |
| index/scatter_nd_add/framework/scatter_nd_add_tf_plugin.cpp | 无问题 |
| index/scatter_nd_max/framework/scatter_nd_max_tf_plugin.cpp | 无问题 |
| index/scatter_nd_min/framework/scatter_nd_min_tf_plugin.cpp | 无问题 |
| index/scatter_nd_sub/framework/scatter_nd_sub_tf_plugin.cpp | 无问题 |
| index/scatter_nd_update/framework/scatter_nd_update_tf_plugin.cpp | 无问题 |
| index/scatter_non_aliasing_add/framework/scatter_non_aliasing_add_tf_plugin.cpp | P3:注释声称同名映射,实际 OriginOpType 不同 |
| index/scatter_sub/framework/scatter_sub_tf_plugin.cpp | 无问题 |
| index/scatter_update/framework/scatter_update_tf_plugin.cpp | 无问题 |
| index/segment_sum/framework/segment_sum_tf_plugin.cpp | 无问题 |
| index/sparse_segment_sum_grad/framework/sparse_segment_sum_grad_tf_plugin.cpp | 无问题 |
| index/tensor_scatter_update/framework/tensor_scatter_update_tf_plugin.cpp | 无问题 |
| index/unsorted_segment_max/framework/unsorted_segment_max_tf_plugin.cpp | 无问题 |
| index/unsorted_segment_min/framework/unsorted_segment_min_tf_plugin.cpp | 无问题 |
| index/unsorted_segment_prod/framework/unsorted_segment_prod_tf_plugin.cpp | 无问题 |
| index/unsorted_segment_sum/framework/unsorted_segment_sum_tf_plugin.cpp | 无问题 |
| index/where/framework/where_tf_plugin.cpp | 无问题 |
⚠️ 已识别出整体风险,但无法提取行内评论,请参考整体评估。


The MR can not be merged, because of CodeReview discussion not resolved
If you want to solve this problem, you can click here to do it in the FAQs.


Pull Request 已合并或已关闭。
If you want to solve this problem, you can click here to do it in the FAQs.


描述
TF plugin修改
关联的Issue
https://gitcode.com/cann/ops-nn/issues/4196
测试
文档更新
类型标签
AI/Agent生成声明