From 2c0efbb1bf884bece4abdbe7693b12b3d3fe42e5 Mon Sep 17 00:00:00 2001
From: Wang Wangwang <wangwang.wang@intel.com>
Date: Wed, 6 Mar 2024 19:20:49 +0800
Subject: [PATCH] [HETERO] Remove ConstantFolding of model in hetero (#22626)

### Details:
constant folding been introduced to hetero, but this may bring other
issue in quantized models , for example, in gpu plugin, precision
conversion will be applied before constant folding, some ops can be
operated in fp16 if hardware support, while direct constant folding in
hetero doesn't take hardware capability into account and fold the
constants directly, this will result in memory concerns, for example,
OOM issue, or the performance issue because of different execution
precision
### ticket-id:
CVS-127934
---
 src/plugins/hetero/src/compiled_model.cpp     |  8 ---
 .../hetero/tests/unit/subgraph_collector.cpp  | 71 +++++++++++++++++++
 2 files changed, 71 insertions(+), 8 deletions(-)

diff --git a/src/plugins/hetero/src/compiled_model.cpp b/src/plugins/hetero/src/compiled_model.cpp
index 3d84d947f81..3032a32d4c7 100644
--- a/src/plugins/hetero/src/compiled_model.cpp
+++ b/src/plugins/hetero/src/compiled_model.cpp
@@ -37,14 +37,6 @@ ov::hetero::CompiledModel::CompiledModel(const std::shared_ptr<ov::Model>& model
 }
 
 void ov::hetero::CompiledModel::compile_model(const std::shared_ptr<ov::Model>& model) {
-    // Calling of ConstantFolding in HETERO plugin is required because
-    // in some cases topology split is happening after constant subgraph.
-    // It may cause replacement of Constant by Parameter in such operations
-    // like Reshape/Transpose/Gather and lead to unexpected dynamism or exception
-    ov::pass::Manager manager;
-    manager.register_pass<ov::pass::ConstantFolding>();
-    manager.run_passes(model);
-
     ov::SupportedOpsMap query_model_result;
     bool user_set_affinities = false;
     // Get user defined affinity
diff --git a/src/plugins/hetero/tests/unit/subgraph_collector.cpp b/src/plugins/hetero/tests/unit/subgraph_collector.cpp
index 7eb3151cde7..8136b0909f9 100644
--- a/src/plugins/hetero/tests/unit/subgraph_collector.cpp
+++ b/src/plugins/hetero/tests/unit/subgraph_collector.cpp
@@ -491,4 +491,75 @@ TEST_F(SubgraphCollectorTest, submodel_with_different_affinity_parameter) {
     };
     auto supported_ops = supported_ops_with_affinity;
     ASSERT_NO_THROW(ov::hetero::mask_model_subgraphs_by_ops(m_model, supported_ops, false, "TEST"));
+}
+
+TEST_F(SubgraphCollectorTest, submodel_with_constant_subgraphs) {
+    auto input = std::make_shared<ov::op::v0::Parameter>(ov::element::u8, ov::PartialShape{1, 3, 16, 16});
+    input->set_friendly_name("input");
+    auto convert = std::make_shared<ov::op::v0::Convert>(input, ov::element::f32);
+    convert->set_friendly_name("convert");
+
+    auto constant1 = ov::op::v0::Constant::create(ov::element::f32, {}, {2.f});
+    constant1->set_friendly_name("constant1");
+
+    auto mul = std::make_shared<ov::op::v1::Multiply>(convert, constant1);
+    mul->set_friendly_name("mul");
+
+    auto constant2 = ov::op::v0::Constant::create(ov::element::i64, ov::Shape{4}, {0, 1, 3, 2});
+    constant2->set_friendly_name("constant2");
+    auto transpose = std::make_shared<ov::op::v1::Transpose>(mul, constant2);
+    transpose->set_friendly_name("transpose");
+    auto shapeOf = std::make_shared<ov::op::v0::ShapeOf>(transpose);
+    shapeOf->set_friendly_name("shapeOf");
+
+    auto reshape_val = ov::op::v0::Constant::create(ov::element::i64, ov::Shape{1}, {-1});
+    reshape_val->set_friendly_name("reshape_val");
+    auto reshape = std::make_shared<ov::op::v1::Reshape>(shapeOf, reshape_val, true);
+    reshape->set_friendly_name("reshape");
+
+    auto zero = ov::op::v0::Constant::create(ov::element::i64, ov::Shape{}, {0});
+    zero->set_friendly_name("zero");
+    auto one = ov::op::v0::Constant::create(ov::element::i64, ov::Shape{}, {1});
+    one->set_friendly_name("one");
+    auto gather = std::make_shared<ov::op::v8::Gather>(reshape, one, zero);
+    gather->set_friendly_name("gather");
+
+    auto result = std::make_shared<ov::op::v0::Result>(gather);
+    result->set_friendly_name("result");
+
+    auto model = std::make_shared<ov::Model>(ov::ResultVector{result}, ov::ParameterVector{input});
+    const std::map<std::string, std::string> supported_ops_with_affinity = {
+        {"input", "MOCK.0"},
+        {"convert", "MOCK.0"},
+        {"mul", "MOCK.0"},
+        {"constant1", "MOCK.0"},
+        {"constant2", "MOCK.0"},
+        {"transpose", "MOCK.1"},
+        {"shapeOf", "MOCK.0"},
+        {"reshape_val", "MOCK.0"},
+        {"reshape", "MOCK.1"},
+        {"zero", "MOCK.0"},
+        {"one", "MOCK.0"},
+        {"gather", "MOCK.1"},
+        {"result", "MOCK.0"},
+    };
+
+    auto supported_ops = supported_ops_with_affinity;
+    ov::hetero::SubgraphsVector ordered_subgraphs;
+    ov::hetero::SubgraphsMappingInfo actual_mapping_info;
+    std::tie(ordered_subgraphs, actual_mapping_info) = get_model_subgraphs(model, supported_ops, true, false);
+    for (const auto& subgraph : ordered_subgraphs) {
+        std::set<std::string> node_set;
+        auto sub_model = std::make_shared<ov::Model>(subgraph._results, subgraph._sinks, subgraph._parameters);
+        for (auto& node : sub_model->get_ordered_ops()) {
+            node_set.insert(node->get_friendly_name());
+        }
+        ASSERT_EQ(node_set.count("transpose"), node_set.count("constant2"));
+        ASSERT_EQ(node_set.count("reshape"), node_set.count("reshape_val"));
+        ASSERT_EQ(node_set.count("gather"), node_set.count("zero"));
+        ASSERT_EQ(node_set.count("gather"), node_set.count("one"));
+        if (node_set.count("transpose") || node_set.count("reshape") || node_set.count("gather")) {
+            ASSERT_EQ(subgraph._affinity, "MOCK.1");
+        }
+    }
 }
\ No newline at end of file