PullThroughReduce and ReduceReshapeFusion transformations - handle more than one consumers (#14101)

src/common/transformations/src/transformations/common_optimizations/pull_through_reduce.cpp

@@ -103,7 +103,7 @@ ov::pass::PullUnsqueezeThroughReduce::PullUnsqueezeThroughReduce() {
 
     const auto input = pattern::any_input(pattern::has_static_rank());
     const auto unsqueeze_axes = pattern::wrap_type<opset9::Constant>();
-    const auto unsqueeze = pattern::wrap_type<opset9::Unsqueeze>({input, unsqueeze_axes}, pattern::has_static_rank());
+    const auto unsqueeze = pattern::wrap_type<opset9::Unsqueeze>({input, unsqueeze_axes}, pattern::consumers_count(1));
     const auto reduce_axes = pattern::wrap_type<opset9::Constant>();
     const auto reduce = pattern::wrap_type<op::util::ArithmeticReductionKeepDims, op::util::LogicalReductionKeepDims>(
         {unsqueeze, reduce_axes});
@@ -123,6 +123,10 @@ ov::pass::PullUnsqueezeThroughReduce::PullUnsqueezeThroughReduce() {
             return false;
         }
 
+        if (unsqueeze_node->get_output_partial_shape(0).rank().is_dynamic()) {
+            return false;
+        }
+
         auto unsqueeze_axes_val = unsqueeze_axes_input->cast_vector<int64_t>();
         normalize_axes(unsqueeze_node.get(),
                        unsqueeze_node->get_output_partial_shape(0).rank().get_length(),
@@ -172,7 +176,7 @@ ov::pass::PullReshapeThroughReduce::PullReshapeThroughReduce() {
     const auto input = pattern::any_input(pattern::has_static_shape());
     const auto reshape_target_shape = pattern::wrap_type<opset9::Constant>();
     const auto reshape =
-        pattern::wrap_type<opset9::Reshape>({input, reshape_target_shape}, pattern::has_static_shape());
+        pattern::wrap_type<opset9::Reshape>({input, reshape_target_shape}, pattern::consumers_count(1));
     const auto reduce_axes = pattern::wrap_type<opset9::Constant>();
     const auto reduce = pattern::wrap_type<op::util::ArithmeticReductionKeepDims, op::util::LogicalReductionKeepDims>(
         {reshape, reduce_axes});
@@ -186,6 +190,9 @@ ov::pass::PullReshapeThroughReduce::PullReshapeThroughReduce() {
             return false;
         }
         const auto reshape_node = pattern_map.at(reshape).get_node_shared_ptr();
+        if (reshape_node->get_output_partial_shape(0).is_dynamic()) {
+            return false;
+        }
         const auto unsqueeze_axes =
             try_get_unsqueeze_axes_from_reshape(reshape_node->get_shape(), input_node->get_shape());
         if (unsqueeze_axes.empty()) {

src/common/transformations/src/transformations/common_optimizations/reduce_reshape_fusion.cpp

@@ -20,7 +20,7 @@ ov::pass::ReduceReshapeFusion::ReduceReshapeFusion() {
     const auto reduce_axes = pattern::wrap_type<opset9::Constant>();
     const auto reduce = pattern::wrap_type<op::util::ArithmeticReductionKeepDims, op::util::LogicalReductionKeepDims>(
         {pattern::any_input(), reduce_axes},
-        pattern::has_static_shape());
+        pattern::consumers_count(1));
     const auto reshape =
         pattern::wrap_type<opset9::Reshape>({reduce, pattern::any_input()}, pattern::has_static_shape());
 
@@ -32,6 +32,9 @@ ov::pass::ReduceReshapeFusion::ReduceReshapeFusion() {
         if (!reduce_node) {
             return false;
         }
+        if (reduce_node->get_output_partial_shape(0).is_dynamic()) {
+            return false;
+        }
         const bool keep_dims = reduce_node->get_keep_dims();
 
         if (keep_dims) {
@@ -41,12 +44,12 @@ ov::pass::ReduceReshapeFusion::ReduceReshapeFusion() {
         const auto reduce_axes_val = reduce_node->get_reduction_axes().to_vector();
         const auto& reshape_shape = reshape_node->get_shape();
 
-        auto reduce_shape_if_keep_dims = reduce_node->get_shape();
+        auto reduce_shape = reduce_node->get_shape();
         for (const auto& axis : reduce_axes_val) {
-            reduce_shape_if_keep_dims.insert(std::next(std::begin(reduce_shape_if_keep_dims), axis), 1);
+            reduce_shape.insert(std::next(std::begin(reduce_shape), axis), 1);
         }
 
-        if (reduce_shape_if_keep_dims != reshape_shape) {
+        if (reduce_shape != reshape_shape) {
             return false;
         }
 

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

@@ -172,6 +172,18 @@ TEST_F(TransformationTestsF, PullUnsqueezeThroughReduceSkipIfNotConstAxes) {
     manager.register_pass<pass::PullUnsqueezeThroughReduce>();
 }
 
+TEST_F(TransformationTestsF, PullUnsqueezeThroughReduceMeanSkipIfMoreThanOneUnsqueezeConsumer) {
+    const auto input = std::make_shared<Parameter>(element::f32, PartialShape{5, 10, 15});
+    const auto unsqueeze_axes = Constant::create(element::i64, Shape{1}, {0});
+    const auto unsqueeze = std::make_shared<Unsqueeze>(input, unsqueeze_axes);
+    const auto reduce_axes = Constant::create(element::i64, Shape{}, {1});
+    const auto reduce_mean = std::make_shared<ReduceMean>(unsqueeze, reduce_axes);
+    const auto add = std::make_shared<Add>(unsqueeze, Constant::create(element::f32, Shape{}, {1}));
+
+    model =  std::make_shared<Model>(NodeVector{reduce_mean, add}, ParameterVector{input});
+    manager.register_pass<pass::PullUnsqueezeThroughReduce>();
+}
+
 struct PullReshapeParams {
     element::Type in_type;
     PartialShape in_shape;
@@ -298,3 +310,15 @@ TEST_F(TransformationTestsF, PullReshapeThroughReduceMeanSkipIfDynamicReshapeOut
     model =  std::make_shared<Model>(NodeVector{reduce_mean}, ParameterVector{input, target_shape});
     manager.register_pass<pass::PullReshapeThroughReduce>();
 }
+
+TEST_F(TransformationTestsF, PullReshapeThroughReduceMeanSkipIfMoreThanOneReshapeConsumer) {
+    const auto input = std::make_shared<Parameter>(element::f32, PartialShape{5, 10, 15});
+    const auto target_shape = Constant::create(element::i64, Shape{4}, {1, 5, 10, 15});
+    const auto reshape = std::make_shared<Reshape>(input, target_shape, false);
+    const auto reduce_axes = Constant::create(element::i64, Shape{}, {2});
+    const auto reduce_mean = std::make_shared<ReduceMean>(reshape, reduce_axes);
+    const auto add = std::make_shared<Add>(reshape, Constant::create(element::f32, Shape{}, {1}));
+
+    model = std::make_shared<Model>(NodeVector{reduce_mean, add}, ParameterVector{input});
+    manager.register_pass<pass::PullReshapeThroughReduce>();
+}

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

@@ -145,3 +145,15 @@ TEST_F(TransformationTestsF, ReduceMeanReshapeFusionSkipIfNonConstReshapeTargetS
     model =  std::make_shared<Model>(NodeVector{reshape}, ParameterVector{input, target_shape});
     manager.register_pass<pass::ReduceReshapeFusion>();
 }
+
+TEST_F(TransformationTestsF, ReduceMeanReshapeFusionSkipIfMoreThanOneReduceConsumer) {
+    const auto input = std::make_shared<Parameter>(element::f32, PartialShape{1, 1});
+    const auto reduce_axes = Constant::create(element::i64, Shape{}, {1});
+    const auto reduce_mean = std::make_shared<ReduceMean>(input, reduce_axes);
+    const auto add = std::make_shared<Add>(reduce_mean, Constant::create(element::f32, Shape{1}, {1}));
+    const auto target_shape = Constant::create(element::i64, Shape{2}, {1, 1});
+    const auto reshape = std::make_shared<Reshape>(reduce_mean, target_shape, false);
+
+    model =  std::make_shared<Model>(NodeVector{reshape, add}, ParameterVector{input});
+    manager.register_pass<pass::ReduceReshapeFusion>();
+}