Add transformation to convert adaptive pool to reduce (#17488)

* Add transformation to convert adaptive pool to reduce

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

* Add tests and apply feedback

* Simplify if branches

* Add to common pipeline

* Remove 3d AdaptivePool with out_shape 1

* Skip test instead of remove

---------

Co-authored-by: Andrei Kochin <andrei.kochin@intel.com>

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

@@ -0,0 +1,29 @@
+// Copyright (C) 2018-2023 Intel Corporation
+// SPDX-License-Identifier: Apache-2.0
+//
+
+#pragma once
+
+#include <memory>
+#include <openvino/pass/graph_rewrite.hpp>
+#include <transformations_visibility.hpp>
+#include <vector>
+
+namespace ov {
+namespace pass {
+
+class TRANSFORMATIONS_API AdaptivePoolToReduce;
+
+}  // namespace pass
+}  // namespace ov
+
+/**
+ * @ingroup ie_transformation_common_api
+ * @brief AdaptivePoolToReduce transformation replaces AdaptiveXXXPool with ReduceXXX when possible
+ */
+
+class ov::pass::AdaptivePoolToReduce : public ov::pass::MatcherPass {
+public:
+    OPENVINO_RTTI("AdaptivePoolToReduce", "0");
+    AdaptivePoolToReduce();
+};

src/common/transformations/src/transformations/common_optimizations/adaptive_pool_to_reduce.cpp

@@ -0,0 +1,63 @@
+// Copyright (C) 2018-2023 Intel Corporation
+// SPDX-License-Identifier: Apache-2.0
+//
+
+#include "transformations/common_optimizations/adaptive_pool_to_reduce.hpp"
+
+#include <memory>
+#include <vector>
+
+#include "itt.hpp"
+#include "openvino/core/rt_info.hpp"
+#include "openvino/op/adaptive_avg_pool.hpp"
+#include "openvino/op/adaptive_max_pool.hpp"
+#include "openvino/op/constant.hpp"
+#include "openvino/op/reduce_max.hpp"
+#include "openvino/op/reduce_mean.hpp"
+#include "openvino/pass/pattern/op/wrap_type.hpp"
+#include "transformations/utils/utils.hpp"
+
+using namespace ov::op;
+
+ov::pass::AdaptivePoolToReduce::AdaptivePoolToReduce() {
+    MATCHER_SCOPE(AdaptivePoolToReduce);
+    auto data_pattern = pattern::any_input();
+    auto out_spatial_shape = pattern::wrap_type<v0::Constant>();
+    auto a_pool = pattern::wrap_type<v8::AdaptiveAvgPool, v8::AdaptiveMaxPool>({data_pattern, out_spatial_shape});
+
+    ov::matcher_pass_callback callback = [=](pattern::Matcher& m) {
+        const auto& pattern_map = m.get_pattern_map();
+
+        const auto& spatial_shape_c = std::dynamic_pointer_cast<v0::Constant>(pattern_map.at(out_spatial_shape));
+        auto spatial_shape = spatial_shape_c->cast_vector<int64_t>();
+        // Verify that all dimensions in adaptive pool shape are 1
+        for (auto& s : spatial_shape) {
+            if (s != 1)
+                return false;
+        }
+
+        auto axes = std::vector<int64_t>(spatial_shape.size(), 0);
+        std::iota(axes.begin(), axes.end(), 2);
+        auto axes_const = v0::Constant::create(element::i64, {spatial_shape.size()}, axes);
+        const auto adaptive_pool = pattern_map.at(a_pool);
+        std::shared_ptr<Node> res_node;
+        if (std::dynamic_pointer_cast<v8::AdaptiveAvgPool>(adaptive_pool)) {
+            res_node = std::make_shared<v1::ReduceMean>(adaptive_pool->input_value(0), axes_const, true);
+        } else if (std::dynamic_pointer_cast<v8::AdaptiveMaxPool>(adaptive_pool)) {
+            if (adaptive_pool->outputs().size() > 1 && adaptive_pool->output(1).get_target_inputs().size() != 0) {
+                // If indexes are used we can't replace it
+                return false;
+            }
+            res_node = std::make_shared<v1::ReduceMax>(adaptive_pool->input_value(0), axes_const, true);
+        } else {
+            return false;
+        }
+        adaptive_pool->output(0).replace(res_node);
+        res_node->set_friendly_name(adaptive_pool->get_friendly_name());
+        copy_runtime_info(adaptive_pool, res_node);
+        return true;
+    };
+
+    auto m = std::make_shared<ov::pass::pattern::Matcher>(a_pool, matcher_name);
+    this->register_matcher(m, callback);
+}

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

@@ -5,6 +5,7 @@
 #include <memory>
 #include <ngraph/pass/constant_folding.hpp>
 #include <ngraph/pass/manager.hpp>
+#include <transformations/common_optimizations/adaptive_pool_to_reduce.hpp>
 #include <transformations/common_optimizations/add_fake_quantize_fusion.hpp>
 #include <transformations/common_optimizations/align_eltwise_input_ranks.hpp>
 #include <transformations/common_optimizations/batch_to_space_fusion.hpp>
@@ -205,6 +206,7 @@ bool ov::pass::MOCTransformations::run_on_model(const std::shared_ptr<ngraph::Fu
     ADD_MATCHER(common_fusions, DepthToSpaceFusion)
     ADD_MATCHER(common_fusions, ShuffleChannelsFusion, !m_use_shapes)
     ADD_MATCHER(common_fusions, NonZeroHorizontalFusion)
+    ADD_MATCHER(common_fusions, AdaptivePoolToReduce)
     common_fusions->set_name("ov::pass::CommonFusions");
 
     REGISTER_PASS(manager, BinarizeWeights)

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

@@ -0,0 +1,98 @@
+// Copyright (C) 2018-2023 Intel Corporation
+// SPDX-License-Identifier: Apache-2.0
+//
+
+#include <gtest/gtest.h>
+
+#include <openvino/opsets/opset10.hpp>
+#include <transformations/common_optimizations/adaptive_pool_to_reduce.hpp>
+
+#include "common_test_utils/ngraph_test_utils.hpp"
+
+using namespace testing;
+using namespace ov;
+
+TEST_F(TransformationTestsF, AdaptiveAvgPool2dToReduceMean) {
+    {
+        auto data = std::make_shared<opset10::Parameter>(element::f32, PartialShape{1, 3, 14, 14});
+        auto out_spatial_shape = opset10::Constant::create(element::i32, Shape{2}, {1, 1});
+        auto adaptive_pool = std::make_shared<opset10::AdaptiveAvgPool>(data, out_spatial_shape);
+        auto result = std::make_shared<opset10::Result>(adaptive_pool);
+        model = std::make_shared<Model>(ResultVector{result}, ParameterVector{data});
+        manager.register_pass<pass::AdaptivePoolToReduce>();
+    }
+    {
+        auto data = std::make_shared<opset10::Parameter>(element::f32, PartialShape{1, 3, 14, 14});
+        auto axes = opset10::Constant::create(element::i64, Shape{2}, {2, 3});
+        auto reduce_mean = std::make_shared<opset10::ReduceMean>(data, axes, true);
+        auto result = std::make_shared<opset10::Result>(reduce_mean);
+        model_ref = std::make_shared<Model>(ResultVector{result}, ParameterVector{data});
+    }
+}
+
+TEST_F(TransformationTestsF, AdaptiveMaxPool2dToReduceMax) {
+    {
+        auto data = std::make_shared<opset10::Parameter>(element::f32, PartialShape{1, 3, 14, 14});
+        auto out_spatial_shape = opset10::Constant::create(element::i32, Shape{2}, {1, 1});
+        auto adaptive_pool = std::make_shared<opset10::AdaptiveMaxPool>(data, out_spatial_shape);
+        auto result = std::make_shared<opset10::Result>(adaptive_pool);
+        model = std::make_shared<Model>(ResultVector{result}, ParameterVector{data});
+        manager.register_pass<pass::AdaptivePoolToReduce>();
+    }
+    {
+        auto data = std::make_shared<opset10::Parameter>(element::f32, PartialShape{1, 3, 14, 14});
+        auto axes = opset10::Constant::create(element::i64, Shape{2}, {2, 3});
+        auto reduce_mean = std::make_shared<opset10::ReduceMax>(data, axes, true);
+        auto result = std::make_shared<opset10::Result>(reduce_mean);
+        model_ref = std::make_shared<Model>(ResultVector{result}, ParameterVector{data});
+    }
+}
+
+TEST_F(TransformationTestsF, AdaptiveMaxPool2dToReduceMaxUsedIndexes) {
+    {
+        auto data = std::make_shared<opset10::Parameter>(element::f32, PartialShape{1, 3, 14, 14});
+        auto out_spatial_shape = opset10::Constant::create(element::i32, Shape{2}, {1, 1});
+        auto adaptive_pool = std::make_shared<opset10::AdaptiveMaxPool>(data, out_spatial_shape);
+        auto result1 = std::make_shared<opset10::Result>(adaptive_pool->output(0));
+        auto result2 = std::make_shared<opset10::Result>(adaptive_pool->output(1));
+        model = std::make_shared<Model>(ResultVector{result1, result2}, ParameterVector{data});
+        manager.register_pass<pass::AdaptivePoolToReduce>();
+    }
+    // Reference model equals initial model
+}
+
+TEST_F(TransformationTestsF, AdaptiveAvgPool3dToReduceMean) {
+    {
+        auto data = std::make_shared<opset10::Parameter>(element::f32, PartialShape{1, 3, 14, 14, 14});
+        auto out_spatial_shape = opset10::Constant::create(element::i32, Shape{3}, {1, 1, 1});
+        auto adaptive_pool = std::make_shared<opset10::AdaptiveAvgPool>(data, out_spatial_shape);
+        auto result = std::make_shared<opset10::Result>(adaptive_pool);
+        model = std::make_shared<Model>(ResultVector{result}, ParameterVector{data});
+        manager.register_pass<pass::AdaptivePoolToReduce>();
+    }
+    {
+        auto data = std::make_shared<opset10::Parameter>(element::f32, PartialShape{1, 3, 14, 14, 14});
+        auto axes = opset10::Constant::create(element::i64, Shape{3}, {2, 3, 4});
+        auto reduce_mean = std::make_shared<opset10::ReduceMean>(data, axes, true);
+        auto result = std::make_shared<opset10::Result>(reduce_mean);
+        model_ref = std::make_shared<Model>(ResultVector{result}, ParameterVector{data});
+    }
+}
+
+TEST_F(TransformationTestsF, AdaptiveMaxPool3dToReduceMax) {
+    {
+        auto data = std::make_shared<opset10::Parameter>(element::f32, PartialShape{1, 3, 14, 14, 14});
+        auto out_spatial_shape = opset10::Constant::create(element::i32, Shape{3}, {1, 1, 1});
+        auto adaptive_pool = std::make_shared<opset10::AdaptiveMaxPool>(data, out_spatial_shape);
+        auto result = std::make_shared<opset10::Result>(adaptive_pool);
+        model = std::make_shared<Model>(ResultVector{result}, ParameterVector{data});
+        manager.register_pass<pass::AdaptivePoolToReduce>();
+    }
+    {
+        auto data = std::make_shared<opset10::Parameter>(element::f32, PartialShape{1, 3, 14, 14, 14});
+        auto axes = opset10::Constant::create(element::i64, Shape{3}, {2, 3, 4});
+        auto reduce_mean = std::make_shared<opset10::ReduceMax>(data, axes, true);
+        auto result = std::make_shared<opset10::Result>(reduce_mean);
+        model_ref = std::make_shared<Model>(ResultVector{result}, ParameterVector{data});
+    }
+}

src/plugins/intel_cpu/tests/functional/shared_tests_instances/skip_tests_config.cpp

@@ -170,7 +170,9 @@ std::vector<std::string> disabledTestPatterns() {
         // 98151. Not valid sorting for slices in reference.
         R"(.*UniqueLayerTestCPU.*axis.*True.*)",
         // AUTO does not support import / export
-        R"(.*smoke_Auto_BehaviorTests/OVCompiledGraphImportExportTest.*(mportExport|readFromV10IR).*/targetDevice=(AUTO).*)"
+        R"(.*smoke_Auto_BehaviorTests/OVCompiledGraphImportExportTest.*(mportExport|readFromV10IR).*/targetDevice=(AUTO).*)",
+        // AdaptiveAvgPool is converted into Reduce op for suitable parameters. CPU Reduce impl doesn't support non planar layout for 3D case
+        R"(.*StaticAdaPoolAvg3DLayoutTest.*OS=\(1\).*_inFmts=(nwc|nCw16c|nCw8c).*)"
     };
 
 #if defined(OPENVINO_ARCH_X86)