IVGCVSW-2024: Support NHWC for Pooling2D CpuRef

* Adds implementation to plumb DataLayout parameter
  for Pooling2D on CpuRef.
* Adds unit tests to execute Pooling2D on CpuRef
  using NHWC data layout.
* Refactors original tests to use DataLayoutIndexed
  and removes duplicate code.

Change-Id: Ife7e0861a886cf58a2042e5be20e5b27af4528c9

11 files changed, 125 insertions, 155 deletions

diff --git a/include/armnn/Descriptors.hpp b/include/armnn/Descriptors.hpp
index a7eca43da3..648477e09b 100644
--- a/include/armnn/Descriptors.hpp
+++ b/include/armnn/Descriptors.hpp
@@ -194,7 +194,7 @@ struct Pooling2dDescriptor
     uint32_t            m_StrideY;
     OutputShapeRounding m_OutputShapeRounding;
     PaddingMethod       m_PaddingMethod;
-    DataLayout          m_DataLayout;
+    DataLayoutIndexed   m_DataLayout;
 };
 
 struct FullyConnectedDescriptor
diff --git a/src/armnn/test/CreateWorkload.hpp b/src/armnn/test/CreateWorkload.hpp
index 01c5e9f689..aac0a4ae6d 100644
--- a/src/armnn/test/CreateWorkload.hpp
+++ b/src/armnn/test/CreateWorkload.hpp
@@ -609,6 +609,7 @@ std::unique_ptr<Pooling2dWorkload> CreatePooling2dWorkloadTest(armnn::IWorkloadF
     BOOST_TEST(queueDescriptor.m_Parameters.m_PadRight == 2);
     BOOST_TEST(queueDescriptor.m_Parameters.m_PadTop == 1);
     BOOST_TEST(queueDescriptor.m_Parameters.m_PadBottom == 1);
+    BOOST_TEST((queueDescriptor.m_Parameters.m_DataLayout == dataLayout));
 
     BOOST_TEST(queueDescriptor.m_Inputs.size() == 1);
     BOOST_TEST(queueDescriptor.m_Outputs.size() == 1);
@@ -856,7 +857,8 @@ std::unique_ptr<ResizeBilinearWorkload> CreateResizeBilinearWorkloadTest(armnn::
             inputShape =  { 2, 4, 4, 3 };
             outputShape = { 2, 2, 2, 3 };
             break;
-        default: // NCHW
+        case DataLayout::NCHW:
+        default:
             inputShape =  { 2, 3, 4, 4 };
             outputShape = { 2, 3, 2, 2 };
     }
diff --git a/src/backends/cl/test/ClCreateWorkloadTests.cpp b/src/backends/cl/test/ClCreateWorkloadTests.cpp
index b5fc031461..4f9989405d 100644
--- a/src/backends/cl/test/ClCreateWorkloadTests.cpp
+++ b/src/backends/cl/test/ClCreateWorkloadTests.cpp
@@ -712,7 +712,8 @@ static void ClResizeBilinearWorkloadTest(DataLayout dataLayout)
             BOOST_TEST(CompareIClTensorHandleShape(inputHandle, { 2, 4, 4, 3 }));
             BOOST_TEST(CompareIClTensorHandleShape(outputHandle, { 2, 2, 2, 3 }));
             break;
-        default: // NCHW
+        case DataLayout::NCHW:
+        default:
             BOOST_TEST(CompareIClTensorHandleShape(inputHandle, { 2, 3, 4, 4 }));
             BOOST_TEST(CompareIClTensorHandleShape(outputHandle, { 2, 3, 2, 2 }));
     }
diff --git a/src/backends/cl/workloads/ClPooling2dWorkload.cpp b/src/backends/cl/workloads/ClPooling2dWorkload.cpp
index 255f57341e..68512ff980 100644
--- a/src/backends/cl/workloads/ClPooling2dWorkload.cpp
+++ b/src/backends/cl/workloads/ClPooling2dWorkload.cpp
@@ -19,8 +19,10 @@ arm_compute::Status ClPooling2dWorkloadValidate(const TensorInfo& input,
     const TensorInfo& output,
     const Pooling2dDescriptor& descriptor)
 {
-    const arm_compute::TensorInfo aclInputInfo = BuildArmComputeTensorInfo(input, descriptor.m_DataLayout);
-    const arm_compute::TensorInfo aclOutputInfo = BuildArmComputeTensorInfo(output, descriptor.m_DataLayout);
+    const arm_compute::TensorInfo aclInputInfo =
+            BuildArmComputeTensorInfo(input, descriptor.m_DataLayout.GetDataLayout());
+    const arm_compute::TensorInfo aclOutputInfo =
+            BuildArmComputeTensorInfo(output, descriptor.m_DataLayout.GetDataLayout());
 
     arm_compute::PoolingLayerInfo layerInfo = BuildArmComputePoolingLayerInfo(descriptor);
 
@@ -36,7 +38,7 @@ ClPooling2dWorkload::ClPooling2dWorkload(
     arm_compute::ICLTensor& input = static_cast<IClTensorHandle*>(m_Data.m_Inputs[0])->GetTensor();
     arm_compute::ICLTensor& output = static_cast<IClTensorHandle*>(m_Data.m_Outputs[0])->GetTensor();
 
-    arm_compute::DataLayout aclDataLayout = ConvertDataLayout(m_Data.m_Parameters.m_DataLayout);
+    arm_compute::DataLayout aclDataLayout = ConvertDataLayout(m_Data.m_Parameters.m_DataLayout.GetDataLayout());
     input.info()->set_data_layout(aclDataLayout);
     output.info()->set_data_layout(aclDataLayout);
 
diff --git a/src/backends/neon/workloads/NeonPooling2dWorkload.cpp b/src/backends/neon/workloads/NeonPooling2dWorkload.cpp
index 7892b3a315..a4ed33b3fa 100644
--- a/src/backends/neon/workloads/NeonPooling2dWorkload.cpp
+++ b/src/backends/neon/workloads/NeonPooling2dWorkload.cpp
@@ -17,8 +17,10 @@ arm_compute::Status NeonPooling2dWorkloadValidate(const TensorInfo& input,
     const TensorInfo& output,
     const Pooling2dDescriptor& descriptor)
 {
-    const arm_compute::TensorInfo aclInputInfo = BuildArmComputeTensorInfo(input, descriptor.m_DataLayout);
-    const arm_compute::TensorInfo aclOutputInfo = BuildArmComputeTensorInfo(output, descriptor.m_DataLayout);
+    const arm_compute::TensorInfo aclInputInfo =
+            BuildArmComputeTensorInfo(input, descriptor.m_DataLayout.GetDataLayout());
+    const arm_compute::TensorInfo aclOutputInfo =
+            BuildArmComputeTensorInfo(output, descriptor.m_DataLayout.GetDataLayout());
 
     arm_compute::PoolingLayerInfo layerInfo = BuildArmComputePoolingLayerInfo(descriptor);
 
@@ -34,7 +36,7 @@ NeonPooling2dWorkload::NeonPooling2dWorkload(
     arm_compute::ITensor& input = boost::polymorphic_downcast<INeonTensorHandle*>(m_Data.m_Inputs[0])->GetTensor();
     arm_compute::ITensor& output = boost::polymorphic_downcast<INeonTensorHandle*>(m_Data.m_Outputs[0])->GetTensor();
 
-    arm_compute::DataLayout aclDataLayout = ConvertDataLayout(m_Data.m_Parameters.m_DataLayout);
+    arm_compute::DataLayout aclDataLayout = ConvertDataLayout(m_Data.m_Parameters.m_DataLayout.GetDataLayout());
     input.info()->set_data_layout(aclDataLayout);
     output.info()->set_data_layout(aclDataLayout);
 
diff --git a/src/backends/reference/test/RefCreateWorkloadTests.cpp b/src/backends/reference/test/RefCreateWorkloadTests.cpp
index 8bad5497a2..d9322709b2 100644
--- a/src/backends/reference/test/RefCreateWorkloadTests.cpp
+++ b/src/backends/reference/test/RefCreateWorkloadTests.cpp
@@ -308,27 +308,51 @@ BOOST_AUTO_TEST_CASE(CreateRefNormalizationNhwcWorkload)
 }
 
 template <typename Pooling2dWorkloadType, armnn::DataType DataType>
-static void RefCreatePooling2dWorkloadTest()
+static void RefCreatePooling2dWorkloadTest(DataLayout dataLayout)
 {
     Graph graph;
     RefWorkloadFactory factory;
-    auto workload = CreatePooling2dWorkloadTest<Pooling2dWorkloadType, DataType>(factory, graph);
+    auto workload = CreatePooling2dWorkloadTest<Pooling2dWorkloadType, DataType>(factory, graph, dataLayout);
+
+    TensorShape inputShape;
+    TensorShape outputShape;
+
+    switch (dataLayout)
+    {
+        case DataLayout::NHWC:
+            inputShape  = { 3, 5, 5, 2 };
+            outputShape = { 3, 2, 4, 2 };
+            break;
+        case DataLayout::NCHW:
+        default:
+            inputShape =  { 3, 2, 5, 5 };
+            outputShape = { 3, 2, 2, 4 };
+    }
 
     // Checks that outputs and inputs are as we expect them (see definition of CreatePooling2dWorkloadTest).
-    CheckInputOutput(
-        std::move(workload),
-        TensorInfo({3, 2, 5, 5}, DataType),
-        TensorInfo({3, 2, 2, 4}, DataType));
+    CheckInputOutput(std::move(workload),
+                     TensorInfo(inputShape, DataType),
+                     TensorInfo(outputShape, DataType));
 }
 
 BOOST_AUTO_TEST_CASE(CreatePooling2dFloat32Workload)
 {
-    RefCreatePooling2dWorkloadTest<RefPooling2dFloat32Workload, armnn::DataType::Float32>();
+    RefCreatePooling2dWorkloadTest<RefPooling2dFloat32Workload, armnn::DataType::Float32>(DataLayout::NCHW);
+}
+
+BOOST_AUTO_TEST_CASE(CreatePooling2dFloat32NhwcWorkload)
+{
+    RefCreatePooling2dWorkloadTest<RefPooling2dFloat32Workload, armnn::DataType::Float32>(DataLayout::NHWC);
 }
 
 BOOST_AUTO_TEST_CASE(CreatePooling2dUint8Workload)
 {
-    RefCreatePooling2dWorkloadTest<RefPooling2dUint8Workload, armnn::DataType::QuantisedAsymm8>();
+    RefCreatePooling2dWorkloadTest<RefPooling2dUint8Workload, armnn::DataType::QuantisedAsymm8>(DataLayout::NCHW);
+}
+
+BOOST_AUTO_TEST_CASE(CreatePooling2dUint8NhwcWorkload)
+{
+    RefCreatePooling2dWorkloadTest<RefPooling2dUint8Workload, armnn::DataType::QuantisedAsymm8>(DataLayout::NHWC);
 }
 
 template <typename SoftmaxWorkloadType, armnn::DataType DataType>
@@ -496,16 +520,16 @@ static void RefCreateResizeBilinearTest(DataLayout dataLayout)
             inputShape  = { 2, 4, 4, 3 };
             outputShape = { 2, 2, 2, 3 };
             break;
-        default: // NCHW
+        case DataLayout::NCHW:
+        default:
             inputShape  = { 2, 3, 4, 4 };
             outputShape = { 2, 3, 2, 2 };
     }
 
     // Checks that outputs and inputs are as we expect them (see definition of CreateResizeBilinearWorkloadTest).
-    CheckInputOutput(
-            std::move(workload),
-            TensorInfo(inputShape, DataType),
-            TensorInfo(outputShape, DataType));
+    CheckInputOutput(std::move(workload),
+                     TensorInfo(inputShape, DataType),
+                     TensorInfo(outputShape, DataType));
 }
 
 BOOST_AUTO_TEST_CASE(CreateResizeBilinearFloat32)
diff --git a/src/backends/reference/test/RefLayerTests.cpp b/src/backends/reference/test/RefLayerTests.cpp
index 9f044cdbaf..30f5b10b0e 100644
--- a/src/backends/reference/test/RefLayerTests.cpp
+++ b/src/backends/reference/test/RefLayerTests.cpp
@@ -80,7 +80,9 @@ ARMNN_AUTO_TEST_CASE(IgnorePaddingL2Pooling2dSize3, IgnorePaddingL2Pooling2dSize
 ARMNN_AUTO_TEST_CASE(IgnorePaddingL2Pooling2dSize3Uint8, IgnorePaddingL2Pooling2dSize3Uint8Test)
 
 ARMNN_AUTO_TEST_CASE(SimpleAveragePooling2d, SimpleAveragePooling2dTest)
+ARMNN_AUTO_TEST_CASE(SimpleAveragePooling2dNhwc, SimpleAveragePooling2dNhwcTest)
 ARMNN_AUTO_TEST_CASE(SimpleAveragePooling2dUint8, SimpleAveragePooling2dUint8Test)
+ARMNN_AUTO_TEST_CASE(SimpleAveragePooling2dUint8Nhwc, SimpleAveragePooling2dUint8NhwcTest)
 ARMNN_AUTO_TEST_CASE(IgnorePaddingAveragePooling2dSize3x2Stride2x2,
                      IgnorePaddingAveragePooling2dSize3x2Stride2x2Test, false)
 ARMNN_AUTO_TEST_CASE(IgnorePaddingAveragePooling2dSize3x2Stride2x2NoPadding,
diff --git a/src/backends/reference/workloads/Pooling2d.cpp b/src/backends/reference/workloads/Pooling2d.cpp
index 5812a290e7..9890920113 100644
--- a/src/backends/reference/workloads/Pooling2d.cpp
+++ b/src/backends/reference/workloads/Pooling2d.cpp
@@ -143,12 +143,16 @@ void Pooling2d(const float* in,
                const TensorInfo& outputInfo,
                const Pooling2dDescriptor& params)
 {
+    const unsigned int channelsIndex = params.m_DataLayout.GetChannelsIndex();
+    const unsigned int heightIndex   = params.m_DataLayout.GetHeightIndex();
+    const unsigned int widthIndex    = params.m_DataLayout.GetWidthIndex();
+
     const int batchSize    = boost::numeric_cast<int>(outputInfo.GetShape()[0]);
-    const int channels     = boost::numeric_cast<int>(outputInfo.GetShape()[1]);
-    const int heightOutput = boost::numeric_cast<int>(outputInfo.GetShape()[2]);
-    const int widthOutput  = boost::numeric_cast<int>(outputInfo.GetShape()[3]);
-    const int heightInput  = boost::numeric_cast<int>(inputInfo.GetShape()[2]);
-    const int widthInput   = boost::numeric_cast<int>(inputInfo.GetShape()[3]);
+    const int channels     = boost::numeric_cast<int>(outputInfo.GetShape()[channelsIndex]);
+    const int heightOutput = boost::numeric_cast<int>(outputInfo.GetShape()[heightIndex]);
+    const int widthOutput  = boost::numeric_cast<int>(outputInfo.GetShape()[widthIndex]);
+    const int heightInput  = boost::numeric_cast<int>(inputInfo.GetShape()[heightIndex]);
+    const int widthInput   = boost::numeric_cast<int>(inputInfo.GetShape()[widthIndex]);
     const int padLeft      = boost::numeric_cast<int>(params.m_PadLeft);
     const int padRight     = boost::numeric_cast<int>(params.m_PadRight);
     const int padTop       = boost::numeric_cast<int>(params.m_PadTop);
diff --git a/src/backends/test/LayerTests.cpp b/src/backends/test/LayerTests.cpp
index e762152b67..5c7a887552 100755
--- a/src/backends/test/LayerTests.cpp
+++ b/src/backends/test/LayerTests.cpp
@@ -5585,17 +5585,22 @@ LayerTestResult<uint8_t, 4> SimpleMaxPooling2dSize3x3Stride2x4Uint8Test(armnn::I
 
 LayerTestResult<float, 4> SimpleAveragePooling2dTest(armnn::IWorkloadFactory& workloadFactory)
 {
-    return SimpleAveragePooling2dTestCommon<float>(workloadFactory);
+    return SimpleAveragePooling2dTest<float>(workloadFactory);
 }
 
 LayerTestResult<float, 4> SimpleAveragePooling2dNhwcTest(armnn::IWorkloadFactory& workloadFactory)
 {
-    return SimpleAveragePooling2dNhwcTestCommon<float>(workloadFactory);
+    return SimpleAveragePooling2dNhwcTest<float>(workloadFactory);
 }
 
 LayerTestResult<uint8_t, 4> SimpleAveragePooling2dUint8Test(armnn::IWorkloadFactory& workloadFactory)
 {
-    return SimpleAveragePooling2dTestCommon<uint8_t>(workloadFactory, 0.5, -1);
+    return SimpleAveragePooling2dTest<uint8_t>(workloadFactory, 0.5, -1);
+}
+
+LayerTestResult<uint8_t, 4> SimpleAveragePooling2dUint8NhwcTest(armnn::IWorkloadFactory& workloadFactory)
+{
+    return SimpleAveragePooling2dNhwcTest<uint8_t>(workloadFactory, 0.5, -1);
 }
 
 LayerTestResult<float, 4> IgnorePaddingAveragePooling2dSize3x2Stride2x2Test(armnn::IWorkloadFactory& workloadFactory,
diff --git a/src/backends/test/LayerTests.hpp b/src/backends/test/LayerTests.hpp
index 5790869366..9aae5da6e0 100644
--- a/src/backends/test/LayerTests.hpp
+++ b/src/backends/test/LayerTests.hpp
@@ -93,6 +93,7 @@ LayerTestResult<uint8_t, 4> IgnorePaddingMaxPooling2dSize3Uint8Test(armnn::IWork
 LayerTestResult<float,   4> SimpleAveragePooling2dTest(armnn::IWorkloadFactory& workloadFactory);
 LayerTestResult<float,   4> SimpleAveragePooling2dNhwcTest(armnn::IWorkloadFactory& workloadFactory);
 LayerTestResult<uint8_t, 4> SimpleAveragePooling2dUint8Test(armnn::IWorkloadFactory& workloadFactory);
+LayerTestResult<uint8_t, 4> SimpleAveragePooling2dUint8NhwcTest(armnn::IWorkloadFactory& workloadFactory);
 LayerTestResult<float,   4> IgnorePaddingAveragePooling2dSize3x2Stride2x2Test(armnn::IWorkloadFactory& workloadFactory,
                                                                               bool forceNoPadding);
 LayerTestResult<float,   4> IgnorePaddingSimpleAveragePooling2dTest(armnn::IWorkloadFactory& workloadFactory);
diff --git a/src/backends/test/Pooling2dTestImpl.hpp b/src/backends/test/Pooling2dTestImpl.hpp
index 7366bcaf44..29263af9bc 100644
--- a/src/backends/test/Pooling2dTestImpl.hpp
+++ b/src/backends/test/Pooling2dTestImpl.hpp
@@ -15,100 +15,44 @@
 #include <algorithm>
 
 template<typename T>
-LayerTestResult<T, 4> SimplePooling2dTestImpl(
-                                              armnn::IWorkloadFactory& workloadFactory,
+LayerTestResult<T, 4> SimplePooling2dTestImpl(armnn::IWorkloadFactory& workloadFactory,
                                               armnn::Pooling2dDescriptor descriptor,
                                               float qScale,
                                               int32_t qOffset,
                                               const boost::multi_array<T, 4>& input,
                                               const boost::multi_array<T, 4>& outputExpected)
 {
-    unsigned int inputHeight = boost::numeric_cast<unsigned int>(input.shape()[2]);
-    unsigned int inputWidth = boost::numeric_cast<unsigned int>(input.shape()[3]);
-    unsigned int inputChannels = boost::numeric_cast<unsigned int>(input.shape()[1]);
-    unsigned int inputBatchSize = boost::numeric_cast<unsigned int>(input.shape()[0]);
-
-    unsigned int outputHeight = boost::numeric_cast<unsigned int>(outputExpected.shape()[2]);
-    unsigned int outputWidth = boost::numeric_cast<unsigned int>(outputExpected.shape()[3]);
-    unsigned int outputChannels = boost::numeric_cast<unsigned int>(outputExpected.shape()[1]);
+    const unsigned int channelsIndex = descriptor.m_DataLayout.GetChannelsIndex();
+    const unsigned int heightIndex   = descriptor.m_DataLayout.GetHeightIndex();
+    const unsigned int widthIndex    = descriptor.m_DataLayout.GetWidthIndex();
+
+    unsigned int inputHeight     = boost::numeric_cast<unsigned int>(input.shape()[heightIndex]);
+    unsigned int inputWidth      = boost::numeric_cast<unsigned int>(input.shape()[widthIndex]);
+    unsigned int inputChannels   = boost::numeric_cast<unsigned int>(input.shape()[channelsIndex]);
+    unsigned int inputBatchSize  = boost::numeric_cast<unsigned int>(input.shape()[0]);
+
+    unsigned int outputHeight    = boost::numeric_cast<unsigned int>(outputExpected.shape()[heightIndex]);
+    unsigned int outputWidth     = boost::numeric_cast<unsigned int>(outputExpected.shape()[widthIndex]);
+    unsigned int outputChannels  = boost::numeric_cast<unsigned int>(outputExpected.shape()[channelsIndex]);
     unsigned int outputBatchSize = boost::numeric_cast<unsigned int>(outputExpected.shape()[0]);
 
-    armnn::TensorInfo inputTensorInfo({ inputBatchSize, inputChannels, inputHeight, inputWidth },
-        armnn::GetDataType<T>());
-    armnn::TensorInfo outputTensorInfo({ outputBatchSize, outputChannels, outputHeight, outputWidth },
-        armnn::GetDataType<T>());
+    armnn::TensorShape inputTensorShape;
+    armnn::TensorShape outputTensorShape;
 
-    // Set quantization parameters if the requested type is a quantized type.
-    if(armnn::IsQuantizedType<T>())
-    {
-        inputTensorInfo.SetQuantizationScale(qScale);
-        inputTensorInfo.SetQuantizationOffset(qOffset);
-        outputTensorInfo.SetQuantizationScale(qScale);
-        outputTensorInfo.SetQuantizationOffset(qOffset);
-    }
-
-    LayerTestResult<T, 4> result(outputTensorInfo);
-
-    std::unique_ptr<armnn::ITensorHandle> inputHandle = workloadFactory.CreateTensorHandle(inputTensorInfo);
-    std::unique_ptr<armnn::ITensorHandle> outputHandle = workloadFactory.CreateTensorHandle(outputTensorInfo);
-
-    armnn::Pooling2dQueueDescriptor queueDescriptor;
-    queueDescriptor.m_Parameters = descriptor;
-    armnn::WorkloadInfo workloadInfo;
-    AddInputToWorkload(queueDescriptor, workloadInfo, inputTensorInfo, inputHandle.get());
-    AddOutputToWorkload(queueDescriptor, workloadInfo, outputTensorInfo, outputHandle.get());
-
-    // Don't execute if Pooling is not supported, as an exception will be raised.
-    armnn::Compute compute = workloadFactory.GetCompute();
-    const size_t reasonIfUnsupportedMaxLen = 255;
-    char reasonIfUnsupported[reasonIfUnsupportedMaxLen+1];
-    result.supported = armnn::IsPooling2dSupported(compute, inputTensorInfo, outputTensorInfo,
-                                                   queueDescriptor.m_Parameters,
-                                                   reasonIfUnsupported, reasonIfUnsupportedMaxLen);
-    if (!result.supported)
+    switch (descriptor.m_DataLayout.GetDataLayout())
     {
-        return result;
+        case armnn::DataLayout::NHWC:
+            inputTensorShape = { inputBatchSize, inputHeight, inputWidth, inputChannels };
+            outputTensorShape = { outputBatchSize, outputHeight, outputWidth, outputChannels };
+            break;
+        case armnn::DataLayout::NCHW:
+        default:
+            inputTensorShape = { inputBatchSize, inputChannels, inputHeight, inputWidth };
+            outputTensorShape = { outputBatchSize, outputChannels, outputHeight, outputWidth };
     }
 
-    std::unique_ptr<armnn::IWorkload> workload = workloadFactory.CreatePooling2d(queueDescriptor, workloadInfo);
-
-    inputHandle->Allocate();
-    outputHandle->Allocate();
-
-    CopyDataToITensorHandle(inputHandle.get(), &input[0][0][0][0]);
-
-    workload->Execute();
-
-    CopyDataFromITensorHandle(&result.output[0][0][0][0], outputHandle.get());
-
-    result.outputExpected = outputExpected;
-
-    return result;
-}
-
-template<typename T>
-LayerTestResult<T, 4> SimplePooling2dNhwcTestImpl(
-                                                  armnn::IWorkloadFactory& workloadFactory,
-                                                  armnn::Pooling2dDescriptor descriptor,
-                                                  float qScale,
-                                                  int32_t qOffset,
-                                                  const boost::multi_array<T, 4>& input,
-                                                  const boost::multi_array<T, 4>& outputExpected)
-{
-    unsigned int inputHeight = boost::numeric_cast<unsigned int>(input.shape()[1]);
-    unsigned int inputWidth = boost::numeric_cast<unsigned int>(input.shape()[2]);
-    unsigned int inputChannels = boost::numeric_cast<unsigned int>(input.shape()[3]);
-    unsigned int inputBatchSize = boost::numeric_cast<unsigned int>(input.shape()[0]);
-
-    unsigned int outputHeight = boost::numeric_cast<unsigned int>(outputExpected.shape()[1]);
-    unsigned int outputWidth = boost::numeric_cast<unsigned int>(outputExpected.shape()[2]);
-    unsigned int outputChannels = boost::numeric_cast<unsigned int>(outputExpected.shape()[3]);
-    unsigned int outputBatchSize = boost::numeric_cast<unsigned int>(outputExpected.shape()[0]);
-
-    armnn::TensorInfo inputTensorInfo({ inputBatchSize, inputHeight, inputWidth, inputChannels },
-                                      armnn::GetDataType<T>());
-    armnn::TensorInfo outputTensorInfo({ outputBatchSize, outputHeight, outputWidth, outputChannels },
-                                       armnn::GetDataType<T>());
+    armnn::TensorInfo inputTensorInfo(inputTensorShape, armnn::GetDataType<T>());
+    armnn::TensorInfo outputTensorInfo(outputTensorShape, armnn::GetDataType<T>());
 
     // Set quantization parameters if the requested type is a quantized type.
     if(armnn::IsQuantizedType<T>())
@@ -126,7 +70,7 @@ LayerTestResult<T, 4> SimplePooling2dNhwcTestImpl(
 
     armnn::Pooling2dQueueDescriptor queueDescriptor;
     queueDescriptor.m_Parameters = descriptor;
-    queueDescriptor.m_Parameters.m_DataLayout = armnn::DataLayout::NHWC;
+    queueDescriptor.m_Parameters.m_DataLayout = descriptor.m_DataLayout;
 
     armnn::WorkloadInfo workloadInfo;
     AddInputToWorkload(queueDescriptor, workloadInfo, inputTensorInfo, inputHandle.get());
@@ -291,6 +235,9 @@ LayerTestResult<T, 4> SimpleMaxPooling2dSize3x3Stride2x4TestCommon(armnn::IWorkl
 
 template<typename T>
 LayerTestResult<T, 4> SimpleAveragePooling2dTestCommon(armnn::IWorkloadFactory& workloadFactory,
+                                                       const armnn::TensorShape& inputTensorShape,
+                                                       const armnn::TensorShape& outputTensorShape,
+                                                       armnn::DataLayout dataLayout,
                                                        float qScale = 1.0f,
                                                        int32_t qOffset = 0)
 {
@@ -303,9 +250,10 @@ LayerTestResult<T, 4> SimpleAveragePooling2dTestCommon(armnn::IWorkloadFactory&
     descriptor.m_PadTop = 1;
     descriptor.m_PadBottom = 1;
     descriptor.m_PaddingMethod = armnn::PaddingMethod::Exclude;
+    descriptor.m_DataLayout = dataLayout;
 
-    armnn::TensorInfo inputTensorInfo({ 1, 1, 4, 4 }, armnn::GetDataType<T>());
-    armnn::TensorInfo outputTensorInfo({ 1, 1, 3, 3 }, armnn::GetDataType<T>());
+    armnn::TensorInfo inputTensorInfo(inputTensorShape, armnn::GetDataType<T>());
+    armnn::TensorInfo outputTensorInfo(outputTensorShape, armnn::GetDataType<T>());
 
     // Set quantization parameters if the requested type is a quantized type.
     if(armnn::IsQuantizedType<T>())
@@ -335,54 +283,33 @@ LayerTestResult<T, 4> SimpleAveragePooling2dTestCommon(armnn::IWorkloadFactory&
 }
 
 template<typename T>
-LayerTestResult<T, 4> SimpleAveragePooling2dNhwcTestCommon(armnn::IWorkloadFactory& workloadFactory,
-                                                           float qScale = 1.0f,
-                                                           int32_t qOffset = 0)
+LayerTestResult<T, 4> SimpleAveragePooling2dTest(armnn::IWorkloadFactory& workloadFactory,
+                                                 float qScale = 1.0f,
+                                                 int32_t qOffset = 0)
 {
-    armnn::Pooling2dDescriptor descriptor;
-    descriptor.m_PoolType = armnn::PoolingAlgorithm::Average;
-    descriptor.m_PoolWidth = descriptor.m_PoolHeight = 2;
-    descriptor.m_StrideX = descriptor.m_StrideY = 2;
-    descriptor.m_PadLeft = 1;
-    descriptor.m_PadRight = 1;
-    descriptor.m_PadTop = 1;
-    descriptor.m_PadBottom = 1;
-    descriptor.m_PaddingMethod = armnn::PaddingMethod::Exclude;
-
-    armnn::TensorInfo inputTensorInfo({ 1, 4, 4, 1 }, armnn::GetDataType<T>());
-    armnn::TensorInfo outputTensorInfo({ 1, 3, 3, 1 }, armnn::GetDataType<T>());
-
-    // Set quantization parameters if the requested type is a quantized type.
-    if(armnn::IsQuantizedType<T>())
-    {
-        inputTensorInfo.SetQuantizationScale(qScale);
-        inputTensorInfo.SetQuantizationOffset(qOffset);
-        outputTensorInfo.SetQuantizationScale(qScale);
-        outputTensorInfo.SetQuantizationOffset(qOffset);
-    }
+    const armnn::TensorShape inputTensorShape  { 1, 1, 4, 4 };
+    const armnn::TensorShape outputTensorShape { 1, 1, 3, 3 };
 
-    auto input = MakeTensor<T, 4>(inputTensorInfo,
-                                  QuantizedVector<T>(qScale, qOffset, {
-                                          1.0f, 2.0f, 3.0f, 4.0f,
-                                          1.0f, 2.0f, 3.0f, 4.0f,
-                                          1.0f, 2.0f, 3.0f, 4.0f,
-                                          1.0f, 2.0f, 3.0f, 4.0f,
-                                  }));
+    return SimpleAveragePooling2dTestCommon<T>(workloadFactory, inputTensorShape, outputTensorShape,
+                                               armnn::DataLayout::NCHW, qScale, qOffset);
+}
 
-    auto outputExpected = MakeTensor<T, 4>(outputTensorInfo,
-                                           QuantizedVector<T>(qScale, qOffset, {
-                                                   1.0f, 2.5f, 4.0f,
-                                                   1.0f, 2.5f, 4.0f,
-                                                   1.0f, 2.5f, 4.0f,
-                                           }));
+template<typename T>
+LayerTestResult<T, 4> SimpleAveragePooling2dNhwcTest(armnn::IWorkloadFactory& workloadFactory,
+                                                     float qScale = 1.0f,
+                                                     int32_t qOffset = 0)
+{
+    const armnn::TensorShape inputTensorShape  { 1, 4, 4, 1 };
+    const armnn::TensorShape outputTensorShape { 1, 3, 3, 1 };
 
-    return SimplePooling2dNhwcTestImpl<T>(workloadFactory, descriptor, qScale, qOffset, input, outputExpected);
+    return SimpleAveragePooling2dTestCommon<T>(workloadFactory, inputTensorShape, outputTensorShape,
+                                               armnn::DataLayout::NHWC, qScale, qOffset);
 }
 
 template<typename T>
 LayerTestResult<T, 4> LargeTensorsAveragePooling2dTestCommon(armnn::IWorkloadFactory& workloadFactory,
-                                                       float qScale = 1.0f,
-                                                       int32_t qOffset = 0)
+                                                             float qScale = 1.0f,
+                                                             int32_t qOffset = 0)
 {
     armnn::Pooling2dDescriptor descriptor;
     descriptor.m_PoolType = armnn::PoolingAlgorithm::Average;