IVGCVSW-7155: Fix Slot replacement during UpdateSubgraphViewSlotPointers

 * Only update boundary slots on actual subgraphview
 * Previously all slots from replacement subgraph added even if internal

Signed-off-by: Francis Murtagh <francis.murtagh@arm.com>
Signed-off-by: Matthew Bentham <matthew.bentham@arm.com>
Change-Id: Ic9ef9fc41ad248838d1c019dd0368378c3119648

2 files changed, 132 insertions, 17 deletions

diff --git a/src/armnn/SubgraphView.cpp b/src/armnn/SubgraphView.cpp
index fef6390bf2..e16c5c7b3a 100644
--- a/src/armnn/SubgraphView.cpp
+++ b/src/armnn/SubgraphView.cpp
@@ -11,6 +11,7 @@
 #include <armnn/utility/NumericCast.hpp>
 #include <armnn/utility/PolymorphicDowncast.hpp>
 
+#include <fmt/format.h>
 #include <utility>
 
 namespace armnn
@@ -539,37 +540,33 @@ void SubgraphView::UpdateSubgraphViewSlotPointers(SubgraphView& patternSubgraph,
 {
     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())
+    for (unsigned long idx = 0; idx < patternSubgraph.GetIInputSlots().size(); idx++)
     {
+        IInputSlot *slot = patternSubgraph.GetIInputSlots()[idx];
         inputSlotPosition = std::find(m_IInputSlots.begin(), m_IInputSlots.end(), slot);
         if (inputSlotPosition != m_IInputSlots.end())
         {
             m_IInputSlots.erase(inputSlotPosition);
+
+            // while here, with correct position, add in replacement InputSlot from the substituteSubgraph
+            m_IInputSlots.insert(inputSlotPosition, substituteSubgraph.GetIInputSlots()[idx]);
         }
     }
 
     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())
+    for (unsigned long idx = 0; idx < patternSubgraph.GetIOutputSlots().size(); idx++)
     {
+        IOutputSlot *slot = patternSubgraph.GetIOutputSlots()[idx];
         outputSlotPosition = std::find(m_IOutputSlots.begin(), m_IOutputSlots.end(), slot);
         if (outputSlotPosition != m_IOutputSlots.end())
         {
             m_IOutputSlots.erase(outputSlotPosition);
+
+            // while here, with correct position, add in replacement OutputSlot from the substituteSubgraph
+            m_IOutputSlots.insert(outputSlotPosition, substituteSubgraph.GetIOutputSlots()[idx]);
         }
     }
-
-    // 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)
@@ -580,6 +577,28 @@ void SubgraphView::SubstituteSubgraph(SubgraphView& patternSubgraph, const Subgr
                                    "Call this function on SubgraphView returned from SubgraphView::GetWorkingCopy()");
     }
 
+    auto numPatternInputs = patternSubgraph.GetIInputSlots().size();
+    auto numSubInputs = substituteSubgraph.GetIInputSlots().size();
+    if (numPatternInputs != numSubInputs)
+    {
+        throw armnn::InvalidArgumentException(
+           fmt::format("Number of InputSlots on substitute SubgraphView ({}) must equal the number of"
+                       " InputSlots on pattern SubgraphView ({})",
+                       numSubInputs,
+                       numPatternInputs));
+    }
+
+    auto numPatternOutputs = patternSubgraph.GetIOutputSlots().size();
+    auto numSubOutputs = substituteSubgraph.GetIOutputSlots().size();
+    if (numPatternOutputs != numSubOutputs)
+    {
+        throw armnn::InvalidArgumentException(
+           fmt::format("Number of OutputSlots on substitute SubgraphView ({}) must equal the number of"
+                       " OutputSlots on pattern SubgraphView ({})",
+                       numSubOutputs,
+                       numPatternOutputs));
+    }
+
     // Add substitute layer to the Main graph i.e. graph in p_WorkingCopyImpl
     auto workingCopyGraph = &p_WorkingCopyImpl->m_Graph;
     substituteSubgraph.ForEachIConnectableLayer([workingCopyGraph](IConnectableLayer* iConnectableLayer)
diff --git a/src/armnn/test/SubgraphViewTests.cpp b/src/armnn/test/SubgraphViewTests.cpp
index 48f4a7fc01..e0fd5fe7c1 100644
--- a/src/armnn/test/SubgraphViewTests.cpp
+++ b/src/armnn/test/SubgraphViewTests.cpp
@@ -2603,13 +2603,109 @@ TEST_CASE("MultipleInputMultipleOutputSlots_SubstituteGraph")
     SubgraphView::IOutputSlots expectedOutputSlots = {activationOutputSlot,
                                                       &standInLayer->GetOutputSlot(0)};
 
-    // expecting addition input slot to be changed with standin input slot
     // convInputSlot MUST remain as an expected input slot
-    SubgraphView::IInputSlots expectedInputSlots = {convInputSlot,
-                                                    &standInLayer->GetInputSlot(0)};
+    SubgraphView::IInputSlots expectedInputSlots = {convInputSlot};
 
     CHECK(expectedOutputSlots == viewCopy.GetIOutputSlots());
     CHECK(expectedInputSlots == viewCopy.GetIInputSlots());
 }
 
+
+
+TEST_CASE("MultipleInputMultipleOutputSlots_SubstituteGraphNewSlots")
+{
+    //         Construct graph      //
+    //                              //
+    //               input          //
+    //                 |            //
+    //              conv2d          //
+    //                 |            //
+    //      const    relu           //
+    //         \   /     \          //
+    //          add    output       //
+    //           |                  //
+    //         output               //
+    //                              //
+    // SubgraphView layers: conv2d relu add const
+
+    Graph graph;
+    Layer* inputLayer = graph.AddLayer<InputLayer>(0, "input");
+
+    Layer* convLayer = graph.AddLayer<Convolution2dLayer>(Convolution2dDescriptor(), "conv");
+    Layer* reluLayer = graph.AddLayer<ActivationLayer>(ActivationDescriptor(), "activation");
+    Layer* constLayer = graph.AddLayer<ConstantLayer>("const");
+    Layer* addLayer = graph.AddLayer<AdditionLayer>("add");
+
+    Layer* outputLayer1 = graph.AddLayer<OutputLayer>(0, "output1");
+    Layer* outputLayer2 = graph.AddLayer<OutputLayer>(1, "output2");
+
+    inputLayer->GetOutputSlot(0).Connect(convLayer->GetInputSlot(0));
+    convLayer->GetOutputSlot(0).Connect(reluLayer->GetInputSlot(0));
+    constLayer->GetOutputSlot(0).Connect(addLayer->GetInputSlot(0));
+    reluLayer->GetOutputSlot(0).Connect(addLayer->GetInputSlot(1));
+    reluLayer->GetOutputSlot(0).Connect(outputLayer1->GetInputSlot(0));
+    addLayer->GetOutputSlot(0).Connect(outputLayer2->GetInputSlot(0));
+
+    // main subgraph creation
+    SubgraphView::IInputSlots inputSlots = {&convLayer->GetInputSlot(0)};
+    SubgraphView::IOutputSlots outputSlots = {&reluLayer->GetOutputSlot(0), &addLayer->GetOutputSlot(0)};
+    auto view = CreateSubgraphViewFrom({convLayer, reluLayer, addLayer, constLayer},
+                                       std::move(inputSlots),
+                                       std::move(outputSlots));
+
+    // need to call GetWorkingCopy() in order for SubstituteSubgraph() to work later on
+    SubgraphView viewCopy = view->GetWorkingCopy();
+    IConnectableLayer* addCopyLayer = nullptr;
+    for (auto layer : viewCopy.GetIConnectableLayers())
+    {
+        // GetWorkingCopy() has caused address pointer of convolution layer to change.
+        // Finding new address pointer...
+        if (layer->GetType() == LayerType::Addition)
+        {
+            addCopyLayer = layer;
+        }
+    }
+
+    // substitute subgraph creation
+    OptimizationViews optimizationViews;
+    IConnectableLayer* standInLayer = optimizationViews.GetINetwork()->AddStandInLayer(StandInDescriptor(2,2),
+                                                                                       "standin");
+    // Extra inputSlot (needed explicit use of vector to prevent ambiguity)
+    auto substituteSubgraph1 = CreateSubgraphViewFrom({standInLayer},
+                                                      {&standInLayer->GetInputSlot(0),
+                                                       &standInLayer->GetInputSlot(1)},
+                                                      std::vector<IOutputSlot*>{&standInLayer->GetOutputSlot(0)});
+    // Extra outputSlot
+    auto substituteSubgraph2 = CreateSubgraphViewFrom({standInLayer},
+                                                      {&standInLayer->GetInputSlot(0)},
+                                                      std::vector<IOutputSlot*>{&standInLayer->GetOutputSlot(0),
+                                                       &standInLayer->GetOutputSlot(1)});
+
+    // pattern subgraph creation
+    IConnectableLayer* constCopyLayer = &addCopyLayer->GetInputSlot(0).GetConnection()->GetOwningIConnectableLayer();
+
+    // Mismatched number of input slots (needed explicit use of vector to prevent ambiguity)
+    SubgraphView::SubgraphViewPtr patternSubgraph1 =
+            CreateSubgraphViewFrom({addCopyLayer, constCopyLayer},
+                                   {&addCopyLayer->GetInputSlot(0)},
+                                   std::vector<IOutputSlot*>{&addCopyLayer->GetOutputSlot(0)});
+
+    // Mismatched number of output slots
+    SubgraphView::SubgraphViewPtr patternSubgraph2 = CreateSubgraphViewFrom({addCopyLayer, constCopyLayer},
+                                                                            {&addCopyLayer->GetInputSlot(0)},
+                                                                            {&addCopyLayer->GetOutputSlot(0)});
+
+
+
+
+    // Ensure that a substitute subgraphView has same number of InputSlots as the pattern subgraphView
+    CHECK_THROWS_AS(viewCopy.SubstituteSubgraph(*patternSubgraph1, *substituteSubgraph1),
+                    armnn::InvalidArgumentException);
+
+    // Ensure that a substitute subgraphView has same number of OutputSlots as the pattern subgraphView
+    CHECK_THROWS_AS(viewCopy.SubstituteSubgraph(*patternSubgraph2, *substituteSubgraph2),
+                    armnn::InvalidArgumentException);
+
+}
+
 }