diff options
Diffstat (limited to 'src/armnn/SubGraph.hpp')
-rw-r--r-- | src/armnn/SubGraph.hpp | 53 |
1 files changed, 53 insertions, 0 deletions
diff --git a/src/armnn/SubGraph.hpp b/src/armnn/SubGraph.hpp new file mode 100644 index 0000000000..312bb115eb --- /dev/null +++ b/src/armnn/SubGraph.hpp @@ -0,0 +1,53 @@ +// +// Copyright © 2017 Arm Ltd. All rights reserved. +// SPDX-License-Identifier: MIT +// + +#pragma once + +#include "Layer.hpp" + +#include <vector> +#include <unordered_set> + +namespace armnn +{ + +/// +/// The SubGraph class represents a subgraph of a Graph. +/// The data it holds, points to data held by layers of the Graph, so the +/// the contents of the SubGraph becomes invalid when the Layers are destroyed +/// or changed. +/// +class SubGraph final +{ +public: + using InputSlots = std::vector<InputSlot *>; + using OutputSlots = std::vector<OutputSlot *>; + using Layers = std::unordered_set<Layer *>; + + SubGraph(); + SubGraph(InputSlots && inputs, + OutputSlots && outputs, + Layers && layers); + + const InputSlots & GetInputSlots() const; + const OutputSlots & GetOutputSlots() const; + const Layers & GetLayers() const; + + const InputSlot* GetInputSlot(unsigned int index) const; + InputSlot* GetInputSlot(unsigned int index); + + const OutputSlot* GetOutputSlot(unsigned int index) const; + OutputSlot* GetOutputSlot(unsigned int index); + + unsigned int GetNumInputSlots() const; + unsigned int GetNumOutputSlots() const; + +private: + InputSlots m_InputSlots; + OutputSlots m_OutputSlots; + Layers m_Layers; +}; + +} // namespace armnn |