[GPU] Fixed removal of reorder connected to non-default out port (#18826)

src/plugins/intel_gpu/include/intel_gpu/graph/program.hpp

@@ -351,6 +351,7 @@ private:
     void rename(program_node& node, primitive_id const& new_id);
     void swap_names(program_node& node1, program_node& node2);
     void replace_all_usages(program_node& old_node, program_node& new_node, bool remove_if_dangling = true);
+    void replace_all_usages(program_node& old_node, std::pair<program_node*, int32_t> new_node, bool remove_if_dangling = true);
 
     // old_node - node which will be replaced
     // new_node - node which will replace the old one

src/plugins/intel_gpu/src/graph/include/program_node.h

@@ -194,9 +194,11 @@ public:
 
     // replaces idx-th dependency of 'this' with 'new_dep', calls program::remove_if_dangling(old_dep)
     void replace_dependency(size_t idx, program_node& new_dep, bool remove_if_dangling = true);
+    void replace_dependency(size_t idx, std::pair<program_node*, int32_t> new_dep, bool remove_if_dangling = true);
     // searches for 'old_dep' in dependencies list of 'this' and replaces it with 'new_dep', calls
     // program::remove_if_dangling(old_dep)
     void replace_dependency(program_node const& old_dep, program_node& new_dep, bool remove_if_dangling = true);
+    void replace_dependency(program_node const& old_dep, std::pair<program_node*, int32_t> new_dep, bool remove_if_dangling = true);
 
     std::vector<primitive_id> get_dependencies_ids() const;
 

src/plugins/intel_gpu/src/graph/network.cpp

@@ -758,6 +758,7 @@ void network::reset_execution(bool wait) {
 }
 
 event::ptr network::set_input_data(const primitive_id& id, memory::ptr data) {
+    GPU_DEBUG_TRACE_DETAIL << "Set input " << id << " " << data->get_layout().to_short_string() << std::endl;
     std::shared_ptr<primitive_inst> primitive_inst;
 
     primitive_inst = find_primitive(id);
@@ -900,6 +901,7 @@ network::output_chains_map::iterator network::add_output_chain(std::shared_ptr<p
 }
 
 std::vector<event::ptr> network::set_output_memory(const primitive_id& id, memory::ptr mem_new) {
+    GPU_DEBUG_TRACE_DETAIL << "Set output " << id << " " << mem_new->get_layout().to_short_string() << std::endl;
     std::shared_ptr<primitive_inst> p_inst;
     std::vector<event::ptr> ret_ev;
     p_inst = find_primitive(id);
@@ -970,7 +972,7 @@ std::string network::get_primitive_info(const primitive_id& id) const {
 bool network::is_cpu_impl(const primitive_id& id) const {
     auto prim_inst = find_primitive(id);
 
-    OPENVINO_ASSERT(prim_inst, "[GPU] Can't get implementation type, since topology",
+    OPENVINO_ASSERT(prim_inst, "[GPU] Can't get implementation type, since topology ",
                                "doesn't contain primitive with requested id: ", id);
 
     return prim_inst->get_impl() ? prim_inst->get_impl()->is_cpu() : true;

src/plugins/intel_gpu/src/graph/program.cpp

@@ -901,6 +901,10 @@ void program::swap_names(program_node& node1, program_node& node2) {
 }
 
 void program::replace_all_usages(program_node& old_node, program_node& new_node, bool remove_if_dangling) {
+    return replace_all_usages(old_node, std::make_pair(&new_node, 0), remove_if_dangling);
+}
+
+void program::replace_all_usages(program_node& old_node, std::pair<program_node*, int32_t> new_node, bool remove_if_dangling) {
     // We need a copy of users of old_node because old_node may be removed when doing replace_dependency()
     const std::list<program_node*> users(old_node.users);
     auto itr = users.begin();
@@ -1013,7 +1017,8 @@ bool program::extract(program_node& node) {
         outputs.push_back(&prev);
     }
 
-    auto& input = node.get_dependency(0);
+    auto input_with_port = node.get_dependency_with_port(0);
+    auto& input = *input_with_port.first;
 
     // update primitive_map of loop primitive,
     // if extracted node is input of loop
@@ -1040,7 +1045,7 @@ bool program::extract(program_node& node) {
     node.dependencies.clear();
 
     if (!node.is_endpoint())
-        replace_all_usages(node, input, false);
+        replace_all_usages(node, input_with_port, false);
 
     if (std::find(processing_order.begin(), processing_order.end(), &node) != processing_order.end())
         processing_order.erase(&node);

src/plugins/intel_gpu/src/graph/program_dump_graph.cpp

@@ -251,7 +251,9 @@ void dump_graph_init(std::ofstream& graph,
             }
             if (it == user->get_dependencies().end())
                 doubled = false;
-            graph << "    " << get_node_id(node) << " -> " << get_node_id(user);
+            graph << "    " << get_node_id(node) << " -> " << get_node_id(user)
+                  << " [label=\"" << it->second << " -> " << std::distance(user->get_dependencies().begin(), it) << "\"]";
+
 
             bool data_flow = node->is_in_data_flow() && user->is_in_data_flow();
             if (data_flow) {

src/plugins/intel_gpu/src/graph/program_node.cpp

@@ -45,15 +45,15 @@ program_node::program_node(std::shared_ptr<primitive> prim, program& prog)
     }
 }
 
-void program_node::replace_dependency(size_t idx, program_node& new_dep, bool remove_if_dangling) {
+void program_node::replace_dependency(size_t idx, std::pair<program_node*, int32_t> new_dep, bool remove_if_dangling) {
     if (idx >= dependencies.size())
         return;
-    if (dependencies[idx].first == &new_dep)
+    if (dependencies[idx].first == new_dep.first)
         return;
 
     if (is_type<loop>()) {
         loop_node& loop = *this;
-        loop.update_primitive_map(dependencies[idx].first->id(), new_dep.id(), true);
+        loop.update_primitive_map(dependencies[idx].first->id(), new_dep.first->id(), true);
     }
 
     auto it = std::find(dependencies[idx].first->users.begin(), dependencies[idx].first->users.end(), this);
@@ -64,16 +64,25 @@ void program_node::replace_dependency(size_t idx, program_node& new_dep, bool re
     if (remove_if_dangling)
         myprog.remove_if_dangling(*dependencies[idx].first);
 
-    dependencies[idx].first = &new_dep;
-    new_dep.users.push_back(this);
+    dependencies[idx].first = new_dep.first;
+    dependencies[idx].second = new_dep.second;
+    new_dep.first->users.push_back(this);
 }
 
-void program_node::replace_dependency(program_node const& old_dep, program_node& new_dep, bool remove_if_dangling) {
+void program_node::replace_dependency(size_t idx, program_node& new_dep, bool remove_if_dangling) {
+    return replace_dependency(idx, std::make_pair(&new_dep, 0), remove_if_dangling);
+}
+
+void program_node::replace_dependency(program_node const& old_dep, std::pair<program_node*, int32_t> new_dep, bool remove_if_dangling) {
     for (size_t i = 0; i < dependencies.size(); ++i)
         if (dependencies[i].first == &old_dep)
             return replace_dependency(i, new_dep, remove_if_dangling);
 }
 
+void program_node::replace_dependency(program_node const& old_dep, program_node& new_dep, bool remove_if_dangling) {
+    return replace_dependency(old_dep, std::make_pair(&new_dep, 0), remove_if_dangling);
+}
+
 std::vector<primitive_id> program_node::get_dependencies_ids() const {
     std::vector<primitive_id> dep_ids;
     for (auto& dependency : dependencies) dep_ids.push_back(dependency.first->get_primitive()->id);

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

@@ -21,6 +21,7 @@
 #include "mvn_inst.h"
 #include "concatenation_inst.h"
 #include "shape_of_inst.h"
+#include "arg_max_min_inst.h"
 #include "gather_inst.h"
 #include "pass_manager.h"
 #include "to_string_utils.h"
@@ -379,3 +380,95 @@ TEST(remove_redundant_reorders, not_to_fuse_concat_with_reorder_inside_shape_of_
 
     ASSERT_EQ(concat_layout.data_type, data_types::f32);
 }
+
+TEST(remove_redundant_reorders, reorder_of_non_default_port) {
+    static const int32_t x_size = 2, y_size = 2, feature_num = 4, batch_num = 2;
+    auto& engine = get_test_engine();
+    const int top_k = 2;
+    auto input = engine.allocate_memory({ data_types::f32, format::bfyx,{ batch_num, feature_num, x_size , y_size } });
+    auto top_k_input = engine.allocate_memory({ data_types::f32, format::bfyx,{ 1, 1, 1 , 1 } });
+
+    topology topology;
+    topology.add(input_layout("input", input->get_layout()));
+    topology.add(cldnn::data("const", {top_k_input}));
+    auto arg_max_min_prim = arg_max_min("arg_max",
+                                        { input_info("input"), input_info("const") },
+                                        ov::op::TopKMode::MAX, top_k,
+                                        0,
+                                        ov::op::TopKSortType::SORT_VALUES,
+                                        false,
+                                        false,
+                                        padding(),
+                                        data_types::f32,
+                                        2);
+    arg_max_min_prim.output_paddings = {padding(), padding()};
+    arg_max_min_prim.output_data_types = {optional_data_type{data_types::f32}, optional_data_type{data_types::f32}};
+    topology.add(arg_max_min_prim);
+    topology.add(reorder("reorder_1", input_info("arg_max", 0), format::bfyx, data_types::f32));
+    topology.add(reorder("reorder_2", input_info("arg_max", 1), format::bfyx, data_types::f32));
+    topology.add(permute("permute_1", input_info("reorder_1", 0), {0, 1, 2, 3}, padding()));
+    topology.add(permute("permute_2", input_info("reorder_2", 0), {0, 1, 2, 3}, padding()));
+    topology.add(concatenation("concat", { input_info("permute_1"), input_info("permute_2") }, 0));
+
+    std::vector<float> input_vec = {
+            //y0x0 y0x1 y1x0 y1x1
+            /*b0f0*/0.1f, 0.2f, 0.3f,  0.4f,
+            /*b0f1*/0.5f, 0.6f,  0.7f, 0.8f,
+            /*b0f2*/0.9f, 1.0f,  1.1f, 1.2f,
+            /*b0f3*/1.3f, 1.4f,  1.5f, 1.6f,
+
+            /*b1f0*/2.1f, 2.2f, 2.3f, 2.4f,
+            /*b1f1*/2.5f, 2.6f, 2.7f, 2.8f,
+            /*b1f2*/2.9f, 3.0f, 3.1f, 3.2f,
+            /*b1f3*/3.3f, 3.4f, 3.5f, 3.6f,
+    };
+
+    std::vector<float> ref_result = {
+            /*indexes*/
+            /*b0*/
+            1, 1, 1, 1, 1, 1, 1, 1,
+            1, 1, 1, 1, 1, 1, 1, 1,
+            /*b1*/
+            0, 0, 0, 0, 0, 0, 0, 0,
+            0, 0, 0, 0, 0, 0, 0, 0,
+            /**values*/
+            /*b0*/
+            2.1f, 2.2f, 2.3f, 2.4f,
+            2.5f, 2.6f, 2.7f, 2.8f,
+            2.9f, 3.0f, 3.1f, 3.2f,
+            3.3f, 3.4f, 3.5f, 3.6f,
+            /*b1*/
+            0.1f, 0.2f, 0.3f,  0.4f,
+            0.5f, 0.6f,  0.7f, 0.8f,
+            0.9f, 1.0f,  1.1f, 1.2f,
+            1.3f, 1.4f,  1.5f, 1.6f,
+    };
+
+    set_values(input, input_vec);
+
+    ExecutionConfig config = get_test_default_config(engine);
+    config.set_property(ov::intel_gpu::optimize_data(true));
+    config.set_property(ov::intel_gpu::allow_new_shape_infer(true));
+
+    network network(engine, topology, config);
+
+    ASSERT_NE(network.get_program(), nullptr);
+    ASSERT_FALSE(has_node_with_type<reorder>(*network.get_program()));
+
+    network.set_input_data("input", input);
+    auto outputs = network.execute();
+
+    ASSERT_EQ(outputs.size(), size_t(1));
+    ASSERT_EQ(outputs.begin()->first, "concat");
+    const int out_size = y_size * feature_num * x_size * top_k * 2;
+    auto output = outputs.at("concat").get_memory();
+    cldnn::mem_lock<float> output_ptr(output, get_test_stream());
+
+    float out_buffer[out_size];
+    for (uint32_t i = 0; i < out_size; i++) {
+        out_buffer[i] = get_value<float>(output_ptr.data(), i);
+    }
+    for (int i = 0; i < out_size; i++) {
+        ASSERT_EQ(out_buffer[i], ref_result[i]);
+    }
+}