[CPU] Ensure uniquness of a parent edge child port (#21965)

This commit is contained in:
Egor Duplenskii 2024-01-05 08:55:13 +01:00 committed by GitHub
parent 058f0206e2
commit 90ee1004a4
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23
2 changed files with 11 additions and 7 deletions

View File

@ -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);

View File

@ -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());