| // |
| // Copyright © 2017 Arm Ltd. All rights reserved. |
| // SPDX-License-Identifier: MIT |
| // |
| |
| #include "SubgraphViewSelector.hpp" |
| #include "Graph.hpp" |
| |
| #include <armnn/utility/IgnoreUnused.hpp> |
| |
| #include <boost/assert.hpp> |
| #include <algorithm> |
| #include <map> |
| #include <queue> |
| #include <unordered_set> |
| |
| namespace armnn |
| { |
| |
| namespace |
| { |
| |
| /// Intermediate data-structure to store the subgraph that a layer has been assigned to. |
| /// This is a "disjoint set" data structure that allows efficient merging of subgraphs, |
| /// which is a key part of the algorithm. Subgraphs are arranged in singly-linked trees |
| /// (with each node storing a pointer to its parent). Subgraphs in the same tree are considered |
| /// to have been merged. Merging subgraphs is performed by attaching one tree to another, |
| /// which is a simple pointer update. |
| /// |
| /// NOTE: Due to the way this is stored, it is almost never correct to directly compare pointers |
| /// to two PartialSubgraphs to check if two layers belong in the same subgraph. Instead you |
| /// should use IsMergedWith(). |
| /// |
| /// This structure also stores information about the dependencies of each subgraph, which is needed |
| /// to determine whether certain subgraphs can be merged. Checking whether a subgraph |
| /// depends on another subgraph is a frequent operation in the algorithm (see AssignSplitId) and so this is optimized |
| /// in preference to the merging of subgraphs. This leads to an approach where each subgraph stores |
| /// a set of all the subgraphs it depends on (for a fast lookup). In order to efficiently update this |
| /// set as subgraphs are merged means we also store a set of subgraphs which *depend on us* (i.e. the |
| /// complement of our dependencies). |
| class PartialSubgraph |
| { |
| public: |
| /// If this subgraph has been merged with another then there is an agreed "representative" for the combined |
| /// subgraph, which uniquely identifies the subgraph. |
| PartialSubgraph* GetRepresentative() |
| { |
| // Recurse up the tree to find the root node. |
| if (m_Parent == nullptr) |
| { |
| return this; |
| } |
| else |
| { |
| PartialSubgraph* result = m_Parent->GetRepresentative(); |
| // Update our parent pointer to point directly to the root in order to speed up future calls to this method. |
| // This essentially "flattens" the tree. |
| m_Parent = result; |
| return result; |
| } |
| } |
| |
| /// Merges this subgraph with another. |
| void MergeWith(PartialSubgraph* other) |
| { |
| if (m_Parent == nullptr) |
| { |
| other = other->GetRepresentative(); |
| if (this == other) |
| { |
| // Already merged - no-op |
| return; |
| } |
| m_Parent = other; |
| |
| // Update others' dependency sets to point to the new representative rather than us. |
| // Keeping these up-to-date means we can rely on these sets containing representatives when |
| // we perform a lookup in HasAntecedent() and so don't need to resolve the representative for each element |
| // of the set. See description at the top of this class for more rationale. |
| for (PartialSubgraph* a : m_Antecedents) |
| { |
| size_t numErased = a->m_Dependants.erase(this); |
| BOOST_ASSERT(numErased == 1); |
| IgnoreUnused(numErased); |
| a->m_Dependants.insert(m_Parent); |
| } |
| for (PartialSubgraph* a : m_Dependants) |
| { |
| size_t numErased = a->m_Antecedents.erase(this); |
| BOOST_ASSERT(numErased == 1); |
| IgnoreUnused(numErased); |
| a->m_Antecedents.insert(m_Parent); |
| } |
| |
| // Merge our dependency sets into our new representative. |
| // We no longer need to maintain our own sets, as requests will always be forwarded to the representative. |
| m_Parent->m_Antecedents.insert(m_Antecedents.begin(), m_Antecedents.end()); |
| m_Antecedents.clear(); |
| m_Parent->m_Dependants.insert(m_Dependants.begin(), m_Dependants.end()); |
| m_Dependants.clear(); |
| } |
| else |
| { |
| // Defer request to the representative |
| GetRepresentative()->MergeWith(other); |
| } |
| } |
| |
| /// Checks if this subgraph has been merged with the given subgraph. |
| bool IsMergedWith(PartialSubgraph* other) |
| { |
| return GetRepresentative() == other->GetRepresentative(); |
| } |
| |
| /// Marks the given subgraph as a direct antecedent (dependency) of this one. |
| void AddDirectAntecedent(PartialSubgraph* antecedent) |
| { |
| if (m_Parent == nullptr) |
| { |
| antecedent = antecedent->GetRepresentative(); |
| |
| m_Antecedents.insert(antecedent); |
| // Also record all of its antecedents, so that we end up with direct and indirect antecedents. |
| // This makes the lookup in HasAntecedent() faster. |
| m_Antecedents.insert(antecedent->m_Antecedents.begin(), antecedent->m_Antecedents.end()); |
| // All of our dependents also need to include the new antecedents |
| for (PartialSubgraph* d : m_Dependants) |
| { |
| d->m_Antecedents.insert(antecedent); |
| d->m_Antecedents.insert(antecedent->m_Antecedents.begin(), antecedent->m_Antecedents.end()); |
| } |
| |
| // Store reverse dependencies as well, required so that we can efficiently navigate the graph |
| // when making updates. |
| antecedent->m_Dependants.insert(this); |
| antecedent->m_Dependants.insert(m_Dependants.begin(), m_Dependants.end()); |
| for (PartialSubgraph* a : antecedent->m_Antecedents) |
| { |
| a->m_Dependants.insert(this); |
| a->m_Dependants.insert(m_Dependants.begin(), m_Dependants.end()); |
| } |
| } |
| else |
| { |
| // Defer request to the representative |
| GetRepresentative()->AddDirectAntecedent(antecedent); |
| } |
| } |
| |
| /// Checks if this subgraph is dependent on the given subgraph, either directly or indirectly. |
| bool HasAntecedent(PartialSubgraph* antecedent) |
| { |
| if (m_Parent == nullptr) |
| { |
| antecedent = antecedent->GetRepresentative(); |
| // Thanks to keeping this set updated in MergeWith and AddDirectAntecedent, we can do an efficient lookup. |
| return m_Antecedents.count(antecedent) > 0; |
| } |
| else |
| { |
| // Defer request to the representative |
| return GetRepresentative()->HasAntecedent(antecedent); |
| } |
| } |
| |
| private: |
| /// Pointer to the parent node in the tree. If this is null then we are the representative for our merged subgraph. |
| PartialSubgraph* m_Parent; |
| /// The representatives of all the subgraphs which we depend on, either directly or indirectly. |
| std::unordered_set<PartialSubgraph*> m_Antecedents; |
| /// The representatives of all the subgraphs which depend on us, either directly or indirectly. |
| std::unordered_set<PartialSubgraph*> m_Dependants; |
| }; |
| |
| /// Intermediate data structure to store information associated with a particular layer. |
| struct LayerSelectionInfo |
| { |
| using LayerInfoContainer = std::map<Layer*, LayerSelectionInfo>; |
| using LayerInfoQueue = std::queue<LayerSelectionInfo*>; |
| |
| LayerSelectionInfo(Layer* layer, const SubgraphViewSelector::LayerSelectorFunction& selector) |
| : m_Layer{layer} |
| , m_Subgraph{nullptr} |
| , m_IsSelected{selector(*layer)} |
| , m_IsProcessed(false) |
| { |
| } |
| |
| bool IsInputLayer() const |
| { |
| return m_Layer->GetType() == armnn::LayerType::Input || m_Layer->GetType() == armnn::LayerType::Constant; |
| } |
| |
| void CollectNonSelectedInputs(LayerSelectionInfo::LayerInfoContainer& layerInfos, |
| SubgraphView::InputSlots& inputSlots) |
| { |
| for (auto&& slot = m_Layer->BeginInputSlots(); slot != m_Layer->EndInputSlots(); ++slot) |
| { |
| OutputSlot* parentLayerOutputSlot = slot->GetConnectedOutputSlot(); |
| BOOST_ASSERT_MSG(parentLayerOutputSlot != nullptr, "The input slots must be connected here."); |
| if (parentLayerOutputSlot) |
| { |
| Layer& parentLayer = parentLayerOutputSlot->GetOwningLayer(); |
| auto parentInfo = layerInfos.find(&parentLayer); |
| if (parentInfo == layerInfos.end() || |
| !m_Subgraph->IsMergedWith(parentInfo->second.m_Subgraph.get())) |
| { |
| // Avoid collecting duplicate input slots |
| InputSlot* inputSlot = &(*slot); |
| if (std::find(inputSlots.begin(), inputSlots.end(), inputSlot) == inputSlots.end()) |
| { |
| inputSlots.push_back(inputSlot); |
| } |
| } |
| } |
| } |
| } |
| |
| 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(); |
| auto childInfo = layerInfos.find(&childLayer); |
| if (childInfo == layerInfos.end() || |
| !m_Subgraph->IsMergedWith(childInfo->second.m_Subgraph.get())) |
| { |
| // Avoid collecting duplicate output slots |
| OutputSlot* outputSlot = &(*slot); |
| if (std::find(outputSlots.begin(), outputSlots.end(), outputSlot) == outputSlots.end()) |
| { |
| outputSlots.push_back(outputSlot); |
| } |
| } |
| } |
| } |
| } |
| |
| Layer* m_Layer; |
| /// Which subgraph this layer has been assigned to. Only valid once m_IsProcessed is true. |
| /// Two layers with different m_Subgraph pointers may in fact have been merged into the same subgraph - |
| /// see the description of the PartialSubgraph class. |
| std::shared_ptr<PartialSubgraph> m_Subgraph; |
| bool m_IsSelected; |
| bool m_IsProcessed; |
| }; |
| |
| } // namespace <anonymous> |
| |
| SubgraphViewSelector::Subgraphs |
| SubgraphViewSelector::SelectSubgraphs(Graph& graph, const LayerSelectorFunction& selector) |
| { |
| SubgraphView subgraph(graph); |
| 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); |
| if (parentInfo != layerInfos.end()) |
| { |
| 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); |
| if (childInfo != layerInfos.end()) |
| { |
| function(childInfo->second); |
| } |
| } |
| } |
| } |
| |
| void AssignSplitId(LayerSelectionInfo::LayerInfoContainer& layerInfos, LayerSelectionInfo& layerInfo) |
| { |
| // Check each input to see if we can attach ourselves to any of the subgraphs that have already been assigned. |
| ForEachLayerInput(layerInfos, layerInfo, [&](LayerSelectionInfo& parentInfo) |
| { |
| // We can only attach ourselves to the subgraph from this input if there isn't a cut here. |
| if (layerInfo.m_IsSelected == parentInfo.m_IsSelected) |
| { |
| // We also need to check that merging into this subgraph won't cause a dependency cycle between subgraphs. |
| // This will be the case if the subgraph that we will become part of is already a dependency |
| // of one of the subgraphs that are input to this layer, e.g: |
| // |
| // 0 | The numbers (0, 1) are the subgraph IDs of each layer and we are looking at layer X. |
| // / \ | |
| // 1 0 | We can't merge X into subgraph 0, because the left-hand input already depends on subgraph 0. |
| // \ / | We can however merge X into subgraph 1. |
| // X | |
| // |
| bool dependenciesOk = true; |
| ForEachLayerInput(layerInfos, layerInfo, [&](LayerSelectionInfo& otherParentInfo) |
| { |
| // We call HasAntecedent() ~ n^2 times, where n is the number of inputs to this layer. |
| // Hence it is important that this is efficient - see PartialSubgraph class description. |
| if (otherParentInfo.m_Subgraph->HasAntecedent(parentInfo.m_Subgraph.get())) |
| { |
| dependenciesOk = false; |
| } |
| }); |
| |
| if (dependenciesOk) |
| { |
| // Merge into the subgraph of this input. If we have already been merged into another subgraph |
| // (from another input of this layer), then merge both of them together. |
| if (layerInfo.m_Subgraph == nullptr) |
| { |
| layerInfo.m_Subgraph = parentInfo.m_Subgraph; |
| } |
| else |
| { |
| // We call MergeWith() ~ n times, where n is the number of inputs to this layer. |
| // Therefore it does not need to be as performant as HasAntecedent(). |
| layerInfo.m_Subgraph->MergeWith(parentInfo.m_Subgraph.get()); |
| } |
| } |
| } |
| }); |
| |
| // If we weren't able to merge into an existing subgraph then we need to make a new one |
| if (layerInfo.m_Subgraph == nullptr) |
| { |
| layerInfo.m_Subgraph = std::make_shared<PartialSubgraph>(); |
| } |
| |
| // Record dependencies of the chosen subgraph based on the inputs of this layer. |
| ForEachLayerInput(layerInfos, layerInfo, [&](LayerSelectionInfo& parentInfo) |
| { |
| // These functions are called ~n times, where n is the number of inputs to this layer. |
| // Therefore it does not need to be as performant as HasAntecedent(). |
| if (!layerInfo.m_Subgraph->IsMergedWith(parentInfo.m_Subgraph.get())) |
| { |
| layerInfo.m_Subgraph->AddDirectAntecedent(parentInfo.m_Subgraph.get()); |
| } |
| }); |
| } |
| |
| 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 layerInfos; |
| |
| LayerSelectionInfo::LayerInfoQueue processQueue; |
| for (auto& layer : subgraph) |
| { |
| auto emplaced = layerInfos.emplace(layer, LayerSelectionInfo{layer, selector}); |
| LayerSelectionInfo& layerInfo = emplaced.first->second; |
| |
| // Start with Input type layers |
| if (layerInfo.IsInputLayer()) |
| { |
| processQueue.push(&layerInfo); |
| } |
| } |
| |
| const SubgraphView::InputSlots& subgraphInputSlots = subgraph.GetInputSlots(); |
| for (auto& inputSlot : subgraphInputSlots) |
| { |
| Layer& layer = inputSlot->GetOwningLayer(); |
| auto emplaced = layerInfos.emplace(&layer, LayerSelectionInfo{&layer, selector}); |
| LayerSelectionInfo& layerInfo = emplaced.first->second; |
| |
| processQueue.push(&layerInfo); |
| } |
| |
| while (!processQueue.empty()) |
| { |
| 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) |
| { |
| // 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 subgraph representative into a map |
| using SelectionInfoPtrs = std::vector<LayerSelectionInfo*>; |
| std::map<PartialSubgraph*, SelectionInfoPtrs> splitMap; |
| for (auto& info : layerInfos) |
| { |
| if (info.second.m_IsSelected) |
| { |
| auto it = splitMap.find(info.second.m_Subgraph->GetRepresentative()); |
| if (it == splitMap.end()) |
| { |
| splitMap.insert( |
| std::make_pair(info.second.m_Subgraph->GetRepresentative(), SelectionInfoPtrs{&info.second})); |
| } |
| else |
| { |
| it->second.push_back(&info.second); |
| } |
| } |
| } |
| |
| // Now each entry in splitMap represents a subgraph |
| Subgraphs result; |
| for (auto& splitGraph : splitMap) |
| { |
| SubgraphView::InputSlots inputs; |
| SubgraphView::OutputSlots outputs; |
| SubgraphView::Layers layers; |
| for (auto&& infoPtr : splitGraph.second) |
| { |
| 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 |
| result.emplace_back(std::make_unique<SubgraphView>(std::move(inputs), |
| std::move(outputs), |
| std::move(layers))); |
| } |
| |
| return result; |
| } |
| |
| } // namespace armnn |