[TF FE] Refactor ops translators to check complex type and speed-up c…

…ompilation (openvinotoolkit#21025) * [TF FE] Refactor AddN translator Signed-off-by: Kazantsev, Roman <[email protected]> * Refactor ArgMin/ArgMax Signed-off-by: Kazantsev, Roman <[email protected]> * Add default check for binary operations Signed-off-by: Kazantsev, Roman <[email protected]> * Refactor Cast translator Signed-off-by: Kazantsev, Roman <[email protected]> * Refactor Const translator Signed-off-by: Kazantsev, Roman <[email protected]> * Refactor Einsum, Elu, and Fill operations Signed-off-by: Kazantsev, Roman <[email protected]> * Refactor MatMul translator Signed-off-by: Kazantsev, Roman <[email protected]> * Refactor Relu6 and Placeholder Signed-off-by: Kazantsev, Roman <[email protected]> * Refactor Round, Square, and unary operation translators Signed-off-by: Kazantsev, Roman <[email protected]> * Update src/frontends/tensorflow_common/src/op/relu_6.cpp * Update src/frontends/tensorflow_common/src/op/const.cpp --------- Signed-off-by: Kazantsev, Roman <[email protected]>
xufang-lisa · Nov 13, 2023 · 7c595f8 · 7c595f8
1 parent 6650b87
commit 7c595f8
Show file tree

Hide file tree

Showing 15 changed files with 114 additions and 76 deletions.
diff --git a/src/frontends/tensorflow_common/src/op/addN.cpp b/src/frontends/tensorflow_common/src/op/addN.cpp
@@ -2,32 +2,29 @@
 // SPDX-License-Identifier: Apache-2.0
 //
 
-#include <numeric>
-
 #include "common_op_table.hpp"
-#include "openvino/opsets/opset8.hpp"
+#include "openvino/op/add.hpp"
 
 using namespace std;
-using namespace ov::opset8;
+using namespace ov;
+using namespace ov::op;
 
 namespace ov {
 namespace frontend {
 namespace tensorflow {
 namespace op {
 
 OutputVector translate_add_n_op(const NodeContext& node) {
-    OutputVector ng_arg_vec;
-    for (size_t i = 0; i < node.get_input_size(); i++) {
-        ng_arg_vec.push_back(node.get_input(static_cast<int>(i)));
+    default_op_checks(node, 1, {"AddN", "ADD_N"});
+    int num_size = static_cast<int>(node.get_input_size());
+
+    Output<Node> result = node.get_input(0);
+    for (int ind = 1; ind < num_size; ++ind) {
+        result = make_shared<v1::Add>(result, node.get_input(ind));
     }
-    auto res = std::accumulate(std::next(ng_arg_vec.begin()),
-                               ng_arg_vec.end(),
-                               ng_arg_vec.at(0),
-                               [](const Output<Node>& a, const Output<Node>& b) -> shared_ptr<Node> {
-                                   return make_shared<Add>(a, b);
-                               });
-    set_node_name(node.get_name(), res.get_node_shared_ptr());
-    return {res};
+
+    set_node_name(node.get_name(), result.get_node_shared_ptr());
+    return {result};
 }
 }  // namespace op
 }  // namespace tensorflow

diff --git a/src/frontends/tensorflow_common/src/op/arg_min_max.cpp b/src/frontends/tensorflow_common/src/op/arg_min_max.cpp
@@ -3,24 +3,27 @@
 //
 
 #include "common_op_table.hpp"
-#include "openvino/opsets/opset8.hpp"
+#include "openvino/op/constant.hpp"
+#include "openvino/op/squeeze.hpp"
+#include "openvino/op/topk.hpp"
 
 using namespace std;
-using namespace ov::opset8;
+using namespace ov::op;
 
 namespace ov {
 namespace frontend {
 namespace tensorflow {
 namespace op {
 
 OutputVector translate_arg_min_max(const NodeContext& node, std::string mode) {
+    default_op_checks(node, 1, {"ArgMax", "ArgMin", "ARG_MAX", "ARG_MIN"});
     auto input = node.get_input(0);
 
     // TensorFlow uses axis with default value equal to zero
     int64_t axis = 0;
     if (node.get_input_size() > 1) {
         TENSORFLOW_OP_VALIDATION(node,
-                                 std::dynamic_pointer_cast<opset8::Constant>(node.get_input(1).get_node_shared_ptr()),
+                                 as_type_ptr<v0::Constant>(node.get_input(1).get_node_shared_ptr()),
                                  "ArgMax/ArgMin is not supported with non-constant axis input");
         std::vector<int64_t> axes;
         get_const_input(node, 1, &axes);
@@ -30,12 +33,12 @@ OutputVector translate_arg_min_max(const NodeContext& node, std::string mode) {
     auto output_type = node.get_attribute<element::Type>("output_type", element::i64);
 
     // compute indices of max/min values using TopK
-    auto k = make_shared<Constant>(element::i64, Shape{}, 1);
+    auto k = make_shared<v0::Constant>(element::i64, Shape{}, 1);
     // TODO: define sort attribute for TensorFlow case
-    auto top_k = std::make_shared<TopK>(input, k, axis, mode, "none", output_type);
+    auto top_k = std::make_shared<v11::TopK>(input, k, axis, mode, "none", output_type);
 
-    auto axis_to_remove = make_shared<Constant>(element::i64, Shape{1}, std::vector<int64_t>({axis}));
-    auto res = make_shared<Squeeze>(top_k->output(1), axis_to_remove);
+    auto axis_to_remove = make_shared<v0::Constant>(element::i64, Shape{1}, vector<int64_t>({axis}));
+    auto res = make_shared<v0::Squeeze>(top_k->output(1), axis_to_remove);
     set_node_name(node.get_name(), res);
     return {res};
 }

diff --git a/src/frontends/tensorflow_common/src/op/binary_op.cpp b/src/frontends/tensorflow_common/src/op/binary_op.cpp
@@ -15,11 +15,12 @@ namespace op {
 
 OutputVector translate_binary_op(const NodeContext& node,
                                  const std::function<Output<Node>(Output<Node>&, Output<Node>&)>& create_binary_op) {
-    auto ng_lhs = node.get_input(0);
-    auto ng_rhs = node.get_input(1);
-    auto ng_node = create_binary_op(ng_lhs, ng_rhs);
-    set_node_name(node.get_name(), ng_node.get_node_shared_ptr());
-    return {ng_node};
+    default_op_checks(node, 2, {});
+    auto lhs = node.get_input(0);
+    auto rhs = node.get_input(1);
+    auto result = create_binary_op(lhs, rhs);
+    set_node_name(node.get_name(), result.get_node_shared_ptr());
+    return {result};
 }
 
 OutputVector translate_floor_div_op(const NodeContext& node) {

diff --git a/src/frontends/tensorflow_common/src/op/cast.cpp b/src/frontends/tensorflow_common/src/op/cast.cpp
@@ -3,21 +3,23 @@
 //
 
 #include "common_op_table.hpp"
-#include "openvino/opsets/opset8.hpp"
+#include "openvino/op/convert.hpp"
 
 using namespace std;
-using namespace ov::opset8;
+using namespace ov::op;
 
 namespace ov {
 namespace frontend {
 namespace tensorflow {
 namespace op {
 
 OutputVector translate_cast_op(const NodeContext& node) {
-    auto ng_input = node.get_input(0);
+    default_op_checks(node, 1, {"Cast", "CAST"});
+    auto x = node.get_input(0);
+
+    auto dst_type = node.get_attribute<element::Type>("DstT");
+    auto res = make_shared<v0::Convert>(x, dst_type);
 
-    auto ng_et = node.get_attribute<element::Type>("DstT");
-    auto res = make_shared<Convert>(ng_input, ng_et);
     set_node_name(node.get_name(), res);
     return res->outputs();
 }

diff --git a/src/frontends/tensorflow_common/src/op/const.cpp b/src/frontends/tensorflow_common/src/op/const.cpp
@@ -5,10 +5,10 @@
 #include "common_op_table.hpp"
 #include "helper_ops/string_constant.hpp"
 #include "helper_ops/unsupported_constant.hpp"
-#include "openvino/opsets/opset8.hpp"
+#include "openvino/op/constant.hpp"
 
 using namespace std;
-using namespace ov::opset8;
+using namespace ov::op;
 using namespace ov;
 
 namespace ov {
@@ -17,6 +17,8 @@ namespace tensorflow {
 namespace op {
 
 OutputVector translate_const_op(const NodeContext& node) {
+    default_op_checks(node, 0, {"Const"});
+
     auto ov_type = node.get_attribute_as_any("dtype");
     std::shared_ptr<Node> const_node;
     if (!ov_type.is<ov::element::Type>() || ov_type.as<ov::element::Type>() == ov::element::dynamic ||
@@ -28,7 +30,7 @@ OutputVector translate_const_op(const NodeContext& node) {
         }
     } else {
         auto tensor = node.get_attribute<Tensor>("value");
-        const_node = std::make_shared<Constant>(tensor);
+        const_node = std::make_shared<v0::Constant>(tensor);
     }
     set_node_name(node.get_name(), const_node);
     return {const_node};

diff --git a/src/frontends/tensorflow_common/src/op/einsum.cpp b/src/frontends/tensorflow_common/src/op/einsum.cpp
@@ -2,28 +2,31 @@
 // SPDX-License-Identifier: Apache-2.0
 //
 
+#include "openvino/op/einsum.hpp"
+
 #include "common_op_table.hpp"
-#include "openvino/opsets/opset8.hpp"
 
 using namespace std;
-using namespace ov::opset8;
+using namespace ov::op;
 
 namespace ov {
 namespace frontend {
 namespace tensorflow {
 namespace op {
 OutputVector translate_einsum_op(const NodeContext& node) {
+    default_op_checks(node, 1, {"Einsum"});
+
     auto op_type = node.get_op_type();
     TENSORFLOW_OP_VALIDATION(node, op_type == "Einsum", "Internal error: incorrect usage of translate_einsum_op.");
-    auto equation = node.get_attribute<std::string>("equation");
+    auto equation = node.get_attribute<string>("equation");
     int input_size = static_cast<int>(node.get_input_size());
 
     OutputVector inputs;
     for (int input_ind = 0; input_ind < input_size; ++input_ind) {
         inputs.push_back(node.get_input(input_ind));
     }
 
-    auto einsum = make_shared<Einsum>(inputs, equation);
+    auto einsum = make_shared<v7::Einsum>(inputs, equation);
     set_node_name(node.get_name(), einsum);
     return {einsum};
 }

diff --git a/src/frontends/tensorflow_common/src/op/elu.cpp b/src/frontends/tensorflow_common/src/op/elu.cpp
@@ -2,22 +2,25 @@
 // SPDX-License-Identifier: Apache-2.0
 //
 
+#include "openvino/op/elu.hpp"
+
 #include "common_op_table.hpp"
-#include "openvino/opsets/opset8.hpp"
 
 using namespace std;
 using namespace ov;
-using namespace ov::opset8;
+using namespace ov::op;
 
 namespace ov {
 namespace frontend {
 namespace tensorflow {
 namespace op {
 
 OutputVector translate_elu_op(const NodeContext& node) {
+    default_op_checks(node, 1, {"Elu", "ELU"});
     auto input = node.get_input(0);
     auto alpha = node.get_attribute<float>("alpha", 1.0);
-    auto res = make_shared<Elu>(input, alpha);
+    auto res = make_shared<v0::Elu>(input, alpha);
+
     set_node_name(node.get_name(), res);
     return res->outputs();
 }

diff --git a/src/frontends/tensorflow_common/src/op/fill.cpp b/src/frontends/tensorflow_common/src/op/fill.cpp
@@ -3,20 +3,22 @@
 //
 
 #include "common_op_table.hpp"
-#include "openvino/opsets/opset8.hpp"
+#include "openvino/op/broadcast.hpp"
 
 using namespace std;
-using namespace ov::opset8;
+using namespace ov::op;
 
 namespace ov {
 namespace frontend {
 namespace tensorflow {
 namespace op {
 
 OutputVector translate_fill_op(const NodeContext& node) {
-    auto ng_dims = node.get_input(0);
-    auto ng_value = node.get_input(1);
-    auto res = make_shared<Broadcast>(ng_value, ng_dims);
+    default_op_checks(node, 2, {"Fill", "FILL"});
+    auto dims = node.get_input(0);
+    auto value = node.get_input(1);
+
+    auto res = make_shared<v3::Broadcast>(value, dims);
     set_node_name(node.get_name(), res);
     return res->outputs();
 }

diff --git a/src/frontends/tensorflow_common/src/op/matmul.cpp b/src/frontends/tensorflow_common/src/op/matmul.cpp
@@ -2,41 +2,48 @@
 // SPDX-License-Identifier: Apache-2.0
 //
 
+#include "openvino/op/matmul.hpp"
+
 #include "common_op_table.hpp"
-#include "openvino/opsets/opset8.hpp"
 
 using namespace std;
-using namespace ov::opset8;
+using namespace ov::op;
 
 namespace ov {
 namespace frontend {
 namespace tensorflow {
 namespace op {
 
 OutputVector translate_mat_mul_op(const NodeContext& node) {
+    default_op_checks(node, 2, {"MatMul"});
+
     auto a = node.get_input(0);
     auto b = node.get_input(1);
     auto transpose_a = node.get_attribute<bool>("transpose_a", false);
     auto transpose_b = node.get_attribute<bool>("transpose_b", false);
 
-    auto res = make_shared<MatMul>(a, b, transpose_a, transpose_b);
+    auto res = make_shared<v0::MatMul>(a, b, transpose_a, transpose_b);
     set_node_name(node.get_name(), res);
     return res->outputs();
 }
 
 OutputVector translate_batch_mat_mul_op(const NodeContext& node) {
+    default_op_checks(node, 2, {"BatchMatMul", "BatchMatMulV2", "BATCH_MATMUL"});
+
     auto x = node.get_input(0);
     auto y = node.get_input(1);
 
     auto adj_x = node.get_attribute<bool>("adj_x", false);
     auto adj_y = node.get_attribute<bool>("adj_y", false);
 
-    auto result = make_shared<MatMul>(x, y, adj_x, adj_y);
+    auto result = make_shared<v0::MatMul>(x, y, adj_x, adj_y);
     set_node_name(node.get_name(), result);
     return result->outputs();
 }
 
 OutputVector translate_batch_mat_mul_with_type_op(const NodeContext& node) {
+    default_op_checks(node, 2, {"BatchMatMulV3"});
+
     auto x = node.get_input(0);
     auto y = node.get_input(1);
 
@@ -46,10 +53,10 @@ OutputVector translate_batch_mat_mul_with_type_op(const NodeContext& node) {
     auto adj_y = node.get_attribute<bool>("adj_y", false);
     auto t_out = node.get_attribute<ov::element::Type>("Tout", input_type);
 
-    auto result = make_shared<MatMul>(x, y, adj_x, adj_y)->output(0);
+    auto result = make_shared<v0::MatMul>(x, y, adj_x, adj_y)->output(0);
 
     if (t_out != input_type) {
-        result = make_shared<Convert>(result, t_out);
+        result = make_shared<v0::Convert>(result, t_out);
     }
 
     set_node_name(node.get_name(), result.get_node_shared_ptr());

diff --git a/src/frontends/tensorflow_common/src/op/placeholder.cpp b/src/frontends/tensorflow_common/src/op/placeholder.cpp
@@ -3,35 +3,39 @@
 //
 
 #include "common_op_table.hpp"
-#include "openvino/opsets/opset8.hpp"
+#include "openvino/op/parameter.hpp"
 
 using namespace std;
-using namespace ov::opset8;
+using namespace ov::op;
 
 namespace ov {
 namespace frontend {
 namespace tensorflow {
 namespace op {
 
 OutputVector translate_placeholder_op(const NodeContext& node) {
-    auto dtype = node.get_attribute<ov::element::Type>("dtype");
-    auto shape = node.get_attribute<ov::PartialShape>("shape", ov::PartialShape::dynamic());
+    default_op_checks(node, 0, {});
+
+    auto dtype = node.get_attribute<element::Type>("dtype");
+    auto shape = node.get_attribute<PartialShape>("shape", PartialShape::dynamic());
     if (shape.rank().is_static() && shape.rank().get_length() == 0 && node.has_attribute("_output_shapes")) {
         // we know some cases when Placeholder operation has empty scalar `shape` attribute value
         // and non-empty `_output_shapes` attribute value.
         // `_output_shapes` attribute value turns to be correct in this case
-        auto output_shapes = node.get_attribute<std::vector<ov::PartialShape>>("_output_shapes");
+        auto output_shapes = node.get_attribute<vector<PartialShape>>("_output_shapes");
         if (output_shapes.size() == 1 && output_shapes[0].rank().is_static()) {
             shape = output_shapes[0];
         }
     }
 
-    auto res = std::make_shared<Parameter>(dtype, shape);
+    auto res = make_shared<v0::Parameter>(dtype, shape);
     set_node_name(node.get_name(), res);
     return res->outputs();
 }
 
 OutputVector translate_placeholder_with_default_op(const NodeContext& node) {
+    default_op_checks(node, 0, {});
+
     // For parity with legacy frontend, it creates a constant node with the default value
     // As a rule, PlaceholderWithDefault is mainly used for is_training variables in the model
     TENSORFLOW_OP_VALIDATION(node,