diff options
Diffstat (limited to 'src')
-rw-r--r-- | src/armnn/Graph.cpp | 18 | ||||
-rw-r--r-- | src/armnn/test/SubgraphViewTests.cpp | 29 |
2 files changed, 41 insertions, 6 deletions
diff --git a/src/armnn/Graph.cpp b/src/armnn/Graph.cpp index b5769f75f3..e5d123830c 100644 --- a/src/armnn/Graph.cpp +++ b/src/armnn/Graph.cpp @@ -497,13 +497,19 @@ void Graph::ReplaceSubgraphConnections(const SubgraphView& subgraph, const Subgr IInputSlot* subgraphInputSlot = subgraphInputSlots.at(inputSlotIdx); ARMNN_ASSERT(subgraphInputSlot); - IOutputSlot* connectedOutputSlot = subgraphInputSlot->GetConnection(); - ARMNN_ASSERT(connectedOutputSlot); - connectedOutputSlot->Disconnect(*subgraphInputSlot); + // Only disconnect if the InputSlot has a connection, this might not be the case when + // dealing with working copies of SubgraphViews + // Note: we don't need this check for OutputSlot as it iterates over a vector of valid connections + if (subgraphInputSlot->GetConnection()) + { + IOutputSlot* connectedOutputSlot = subgraphInputSlot->GetConnection(); + ARMNN_ASSERT(connectedOutputSlot); + connectedOutputSlot->Disconnect(*subgraphInputSlot); - IInputSlot* substituteInputSlot = substituteSubgraphInputSlots.at(inputSlotIdx); - ARMNN_ASSERT(substituteInputSlot); - connectedOutputSlot->Connect(*substituteInputSlot); + IInputSlot* substituteInputSlot = substituteSubgraphInputSlots.at(inputSlotIdx); + ARMNN_ASSERT(substituteInputSlot); + connectedOutputSlot->Connect(*substituteInputSlot); + } } // Step 2: process output slots diff --git a/src/armnn/test/SubgraphViewTests.cpp b/src/armnn/test/SubgraphViewTests.cpp index 4ce67b0fec..9bb5e69bbb 100644 --- a/src/armnn/test/SubgraphViewTests.cpp +++ b/src/armnn/test/SubgraphViewTests.cpp @@ -2063,6 +2063,35 @@ TEST_CASE("SubgraphViewWorkingCopySubstituteSubgraph") CHECK_THROWS_AS(workingCopy.GetWorkingCopy(), Exception); } +TEST_CASE("SubgraphViewPartialWorkingCopySubstituteSubgraph") +{ + Graph graph; + + auto input = graph.AddLayer<InputLayer>(0, "Input"); + auto activation = graph.AddLayer<ActivationLayer>(ActivationDescriptor{}, "Activation"); + auto output = graph.AddLayer<OutputLayer>(1, "Output"); + + input->GetOutputSlot(0).Connect(activation->GetInputSlot(0)); + activation->GetOutputSlot(0).Connect(output->GetInputSlot(0)); + + //Add in out of order + auto view = CreateSubgraphViewFrom({activation}, + {&activation->GetInputSlot(0)}, + {&activation->GetOutputSlot(0)}); + + auto workingCopy = view->GetWorkingCopy(); + + // First (and only) layer in the subgraph is the Activation + CHECK(std::string((*workingCopy.beginIConnectable())->GetName()) == "Activation"); + + // Substitute the "Activation" layer for an equivalent layer + auto activation2 = graph.AddLayer<ActivationLayer>(ActivationDescriptor{}, "Activation2"); + SubgraphView pattern(*workingCopy.beginIConnectable()); + workingCopy.SubstituteSubgraph(pattern, activation2); + + CHECK(std::string((*workingCopy.beginIConnectable())->GetName()) == "Activation2"); +} + TEST_CASE("SubgraphViewWorkingCopyOptimizationViews") { Graph graph; |