IVGCVSW-3418 Add Arm NN front end support for the new Stack layer

 * Added new StackLayer class
 * Made necessary changes to Descriptors, ILayerSupport, ILayerVisitor, etc.
 * Added unit tests

Signed-off-by: Matthew Jackson <matthew.jackson@arm.com>
Change-Id: Ieb97a928a342ffe1901c6058eb895711c358fd3d

8 files changed, 319 insertions, 0 deletions

diff --git a/src/armnn/InternalTypes.hpp b/src/armnn/InternalTypes.hpp
index b097265d81..bf095ac8a2 100644
--- a/src/armnn/InternalTypes.hpp
+++ b/src/armnn/InternalTypes.hpp
@@ -58,6 +58,7 @@ enum class LayerType
     SpaceToBatchNd,
     SpaceToDepth,
     Splitter,
+    Stack,
     StridedSlice,
     Subtraction,
     Switch,
diff --git a/src/armnn/LayersFwd.hpp b/src/armnn/LayersFwd.hpp
index 0f9633a58c..b3f7adc02c 100644
--- a/src/armnn/LayersFwd.hpp
+++ b/src/armnn/LayersFwd.hpp
@@ -50,6 +50,7 @@
 #include "layers/SpaceToBatchNdLayer.hpp"
 #include "layers/SpaceToDepthLayer.hpp"
 #include "layers/SplitterLayer.hpp"
+#include "layers/StackLayer.hpp"
 #include "layers/StridedSliceLayer.hpp"
 #include "layers/SubtractionLayer.hpp"
 #include "layers/SwitchLayer.hpp"
@@ -126,6 +127,7 @@ DECLARE_LAYER(Softmax)
 DECLARE_LAYER(SpaceToBatchNd)
 DECLARE_LAYER(SpaceToDepth)
 DECLARE_LAYER(Splitter)
+DECLARE_LAYER(Stack)
 DECLARE_LAYER(StridedSlice)
 DECLARE_LAYER(Subtraction)
 DECLARE_LAYER(Switch)
diff --git a/src/armnn/Network.cpp b/src/armnn/Network.cpp
index 3b7a1cf2b3..29493816a8 100644
--- a/src/armnn/Network.cpp
+++ b/src/armnn/Network.cpp
@@ -1422,6 +1422,12 @@ IConnectableLayer* Network::AddTransposeConvolution2dLayer(const TransposeConvol
     return layer;
 }
 
+IConnectableLayer* Network::AddStackLayer(const StackDescriptor& stackDescriptor,
+                                          const char* name)
+{
+    return m_Graph->AddLayer<StackLayer>(stackDescriptor, name);
+}
+
 void Network::Accept(ILayerVisitor& visitor) const
 {
     for (auto layer : GetGraph())
diff --git a/src/armnn/Network.hpp b/src/armnn/Network.hpp
index 7fc5b651d0..8a99debb47 100644
--- a/src/armnn/Network.hpp
+++ b/src/armnn/Network.hpp
@@ -197,6 +197,9 @@ public:
                                                       const Optional<ConstTensor>& biases,
                                                       const char* name = nullptr) override;
 
+    IConnectableLayer* AddStackLayer(const StackDescriptor& stackDescriptor,
+                                     const char* name = nullptr) override;
+
     void Accept(ILayerVisitor& visitor) const override;
 
 private:
diff --git a/src/armnn/layers/StackLayer.cpp b/src/armnn/layers/StackLayer.cpp
new file mode 100644
index 0000000000..59bc8d5a13
--- /dev/null
+++ b/src/armnn/layers/StackLayer.cpp
@@ -0,0 +1,98 @@
+//
+// Copyright © 2017 Arm Ltd. All rights reserved.
+// SPDX-License-Identifier: MIT
+//
+#include "StackLayer.hpp"
+#include "LayerCloneBase.hpp"
+
+#include <armnn/TypesUtils.hpp>
+#include <backendsCommon/WorkloadData.hpp>
+#include <backendsCommon/WorkloadFactory.hpp>
+
+#include <queue>
+
+namespace armnn
+{
+
+StackLayer::StackLayer(const StackDescriptor& param, const char* name)
+    : LayerWithParameters(param.m_NumInputs, 1, LayerType::Stack, param, name)
+{
+}
+
+std::unique_ptr<IWorkload> StackLayer::CreateWorkload(const Graph& graph, const IWorkloadFactory& factory) const
+{
+    StackQueueDescriptor descriptor;
+    return factory.CreateStack(descriptor, PrepInfoAndDesc(descriptor, graph));
+}
+
+StackLayer* StackLayer::Clone(Graph& graph) const
+{
+    return CloneBase<StackLayer>(graph, m_Param, GetName());
+}
+
+std::vector<TensorShape> StackLayer::InferOutputShapes(const std::vector<TensorShape>& inputShapes) const
+{
+    const TensorShape& inputShape = m_Param.m_InputShape;
+    const unsigned int inputNumDimensions = inputShape.GetNumDimensions();
+    const unsigned int axis = m_Param.m_Axis;
+
+    BOOST_ASSERT(axis <= inputNumDimensions);
+
+    unsigned int dimensionSizes[inputNumDimensions + 1];
+    for (unsigned int i = 0; i < axis; ++i)
+    {
+        dimensionSizes[i] = inputShape[i];
+    }
+
+    dimensionSizes[axis] = m_Param.m_NumInputs;
+
+    for (unsigned int i = axis + 1; i < inputNumDimensions + 1; ++i)
+    {
+        dimensionSizes[i] = inputShape[i-1];
+    }
+
+    TensorShape targetShape = TensorShape(inputNumDimensions + 1, dimensionSizes);
+
+    return std::vector<TensorShape>({ targetShape });
+}
+
+void StackLayer::ValidateTensorShapesFromInputs()
+{
+    // Validates Stack layer.
+    ConditionalThrowIfNotEqual<LayerValidationException>(
+        "StackLayer: Num Input Slots must match Num Inputs.",
+        m_Param.m_NumInputs,
+        GetNumInputSlots());
+
+    VerifyLayerConnections(m_Param.m_NumInputs, CHECK_LOCATION());
+
+    // Constructs and validates input shapes
+    std::vector<TensorShape> inputShapes;
+    for (unsigned int i = 0; i < GetNumInputSlots(); ++i)
+    {
+        TensorShape inputShape = GetInputSlot(i).GetConnection()->GetTensorInfo().GetShape();
+        if (inputShape != m_Param.m_InputShape)
+        {
+            throw LayerValidationException("ConcatLayer: TensorShape set on InputSlot[" +
+                                           std::to_string(i) +
+                                           "] does not match defined input shape");
+        }
+        inputShapes.push_back(inputShape);
+    }
+
+    auto inferredShapes = InferOutputShapes(inputShapes);
+
+    BOOST_ASSERT(inferredShapes.size() == 1);
+
+    ConditionalThrowIfNotEqual<LayerValidationException>(
+        "StackLayer: TensorShape set on OutputSlot[0] does not match the inferred shape.",
+        GetOutputSlot(0).GetTensorInfo().GetShape(),
+        inferredShapes[0]);
+}
+
+void StackLayer::Accept(ILayerVisitor& visitor) const
+{
+    visitor.VisitStackLayer(this, GetParameters(), GetName());
+}
+
+} // namespace armnn armnn
diff --git a/src/armnn/layers/StackLayer.hpp b/src/armnn/layers/StackLayer.hpp
new file mode 100644
index 0000000000..6c845972d0
--- /dev/null
+++ b/src/armnn/layers/StackLayer.hpp
@@ -0,0 +1,49 @@
+//
+// Copyright © 2017 Arm Ltd. All rights reserved.
+// SPDX-License-Identifier: MIT
+//
+#pragma once
+
+#include "LayerWithParameters.hpp"
+
+namespace armnn
+{
+
+/// This layer represents a stack operation.
+class StackLayer : public LayerWithParameters<StackDescriptor>
+{
+public:
+    /// Makes a workload for the Stack type.
+    /// @param [in] graph The graph where this layer can be found.
+    /// @param [in] factory The workload factory which will create the workload.
+    /// @return A pointer to the created workload, or nullptr if not created.
+    virtual std::unique_ptr<IWorkload>CreateWorkload(const Graph& graph,
+                                                     const IWorkloadFactory& factory) const override;
+
+    /// Creates a dynamically-allocated copy of this layer.
+    /// @param [in] graph The graph into which this layer is being cloned.
+    StackLayer* Clone(Graph& graph) const override;
+
+    /// Check if the input tensor shape(s)
+    /// will lead to a valid configuration of @ref StackLayer.
+    void ValidateTensorShapesFromInputs() override;
+
+    /// By default returns inputShapes if the number of inputs are equal to number of outputs,
+    /// otherwise infers the output shapes from given input shapes and layer properties.
+    /// @param [in] inputShapes The input shapes layer has.
+    /// @return A vector to the inferred output shape.
+    std::vector<TensorShape> InferOutputShapes(const std::vector<TensorShape>& inputShapes) const override;
+
+    void Accept(ILayerVisitor& visitor) const override;
+
+protected:
+    /// Constructor to create a StackLayer.
+    /// @param [in] param StackDescriptor to configure the stack operation.
+    /// @param [in] name Optional name for the layer.
+    StackLayer(const StackDescriptor& param, const char* name);
+
+    /// Default destructor
+    ~StackLayer() = default;
+};
+
+} // namespace
diff --git a/src/armnn/test/InferOutputTests.cpp b/src/armnn/test/InferOutputTests.cpp
index 6ce56e9805..24ae8b263b 100644
--- a/src/armnn/test/InferOutputTests.cpp
+++ b/src/armnn/test/InferOutputTests.cpp
@@ -25,4 +25,10 @@ ARMNN_SIMPLE_TEST_CASE(PreluInferOutputShapeNoMatch,               PreluInferOut
 ARMNN_SIMPLE_TEST_CASE(PreluValidateTensorShapesFromInputsMatch,   PreluValidateTensorShapesFromInputsMatchTest)
 ARMNN_SIMPLE_TEST_CASE(PreluValidateTensorShapesFromInputsNoMatch, PreluValidateTensorShapesFromInputsNoMatchTest)
 
+// Stack
+ARMNN_SIMPLE_TEST_CASE(StackInferOutputShapeFromInputsMatch,       StackInferOutputShapeFromInputsMatchTest)
+ARMNN_SIMPLE_TEST_CASE(StackInferOutputShapeFromInputsNoMatch,     StackInferOutputShapeFromInputsNoMatchTest)
+ARMNN_SIMPLE_TEST_CASE(StackValidateTensorShapesFromInputsMatch,   StackValidateTensorShapesFromInputsMatchTest)
+ARMNN_SIMPLE_TEST_CASE(StackValidateTensorShapesFromInputsNoMatch, StackValidateTensorShapesFromInputsNoMatchTest)
+
 BOOST_AUTO_TEST_SUITE_END()
diff --git a/src/armnn/test/InferOutputTests.hpp b/src/armnn/test/InferOutputTests.hpp
index 6e5602a296..47eabd3cb0 100644
--- a/src/armnn/test/InferOutputTests.hpp
+++ b/src/armnn/test/InferOutputTests.hpp
@@ -13,6 +13,7 @@
 #include <layers/BatchToSpaceNdLayer.hpp>
 #include <layers/SpaceToDepthLayer.hpp>
 #include <layers/PreluLayer.hpp>
+#include <layers/StackLayer.hpp>
 
 #include <boost/algorithm/string.hpp>
 #include <boost/test/unit_test.hpp>
@@ -193,3 +194,156 @@ void PreluValidateTensorShapesFromInputsNoMatchTest()
     // Graph::InferTensorInfos calls Layer::ValidateTensorShapesFromInputs
     BOOST_CHECK_THROW(graph.InferTensorInfos(), armnn::LayerValidationException);
 }
+
+void StackInferOutputShapeImpl(const armnn::StackDescriptor           descriptor,
+                               const std::vector<armnn::TensorShape>& inputShapes,
+                               std::vector<armnn::TensorShape>&       outputShapes)
+{
+    armnn::Graph graph;
+    armnn::StackLayer* const stackLayer = graph.AddLayer<armnn::StackLayer>(descriptor, "stack");
+    outputShapes = stackLayer->InferOutputShapes(inputShapes);
+}
+
+void StackInferOutputShapeFromInputsMatchTest()
+{
+    armnn::Graph graph;
+
+    armnn::StackDescriptor descriptor;
+    descriptor.m_Axis = 1;
+    descriptor.m_NumInputs = 3;
+    descriptor.m_InputShape = armnn::TensorShape
+    (
+        { 4, 2 }  // Defined input shape
+    );
+
+    const std::vector<armnn::TensorShape> inputShapes
+    {
+        { 4, 2 }, // Actual input shapes
+        { 4, 2 },
+        { 4, 2 }
+    };
+
+    std::vector<armnn::TensorShape> outputShapes;
+    BOOST_CHECK_NO_THROW(StackInferOutputShapeImpl(descriptor, inputShapes, outputShapes));
+
+    armnn::TensorShape expectedOutputShape
+    (
+        { 4, 3, 2 }
+    );
+    BOOST_CHECK(outputShapes.size() == 1);
+    BOOST_CHECK(outputShapes[0] == expectedOutputShape);
+}
+
+void StackInferOutputShapeFromInputsNoMatchTest()
+{
+    armnn::Graph graph;
+
+    armnn::StackDescriptor descriptor;
+    descriptor.m_Axis = 1;
+    descriptor.m_NumInputs = 3;
+    descriptor.m_InputShape = armnn::TensorShape
+    (
+        { 4, 2 }  // Defined input shape
+    );
+
+    const std::vector<armnn::TensorShape> inputShapes
+    {
+        { 4, 2 }, // Actual input shapes
+        { 4, 5 }, // Incorrectly shaped input tensor
+        { 4, 2 }
+    };
+
+    // Output shape is inferred from the descriptor, so should still be correct despite mismatching input shapes
+    std::vector<armnn::TensorShape> outputShapes;
+    BOOST_CHECK_NO_THROW(StackInferOutputShapeImpl(descriptor, inputShapes, outputShapes));
+
+    armnn::TensorShape expectedOutputShape
+    (
+        { 4, 3, 2 }
+    );
+    BOOST_CHECK(outputShapes.size() == 1);
+    BOOST_CHECK(outputShapes[0] == expectedOutputShape);
+}
+
+void CreateStackLayerHelper(armnn::Graph& graph,
+                            const armnn::StackDescriptor& descriptor,
+                            const std::vector<armnn::TensorShape>& inputShapes,
+                            const armnn::TensorShape& outputShape)
+{
+    // Creates the Stack layer
+    armnn::Layer* const stackLayer = graph.AddLayer<armnn::StackLayer>(descriptor, "stack");
+
+    // Creates extra layers
+    std::vector<armnn::Layer*> inputs;
+    for (unsigned int i=0; i<inputShapes.size(); ++i)
+    {
+        inputs.push_back(graph.AddLayer<armnn::InputLayer>(static_cast<int>(i), "input"));
+    }
+    armnn::Layer* const output = graph.AddLayer<armnn::OutputLayer>(0, "output");
+
+    // Connects up
+    std::vector<armnn::TensorInfo> inputTensorInfos;
+    for (unsigned int i=0; i<inputs.size(); ++i)
+    {
+        inputTensorInfos.push_back(armnn::TensorInfo(inputShapes[i], armnn::DataType::Float32));
+    }
+    armnn::TensorInfo outputTensorInfo(outputShape, armnn::DataType::Float32);
+
+    for (unsigned int i=0; i<inputs.size(); ++i)
+    {
+        Connect(inputs[i], stackLayer, inputTensorInfos[i], 0, i);
+    }
+    Connect(stackLayer, output, outputTensorInfo, 0, 0);
+}
+
+void StackValidateTensorShapesFromInputsMatchTest()
+{
+    armnn::Graph graph;
+
+    armnn::StackDescriptor descriptor;
+    descriptor.m_Axis = 0;
+    descriptor.m_NumInputs = 3;
+    descriptor.m_InputShape = armnn::TensorShape
+    (
+        { 2, 5 }  // Defined input shape
+    );
+
+    const std::vector<armnn::TensorShape> inputShapes
+    {
+        { 2, 5 }, // Actual input shapes
+        { 2, 5 },
+        { 2, 5 }
+    };
+
+    // Creates the Stack layer
+    CreateStackLayerHelper(graph, descriptor, inputShapes, { 3, 2, 5 });
+
+    // Graph::InferTensorInfos calls Layer::ValidateTensorShapesFromInputs
+    BOOST_CHECK_NO_THROW(graph.InferTensorInfos());
+}
+
+void StackValidateTensorShapesFromInputsNoMatchTest()
+{
+    armnn::Graph graph;
+
+    armnn::StackDescriptor descriptor;
+    descriptor.m_Axis = 0;
+    descriptor.m_NumInputs = 3;
+    descriptor.m_InputShape = armnn::TensorShape
+    (
+        { 2, 5 }  // Defined input shape
+    );
+
+    const std::vector<armnn::TensorShape> inputShapes
+    {
+        { 2, 5 }, // Actual input shapes
+        { 2, 2 }, // Incorrectly shaped input tensor
+        { 2, 5 }
+    };
+
+    // Creates the Stack layer
+    CreateStackLayerHelper(graph, descriptor, inputShapes, { 3, 2, 5 });
+
+    // Graph::InferTensorInfos calls Layer::ValidateTensorShapesFromInputs
+    BOOST_CHECK_THROW(graph.InferTensorInfos(), armnn::LayerValidationException);
+}