IVGCVSW-2989 Generate subgraphs without cyclic dependencies

Change-Id: I45f81aa4ca8a964e423594fe271825c4a52b21f4
Signed-off-by: Derek Lamberti <derek.lamberti@arm.com>

2 files changed, 239 insertions, 56 deletions

diff --git a/src/armnn/SubgraphViewSelector.cpp b/src/armnn/SubgraphViewSelector.cpp
index e8919c1e12..cc821ec956 100644
--- a/src/armnn/SubgraphViewSelector.cpp
+++ b/src/armnn/SubgraphViewSelector.cpp
@@ -8,6 +8,7 @@
 #include <boost/assert.hpp>
 #include <algorithm>
 #include <unordered_map>
+#include <queue>
 
 namespace armnn
 {
@@ -17,14 +18,16 @@ namespace
 
 struct LayerSelectionInfo
 {
+    using SplitId = uint32_t;
     using LayerInfoContainer = std::unordered_map<Layer*, LayerSelectionInfo>;
+    using LayerInfoQueue = std::queue<LayerSelectionInfo*>;
     static constexpr uint32_t InitialSplitId() { return 1; }
 
     LayerSelectionInfo(Layer* layer, const SubgraphViewSelector::LayerSelectorFunction& selector)
     : m_Layer{layer}
     , m_SplitId{0}
     , m_IsSelected{selector(*layer)}
-    , m_IsVisited(false)
+    , m_IsProcessed(false)
     {
         // fill topology information by storing direct children
         for (auto&& slot = m_Layer->BeginOutputSlots(); slot != m_Layer->EndOutputSlots(); ++slot)
@@ -37,47 +40,13 @@ struct LayerSelectionInfo
         }
     }
 
-    void MarkChildrenSplits(LayerInfoContainer& network,
-                            uint32_t splitId,
-                            bool prevSelected)
-    {
-        if (m_IsVisited)
-        {
-            return;
-        }
-        m_IsVisited = true;
-
-        if (m_SplitId < splitId)
-        {
-            m_SplitId = splitId;
-        }
-
-        // introduce a new split point at all non-selected points, but only if the
-        // previous point was selected. this prevents creating a new subgraph at
-        // every non-selected layer
-        if (!m_IsSelected && prevSelected)
-        {
-            ++m_SplitId;
-        }
-
-        for (auto& layer : m_DirectChildren)
-        {
-            auto it = network.find(layer);
-            BOOST_ASSERT_MSG(it != network.end(), "All layers must be part of the topology.");
-            if (it != network.end())
-            {
-                it->second.MarkChildrenSplits(network, m_SplitId, m_IsSelected);
-            }
-        }
-    }
-
     bool IsInputLayer() const
     {
-        return m_Layer->GetType() == armnn::LayerType::Input;
+        return m_Layer->GetType() == armnn::LayerType::Input || m_Layer->GetType() == armnn::LayerType::Constant;
     }
 
-    void CollectNonSelectedInputs(SubgraphView::InputSlots& inputSlots,
-                                  const SubgraphViewSelector::LayerSelectorFunction& selector)
+    void CollectNonSelectedInputs(LayerSelectionInfo::LayerInfoContainer& layerInfos,
+                                  SubgraphView::InputSlots& inputSlots)
     {
         for (auto&& slot = m_Layer->BeginInputSlots(); slot != m_Layer->EndInputSlots(); ++slot)
         {
@@ -86,7 +55,8 @@ struct LayerSelectionInfo
             if (parentLayerOutputSlot)
             {
                 Layer& parentLayer = parentLayerOutputSlot->GetOwningLayer();
-                if (selector(parentLayer) == false)
+                auto parentInfo = layerInfos.find(&parentLayer);
+                if (m_SplitId != parentInfo->second.m_SplitId)
                 {
                     inputSlots.push_back(&(*slot));
                 }
@@ -94,15 +64,16 @@ struct LayerSelectionInfo
         }
     }
 
-    void CollectNonSelectedOutputSlots(SubgraphView::OutputSlots& outputSlots,
-                                       const SubgraphViewSelector::LayerSelectorFunction& selector)
+    void CollectNonSelectedOutputSlots(LayerSelectionInfo::LayerInfoContainer& layerInfos,
+                                       SubgraphView::OutputSlots& outputSlots)
     {
         for (auto&& slot = m_Layer->BeginOutputSlots(); slot != m_Layer->EndOutputSlots(); ++slot)
         {
             for (InputSlot* childLayerInputSlot : slot->GetConnections())
             {
                 Layer& childLayer = childLayerInputSlot->GetOwningLayer();
-                if (selector(childLayer) == false)
+                auto childInfo = layerInfos.find(&childLayer);
+                if (m_SplitId != childInfo->second.m_SplitId)
                 {
                     outputSlots.push_back(&(*slot));
                 }
@@ -112,9 +83,9 @@ struct LayerSelectionInfo
 
     std::vector<Layer*> m_DirectChildren;
     Layer* m_Layer;
-    uint32_t m_SplitId;
+    SplitId m_SplitId;
     bool m_IsSelected;
-    bool m_IsVisited;
+    bool m_IsProcessed;
 };
 
 } // namespace <anonymous>
@@ -126,32 +97,157 @@ SubgraphViewSelector::SelectSubgraphs(Graph& graph, const LayerSelectorFunction&
     return SubgraphViewSelector::SelectSubgraphs(subgraph, selector);
 }
 
+
+template<typename Delegate>
+void ForEachLayerInput(LayerSelectionInfo::LayerInfoContainer& layerInfos,
+                       LayerSelectionInfo& layerInfo,
+                       Delegate function)
+{
+    Layer& layer = *layerInfo.m_Layer;
+
+    for (auto inputSlot : layer.GetInputSlots())
+    {
+        auto connectedInput = boost::polymorphic_downcast<OutputSlot*>(inputSlot.GetConnection());
+        BOOST_ASSERT_MSG(connectedInput, "Dangling input slot detected.");
+        Layer& inputLayer = connectedInput->GetOwningLayer();
+
+        auto parentInfo = layerInfos.find(&inputLayer);
+        function(parentInfo->second);
+    }
+}
+
+template<typename Delegate>
+void ForEachLayerOutput(LayerSelectionInfo::LayerInfoContainer& layerInfos,
+                        LayerSelectionInfo& layerInfo,
+                        Delegate function)
+{
+    Layer& layer= *layerInfo.m_Layer;
+
+    for (auto& outputSlot : layer.GetOutputSlots())
+    {
+        for (auto& output : outputSlot.GetConnections())
+        {
+            Layer& childLayer = output->GetOwningLayer();
+
+            auto childInfo = layerInfos.find(&childLayer);
+            function(childInfo->second);
+        }
+    }
+}
+
+void AssignSplitId(LayerSelectionInfo::LayerInfoContainer& layerInfos, LayerSelectionInfo& layerInfo)
+{
+    bool newSplit = false;
+    LayerSelectionInfo::SplitId minSplitId = std::numeric_limits<LayerSelectionInfo::SplitId>::max();
+    LayerSelectionInfo::SplitId maxSplitId = std::numeric_limits<LayerSelectionInfo::SplitId>::lowest();
+    LayerSelectionInfo::SplitId maxSelectableId = std::numeric_limits<LayerSelectionInfo::SplitId>::lowest();
+
+    ForEachLayerInput(layerInfos, layerInfo, [&newSplit, &minSplitId, &maxSplitId, &maxSelectableId, &layerInfo](
+        LayerSelectionInfo& parentInfo)
+        {
+            minSplitId = std::min(minSplitId, parentInfo.m_SplitId);
+            maxSplitId = std::max(maxSplitId, parentInfo.m_SplitId);
+            if (parentInfo.m_IsSelected && layerInfo.m_IsSelected)
+            {
+                maxSelectableId = std::max(maxSelectableId, parentInfo.m_SplitId);
+            }
+
+            if (layerInfo.m_IsSelected != parentInfo.m_IsSelected)
+            {
+                newSplit = true;
+            }
+
+        });
+
+    // Assign the split Id for the current layerInfo
+    if (newSplit)
+    {
+        if (maxSelectableId > minSplitId)
+        {
+            // We can be overly aggressive when choosing to create a new split so
+            // here we determine if one of the parent branches are suitable candidates for continuation instead.
+            // Any splitId > minSplitId will come from a shorter branch...and therefore should not be from
+            // the split containing the original fork and thus we avoid the execution dependency.
+            layerInfo.m_SplitId = maxSelectableId;
+        }
+        else
+        {
+            layerInfo.m_SplitId = ++maxSplitId;
+        }
+    } else
+    {
+        // The branch with the highest splitId represents the shortest path of selected nodes.
+        layerInfo.m_SplitId = maxSplitId;
+    }
+}
+
+bool IsReadyForSplitAssignment(LayerSelectionInfo::LayerInfoContainer& layerInfos, LayerSelectionInfo& layerInfo)
+{
+    bool ready = true;
+    ForEachLayerInput(layerInfos, layerInfo,
+                      [&ready](LayerSelectionInfo& parentInfo)
+                          {
+                              if (!parentInfo.m_IsProcessed)
+                              {
+                                  ready = false;
+                              }
+                          });
+    return ready;
+}
+
 SubgraphViewSelector::Subgraphs
 SubgraphViewSelector::SelectSubgraphs(SubgraphView& subgraph, const LayerSelectorFunction& selector)
 {
-    LayerSelectionInfo::LayerInfoContainer layerInfo;
+    LayerSelectionInfo::LayerInfoContainer layerInfos;
 
+    LayerSelectionInfo::LayerInfoQueue processQueue;
     for (auto& layer : subgraph)
     {
-        layerInfo.emplace(layer, LayerSelectionInfo{layer, selector});
+        auto emplaced = layerInfos.emplace(layer, LayerSelectionInfo{layer, selector});
+        LayerSelectionInfo& layerInfo = emplaced.first->second;
+
+        // Start with Input type layers
+        if (layerInfo.IsInputLayer())
+        {
+            processQueue.push(&layerInfo);
+        }
     }
 
-    uint32_t splitNo = LayerSelectionInfo::InitialSplitId();
-    for (auto& info : layerInfo)
+    while (!processQueue.empty())
     {
-        if (info.second.IsInputLayer())
+        LayerSelectionInfo& layerInfo = *processQueue.front();
+        processQueue.pop(); // remove front from queue
+
+        // This layerInfo may have been added to the queue multiple times, so skip if we have already processed it
+        if (!layerInfo.m_IsProcessed)
         {
-            // For each input layer we mark the graph where subgraph
-            // splits need to happen because of the dependency between
-            // the selected and non-selected nodes
-            info.second.MarkChildrenSplits(layerInfo, splitNo, false);
+
+            // Only process this layerInfo if all inputs have been processed
+            if (!IsReadyForSplitAssignment(layerInfos, layerInfo))
+            {
+                // Put back of the process queue if we can't process it just yet
+                processQueue.push(&layerInfo);
+                continue; // Skip to next iteration
+            }
+
+            // Now we do the processing
+            AssignSplitId(layerInfos, layerInfo);
+
+            // Queue any child nodes for processing
+            ForEachLayerOutput(layerInfos, layerInfo, [&processQueue](LayerSelectionInfo& childInfo)
+                {
+                    processQueue.push(&childInfo);
+                });
+
+            // We don't need to process this node again
+            layerInfo.m_IsProcessed = true;
         }
     }
 
     // Collect all selected layers keyed by split id into a map
     using SelectionInfoPtrs = std::vector<LayerSelectionInfo*>;
     std::unordered_map<uint32_t, SelectionInfoPtrs> splitMap;
-    for (auto& info : layerInfo)
+    for (auto& info : layerInfos)
     {
         if (info.second.m_IsSelected)
         {
@@ -178,8 +274,8 @@ SubgraphViewSelector::SelectSubgraphs(SubgraphView& subgraph, const LayerSelecto
             SubgraphView::Layers layers;
             for (auto&& infoPtr : splitGraph.second)
             {
-                infoPtr->CollectNonSelectedInputs(inputs, selector);
-                infoPtr->CollectNonSelectedOutputSlots(outputs, selector);
+                infoPtr->CollectNonSelectedInputs(layerInfos, inputs);
+                infoPtr->CollectNonSelectedOutputSlots(layerInfos, outputs);
                 layers.push_back(infoPtr->m_Layer);
             }
             // Create a new sub-graph with the new lists of input/output slots and layer
diff --git a/src/armnn/test/SubgraphViewTests.cpp b/src/armnn/test/SubgraphViewTests.cpp
index 4a9316428b..d580385797 100644
--- a/src/armnn/test/SubgraphViewTests.cpp
+++ b/src/armnn/test/SubgraphViewTests.cpp
@@ -1028,4 +1028,91 @@ BOOST_AUTO_TEST_CASE(MultipleSubgraphs)
     }
 }
 
+BOOST_AUTO_TEST_CASE(SubgraphCycles)
+{
+    // This case represent the scenario when a naive split could lead to a cyclic dependency between two subgraphs
+    //
+    // X0 -> M0 -> X1 -> M2 -> X2
+    // X0 -> M0 -> M1 -> M2 -> X2
+    //
+    /*
+          X0
+          |
+          |
+          M0
+         / |
+        /  |
+       X1  M1
+        \ /
+        M2
+        |
+        X2
+    */
+    // The expected result for this is that M0,M1 will be part of one subgraph and M2 in another and the
+    // input and output slots in the subgraphs will be set accordingly.
+    //
+    Graph graph;
+
+    OriginsDescriptor mergerDescriptor(2);
+    auto x0 = graph.AddLayer<InputLayer>(0, "x0");
+    auto m0 = graph.AddLayer<ActivationLayer>(ActivationDescriptor{}, "m0");
+    auto x1 = graph.AddLayer<ActivationLayer>(ActivationDescriptor{}, "x1");
+    auto m1 = graph.AddLayer<ActivationLayer>(ActivationDescriptor{}, "m1");
+    auto m2 = graph.AddLayer<AdditionLayer>("m2");
+    auto x2 = graph.AddLayer<ActivationLayer>(ActivationDescriptor{}, "x2");
+
+    x0->GetOutputSlot(0).Connect(m0->GetInputSlot(0));
+    m0->GetOutputSlot(0).Connect(x1->GetInputSlot(0));
+    m0->GetOutputSlot(0).Connect(m1->GetInputSlot(0));
+    x1->GetOutputSlot(0).Connect(m2->GetInputSlot(0));
+    m1->GetOutputSlot(0).Connect(m2->GetInputSlot(1));
+    m2->GetOutputSlot(0).Connect(x2->GetInputSlot(0));
+
+    // All selected 'M*' layers will be have 'm' in the name
+    SubgraphViewSelector::Subgraphs subgraphs =
+        SubgraphViewSelector::SelectSubgraphs(
+            graph,
+            // select the middle layers only
+            [](const Layer & l)
+                {
+                    bool toSelect = (l.GetNameStr().find('m') != std::string::npos);
+                    return toSelect;
+                });
+
+    // expected results to test against
+    auto inputSubgraph = CreateSubgraphViewFrom(CreateInputsFrom({m0}),
+                                                CreateOutputsFrom({m0, m1}),
+                                                {m0, m1});
+
+    auto outputSubgraph = CreateSubgraphViewFrom(CreateInputsFrom({m2}),
+                                                 CreateOutputsFrom({m2}),
+                                                 {m2});
+
+    BOOST_TEST(subgraphs.size() == 2);
+    if (subgraphs.size() == 2)
+    {
+        // we need to have valid subgraph pointers here
+        BOOST_TEST((subgraphs[0] != nullptr));
+        BOOST_TEST((subgraphs[1] != nullptr));
+
+        if (subgraphs[0].get() != nullptr && subgraphs[1].get() != nullptr)
+        {
+            // sort the subgraphs by layer size, so it is simpler to test
+            std::sort(subgraphs.begin(), subgraphs.end(),
+                      [](SubgraphViewSelector::SubgraphViewPtr & lhs, SubgraphViewSelector::SubgraphViewPtr & rhs)
+                          {
+                              return (lhs->GetLayers().size() < rhs->GetLayers().size());
+                          }
+            );
+
+            // one subgraph needs to be size=1 and the other one is 4
+            BOOST_TEST(subgraphs[0]->GetLayers().size() == 1);
+            BOOST_TEST(subgraphs[1]->GetLayers().size() == 2);
+
+            CompareSubgraphViews(subgraphs[0], outputSubgraph);
+            CompareSubgraphViews(subgraphs[1], inputSubgraph);
+        }
+    }
+}
+
 BOOST_AUTO_TEST_SUITE_END()