diff options
Diffstat (limited to 'src/armnn/Graph.hpp')
-rw-r--r-- | src/armnn/Graph.hpp | 315 |
1 files changed, 315 insertions, 0 deletions
diff --git a/src/armnn/Graph.hpp b/src/armnn/Graph.hpp new file mode 100644 index 0000000000..8888034197 --- /dev/null +++ b/src/armnn/Graph.hpp @@ -0,0 +1,315 @@ +// +// Copyright © 2017 Arm Ltd. All rights reserved. +// See LICENSE file in the project root for full license information. +// +#pragma once + +#include "Layers.hpp" + +#include <armnn/Types.hpp> +#include <armnn/TensorFwd.hpp> +#include <armnn/NetworkFwd.hpp> +#include <armnn/Exceptions.hpp> + +#include <list> +#include <unordered_map> +#include <unordered_set> +#include <vector> + +#include <boost/assert.hpp> +#include <boost/iterator/transform_iterator.hpp> + +namespace armnn +{ +class Graph +{ +public: + template <typename CVLayerT> + static CVLayerT* PtrCast(Layer* const layer) + { + return boost::polymorphic_downcast<CVLayerT*>(layer); + } + + using LayersList = std::list<Layer*>; + using Iterator = LayersList::const_iterator; // const so pointers in the list can't be modified externally + using ConstIterator = boost::transform_iterator<decltype(&PtrCast<const Layer>), Iterator>; + using IteratorDifference = Iterator::difference_type; + + using ConstIteratorInputs = boost::transform_iterator<decltype(&PtrCast<const InputLayer>), Iterator>; + using ConstIteratorOutputs = boost::transform_iterator<decltype(&PtrCast<const OutputLayer>), Iterator>; + + /// Wrapper class returned by Graph::GetInputLayers() + struct InputLayersAccessor + { + explicit InputLayersAccessor(const Graph& graph) : m_Graph(graph) {} + + ConstIteratorInputs begin() const + { + return { m_Graph.m_Layers.begin(), &PtrCast<const InputLayer> }; + } + + ConstIteratorInputs end() const + { + return { std::next(m_Graph.m_Layers.begin(), static_cast<IteratorDifference>(m_Graph.GetNumInputs())), + &PtrCast<const InputLayer> }; + } + + const Graph& m_Graph; + }; + + /// Wrapper class returned by Graph::GetOutputLayers() + struct OutputLayersAccessor + { + explicit OutputLayersAccessor(const Graph& graph) : m_Graph(graph) {} + + ConstIteratorOutputs begin() const + { + return { std::prev(m_Graph.m_Layers.end(), static_cast<IteratorDifference>(m_Graph.GetNumOutputs())), + &PtrCast<const OutputLayer> }; + } + + ConstIteratorOutputs end() const + { + return { m_Graph.m_Layers.end(), &PtrCast<const OutputLayer> }; + } + + const Graph& m_Graph; + }; + + Graph() : m_LayersInOrder(true) {} + + Graph(const Graph& other); + + Graph& operator=(const Graph& other) = delete; + + ~Graph() + { + for (auto&& layer : m_Layers) + { + delete layer; + } + } + + Status Print() const; + + /// Adds a new layer of type LaterType to the graph constructed with the arguments passed. + template <typename LayerT, typename... Args> + LayerT* AddLayer(Args&&... args); + + /// Inserts a new layer between the output slot currently connected to insertBefore + /// and insertBefore itself. + template <typename LayerT, typename... Args> + LayerT* InsertNewLayer(InputSlot& insertBefore, Args&&... args); + + /// Deletes the layer at the specified position and returns an iterator pointing + /// to the next element after the one being deleted. + Iterator EraseLayer(Iterator pos); + + /// Deletes the layer and returns an iterator pointing to the next layer in the graph + /// (next in the list, after the one being deleted). Sets @a layer to nullptr on return. + /// Templated to support pointers to any layer type. + template <typename LayerT> + Iterator EraseLayer(LayerT*& layer); + + /// Return iterator pointing to begin of list. Lowercase for range-based for loops. + Iterator begin() { return m_Layers.begin(); } + /// Return iterator pointing to end of list. Lowercase for range-based for loops. + Iterator end() { return m_Layers.end(); } + + /// Return const iterator pointing to begin of list. Lowercase for range-based for loops. + ConstIterator begin() const { return {m_Layers.begin(), &PtrCast<const Layer>}; } + /// Return const iterator pointing to end of list. Lowercase for range-based for loops. + ConstIterator end() const { return {m_Layers.end(), &PtrCast<const Layer>}; } + + /// Sort layers in topological order and return this. + Graph& TopologicalSort() { const_cast<const Graph*>(this)->TopologicalSort(); return *this; } + const Graph& TopologicalSort() const; + + size_t GetNumInputs() const { return m_InputIds.size(); } + size_t GetNumOutputs() const { return m_OutputIds.size(); } + + /// Returns a wrapper object with begin(), end() methods to iterate over the input layers + /// in a range-based for loop + InputLayersAccessor GetInputLayers() const { return InputLayersAccessor(*this); } + + /// Returns a wrapper object with begin(), end() methods to iterate over the output layers + /// in a range-based for loop + OutputLayersAccessor GetOutputLayers() const { return OutputLayersAccessor(*this); } + + size_t GetNumLayers() const { return m_Layers.size(); } + + /// Allocate memory for all tensors under output tensor handers of each layer + Status AllocateDynamicBuffers(); + + /// Modifies the graph in-place, removing edges connecting layers using different compute devices, + /// and relinking them via an intermediary copy layers. + void AddCopyLayers(); + + void InferTensorInfos(); + +private: + template <typename LayerT> + class LayerInGraphBase; + + template <typename LayerT> + class LayerInGraph; + + /// Get the position of a layer in the graph. + Iterator GetPosInGraph(Layer& layer); + + /// Adds a new layer of type LaterType to the graph constructed with the arguments passed. + template <typename LayerT, typename... Args> + LayerInGraph<LayerT>* AddLayerImpl(Iterator insertBefore, Args&&... args); + + std::unordered_set<LayerBindingId> m_InputIds; + std::unordered_set<LayerBindingId> m_OutputIds; + std::unordered_map<const Layer*, Iterator> m_PosInGraphMap; + + /// Mutable to allow sorting on const object. + mutable LayersList m_Layers; + mutable bool m_LayersInOrder; +}; + +/// Common base class for layers in the graph +template <typename LayerT> +class Graph::LayerInGraphBase : public LayerT +{ +protected: + template <typename... Args> + LayerInGraphBase(Graph& graph, Iterator insertBefore, Args&&... args) + : LayerT(std::forward<Args>(args)...), m_Graph(graph) + { + m_Graph.m_PosInGraphMap.emplace(this, m_Graph.m_Layers.emplace(insertBefore, this)); + } + ~LayerInGraphBase() + { + const size_t numErased = m_Graph.m_PosInGraphMap.erase(this); + boost::ignore_unused(numErased); + BOOST_ASSERT(numErased == 1); + } + + Graph& m_Graph; +}; + +/// Input/Output layers specialize this template +template <typename LayerT> +class Graph::LayerInGraph final : public LayerInGraphBase<LayerT> +{ +public: + template <typename... Args> + LayerInGraph(Graph& graph, Iterator insertBefore, Args&&... args) + : LayerInGraphBase<LayerT>(graph, insertBefore, std::forward<Args>(args)...) + { + } +}; + +/// Inputs add/remove their binding id to m_InputIds in the graph. +template <> +class Graph::LayerInGraph<InputLayer> final : public LayerInGraphBase<InputLayer> +{ +public: + template <typename... Args> + LayerInGraph(Graph& graph, Iterator insertBefore, Args&&... args) + : LayerInGraphBase<InputLayer>(graph, insertBefore, std::forward<Args>(args)...) + { + const bool isNewId = m_Graph.m_InputIds.emplace(GetBindingId()).second; + if (!isNewId) + { + throw InvalidArgumentException("A layer already exists with the specified id"); + } + } + ~LayerInGraph() override + { + const size_t numErased = m_Graph.m_InputIds.erase(GetBindingId()); + boost::ignore_unused(numErased); + BOOST_ASSERT(numErased == 1); + } +}; + +/// Outputs add/remove their binding id to m_OutputIds in the graph. +template <> +class Graph::LayerInGraph<OutputLayer> final : public LayerInGraphBase<OutputLayer> +{ +public: + template <typename... Args> + LayerInGraph(Graph& graph, Iterator insertBefore, Args&&... args) + : LayerInGraphBase<OutputLayer>(graph, insertBefore, std::forward<Args>(args)...) + { + const bool isNewId = m_Graph.m_OutputIds.emplace(GetBindingId()).second; + if (!isNewId) + { + throw InvalidArgumentException("A layer already exists with the specified id"); + } + } + ~LayerInGraph() override + { + const size_t numErased = m_Graph.m_OutputIds.erase(GetBindingId()); + boost::ignore_unused(numErased); + BOOST_ASSERT(numErased == 1); + } +}; + +inline Graph::Iterator Graph::GetPosInGraph(Layer& layer) +{ + auto it = m_PosInGraphMap.find(&layer); + BOOST_ASSERT(it != m_PosInGraphMap.end()); + return it->second; +} + +template <typename LayerT, typename... Args> +inline Graph::LayerInGraph<LayerT>* Graph::AddLayerImpl(Iterator insertBefore, Args&&... args) +{ + return new LayerInGraph<LayerT>(*this, insertBefore, std::forward<Args>(args)...); +} + +/// Inputs are inserted at the front of the list, to keep the order correct if the list is sorted. +/// Outputs are inserted at the back of the list, to keep the order correct if the list is sorted. +/// Other layers are inserted before existing outputs, so the latter remain at the back of the list. +template <typename LayerT, typename... Args> +inline LayerT* Graph::AddLayer(Args&&... args) +{ + switch (LayerEnumOf<LayerT>()) + { + case LayerType::Input: + { + return AddLayerImpl<LayerT>(begin(), std::forward<Args>(args)...); + } + case LayerType::Output: + { + return AddLayerImpl<LayerT>(end(), std::forward<Args>(args)...); + } + default: + { + m_LayersInOrder = false; + const auto pos = std::prev(end(), IteratorDifference(GetNumOutputs())); + return AddLayerImpl<LayerT>(pos, std::forward<Args>(args)...); + } + } +} + +template <typename LayerT, typename... Args> +inline LayerT* Graph::InsertNewLayer(InputSlot& insertBefore, Args&&... args) +{ + // Insert before the child layer so topological order is kept. + const Iterator pos = GetPosInGraph(insertBefore.GetOwningLayer()); + LayerT* const layer = AddLayerImpl<LayerT>(pos, std::forward<Args>(args)...); + insertBefore.Insert(*layer); + return layer; +} + +inline Graph::Iterator Graph::EraseLayer(Iterator pos) +{ + delete *pos; + return m_Layers.erase(pos); +} + +template <typename LayerT> +inline Graph::Iterator Graph::EraseLayer(LayerT*& layer) +{ + BOOST_ASSERT(layer != nullptr); + Iterator next = EraseLayer(GetPosInGraph(*layer)); + layer = nullptr; + return next; +} + +} // namespace armnn |