[GPU] Fix reshape optimization (#19270)

src/plugins/intel_gpu/src/graph/graph_optimizer/handle_reshape.cpp

@@ -2,6 +2,7 @@
 // SPDX-License-Identifier: Apache-2.0
 //
 
+#include "intel_gpu/runtime/internal_properties.hpp"
 #include "pass_manager.h"
 #include "program_helpers.h"
 #include "reshape_inst.h"
@@ -63,8 +64,20 @@ void handle_reshape::run(program& p) {
                 return;
 
             auto& out_node = node.get_users().front();
-            if (out_node->is_type<reshape>())
-                p.extract_and_remove(node);
+
+            if (!out_node->is_type<reshape>())
+                return;
+
+            const auto& out_reshape = out_node->as<reshape>();
+            // In case of new shape infer we should not shrink reshapes chain if first reshape changes input rank, e.g.
+            // [a, b] -> reshape1 -> [a1, b1, c1] -> reshape2 -> [a2, b2, 0] and any of the reshapes has special_zero=true
+            // Configuration above will fail if we remove reshape1 node as attempt to handle special zero will fail due to small rank of input
+            if (p.get_config().get_property(ov::intel_gpu::allow_new_shape_infer) &&
+                out_node->get_output_pshape().size() != node.get_input_pshape().size() &&
+                (out_reshape.get_primitive()->special_zero || node.get_primitive()->special_zero))
+                return;
+
+            p.extract_and_remove(node);
         });
     }
 

src/plugins/intel_gpu/src/graph/graph_optimizer/remove_redundant_reorders.cpp

@@ -651,6 +651,14 @@ void remove_redundant_reorders::run(program& p) {
 
         auto& reshape_input_node = dep_node.as<reshape>();
 
+        // In case of new shape infer we should not shrink reshapes chain if first reshape changes input rank, e.g.
+        // [a, b] -> reshape1 -> [a1, b1, c1] -> reshape2 -> [a2, b2, 0] and any of the reshapes has special_zero=true
+        // Configuration above will fail if we remove reshape1 node as attempt to handle special zero will fail due to small rank of input
+        if (p.get_config().get_property(ov::intel_gpu::allow_new_shape_infer) &&
+            reshape_node.get_output_pshape().size() != dep_node.get_input_pshape().size() &&
+            (reshape_node.get_primitive()->special_zero || reshape_input_node.get_primitive()->special_zero))
+            continue;
+
         if (reshape_node.is_dynamic())
             continue;
 

src/plugins/intel_gpu/src/graph/reshape.cpp

@@ -148,6 +148,8 @@ std::string reshape_inst::to_string(reshape_node const& node) {
     reshape_info.add("input id", input.id());
     reshape_info.add("output shape", desc->output_shape);
     reshape_info.add("output pshape", desc->output_partial_shape);
+    reshape_info.add("output pattern", desc->output_pattern);
+    reshape_info.add("special zero", desc->special_zero);
 
     node_info->add("reshape info", reshape_info);
     node_info->dump(primitive_description);

src/plugins/intel_gpu/tests/unit/passes/handle_reshape.cpp

@@ -55,6 +55,44 @@ TEST(handle_reshape, dont_remove_reshape_that_changes_rank) {
     ASSERT_TRUE(prog->get_node("reshape").can_be_optimized());
 }
 
+TEST(handle_reshape, dont_remove_reshape_that_changes_rank_chain) {
+    auto& engine = get_test_engine();
+    auto data0_layout = engine.allocate_memory({ ov::PartialShape{1, 1}, data_types::f16, format::bfyx });
+    auto data1_layout = engine.allocate_memory({ ov::PartialShape{1, 1, 1}, data_types::f16, format::bfyx });
+    auto in_layout = layout{ ov::PartialShape{2, 4}, data_types::f16, format::bfyx };
+
+    topology topology;
+    topology.add(input_layout("input", in_layout));
+    topology.add(data("data0", data0_layout));
+    topology.add(data("data1", data1_layout));
+    topology.add(eltwise("e1", input_info("input"), input_info("data0"), eltwise_mode::sum));
+    topology.add(reshape("reshape1", input_info("e1"), false, {2, 1, 4}, {2, 1, 4})); // 2d -> 3d
+    topology.add(reshape("reshape2", input_info("reshape1"), true, {1, 2, 0}, {1, 2, 0})); // swap 0 and 1 axis and keep 2nd axis as is
+    topology.add(eltwise("e2", input_info("reshape2"), input_info("data1"), eltwise_mode::sum));
+
+    ExecutionConfig config = get_test_default_config(engine);
+    config.set_property(ov::intel_gpu::allow_new_shape_infer(true));
+    config.set_property(ov::intel_gpu::optimize_data(true));
+    auto prog = program::build_program(engine, topology, config, false, true);
+
+    layout_optimizer lo(true);
+
+    program_wrapper::apply_opt_pass<prepare_primitive_fusing>(*prog, lo);
+    program_wrapper::apply_opt_pass<handle_reshape>(*prog);
+    program_wrapper::apply_opt_pass<remove_redundant_reorders>(*prog, lo);
+
+    for (auto& n : prog->get_processing_order()) {
+        n->recalc_output_layout(true);
+    }
+
+    ASSERT_NE(prog, nullptr);
+    ASSERT_TRUE(has_node(*prog, "reshape1"));
+    ASSERT_TRUE(has_node(*prog, "reshape2"));
+
+    ASSERT_TRUE(prog->get_node("reshape1").can_be_optimized());
+    ASSERT_TRUE(prog->get_node("reshape2").can_be_optimized());
+}
+
 TEST(handle_reshape, skip_reorder_node_to_split_when_onndnn_not_support) {
     // Onednn FC does not support fp32 input, fp16 weight. In such case, we need to ignore reorder_split from handle_reshape pass
     auto& engine = get_test_engine();