Introduce PullThroughReduce and ReduceReshapeFusion transformations (#13047)

* Initial version of PullThroughReduce and ReduceReshapeFusion transformation

* remove std namespace

* headers + namespaces clean-up

* resolved problem with ambiguous namespace

* clang format applied

* tests refactor. part 1

* tests refactor. part.2

* changed includes

* review remarks

* tests refactor. part.3

* review remarks

* fix build problem

* remarks

* fix reserving memory

* review remarks

* styles applied

* handle axes in the middle in try_get_unsqueeze_axes_from_reshape

* handle case of scalar input shape

Co-authored-by: Michal Lukaszewski <michal.lukaszewski@intel.com>

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

@@ -0,0 +1,57 @@
+// Copyright (C) 2022 Intel Corporation
+// SPDX-License-Identifier: Apache-2.0
+//
+
+#pragma once
+
+#include "openvino/pass/graph_rewrite.hpp"
+#include "transformations_visibility.hpp"
+
+namespace ov {
+namespace pass {
+
+class TRANSFORMATIONS_API PullThroughReduce;
+class TRANSFORMATIONS_API PullUnsqueezeThroughReduce;
+class TRANSFORMATIONS_API PullReshapeThroughReduce;
+
+}  // namespace pass
+}  // namespace ov
+
+/**
+ * @ingroup ie_transformation_common_api
+ * @brief PullUnsqueezeThroughReduce transformation
+ * The transformation pulls Unsqueeze operator though Reduce ops if possible.
+ * In the further processing such Unsqueeze can be often skipped as nop.
+ */
+class ov::pass::PullUnsqueezeThroughReduce : public ov::pass::MatcherPass {
+public:
+    OPENVINO_RTTI("PullUnsqueezeThroughReduce", "0");
+    PullUnsqueezeThroughReduce();
+};
+
+/**
+ * @ingroup ie_transformation_common_api
+ * @brief PullReshapeThroughReduce transformation
+ * The transformation pulls Reshape operator though Reduce ops if possible.
+ * In the further processing such Reshape can be often skipped as nop.
+ */
+class ov::pass::PullReshapeThroughReduce : public ov::pass::MatcherPass {
+public:
+    OPENVINO_RTTI("PullReshapeThroughReduce", "0");
+    PullReshapeThroughReduce();
+};
+
+/**
+ * @ingroup ie_transformation_common_api
+ * @brief PullThroughReduce transformation
+ * The transformation pulls Reshape or Unsqueeze operators though Reduce ops if possible.
+ * In the further processing such Reshape/Unsqueeze can be often skipped as nop.
+ */
+class ov::pass::PullThroughReduce : public ov::pass::GraphRewrite {
+public:
+    OPENVINO_RTTI("PullThroughReduce", "0");
+    PullThroughReduce() {
+        add_matcher<ov::pass::PullUnsqueezeThroughReduce>();
+        add_matcher<ov::pass::PullReshapeThroughReduce>();
+    }
+};

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

@@ -0,0 +1,27 @@
+// Copyright (C) 2022 Intel Corporation
+// SPDX-License-Identifier: Apache-2.0
+//
+
+#pragma once
+
+#include "openvino/pass/graph_rewrite.hpp"
+#include "transformations_visibility.hpp"
+
+namespace ov {
+namespace pass {
+
+class TRANSFORMATIONS_API ReduceReshapeFusion;
+
+}  // namespace pass
+}  // namespace ov
+
+/**
+ * @ingroup ie_transformation_common_api
+ * @brief ReduceReshapeFusion transformation
+ * Fuse ReduceOp(keep_dims=false)+Reshape to ReduceOp(keep_dims=true)
+ */
+class ov::pass::ReduceReshapeFusion : public ov::pass::MatcherPass {
+public:
+    OPENVINO_RTTI("ReduceReshapeFusion", "0");
+    ReduceReshapeFusion();
+};

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

@@ -42,8 +42,10 @@
 #include <transformations/common_optimizations/optimize_strided_slice.hpp>
 #include <transformations/common_optimizations/pad_fusion.hpp>
 #include <transformations/common_optimizations/prelu_fusion.hpp>
+#include <transformations/common_optimizations/pull_through_reduce.hpp>
 #include <transformations/common_optimizations/pull_transpose_through_fq.hpp>
 #include <transformations/common_optimizations/random_uniform_fusion.hpp>
+#include <transformations/common_optimizations/reduce_reshape_fusion.hpp>
 #include <transformations/common_optimizations/relu_fake_quantize_fusion.hpp>
 #include <transformations/common_optimizations/remove_concat_zero_dim_input.hpp>
 #include <transformations/common_optimizations/remove_filtering_boxes_by_size.hpp>
@@ -132,10 +134,9 @@ bool ngraph::pass::MOCTransformations::run_on_model(const std::shared_ptr<ngraph
     }
     // workaround until dynamism in NMS is not supported
     manager.register_pass<ngraph::pass::ConvertNmsGatherPathToUnsigned>();
-
     manager.register_pass<ngraph::pass::StridedSliceOptimization>(m_use_shapes);
-
     manager.register_pass<ngraph::pass::BroadcastElementwiseFusion>();
+    manager.register_pass<ov::pass::PullThroughReduce>();
 
     auto transpose_sinking = manager.register_pass<ngraph::pass::GraphRewrite>();
     transpose_sinking->add_matcher<ngraph::pass::TransposeSinking>();
@@ -162,6 +163,7 @@ bool ngraph::pass::MOCTransformations::run_on_model(const std::shared_ptr<ngraph
     common_fusions->add_matcher<ngraph::pass::ClampFusion>();
     common_fusions->add_matcher<ngraph::pass::PadFusion>();
     common_fusions->add_matcher<ngraph::pass::SoftmaxFusion>();
+    common_fusions->add_matcher<ov::pass::ReduceReshapeFusion>();
     common_fusions->add_matcher<ngraph::pass::MVNFusion>();
     common_fusions->add_matcher<ngraph::pass::DilatedConvolutionConverter>();
     common_fusions->add_matcher<ngraph::pass::GeluFusion>();

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

@@ -0,0 +1,231 @@
+// Copyright (C) 2022 Intel Corporation
+// SPDX-License-Identifier: Apache-2.0
+//
+
+#include "transformations/common_optimizations/pull_through_reduce.hpp"
+
+#include <memory>
+#include <vector>
+
+#include "itt.hpp"
+#include "openvino/core/validation_util.hpp"
+#include "openvino/op/util/reduction_base.hpp"
+#include "openvino/opsets/opset9.hpp"
+#include "openvino/pass/pattern/op/or.hpp"
+#include "openvino/pass/pattern/op/wrap_type.hpp"
+#include "sequnce_generator.hpp"
+#include "transformations/utils/utils.hpp"
+
+namespace {
+// Adjust axes of Unsqueeze/Reduce ops after Unsqueeze pulling
+// For example if we have:
+//          input(shape={5,10,15})
+//                 |
+//          Unsqueeze(axes=[0,3]) -> output_shape = {1,5,10,1,15}
+//                 |
+// ReduceOp(axes=[2,4], keep_dims=false) -> output_shape = {1,5,1}
+// after pulling it will be:
+//          input(shape={5,10,15})
+//                 |
+// ReduceOp(axes=[1,2], keep_dims=false) -> output_shape = {5}
+//                 |
+//          Unsqueeze(axes=[0,2]) -> output_shape = {1,5,1}
+const std::vector<int64_t> adjust_axes(const std::vector<int64_t>& axes_to_align,
+                                       const std::vector<int64_t>& offset_axes) {
+    auto number_of_axes_less_than = [&offset_axes](const int64_t current_axis) {
+        return std::count_if(std::begin(offset_axes),
+                             std::end(offset_axes),
+                             [&current_axis](const int64_t excluded_axis) {
+                                 return excluded_axis < current_axis;
+                             });
+    };
+    std::vector<int64_t> result;
+    for (const auto& axis : axes_to_align) {
+        result.push_back(axis - number_of_axes_less_than(axis));
+    }
+    return result;
+}
+
+// Try to represent given Reshape node via Unsqueeze and calculate axes of such Unsqueeze
+// - Reshape(input_shape={5,10,15}, target_shape={5,10,15,1}), 3 axis returned
+// - Reshape(input_shape={5,10,15}, target_shape={1,5,10,15}), 0 axis returned
+// - Reshape(input_shape={5,10,15}, target_shape={1,5,10,15,1}), 0 and 3 axes returned
+// - Reshape(input_shape={5,10,15}, target_shape={5,10,1,15}), 2 axis is returned
+std::vector<int64_t> try_get_unsqueeze_axes_from_reshape(const ov::Shape& target_shape, const ov::Shape& input_shape) {
+    std::vector<int64_t> result;
+    if (input_shape.size() == 0) {  // scalar case - can be reshaped only to [1,..,1] shape
+        result.resize(target_shape.size());
+        std::iota(std::begin(result), std::end(result), 0);
+        return result;
+    }
+    auto cur_input_shape_elem_idx = 0;
+    auto cur_input_shape_elem = input_shape[cur_input_shape_elem_idx];
+    auto target_shape_idx = 0;
+    for (; target_shape_idx < target_shape.size(); ++target_shape_idx) {
+        if (cur_input_shape_elem == target_shape[target_shape_idx] &&
+            cur_input_shape_elem_idx + 1 < input_shape.size()) {
+            ++cur_input_shape_elem_idx;
+            cur_input_shape_elem = input_shape[cur_input_shape_elem_idx];
+        } else if (target_shape[target_shape_idx] == 1 &&
+                   (target_shape_idx >= input_shape.size() + result.size() || cur_input_shape_elem != 1)) {
+            result.push_back(target_shape_idx);
+        }
+    }
+    if (cur_input_shape_elem_idx == input_shape.size() - 1 && target_shape_idx == target_shape.size()) {
+        return result;
+    } else {
+        return {};
+    }
+    return result;
+}
+
+// Update given reshape_input_shape by inserting "1" dimension on the postion represented by axes_to_insert
+std::shared_ptr<ov::opset9::Constant> update_reshape_target_shape(const ov::Shape& reshape_input_shape,
+                                                                  const std::vector<int64_t>& axes_to_insert) {
+    auto result = std::vector<int64_t>(std::begin(reshape_input_shape), std::end(reshape_input_shape));
+    for (const auto& axis : axes_to_insert) {
+        result.insert(std::next(std::begin(result), axis), 1);
+    }
+    return ov::opset9::Constant::create(ov::element::i64, ov::Shape{result.size()}, result);
+}
+
+// Return true if given inputs have some common elements, otherwise return false.
+bool have_same_axes(const std::vector<int64_t>& unsqueeze_axes, const std::vector<int64_t>& reduce_op_axes) {
+    return std::find_first_of(std::begin(unsqueeze_axes),
+                              std::end(unsqueeze_axes),
+                              std::begin(reduce_op_axes),
+                              std::end(reduce_op_axes)) != std::end(unsqueeze_axes);
+}
+}  // namespace
+
+ov::pass::PullUnsqueezeThroughReduce::PullUnsqueezeThroughReduce() {
+    MATCHER_SCOPE(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 reduce_axes = pattern::wrap_type<opset9::Constant>();
+    const auto reduce = pattern::wrap_type<op::util::ArithmeticReductionKeepDims, op::util::LogicalReductionKeepDims>(
+        {unsqueeze, reduce_axes});
+
+    matcher_pass_callback callback = [=](pattern::Matcher& m) {
+        auto& pattern_map = m.get_pattern_value_map();
+        const auto input_node = pattern_map.at(input);
+        const auto reduce_node =
+            std::dynamic_pointer_cast<op::util::ReductionBase>(pattern_map.at(reduce).get_node_shared_ptr());
+        const auto unsqueeze_node = pattern_map.at(unsqueeze).get_node_shared_ptr();
+        auto unsqueeze_axes_input =
+            std::dynamic_pointer_cast<opset9::Constant>(pattern_map.at(unsqueeze_axes).get_node_shared_ptr());
+        auto reduce_axes_input =
+            std::dynamic_pointer_cast<opset9::Constant>(pattern_map.at(reduce_axes).get_node_shared_ptr());
+
+        if (!unsqueeze_axes_input || !reduce_axes_input || !reduce_node) {
+            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(),
+                       unsqueeze_axes_val);
+        const auto reduce_axes_val = reduce_node->get_reduction_axes().to_vector();
+
+        if (have_same_axes(unsqueeze_axes_val, reduce_axes_val)) {
+            return false;
+        }
+
+        const bool keep_dims = reduce_node->get_keep_dims();
+
+        if (!keep_dims) {
+            const auto unsqueeze_adjusted_axes = adjust_axes(unsqueeze_axes_val, reduce_axes_val);
+            if (unsqueeze_adjusted_axes != unsqueeze_axes_val) {
+                unsqueeze_axes_input = opset9::Constant::create(unsqueeze_axes_input->get_element_type(),
+                                                                unsqueeze_axes_input->get_shape(),
+                                                                unsqueeze_adjusted_axes);
+            }
+        }
+
+        const auto reduce_adjusted_axes = adjust_axes(reduce_axes_val, unsqueeze_axes_val);
+        if (reduce_adjusted_axes != reduce_axes_val) {
+            reduce_axes_input = opset9::Constant::create(reduce_axes_input->get_element_type(),
+                                                         reduce_axes_input->get_shape(),
+                                                         reduce_adjusted_axes);
+        }
+
+        const auto new_reduce_node = reduce_node->clone_with_new_inputs({input_node, reduce_axes_input});
+        new_reduce_node->set_friendly_name(unsqueeze_node->get_friendly_name());
+        const auto new_unsqueeze_node = unsqueeze_node->clone_with_new_inputs({new_reduce_node, unsqueeze_axes_input});
+        new_unsqueeze_node->set_friendly_name(reduce_node->get_friendly_name());
+
+        copy_runtime_info({reduce_node, unsqueeze_node}, {new_reduce_node, new_unsqueeze_node});
+        replace_node(m.get_match_root(), new_unsqueeze_node);
+
+        return true;
+    };
+
+    auto m = std::make_shared<pattern::Matcher>(reduce, matcher_name);
+    register_matcher(m, callback);
+}
+
+ov::pass::PullReshapeThroughReduce::PullReshapeThroughReduce() {
+    MATCHER_SCOPE(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());
+    const auto reduce_axes = pattern::wrap_type<opset9::Constant>();
+    const auto reduce = pattern::wrap_type<op::util::ArithmeticReductionKeepDims, op::util::LogicalReductionKeepDims>(
+        {reshape, reduce_axes});
+
+    matcher_pass_callback callback = [=](pattern::Matcher& m) {
+        auto& pattern_map = m.get_pattern_value_map();
+        const auto input_node = pattern_map.at(input).get_node_shared_ptr();
+        const auto reduce_node =
+            std::dynamic_pointer_cast<op::util::ReductionBase>(pattern_map.at(reduce).get_node_shared_ptr());
+        if (!reduce_node) {
+            return false;
+        }
+        const auto reshape_node = pattern_map.at(reshape).get_node_shared_ptr();
+        const auto unsqueeze_axes =
+            try_get_unsqueeze_axes_from_reshape(reshape_node->get_shape(), input_node->get_shape());
+        if (unsqueeze_axes.empty()) {
+            return false;
+        }
+
+        const auto reduce_axes_val = reduce_node->get_reduction_axes().to_vector();
+
+        if (have_same_axes(unsqueeze_axes, reduce_axes_val)) {
+            return false;
+        }
+
+        const auto unsqueeze_adjusted_axes = adjust_axes(unsqueeze_axes, reduce_axes_val);
+        const auto reduce_adjusted_axes = adjust_axes(reduce_axes_val, unsqueeze_axes);
+
+        auto reduce_axes_input =
+            std::dynamic_pointer_cast<opset9::Constant>(pattern_map.at(reduce_axes).get_node_shared_ptr());
+
+        if (!reduce_axes_input) {
+            return false;
+        }
+
+        if (reduce_adjusted_axes != reduce_axes_val) {
+            reduce_axes_input = opset9::Constant::create(reduce_axes_input->get_element_type(),
+                                                         reduce_axes_input->get_shape(),
+                                                         reduce_adjusted_axes);
+        }
+
+        const auto new_reduce_node = reduce_node->clone_with_new_inputs({input_node, reduce_axes_input});
+        new_reduce_node->set_friendly_name(reshape_node->get_friendly_name());
+        const auto new_reshape_node = reshape_node->clone_with_new_inputs(
+            {new_reduce_node, update_reshape_target_shape(new_reduce_node->get_shape(), unsqueeze_adjusted_axes)});
+        new_reshape_node->set_friendly_name(reduce_node->get_friendly_name());
+
+        copy_runtime_info({reduce_node, reshape_node}, {new_reduce_node, new_reshape_node});
+        replace_node(m.get_match_root(), new_reshape_node);
+
+        return true;
+    };
+
+    auto m = std::make_shared<pattern::Matcher>(reduce, matcher_name);
+    register_matcher(m, callback);
+}

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

@@ -0,0 +1,69 @@
+// Copyright (C) 2022 Intel Corporation
+// SPDX-License-Identifier: Apache-2.0
+//
+
+#include "transformations/common_optimizations/reduce_reshape_fusion.hpp"
+
+#include <memory>
+#include <vector>
+
+#include "itt.hpp"
+#include "openvino/op/util/reduction_base.hpp"
+#include "openvino/opsets/opset9.hpp"
+#include "openvino/pass/pattern/op/or.hpp"
+#include "openvino/pass/pattern/op/wrap_type.hpp"
+#include "transformations/utils/utils.hpp"
+
+ov::pass::ReduceReshapeFusion::ReduceReshapeFusion() {
+    MATCHER_SCOPE(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());
+    const auto reshape =
+        pattern::wrap_type<opset9::Reshape>({reduce, pattern::any_input()}, pattern::has_static_shape());
+
+    matcher_pass_callback callback = [=](pattern::Matcher& m) {
+        auto& pattern_map = m.get_pattern_value_map();
+        auto reshape_node = pattern_map.at(reshape).get_node_shared_ptr();
+        const auto reduce_node =
+            std::dynamic_pointer_cast<op::util::ReductionBase>(pattern_map.at(reduce).get_node_shared_ptr());
+        if (!reduce_node) {
+            return false;
+        }
+        const bool keep_dims = reduce_node->get_keep_dims();
+
+        if (keep_dims) {
+            return false;
+        }
+
+        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();
+        for (const auto& axis : reduce_axes_val) {
+            reduce_shape_if_keep_dims.insert(std::next(std::begin(reduce_shape_if_keep_dims), axis), 1);
+        }
+
+        if (reduce_shape_if_keep_dims != reshape_shape) {
+            return false;
+        }
+
+        if (auto arithmetic_reduce_node =
+                std::dynamic_pointer_cast<op::util::ArithmeticReductionKeepDims>(reduce_node)) {
+            arithmetic_reduce_node->set_keep_dims(true);
+        } else if (auto logical_reduce_node =
+                       std::dynamic_pointer_cast<op::util::LogicalReductionKeepDims>(reduce_node)) {
+            logical_reduce_node->set_keep_dims(true);
+        }
+        reduce_node->set_friendly_name(reshape_node->get_friendly_name());
+        copy_runtime_info(reshape_node, reduce_node);
+        replace_node(m.get_match_root(), reduce_node);
+
+        return true;
+    };
+
+    auto m = std::make_shared<pattern::Matcher>(reshape, matcher_name);
+    register_matcher(m, callback);
+}