已合并
refactor: 迁移 TF 插件注册至新框架 API(part 3) #9758
refactor: 迁移 TF 插件注册至新框架 API(part 3) #9758
已合并
Nice try创建于 24 天前
共 22 个文件变更+613-8
@@ -0,0 +1,44 @@
1+/**
2+ * Copyright (c) 2026 Huawei Technologies Co., Ltd.
3+ * This program is free software, you can redistribute it and/or modify it under the terms and conditions of
4+ * CANN Open Software License Agreement Version 2.0 (the "License").
5+ * Please refer to the License for details. You may not use this file except in compliance with the License.
6+ * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED,
7+ * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE.
8+ * See LICENSE in the root of the software repository for the full text of the License.
9+ */
10+ 
11+/*!
12+ * \file ascend_weightquant_tf_plugin.cpp
13+ * \brief
14+ */
15+#include <string>
16+ 
17+#include "graph/types.h"
18+#include "log/log.h"
19+#include "register/register.h"
20+ 
21+namespace domi {
22+static Status ParseParamsAscendWeightQuant(const ge::Operator& op_src, ge::Operator& op)
23+{
24+ AutoMappingByOpFn(op_src, op);
25+ 
26+ std::string dst_type_str;
27+ if (op.GetAttr("dst_type", dst_type_str) == ge::GRAPH_SUCCESS) {
28+ int dst_type = ge::DT_INT8;
29+ if (dst_type_str == "INT4") {
30+ dst_type = ge::DT_INT4;
31+ }
32+ op.SetAttr("dst_type", dst_type);
33+ }
34+ 
35+ OP_LOGI("AscendWeightQuant", "op[AscendWeightQuant] tensowflow plugin parser [AutoMapping] success.");
36+ return SUCCESS;
37+}
38+ 
39+REGISTER_CUSTOM_OP("AscendWeightQuant")
40+ .FrameworkType(TENSORFLOW)
41+ .OriginOpType("AscendWeightQuant")
42+ .ParseParamsByOperatorFn(ParseParamsAscendWeightQuant)
43+ .ImplyType(ImplyType::TVM);
44+} // namespace domi
@@ -0,0 +1,73 @@
1+/**
2+ * Copyright (c) 2026 Huawei Technologies Co., Ltd.
3+ * This program is free software, you can redistribute it and/or modify it under the terms and conditions of
4+ * CANN Open Software License Agreement Version 2.0 (the "License").
5+ * Please refer to the License for details. You may not use this file except in compliance with the License.
6+ * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED,
7+ * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE.
8+ * See LICENSE in the root of the software repository for the full text of the License.
9+ */
10+ 
11+/*!
12+ * \file gather_point_grad_tf_plugin.cpp
13+ * \brief
14+ */
15+#include <string>
16+#include <vector>
17+#include <map>
18+#include "register/register.h"
19+#include "framework/plugin_util.h"
20+#include "graph/operator.h"
21+#include "stub_ops.h"
22+#include "index/scatter_update/op_graph/scatter_update_proto.h"
23+#include "log/log.h"
24+ 
25+namespace domi {
26+using namespace ge;
27+ 
28+static Status ParseParamsGatherPointGrad(const ge::Operator& op_src, ge::Operator& op_dest)
29+{
30+ int n = static_cast<int>(op_dest.GetInputsSize());
31+ OP_LOGI(GetOpName(op_dest).c_str(), "ParseParamsGatherPointGrad input_size = %d", n);
32+ // 2.set original_type
33+ op_dest.SetAttr("original_type", std::string("GatherPointGrad"));
34+ // 3.set attr if needed
35+ op_dest.SetAttr("name", GetOpName(op_dest));
36+ 
37+ return SUCCESS;
38+}
39+ 
40+static Status ParseOpToGraphGatherPointGrad(const ge::Operator& op, ge::Graph& graph)
41+{
42+ std::string ori_name;
43+ if (op.GetAttr("name", ori_name) != SUCCESS) {
44+ OP_LOGE(GetOpName(op).c_str(), "get name from op failed");
45+ return FAILED;
46+ }
47+ 
48+ ge::Operator data_0 = op::Data("inp").set_attr_index(0);
49+ ge::Operator data_1 = op::Data("idx").set_attr_index(1);
50+ ge::Operator data_2 = op::Data("out_g").set_attr_index(2);
51+ auto use_locking = false;
52+ 
53+ auto ScatterUpdate = op::ScatterUpdate(ori_name.c_str())
54+ .set_input_var(data_0)
55+ .set_input_indices(data_1)
56+ .set_input_updates(data_2)
57+ .set_attr_use_locking(use_locking);
58+ std::vector<ge::Operator> inputs{data_0, data_1, data_2};
59+ std::vector<std::pair<ge::Operator, std::vector<size_t>>> output_indexs;
60+ output_indexs.emplace_back(ScatterUpdate, vector<std::size_t>{0});
61+ graph.SetInputs(inputs).SetOutputs(output_indexs);
62+ 
63+ return SUCCESS;
64+}
65+ 
66+// register GatherPointGrad op info to GE
67+REGISTER_CUSTOM_OP("PartitionedCall")
68+ .FrameworkType(TENSORFLOW)
69+ .OriginOpType("GatherPointGrad")
70+ .ParseParamsByOperatorFn(ParseParamsGatherPointGrad)
71+ .ParseOpToGraphFn(ParseOpToGraphGatherPointGrad)
72+ .ImplyType(ImplyType::TVM);
73+} // namespace domi
@@ -0,0 +1,76 @@
1+/**
2+ * Copyright (c) 2026 Huawei Technologies Co., Ltd.
3+ * This program is free software, you can redistribute it and/or modify it under the terms and conditions of
4+ * CANN Open Software License Agreement Version 2.0 (the "License").
5+ * Please refer to the License for details. You may not use this file except in compliance with the License.
6+ * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED,
7+ * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE.
8+ * See LICENSE in the root of the software repository for the full text of the License.
9+ */
10+ 
11+/*!
12+ * \file gather_point_tf_plugin.cpp
13+ * \brief
14+ */
15+#include <string>
16+#include <vector>
17+#include <map>
18+#include "register/register.h"
19+#include "framework/plugin_util.h"
20+#include "graph/operator.h"
21+#include "stub_ops.h"
22+#include "index/gather_v2/op_graph/gather_v2_proto.h"
23+#include "log/log.h"
24+ 
25+namespace domi {
26+using namespace ge;
27+ 
28+static Status ParseParamsGatherPoint(const ge::Operator& op_src, ge::Operator& op_dest)
29+{
30+ int n = static_cast<int>(op_dest.GetInputsSize());
31+ OP_LOGI(GetOpName(op_dest).c_str(), "ParseParamsGatherPoint input_size = %d", n);
32+ // 2.set original_type
33+ op_dest.SetAttr("original_type", std::string("GatherPoint"));
34+ // 3.set attr if needed
35+ op_dest.SetAttr("name", GetOpName(op_dest));
36+ 
37+ return SUCCESS;
38+}
39+ 
40+static Status ParseOpToGraphGatherPoint(const ge::Operator& op, ge::Graph& graph)
41+{
42+ std::string ori_name;
43+ if (op.GetAttr("name", ori_name) != SUCCESS) {
44+ OP_LOGE(GetOpName(op).c_str(), "get name from op failed");
45+ return FAILED;
46+ }
47+ 
48+ ge::Operator data_0 = op::Data("inp").set_attr_index(0);
49+ ge::Operator data_1 = op::Data("idx").set_attr_index(1);
50+ int32_t axis_val = 1;
51+ 
52+ TensorDesc tensor1_desc(ge::Shape(), FORMAT_ND, DT_INT32);
53+ ge::Tensor const_value(tensor1_desc, (uint8_t*)&axis_val, sizeof(axis_val));
54+ auto const_op = op::Const("const_data").set_attr_value(const_value);
55+ int batch_dims = 1;
56+ auto GatherV2 = op::GatherV2(ori_name.c_str())
57+ .set_input_x(data_0)
58+ .set_input_indices(data_1)
59+ .set_input_axis(const_op)
60+ .set_attr_batch_dims(batch_dims);
61+ std::vector<ge::Operator> inputs{data_0, data_1, const_op};
62+ std::vector<std::pair<ge::Operator, std::vector<size_t>>> output_indexs;
63+ output_indexs.emplace_back(GatherV2, vector<std::size_t>{0});
64+ graph.SetInputs(inputs).SetOutputs(output_indexs);
65+ 
66+ return SUCCESS;
67+}
68+ 
69+// register GatherPoint op info to GE
70+REGISTER_CUSTOM_OP("PartitionedCall")
71+ .FrameworkType(TENSORFLOW)
72+ .OriginOpType("GatherPoint")
73+ .ParseParamsByOperatorFn(ParseParamsGatherPoint)
74+ .ParseOpToGraphFn(ParseOpToGraphGatherPoint)
75+ .ImplyType(ImplyType::TVM);
76+} // namespace domi
@@ -0,0 +1,23 @@
1+/**
2+ * Copyright (c) 2026 Huawei Technologies Co., Ltd.
3+ * This program is free software, you can redistribute it and/or modify it under the terms and conditions of
4+ * CANN Open Software License Agreement Version 2.0 (the "License").
5+ * Please refer to the License for details. You may not use this file except in compliance with the License.
6+ * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED,
7+ * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE.
8+ * See LICENSE in the root of the software repository for the full text of the License.
9+ */
10+ 
11+/*!
12+ * \file broadcast_gradient_args_tf_plugin.cpp
13+ * \brief BroadcastGradientArgs TensorFlow plugin mapping.
14+ */
15+#include "register/register.h"
16+ 
17+namespace domi {
18+REGISTER_CUSTOM_OP("BroadcastGradientArgs")
19+ .FrameworkType(TENSORFLOW)
20+ .OriginOpType("BroadcastGradientArgs")
21+ .ParseParamsByOperatorFn(AutoMappingByOpFn)
22+ .ImplyType(ImplyType::TVM);
23+} // namespace domi
@@ -21,4 +21,10 @@ REGISTER_CUSTOM_OP("GatherV2")
21 .OriginOpType("GatherV2")21 .OriginOpType("GatherV2")
22 .ParseParamsByOperatorFn(AutoMappingByOpFn)22 .ParseParamsByOperatorFn(AutoMappingByOpFn)
23 .ImplyType(ImplyType::TVM);23 .ImplyType(ImplyType::TVM);
24+ 
25+REGISTER_CUSTOM_OP("Gather")
26+ .FrameworkType(TENSORFLOW)
27+ .OriginOpType(std::vector<ge::AscendString>{"Gather", "ResourceGather"})
28+ .ParseParamsByOperatorFn(AutoMappingByOpFn)
29+ .ImplyType(ImplyType::TVM);
24} // namespace domi30} // namespace domi
@@ -21,4 +21,10 @@ REGISTER_CUSTOM_OP("ScatterAdd")
21 .OriginOpType("ScatterAdd")21 .OriginOpType("ScatterAdd")
22 .ParseParamsByOperatorFn(AutoMappingByOpFn)22 .ParseParamsByOperatorFn(AutoMappingByOpFn)
23 .ImplyType(ImplyType::TVM);23 .ImplyType(ImplyType::TVM);
24+ 
25+REGISTER_CUSTOM_OP("TensorScatterAdd")
26+ .FrameworkType(TENSORFLOW)
27+ .OriginOpType("TensorScatterAdd")
28+ .ParseParamsByOperatorFn(AutoMappingByOpFn)
29+ .ImplyType(ImplyType::TVM);
24} // namespace domi30} // namespace domi
@@ -21,4 +21,10 @@ REGISTER_CUSTOM_OP("ScatterSub")
21 .OriginOpType("ScatterSub")21 .OriginOpType("ScatterSub")
22 .ParseParamsByOperatorFn(AutoMappingByOpFn)22 .ParseParamsByOperatorFn(AutoMappingByOpFn)
23 .ImplyType(ImplyType::TVM);23 .ImplyType(ImplyType::TVM);
24+ 
25+REGISTER_CUSTOM_OP("TensorScatterSub")
26+ .FrameworkType(TENSORFLOW)
27+ .OriginOpType("TensorScatterSub")
28+ .ParseParamsByOperatorFn(AutoMappingByOpFn)
29+ .ImplyType(ImplyType::TVM);
24} // namespace domi30} // namespace domi
@@ -0,0 +1,23 @@
1+/**
2+ * Copyright (c) 2026 Huawei Technologies Co., Ltd.
3+ * This program is free software, you can redistribute it and/or modify it under the terms and conditions of
4+ * CANN Open Software License Agreement Version 2.0 (the "License").
5+ * Please refer to the License for details. You may not use this file except in compliance with the License.
6+ * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED,
7+ * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE.
8+ * See LICENSE in the root of the software repository for the full text of the License.
9+ */
10+ 
11+/*!
12+ * \file sparse_segment_mean_tf_plugin.cpp
13+ * \brief SparseSegmentMean TensorFlow plugin mapping.
14+ */
15+#include "register/register.h"
16+ 
17+namespace domi {
18+REGISTER_CUSTOM_OP("SparseSegmentMean")
19+ .FrameworkType(TENSORFLOW)
20+ .OriginOpType("SparseSegmentMean")
21+ .ParseParamsByOperatorFn(AutoMappingByOpFn)
22+ .ImplyType(ImplyType::TVM);
23+} // namespace domi
@@ -0,0 +1,23 @@
1+/**
2+ * Copyright (c) 2026 Huawei Technologies Co., Ltd.
3+ * This program is free software, you can redistribute it and/or modify it under the terms and conditions of
4+ * CANN Open Software License Agreement Version 2.0 (the "License").
5+ * Please refer to the License for details. You may not use this file except in compliance with the License.
6+ * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED,
7+ * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE.
8+ * See LICENSE in the root of the software repository for the full text of the License.
9+ */
10+ 
11+/*!
12+ * \file sparse_segment_mean_grad_tf_plugin.cpp
13+ * \brief SparseSegmentMeanGrad TensorFlow plugin mapping.
14+ */
15+#include "register/register.h"
16+ 
17+namespace domi {
18+REGISTER_CUSTOM_OP("SparseSegmentMeanGrad")
19+ .FrameworkType(TENSORFLOW)
20+ .OriginOpType("SparseSegmentMeanGrad")
21+ .ParseParamsByOperatorFn(AutoMappingByOpFn)
22+ .ImplyType(ImplyType::TVM);
23+} // namespace domi
@@ -0,0 +1,23 @@
1+/**
2+ * Copyright (c) 2026 Huawei Technologies Co., Ltd.
3+ * This program is free software, you can redistribute it and/or modify it under the terms and conditions of
4+ * CANN Open Software License Agreement Version 2.0 (the "License").
5+ * Please refer to the License for details. You may not use this file except in compliance with the License.
6+ * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED,
7+ * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE.
8+ * See LICENSE in the root of the software repository for the full text of the License.
9+ */
10+ 
11+/*!
12+ * \file sparse_slice_tf_plugin.cpp
13+ * \brief SparseSlice TensorFlow plugin mapping.
14+ */
15+#include "register/register.h"
16+ 
17+namespace domi {
18+REGISTER_CUSTOM_OP("SparseSlice")
19+ .FrameworkType(TENSORFLOW)
20+ .OriginOpType("SparseSlice")
21+ .ParseParamsByOperatorFn(AutoMappingByOpFn)
22+ .ImplyType(ImplyType::TVM);
23+} // namespace domi
@@ -0,0 +1,23 @@
1+/**
2+ * Copyright (c) 2026 Huawei Technologies Co., Ltd.
3+ * This program is free software, you can redistribute it and/or modify it under the terms and conditions of
4+ * CANN Open Software License Agreement Version 2.0 (the "License").
5+ * Please refer to the License for details. You may not use this file except in compliance with the License.
6+ * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED,
7+ * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE.
8+ * See LICENSE in the root of the software repository for the full text of the License.
9+ */
10+ 
11+/*!
12+ * \file sparse_to_dense_tf_plugin.cpp
13+ * \brief SparseToDense TensorFlow plugin mapping.
14+ */
15+#include "register/register.h"
16+ 
17+namespace domi {
18+REGISTER_CUSTOM_OP("SparseToDense")
19+ .FrameworkType(TENSORFLOW)
20+ .OriginOpType("SparseToDense")
21+ .ParseParamsByOperatorFn(AutoMappingByOpFn)
22+ .ImplyType(ImplyType::TVM);
23+} // namespace domi
@@ -0,0 +1,29 @@
1+/**
2+ * Copyright (c) 2026 Huawei Technologies Co., Ltd.
3+ * This program is free software, you can redistribute it and/or modify it under the terms and conditions of
4+ * CANN Open Software License Agreement Version 2.0 (the "License").
5+ * Please refer to the License for details. You may not use this file except in compliance with the License.
6+ * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED,
7+ * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE.
8+ * See LICENSE in the root of the software repository for the full text of the License.
9+ */
10+ 
11+/*!
12+ * \file einsum_tf_plugin.cpp
13+ * \brief
14+ */
15+#include "register/register.h"
16+ 
17+namespace domi {
18+static Status AutoMappingFnEinSum(const ge::Operator& op_src, ge::Operator& op)
19+{
20+ AutoMappingByOpFn(op_src, op);
21+ return SUCCESS;
22+}
23+ 
24+REGISTER_CUSTOM_OP("Einsum")
25+ .FrameworkType(TENSORFLOW)
26+ .OriginOpType("Einsum")
27+ .ParseParamsByOperatorFn(AutoMappingFnEinSum)
28+ .ImplyType(ImplyType::TVM);
29+} // namespace domi
@@ -19,6 +19,6 @@ namespace domi {
19REGISTER_CUSTOM_OP("BNInfer")19REGISTER_CUSTOM_OP("BNInfer")
20 .FrameworkType(TENSORFLOW)20 .FrameworkType(TENSORFLOW)
21 .OriginOpType("BNInfer")21 .OriginOpType("BNInfer")
22- .ParseParamsFn(AutoMappingFn)22+ .ParseParamsByOperatorFn(AutoMappingByOpFn)
23 .ImplyType(ImplyType::TVM);23 .ImplyType(ImplyType::TVM);
24} // namespace domi24} // namespace domi
@@ -19,6 +19,6 @@ namespace domi {
19REGISTER_CUSTOM_OP("BNTrainingReduceGrad")19REGISTER_CUSTOM_OP("BNTrainingReduceGrad")
20 .FrameworkType(TENSORFLOW)20 .FrameworkType(TENSORFLOW)
21 .OriginOpType("BNTrainingReduceGrad")21 .OriginOpType("BNTrainingReduceGrad")
22- .ParseParamsFn(AutoMappingFn)22+ .ParseParamsByOperatorFn(AutoMappingByOpFn)
23 .ImplyType(ImplyType::TVM);23 .ImplyType(ImplyType::TVM);
24} // namespace domi24} // namespace domi
@@ -20,6 +20,6 @@ namespace domi {
20REGISTER_CUSTOM_OP("BNTrainingUpdateGrad")20REGISTER_CUSTOM_OP("BNTrainingUpdateGrad")
21 .FrameworkType(TENSORFLOW)21 .FrameworkType(TENSORFLOW)
22 .OriginOpType("BNTrainingUpdateGrad")22 .OriginOpType("BNTrainingUpdateGrad")
23- .ParseParamsFn(AutoMappingFn)23+ .ParseParamsByOperatorFn(AutoMappingByOpFn)
24 .ImplyType(ImplyType::TVM);24 .ImplyType(ImplyType::TVM);
25} // namespace domi25} // namespace domi
@@ -20,6 +20,12 @@ namespace domi {
20REGISTER_CUSTOM_OP("BNTrainingUpdateV2")20REGISTER_CUSTOM_OP("BNTrainingUpdateV2")
21 .FrameworkType(TENSORFLOW)21 .FrameworkType(TENSORFLOW)
22 .OriginOpType("BNTrainingUpdateV2")22 .OriginOpType("BNTrainingUpdateV2")
23- .ParseParamsFn(AutoMappingFn)23+ .ParseParamsByOperatorFn(AutoMappingByOpFn)
24+ .ImplyType(ImplyType::TVM);
25+ 
26+REGISTER_CUSTOM_OP("BNTrainingUpdate")
27+ .FrameworkType(TENSORFLOW)
28+ .OriginOpType("BNTrainingUpdate")
29+ .ParseParamsByOperatorFn(AutoMappingByOpFn)
24 .ImplyType(ImplyType::TVM);30 .ImplyType(ImplyType::TVM);
25} // namespace domi31} // namespace domi
@@ -19,7 +19,7 @@ namespace domi {
19// The attributes `begin_norm_axis`, `begin_params_axis`, `epsilon` share the same name in TF and GE.19// The attributes `begin_norm_axis`, `begin_params_axis`, `epsilon` share the same name in TF and GE.
20REGISTER_CUSTOM_OP("LayerNorm")20REGISTER_CUSTOM_OP("LayerNorm")
21 .FrameworkType(TENSORFLOW)21 .FrameworkType(TENSORFLOW)
22- .OriginOpType("LayerNorm")22+ .OriginOpType(std::vector<ge::AscendString>{"LayerNorm", "FusedLayerNorm"})
23 .ParseParamsByOperatorFn(AutoMappingByOpFn)23 .ParseParamsByOperatorFn(AutoMappingByOpFn)
24 .ImplyType(ImplyType::TVM);24 .ImplyType(ImplyType::TVM);
25} // namespace domi25} // namespace domi
@@ -18,7 +18,7 @@ namespace domi {
18// TensorFlow LayerNormGrad maps directly from the TensorFlow op of the same name; auto operator mapping suffices.18// TensorFlow LayerNormGrad maps directly from the TensorFlow op of the same name; auto operator mapping suffices.
19REGISTER_CUSTOM_OP("LayerNormGrad")19REGISTER_CUSTOM_OP("LayerNormGrad")
20 .FrameworkType(TENSORFLOW)20 .FrameworkType(TENSORFLOW)
21- .OriginOpType("LayerNormGrad")21+ .OriginOpType(std::vector<ge::AscendString>{"LayerNormGrad", "FusedLayerNormGrad"})
22 .ParseParamsByOperatorFn(AutoMappingByOpFn)22 .ParseParamsByOperatorFn(AutoMappingByOpFn)
23 .ImplyType(ImplyType::TVM);23 .ImplyType(ImplyType::TVM);
24} // namespace domi24} // namespace domi
@@ -17,7 +17,7 @@
17namespace domi {17namespace domi {
18REGISTER_CUSTOM_OP("ApplyAdagrad")18REGISTER_CUSTOM_OP("ApplyAdagrad")
19 .FrameworkType(TENSORFLOW)19 .FrameworkType(TENSORFLOW)
20- .OriginOpType({"ApplyAdagrad", "ResourceApplyAdagrad"})20+ .OriginOpType(std::vector<ge::AscendString>{"ApplyAdagrad", "ResourceApplyAdagrad"})
21- .ParseParamsFn(AutoMappingFn)21+ .ParseParamsByOperatorFn(AutoMappingByOpFn)
22 .ImplyType(ImplyType::TVM);22 .ImplyType(ImplyType::TVM);
23} // namespace domi23} // namespace domi
@@ -0,0 +1,106 @@
1+/**
2+ * Copyright (c) 2026 Huawei Technologies Co., Ltd.
3+ * This program is free software, you can redistribute it and/or modify it under the terms and conditions of
4+ * CANN Open Software License Agreement Version 2.0 (the "License").
5+ * Please refer to the License for details. You may not use this file except in compliance with the License.
6+ * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED,
7+ * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE.
8+ * See LICENSE in the root of the software repository for the full text of the License.
9+ */
10+ 
11+/*!
12+ * \file block_lstm_tf_plugin.cpp
13+ * \brief
14+ */
15+#include <string>
16+#include <vector>
17+#include <map>
18+#include "register/register.h"
19+#include "framework/plugin_util.h"
20+#include "graph/operator.h"
21+#include "stub_ops.h"
22+#include "rnn/dynamic_rnn/op_graph/dynamic_rnn_proto.h"
23+ 
24+#include "log/log.h"
25+#include "graph/types.h"
26+ 
27+namespace domi {
28+using namespace ge;
29+ 
30+static Status ParseParamsRNN(const ge::Operator& op_src, ge::Operator& op)
31+{
32+ // Set original_type
33+ op.SetAttr("original_type", "BlockLSTM");
34+ 
35+ // 未调用 AutoMappingByOpFn,TF 源节点属性不会自动映射到目标算子;
36+ // 因此从 op_src(携带 TF 源节点属性)显式解析后写入目标算子,缺失时保持默认值。
37+ float forget_bias = 0.0f;
38+ (void)op_src.GetAttr("forget_bias", forget_bias);
39+ (void)op.SetAttr("forget_bias", forget_bias);
40+ 
41+ float cell_clip = 3.0f;
42+ (void)op_src.GetAttr("cell_clip", cell_clip);
43+ (void)op.SetAttr("cell_clip", cell_clip);
44+ 
45+ bool use_peephole = false;
46+ (void)op_src.GetAttr("use_peephole", use_peephole);
47+ (void)op.SetAttr("use_peephole", use_peephole);
48+ 
49+ return SUCCESS;
50+}
51+ 
52+static Status ParseOpToGraphRNN(const ge::Operator& op, ge::Graph& graph)
53+{
54+ ge::Operator data_0 = op::Data("seq_len_max").set_attr_index(0);
55+ ge::Operator data_1 = op::Data("x").set_attr_index(1);
56+ ge::Operator data_2 = op::Data("cs_prev").set_attr_index(2);
57+ ge::Operator data_3 = op::Data("h_prev").set_attr_index(3);
58+ ge::Operator data_4 = op::Data("w").set_attr_index(4);
59+ ge::Operator data_8 = op::Data("b").set_attr_index(8);
60+ 
61+ float forget_bias = 0.0;
62+ if (op.GetAttr("forget_bias", forget_bias) != ge::GRAPH_SUCCESS) {
63+ OP_LOGE(GetOpName(op).c_str(), "get attr forget_bias failed.");
64+ return FAILED;
65+ }
66+ 
67+ float cell_clip = 3.0;
68+ if (op.GetAttr("cell_clip", cell_clip) != ge::GRAPH_SUCCESS) {
69+ OP_LOGE(GetOpName(op).c_str(), "get attr cell_clip failed.");
70+ return FAILED;
71+ }
72+ 
73+ bool use_peephole = false;
74+ if (op.GetAttr("use_peephole", use_peephole) != ge::GRAPH_SUCCESS) {
75+ OP_LOGE(GetOpName(op).c_str(), "get attr use_peephole failed.");
76+ return FAILED;
77+ }
78+ 
79+ auto cast = op::Cast().set_input_x(data_0).set_attr_dst_type(3);
80+ auto rnn = op::DynamicRNN()
81+ .set_input_x(data_1)
82+ .set_input_w(data_4)
83+ .set_input_b(data_8)
84+ .set_input_seq_length(cast)
85+ .set_input_init_h(data_3)
86+ .set_input_init_c(data_2)
87+ .set_attr_forget_bias(forget_bias)
88+ .set_attr_cell_clip(cell_clip)
89+ .set_attr_use_peephole(use_peephole)
90+ .set_attr_cell_type("BLOCKLSTM");
91+ 
92+ std::vector<ge::Operator> inputs{data_1, data_4, data_8, data_0, data_3, data_2};
93+ std::vector<std::pair<ge::Operator, std::vector<size_t>>> output_indexs;
94+ output_indexs.emplace_back(rnn, vector<std::size_t>{3, 2, 5, 6, 4, 7, 1});
95+ graph.SetInputs(inputs).SetOutputs(output_indexs);
96+ return SUCCESS;
97+}
98+ 
99+// register BlockLSTM op info to GE
100+REGISTER_CUSTOM_OP("DynamicRNN")
101+ .FrameworkType(TENSORFLOW)
102+ .OriginOpType("BlockLSTM")
103+ .ParseParamsByOperatorFn(ParseParamsRNN)
104+ .ParseOpToGraphFn(ParseOpToGraphRNN)
105+ .ImplyType(ImplyType::TVM);
106+} // namespace domi
@@ -0,0 +1,96 @@
1+/**
2+ * Copyright (c) 2026 Huawei Technologies Co., Ltd.
3+ * This program is free software, you can redistribute it and/or modify it under the terms and conditions of
4+ * CANN Open Software License Agreement Version 2.0 (the "License").
5+ * Please refer to the License for details. You may not use this file except in compliance with the License.
6+ * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED,
7+ * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE.
8+ * See LICENSE in the root of the software repository for the full text of the License.
9+ */
10+ 
11+/*!
12+ * \file dynamic_rnn_tf_plugin.cpp
13+ * \brief
14+ */
15+ 
16+#include <cstring>
17+#include <string>
18+#include <vector>
19+ 
20+#include "register/register.h"
21+#include "framework/plugin_util.h"
22+#include "graph/operator.h"
23+#include "log/log.h"
24+ 
25+namespace domi {
26+uint32_t wRnnInputPosition = 1;
27+static const char* const kForgetBias = "lstm_cell/add/y";
28+static const char* const kTransposeNode = "Transpose";
29+ 
30+static Status DynamicRNNParserParams(const std::vector<ge::Operator>& inside_nodes, ge::Operator& op)
31+{
32+ OP_LOGI(GetOpName(op).c_str(), "Enter DynamicRNN fusion parser.");
33+ 
34+ // 基于融合子图算子恢复旧 NodeDef 解析语义(node_def.pb.h 在新工具链已不存在,改用算子版融合解析):
35+ // - 子图含 Transpose 节点 → time_major=false
36+ // - 名为 lstm_cell/add/y 的 Const 节点 → 从 value 张量解析 forget_bias
37+ bool time_major = true;
38+ float forget_bias = 0.0f;
陈
陈陈展熹24 天前

这里直接固定 time_major=true、forget_bias=0,导致 inside_nodes 中的模型信息被丢弃。旧实现会在融合子图包含 Transpose 时将 time_major 置为 false,并从 lstm_cell/add/y 常量节点解析 forget_bias;当前写法会让所有非默认模型按默认属性执行,输出语义可能变化。迁移需要保留这两个属性的解析,并补充非默认值用例。

likedislike
Nice try
Nice try
24 天前 评论:
39+ for (const auto& node : inside_nodes) {
40+ ge::AscendString node_type;
41+ if (node.GetOpType(node_type) == ge::GRAPH_SUCCESS && std::string(node_type.GetString()) == kTransposeNode) {
42+ time_major = false;
43+ }
44+ 
45+ ge::AscendString node_name;
46+ if (node.GetName(node_name) != ge::GRAPH_SUCCESS) {
47+ continue;
48+ }
49+ if (std::string(node_name.GetString()).find(kForgetBias) == std::string::npos) {
50+ continue;
51+ }
52+ ge::Tensor const_value;
53+ if (node.GetAttr("value", const_value) != ge::GRAPH_SUCCESS) {
54+ OP_LOGE(GetOpName(op).c_str(), "parse forget_bias from const node %s failed", node_name.GetString());
55+ return PARAM_INVALID;
56+ }
57+ const uint8_t* value_data = const_value.GetData();
58+ const size_t value_size = const_value.GetSize();
59+ if (value_data == nullptr || value_size < sizeof(float)) {
60+ OP_LOGE(GetOpName(op).c_str(), "parse forget_bias from const node %s failed", node_name.GetString());
61+ return PARAM_INVALID;
62+ }
63+ (void)memcpy(&forget_bias, value_data, sizeof(float));
64+ }
65+ op.SetAttr("time_major", time_major);
66+ op.SetAttr("forget_bias", forget_bias);
67+ OP_LOGD(GetOpName(op).c_str(), "parser stage set DynamicRNN's attr time_major is %s forget_bias is %.1f",
68+ time_major ? "true" : "false", forget_bias);
69+ 
70+ ge::TensorDesc input_desc = op.GetInputDesc(wRnnInputPosition);
71+ input_desc.SetOriginFormat(ge::FORMAT_HWCN);
72+ input_desc.SetFormat(ge::FORMAT_HWCN);
73+ 
74+ if (op.UpdateInputDesc(wRnnInputPosition, input_desc) != ge::GRAPH_SUCCESS) {
75+ OP_LOGE(GetOpName(op).c_str(), "Update input desc fail, index:%u.", wRnnInputPosition);
76+ return FAILED;
77+ }
78+ 
79+ return SUCCESS;
80+}
81+ 
82+// 来自 dynamic_rnn_plugin.cc:TF 原生 DynamicRNN
83+REGISTER_CUSTOM_OP("DynamicRNN")
84+ .FrameworkType(TENSORFLOW)
85+ .OriginOpType("DynamicRNN")
86+ .ParseParamsByOperatorFn(AutoMappingByOpFn)
87+ .FusionParseParamsFn(DynamicRNNParserParams)
88+ .ImplyType(ImplyType::TVM);
89+ 
90+// 来自 dynamic_rnn_tf_plugin.cc:DynamicRnn(保持原独立注册语义,不附加融合解析)
91+REGISTER_CUSTOM_OP("DynamicRNN")
92+ .FrameworkType(TENSORFLOW)
93+ .OriginOpType("DynamicRnn")
94+ .ParseParamsByOperatorFn(AutoMappingByOpFn)
95+ .ImplyType(ImplyType::TVM);
96+} // namespace domi
@@ -0,0 +1,42 @@
1+/**
2+ * Copyright (c) 2026 Huawei Technologies Co., Ltd.
3+ * This program is free software, you can redistribute it and/or modify it under the terms and conditions of
4+ * CANN Open Software License Agreement Version 2.0 (the "License").
5+ * Please refer to the License for details. You may not use this file except in compliance with the License.
6+ * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED,
7+ * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE.
8+ * See LICENSE in the root of the software repository for the full text of the License.
9+ */
10+ 
11+/*!
12+ * \file dynamic_rnnv2_tf_plugin.cpp
13+ * \brief DynamicRnnV2 / DynamicRNNV2 TensorFlow plugin mapping.
14+ */
15+#include "register/register.h"
16+ 
17+namespace domi {
18+static Status ParseParamsDynamicRNN(const ge::Operator& op_src, ge::Operator& op_dest)
19+{
20+ AutoMappingByOpFn(op_src, op_dest);
21+ op_dest.SetAttr("is_misplaced", true);
22+ return SUCCESS;
23+}
24+ 
25+REGISTER_CUSTOM_OP("DynamicRNN")
26+ .FrameworkType(TENSORFLOW)
27+ .OriginOpType("DynamicRnnV2")
28+ .ParseParamsByOperatorFn(ParseParamsDynamicRNN)
29+ .ImplyType(ImplyType::TVM);
30+ 
31+REGISTER_CUSTOM_OP("DynamicRNNV2")
32+ .FrameworkType(TENSORFLOW)
33+ .OriginOpType("DynamicRnnv2WithoutSeqlength")
34+ .ParseParamsByOperatorFn(ParseParamsDynamicRNN)
35+ .ImplyType(ImplyType::TVM);
36+ 
37+REGISTER_CUSTOM_OP("DynamicRNNV2")
38+ .FrameworkType(TENSORFLOW)
39+ .OriginOpType("DynamicRnnv2WithSeqlength")
40+ .ParseParamsByOperatorFn(AutoMappingByOpFn)
41+ .ImplyType(ImplyType::TVM);
42+} // namespace domi