From 83664e7b3c803248dc18891590b38900e2cf5a01 Mon Sep 17 00:00:00 2001
From: Cathal Corbett <cathal.corbett@arm.com>
Date: Mon, 22 Aug 2022 17:59:32 +0100
Subject: IVGCVSW-7155 SubgraphView::SubstituteSubgraph IOutputSlots
 incorrectly overridden

Signed-off-by: Cathal Corbett <cathal.corbett@arm.com>
Change-Id: If594e291951a5f9ed1957a19a971c498f6e7843f
---
 src/armnn/SubgraphView.cpp | 41 +++++++++++++++++++++++++++++++++++++++--
 1 file changed, 39 insertions(+), 2 deletions(-)

(limited to 'src/armnn/SubgraphView.cpp')

diff --git a/src/armnn/SubgraphView.cpp b/src/armnn/SubgraphView.cpp
index 804ff731fb..b48529c523 100644
--- a/src/armnn/SubgraphView.cpp
+++ b/src/armnn/SubgraphView.cpp
@@ -525,6 +525,44 @@ void SubgraphView::SubstituteSubgraph(SubgraphView& subgraph, IConnectableLayer*
     SubstituteSubgraph(subgraph, substituteSubgraph);
 }
 
+void SubgraphView::UpdateSubgraphViewSlotPointers(SubgraphView& patternSubgraph,
+                                                  const SubgraphView& substituteSubgraph)
+{
+    std::vector<IInputSlot*>::iterator inputSlotPosition;
+    // search for and erase any InputSlots that appear in the WorkingCopy that match those in the PatternSubgraph
+    for (auto slot : patternSubgraph.GetIInputSlots())
+    {
+        inputSlotPosition = std::find(m_IInputSlots.begin(), m_IInputSlots.end(), slot);
+        if (inputSlotPosition != m_IInputSlots.end())
+        {
+            m_IInputSlots.erase(inputSlotPosition);
+        }
+    }
+
+    std::vector<IOutputSlot*>::iterator outputSlotPosition;
+    // search for and erase any OutputSlots that appear in the WorkingCopy that match those in the PatternSubgraph
+    for (auto slot : patternSubgraph.GetIOutputSlots())
+    {
+        outputSlotPosition = std::find(m_IOutputSlots.begin(), m_IOutputSlots.end(), slot);
+        if (outputSlotPosition != m_IOutputSlots.end())
+        {
+            m_IOutputSlots.erase(outputSlotPosition);
+        }
+    }
+
+    // append InputSlots from the SubstituteSubgraph to the WorkingCopy m_IInputSlots vector variable
+    // at the position in the vector where PatternSubgraph InputSlots were last removed
+    m_IInputSlots.insert(inputSlotPosition,
+                         std::make_move_iterator(substituteSubgraph.m_IInputSlots.begin()),
+                         std::make_move_iterator(substituteSubgraph.m_IInputSlots.end()));
+
+    // append OutputSlots from the SubstituteSubgraph to the WorkingCopy m_IOutputSlots vector variable
+    // at the position in the vector where PatternSubgraph OutputSlots were last removed
+    m_IOutputSlots.insert(outputSlotPosition,
+                         std::make_move_iterator(substituteSubgraph.m_OutputSlots.begin()),
+                         std::make_move_iterator(substituteSubgraph.m_OutputSlots.end()));
+}
+
 void SubgraphView::SubstituteSubgraph(SubgraphView& patternSubgraph, const SubgraphView& substituteSubgraph)
 {
     if (!p_WorkingCopyImpl)
@@ -556,8 +594,7 @@ void SubgraphView::SubstituteSubgraph(SubgraphView& patternSubgraph, const Subgr
     workingCopyGraph->ReplaceSubgraphConnections(patternSubgraph, substituteSubgraph);
 
     // Update input/outputSlot pointers
-    m_IInputSlots = std::move(substituteSubgraph.m_IInputSlots);
-    m_IOutputSlots = std::move(substituteSubgraph.m_IOutputSlots);
+    UpdateSubgraphViewSlotPointers(patternSubgraph, substituteSubgraph);
 
     // Delete the old layers.
     workingCopyGraph->EraseSubgraphLayers(patternSubgraph);
-- 
cgit v1.2.1