From 79b267033c24c1cef686a4fc1fba219a8323dcf5 Mon Sep 17 00:00:00 2001
From: Wilson Seok <wilson.seok@intel.com>
Date: Wed, 29 Mar 2023 04:25:22 +0900
Subject: [PATCH] [GPU] Fix program::replace() to copy duplicated connection
 from single constant (#16529)

* fix program::replace() to copy duplicated connection from single constant

* add unit test

* modified with review feedback
---
 src/plugins/intel_gpu/src/graph/program.cpp   |  6 ++-
 .../tests/passes/prepare_quantization.cpp     | 47 +++++++++++++++++++
 2 files changed, 52 insertions(+), 1 deletion(-)
 create mode 100644 src/plugins/intel_gpu/tests/passes/prepare_quantization.cpp

diff --git a/src/plugins/intel_gpu/src/graph/program.cpp b/src/plugins/intel_gpu/src/graph/program.cpp
index 5eff06f4276..bd04286be7e 100644
--- a/src/plugins/intel_gpu/src/graph/program.cpp
+++ b/src/plugins/intel_gpu/src/graph/program.cpp
@@ -931,9 +931,13 @@ void program::replace(program_node& old_node, program_node& new_node) {
     new_node.valid_output_layouts = old_node.valid_output_layouts;
 
     // copy old's dependencies
+    // First copy them from old node to new node
+    for (auto& dependency : old_node.dependencies) {
+        add_connection(*dependency.first, new_node);
+    }
+    // Second delete them from old node
     while (!old_node.dependencies.empty()) {
         auto& dep = old_node.dependencies.front().first;
-        add_connection(*dep, new_node);
         remove_connection(*dep, old_node);
     }
 
diff --git a/src/plugins/intel_gpu/tests/passes/prepare_quantization.cpp b/src/plugins/intel_gpu/tests/passes/prepare_quantization.cpp
new file mode 100644
index 00000000000..bbd918fde11
--- /dev/null
+++ b/src/plugins/intel_gpu/tests/passes/prepare_quantization.cpp
@@ -0,0 +1,47 @@
+// Copyright (C) 2023 Intel Corporation
+// SPDX-License-Identifier: Apache-2.0
+//
+
+#include "test_utils.h"
+
+#include "intel_gpu/runtime/engine.hpp"
+
+#include "intel_gpu/graph/network.hpp"
+#include "intel_gpu/graph/program.hpp"
+#include "data_inst.h"
+#include "quantize_inst.h"
+#include "pass_manager.h"
+#include "to_string_utils.h"
+
+#include "program_wrapper.h"
+
+#include <memory>
+
+using namespace cldnn;
+using namespace ::tests;
+
+TEST(prepare_quantization, program_replace_check_num_of_nodes) {
+    auto& engine = get_test_engine();
+    auto data0_layout = engine.allocate_memory({ ov::PartialShape{1}, data_types::f32, format::bfyx });
+    auto data1_layout = engine.allocate_memory({ ov::PartialShape{1}, data_types::f32, format::bfyx });
+    auto in_layout = layout{ ov::PartialShape::dynamic(0), data_types::f32, format::bfyx };
+
+    topology topology;
+    topology.add(input_layout("input", in_layout));
+    topology.add(data("input_low", data0_layout));
+    topology.add(data("input_high", data1_layout));
+    topology.add(quantize("quantize", input_info("input"), input_info("input_low"), input_info("input_high"), input_info("input_low"), input_info("input_high"), 256, data_types::f32));
+
+    ExecutionConfig config;
+    config.set_property(ov::intel_gpu::optimize_data(true));
+    auto prog = program::build_program(engine, topology, config, false, true);
+
+    ASSERT_NE(prog, nullptr);
+    ASSERT_TRUE(prog->get_node("quantize").get_dependencies().size() == 5);
+
+    layout_optimizer lo(true);
+
+    program_wrapper::apply_opt_pass<prepare_quantization>(*prog);
+
+    ASSERT_TRUE(prog->get_node("quantize").get_dependencies().size() == 9);
+}
\ No newline at end of file