From 90ee1004a4597dd3e00a206ad141d1b622e8dde3 Mon Sep 17 00:00:00 2001
From: Egor Duplenskii <egor.duplensky@gmail.com>
Date: Fri, 5 Jan 2024 08:55:13 +0100
Subject: [PATCH] [CPU] Ensure uniquness of a parent edge child port (#21965)

---
 src/plugins/intel_cpu/src/graph.cpp           |  7 ++++---
 src/plugins/intel_cpu/src/graph_optimizer.cpp | 11 +++++++----
 2 files changed, 11 insertions(+), 7 deletions(-)

diff --git a/src/plugins/intel_cpu/src/graph.cpp b/src/plugins/intel_cpu/src/graph.cpp
index 6a888ecc15f..888e4094c5c 100644
--- a/src/plugins/intel_cpu/src/graph.cpp
+++ b/src/plugins/intel_cpu/src/graph.cpp
@@ -29,6 +29,7 @@
 #include "nodes/input.h"
 #include "nodes/reorder.h"
 #include "nodes/memory.hpp"
+#include "openvino/core/except.hpp"
 #include "openvino/core/model.hpp"
 #include "openvino/core/node.hpp"
 #include "utils/debug_capabilities.h"
@@ -1422,9 +1423,9 @@ void Graph::CreateEdge(const NodePtr& parent,
                        int childPort) {
     assert(parentPort >= 0 && childPort >= 0);
     assert(std::none_of(child->getParentEdges().begin(), child->getParentEdges().end(),
-                       [&childPort](const EdgeWeakPtr& edge){
-                           return edge.lock()->getOutputNum() == childPort;
-                       }));
+                        [&childPort](const EdgeWeakPtr& edge){
+                            return edge.lock()->getOutputNum() == childPort;
+                        }));
 
     auto edge = std::make_shared<Edge>(parent, child, parentPort, childPort);
 
diff --git a/src/plugins/intel_cpu/src/graph_optimizer.cpp b/src/plugins/intel_cpu/src/graph_optimizer.cpp
index 6552d74fb5c..c5a55909572 100644
--- a/src/plugins/intel_cpu/src/graph_optimizer.cpp
+++ b/src/plugins/intel_cpu/src/graph_optimizer.cpp
@@ -2612,7 +2612,7 @@ void GraphOptimizer::reshapeRnnSeq(Graph &graph) {
             unsqueeze->set_friendly_name(parentNode->getName() + "_abc_a1bc_" + std::to_string(j));
 
             const auto cpuUnsqueeze = std::make_shared<Reshape>(unsqueeze, graph.getGraphContext());
-            graph.InsertNode(parentNode, childNode, cpuUnsqueeze, edge->getInputNum(), edge->getOutputNum(), false);
+            graph.InsertNode(edge, cpuUnsqueeze, false);
 
             const auto cpuConstant = std::make_shared<node::Input>(secondInput, graph.getGraphContext());
             graph.AddNode(cpuConstant);
@@ -2790,14 +2790,16 @@ void GraphOptimizer::MatchSdpaKvCache(Graph &graph) {
         if (!memInputNode->getParentEdges().empty()) {
             auto parentEdge = memInputNode->getParentEdgeAt(0);
             auto parent = parentEdge->getParent();
-            graph.CreateEdge(parent, memInputSdpa, parentEdge->getInputNum(), 0);
+            const auto inputNum = parentEdge->getInputNum();
             graph.RemoveEdge(parentEdge);
+            graph.CreateEdge(parent, memInputSdpa, inputNum, 0);
         }
 
         for (auto&& edge : memInputNode->getChildEdgesAtPort(0)) {
             auto child = edge->getChild();
-            graph.CreateEdge(memInputSdpa, child, 0, edge->getOutputNum());
+            const auto outputNum = edge->getOutputNum();
             graph.RemoveEdge(edge);
+            graph.CreateEdge(memInputSdpa, child, 0, outputNum);
         }
 
         //create a stub memory output
@@ -2812,8 +2814,9 @@ void GraphOptimizer::MatchSdpaKvCache(Graph &graph) {
             graph.getGraphContext());
 
         auto memOutputEdge = memOutput.getParentEdgeAt(0);
-        graph.CreateEdge(sdpa, memOutputStub, memOutputEdge->getInputNum(), 0);
+        const auto inputNum = memOutputEdge->getInputNum();
         graph.RemoveEdge(memOutputEdge);
+        graph.CreateEdge(sdpa, memOutputStub, inputNum, 0);
 
         memInputSdpa->registerOutputNode(memOutputStub.get());