[Transforamtions] NonZero horizontal fusion (#16571)

* Added ValuePredicate 'consumers_more_than'

* NonZero fusion

* NonZero fusion tests

src/common/transformations/include/transformations/common_optimizations/nonzero_fusion.hpp

@@ -0,0 +1,28 @@
+// Copyright (C) 2018-2023 Intel Corporation
+// SPDX-License-Identifier: Apache-2.0
+//
+
+#pragma once
+
+#include <memory>
+#include <openvino/pass/graph_rewrite.hpp>
+#include <openvino/pass/pattern/matcher.hpp>
+#include <transformations_visibility.hpp>
+
+namespace ov {
+namespace pass {
+
+class TRANSFORMATIONS_API NonZeroFusion;
+
+}  // namespace pass
+}  // namespace ov
+
+/**
+ * @ingroup ie_transformation_common_api
+ * @brief NonZeroFusion transformation makes horizontal fusion for equal NonZero layers
+ */
+class ov::pass::NonZeroFusion : public ov::pass::MatcherPass {
+public:
+    OPENVINO_RTTI("NonZeroFusion", "0");
+    NonZeroFusion();
+};

src/common/transformations/src/transformations/common_optimizations/moc_transformations.cpp

@@ -38,6 +38,7 @@
 #include <transformations/common_optimizations/mul_fake_quantize_fusion.hpp>
 #include <transformations/common_optimizations/mvn_fusion.hpp>
 #include <transformations/common_optimizations/nearest_neighbor_upsampling_fusion.hpp>
+#include <transformations/common_optimizations/nonzero_fusion.hpp>
 #include <transformations/common_optimizations/nop_elimination.hpp>
 #include <transformations/common_optimizations/normalize_l2_fusion.hpp>
 #include <transformations/common_optimizations/optimize_strided_slice.hpp>
@@ -200,6 +201,7 @@ bool ov::pass::MOCTransformations::run_on_model(const std::shared_ptr<ngraph::Fu
     ADD_MATCHER(common_fusions, PReluFusion)
     ADD_MATCHER(common_fusions, DepthToSpaceFusion)
     ADD_MATCHER(common_fusions, ShuffleChannelsFusion, !m_use_shapes)
+    ADD_MATCHER(common_fusions, NonZeroFusion)
     common_fusions->set_name("ov::pass::CommonFusions");
 
     REGISTER_PASS(manager, BinarizeWeights)

src/common/transformations/src/transformations/common_optimizations/nonzero_fusion.cpp

@@ -0,0 +1,40 @@
+// Copyright (C) 2018-2023 Intel Corporation
+// SPDX-License-Identifier: Apache-2.0
+//
+
+#include "transformations/common_optimizations/nonzero_fusion.hpp"
+
+#include <memory>
+#include <openvino/opsets/opset10.hpp>
+#include <openvino/pass/pattern/op/wrap_type.hpp>
+#include <vector>
+
+#include "itt.hpp"
+#include "transformations/utils/utils.hpp"
+
+ov::pass::NonZeroFusion::NonZeroFusion() {
+    MATCHER_SCOPE(NonZeroFusion);
+    auto input_m = pass::pattern::any_input(ov::pass::pattern::consumers_more_than(1));
+    auto nonzero_m = pass::pattern::wrap_type<ov::opset10::NonZero>({input_m});
+
+    ov::matcher_pass_callback callback = [=](ov::pass::pattern::Matcher& m) {
+        const auto& pattern_map = m.get_pattern_value_map();
+        const auto nonzero = ov::as_type_ptr<ov::opset10::NonZero>(pattern_map.at(nonzero_m).get_node_shared_ptr());
+        const auto out_prc = nonzero->get_output_type();
+
+        bool status = false;
+        auto replace_if_nodes_match = [&](const ov::Input<ov::Node>& in) {
+            auto cur_nonzero = ov::as_type_ptr<ov::opset10::NonZero>(in.get_node()->shared_from_this());
+            if (cur_nonzero && cur_nonzero->get_output_type() == out_prc) {
+                status |= ov::replace_output_update_name(cur_nonzero->output(0), nonzero->output(0));
+            }
+        };
+
+        const auto consumers = pattern_map.at(input_m).get_target_inputs();
+        std::for_each(consumers.begin(), consumers.end(), replace_if_nodes_match);
+        return status;
+    };
+
+    auto m = std::make_shared<ov::pass::pattern::Matcher>(nonzero_m, matcher_name);
+    register_matcher(m, callback);
+}

src/common/transformations/tests/common_optimizations/nonzero_fusion_test.cpp

@@ -0,0 +1,130 @@
+// Copyright (C) 2018-2023 Intel Corporation
+// SPDX-License-Identifier: Apache-2.0
+//
+
+#include <gtest/gtest.h>
+
+#include <memory>
+#include <openvino/opsets/opset10.hpp>
+#include <string>
+#include <transformations/common_optimizations/nonzero_fusion.hpp>
+
+#include "common_test_utils/ngraph_test_utils.hpp"
+
+using namespace testing;
+
+enum NonZeroType { I32, I64, NONE };
+
+struct NonZeroFusionBuilder {
+    NonZeroFusionBuilder() = default;
+    NonZeroFusionBuilder(const std::vector<NonZeroType>& props) : branch_props(props) {}
+
+    std::shared_ptr<ov::Model> getOriginal() {
+        const auto input = std::make_shared<ov::opset10::Parameter>(ov::element::f32, ov::PartialShape::dynamic(4));
+        ov::NodeVector results;
+        for (size_t i = 0; i < branch_props.size(); ++i) {
+            std::shared_ptr<ov::Node> nonzero;
+            switch (branch_props[i]) {
+            case NonZeroType::I32:
+                nonzero = std::make_shared<ov::opset10::NonZero>(input, ov::element::i32);
+                break;
+            case NonZeroType::I64:
+                nonzero = std::make_shared<ov::opset10::NonZero>(input, ov::element::i64);
+                break;
+            default:
+                nonzero = input;
+                break;
+            }
+            auto last_node = std::make_shared<ov::opset10::Relu>(nonzero);
+            last_node->set_friendly_name("last_node_" + std::to_string(i));
+            results.push_back(last_node);
+        }
+        return std::make_shared<ov::Model>(results, ov::ParameterVector{input});
+    };
+
+    std::shared_ptr<ov::Model> getReference() {
+        const auto input = std::make_shared<ov::opset10::Parameter>(ov::element::f32, ov::PartialShape::dynamic(4));
+
+        std::shared_ptr<ov::Node> i32_node;
+        std::shared_ptr<ov::Node> i64_node;
+        ov::NodeVector results;
+        for (size_t i = 0; i < branch_props.size(); ++i) {
+            std::shared_ptr<ov::Node> nonzero;
+            if (branch_props[i] == NonZeroType::I32) {
+                nonzero = i32_node ? i32_node : std::make_shared<ov::opset10::NonZero>(input, ov::element::i32);
+                if (!i32_node)
+                    i32_node = nonzero;
+            } else if (branch_props[i] == NonZeroType::I64) {
+                nonzero = i64_node ? i64_node : std::make_shared<ov::opset10::NonZero>(input, ov::element::i64);
+                if (!i64_node)
+                    i64_node = nonzero;
+            } else {
+                nonzero = input;
+            }
+            auto last_node = std::make_shared<ov::opset10::Relu>(nonzero);
+            last_node->set_friendly_name("last_node_" + std::to_string(i));
+            results.push_back(last_node);
+        }
+        return std::make_shared<ov::Model>(results, ov::ParameterVector{input});
+    }
+
+    std::vector<NonZeroType> branch_props;
+};
+
+class NonZeroFusionTests : public testing::WithParamInterface<std::vector<NonZeroType>>, public TransformationTestsF {
+public:
+    NonZeroFusionTests() : TransformationTestsF() {
+        comparator.enable(FunctionsComparator::CONSUMERS_COUNT);
+    }
+
+    static std::string getTestCaseName(testing::TestParamInfo<std::vector<NonZeroType>> obj) {
+        const std::vector<NonZeroType> testValues = obj.param;
+        std::ostringstream result;
+        result << "branch_props_{";
+        for (const auto& value : testValues) {
+            switch (value) {
+            case NonZeroType::I32:
+                result << "nonzero_i32,";
+                break;
+            case NonZeroType::I64:
+                result << "nonzero_i64,";
+                break;
+            default:
+                result << "wo_nonzero,";
+                break;
+            }
+        }
+        result << "}";
+        return result.str();
+    }
+
+protected:
+    void SetUp() override {
+        TransformationTestsF::SetUp();
+        const auto branch_props = GetParam();
+        builder = NonZeroFusionBuilder(branch_props);
+        manager.register_pass<ov::pass::NonZeroFusion>();
+    }
+
+    NonZeroFusionBuilder builder;
+};
+
+TEST_P(NonZeroFusionTests, NonZeroFusion) {
+    model = builder.getOriginal();
+    model_ref = builder.getReference();
+}
+
+namespace NonZeroFusionTestsInstantiation {
+std::vector<std::vector<NonZeroType>> test_params{std::vector<NonZeroType>(5, I32),
+                                                  std::vector<NonZeroType>(5, I64),
+                                                  std::vector<NonZeroType>(2, NONE),
+                                                  {I32, I64, I32, I64, I32},
+                                                  {I32, I64, NONE, I64, I32},
+                                                  {NONE, I64, NONE, I64, I32}};
+
+INSTANTIATE_TEST_SUITE_P(TransformationTestsF,
+                         NonZeroFusionTests,
+                         ::testing::ValuesIn(test_params),
+                         NonZeroFusionTests::getTestCaseName);
+
+}  // namespace NonZeroFusionTestsInstantiation

src/core/include/openvino/pass/pattern/op/pattern.hpp

@@ -39,6 +39,9 @@ std::function<bool(std::shared_ptr<Node>)> has_class() {
 OPENVINO_API
 std::function<bool(Output<Node>)> consumers_count(size_t n);
 
+OPENVINO_API
+std::function<bool(Output<Node>)> consumers_more_than(size_t n);
+
 OPENVINO_API
 std::function<bool(Output<Node>)> has_static_dim(size_t pos);
 

src/core/src/pattern/op/pattern.cpp

@@ -50,6 +50,12 @@ std::function<bool(Output<Node>)> consumers_count(size_t n) {
     };
 }
 
+std::function<bool(Output<Node>)> consumers_more_than(size_t n) {
+    return [=](Output<Node> output) -> bool {
+        return output.get_target_inputs().size() > n;
+    };
+}
+
 std::function<bool(Output<Node>)> has_static_dim(size_t pos) {
     return [=](Output<Node> output) -> bool {
         const auto& shape = output.get_partial_shape();

src/plugins/intel_gna/src/transformations/gather_sinking_unary.cpp

@@ -156,18 +156,10 @@ GatherSinkingUnaryBackwardSingleConsumer::GatherSinkingUnaryBackwardSingleConsum
     register_matcher(m, matcher_pass_callback);
 }
 
-namespace {
-std::function<bool(Output<Node>)> consumers_more_than(size_t n) {
-    return [=](Output<Node> output) -> bool {
-        return output.get_target_inputs().size() > n;
-    };
-}
-}  // namespace
-
 GatherSinkingUnaryBackwardMultiConsumers::GatherSinkingUnaryBackwardMultiConsumers() {
     MATCHER_SCOPE(GatherSinkingUnaryBackwardMultiConsumers);
     auto unary_restrictions = [](const Output<Node>& output) -> bool {
-        return consumers_more_than(1)(output) && HasSameOutputGatherNodes(output);
+        return ov::pass::pattern::consumers_more_than(1)(output) && HasSameOutputGatherNodes(output);
     };
 
     auto unary_label =