diff options
Diffstat (limited to 'src/armnn/SubgraphViewSelector.cpp')
-rw-r--r-- | src/armnn/SubgraphViewSelector.cpp | 208 |
1 files changed, 152 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 |