IVGCVSW-1883 Add support for different memory layouts

Change-Id: I6e9973bf25acad980fb4e96af8080ac829db0d28

4 files changed, 36 insertions, 2 deletions

diff --git a/include/armnn/Types.hpp b/include/armnn/Types.hpp
index 6d89e0ea00..172df1b454 100644
--- a/include/armnn/Types.hpp
+++ b/include/armnn/Types.hpp
@@ -28,6 +28,12 @@ enum class DataType
     Signed32  = 3
 };
 
+enum class DataLayout
+{
+    NCHW = 1,
+    NHWC = 2
+};
+
 enum class ActivationFunction
 {
     Sigmoid     = 0,
diff --git a/src/armnn/Layer.cpp b/src/armnn/Layer.cpp
index 48ab8b5dbe..b0a3ac543b 100644
--- a/src/armnn/Layer.cpp
+++ b/src/armnn/Layer.cpp
@@ -124,10 +124,15 @@ LayerGuid GenerateLayerGuid()
 }
 } // namespace
 
-Layer::Layer(unsigned int numInputSlots, unsigned int numOutputSlots, LayerType type, const char* name)
+Layer::Layer(unsigned int numInputSlots,
+             unsigned int numOutputSlots,
+             LayerType type,
+             DataLayout layout,
+             const char* name)
 : m_OutputHandlers(numOutputSlots)
 , m_LayerName(name ? name : "")
 , m_Type(type)
+, m_DataLayout(layout)
 , m_ComputeDevice(Compute::Undefined)
 , m_Guid(GenerateLayerGuid())
 {
@@ -144,6 +149,14 @@ Layer::Layer(unsigned int numInputSlots, unsigned int numOutputSlots, LayerType
     }
 }
 
+Layer::Layer(unsigned int numInputSlots,
+             unsigned int numOutputSlots,
+             LayerType type,
+             const char* name)
+: Layer(numInputSlots, numOutputSlots, type, DataLayout::NCHW, name)
+{
+}
+
 void Layer::CollectWorkloadInputs(WorkloadDataCollector& dataCollector, const Graph& graph) const
 {
     for (auto&& inputSlot : GetInputSlots())
diff --git a/src/armnn/Layer.hpp b/src/armnn/Layer.hpp
index 31837a5cc0..fd523cee9e 100644
--- a/src/armnn/Layer.hpp
+++ b/src/armnn/Layer.hpp
@@ -186,6 +186,7 @@ class Layer : public IConnectableLayer
 public:
     /// @param name - Optional name for the layer (may be nullptr).
     Layer(unsigned int numInputSlots, unsigned int numOutputSlots, LayerType type, const char* name);
+    Layer(unsigned int numInputSlots, unsigned int numOutputSlots, LayerType type, DataLayout layout, const char* name);
 
     const std::string& GetNameStr() const
     {
@@ -234,6 +235,8 @@ public:
 
     DataType GetDataType() const;
 
+    DataLayout GetDataLayout() const { return m_DataLayout; }
+
     Compute GetComputeDevice() const { return m_ComputeDevice; }
     void SetComputeDevice(Compute device) { m_ComputeDevice = device; }
 
@@ -341,6 +344,7 @@ private:
     std::vector<OutputSlot> m_OutputSlots;
 
     const LayerType m_Type;
+    const DataLayout m_DataLayout;
     Compute m_ComputeDevice;
 
     /// Used for sorting.
diff --git a/src/armnn/layers/LayerWithParameters.hpp b/src/armnn/layers/LayerWithParameters.hpp
index 0cb970a2f1..6156d6b66c 100644
--- a/src/armnn/layers/LayerWithParameters.hpp
+++ b/src/armnn/layers/LayerWithParameters.hpp
@@ -28,13 +28,24 @@ protected:
     LayerWithParameters(unsigned int numInputSlots,
                         unsigned int numOutputSlots,
                         LayerType type,
+                        DataLayout layout,
                         const Parameters& param,
                         const char* name)
-    :   Layer(numInputSlots, numOutputSlots, type, name)
+    :   Layer(numInputSlots, numOutputSlots, type, layout, name)
     ,   m_Param(param)
     {
     }
 
+    LayerWithParameters(unsigned int numInputSlots,
+                        unsigned int numOutputSlots,
+                        LayerType type,
+                        const Parameters& param,
+                        const char* name)
+        : Layer(numInputSlots, numOutputSlots, type, name)
+        , m_Param(param)
+    {
+    }
+
     ~LayerWithParameters() = default;
 
     /// Helper function to reduce duplication in *Layer::CreateWorkload.