COMPMID-806 Add NHWC data format support format for NEON pooling

Change-Id: I7ab174c72f3d56134fcec259a137739061fd12e9
Reviewed-on: https://eu-gerrit-1.euhpc.arm.com/123065
Reviewed-by: Georgios Pinitas <georgios.pinitas@arm.com>
Tested-by: Jenkins <bsgcomp@arm.com>

1 files changed, 36 insertions, 11 deletions

diff --git a/src/runtime/NEON/functions/NEPoolingLayer.cpp b/src/runtime/NEON/functions/NEPoolingLayer.cpp
index bc0b6f86d3..cbfd68485f 100644
--- a/src/runtime/NEON/functions/NEPoolingLayer.cpp
+++ b/src/runtime/NEON/functions/NEPoolingLayer.cpp
@@ -31,7 +31,7 @@
 using namespace arm_compute;
 
 NEPoolingLayer::NEPoolingLayer()
-    : _pooling_layer_kernel(), _border_handler(), _is_global_pooling_layer(false)
+    : _pooling_layer_kernel(), _border_handler(), _is_global_pooling_layer(false), _data_layout(DataLayout::NCHW)
 {
 }
 
@@ -40,17 +40,31 @@ void NEPoolingLayer::configure(ITensor *input, ITensor *output, const PoolingLay
     // Check if we have Global Pooling Layer
     _is_global_pooling_layer = (input->info()->dimension(0) == pool_info.pool_size().width) && (input->info()->dimension(1) == pool_info.pool_size().height);
 
+    // Get data layout
+    _data_layout = input->info()->data_layout();
+
     // Configure pooling kernel
     _pooling_layer_kernel.configure(input, output, pool_info);
 
-    // Configure border depending on operation required (quantize border in case of asymmetric data_type)
-    BorderMode border_mode = (pool_info.pool_type() == PoolingType::MAX) ? BorderMode::REPLICATE : BorderMode::CONSTANT;
-    PixelValue zero_value(0.f);
-    if(is_data_type_quantized_asymmetric(input->info()->data_type()) && !pool_info.exclude_padding())
+    switch(_data_layout)
     {
-        zero_value = PixelValue(static_cast<uint32_t>(input->info()->quantization_info().offset));
+        case DataLayout::NCHW:
+        {
+            // Configure border depending on operation required (quantize border in case of asymmetric data_type)
+            BorderMode border_mode = (pool_info.pool_type() == PoolingType::MAX) ? BorderMode::REPLICATE : BorderMode::CONSTANT;
+            PixelValue zero_value(0.f);
+            if(is_data_type_quantized_asymmetric(input->info()->data_type()) && !pool_info.exclude_padding())
+            {
+                zero_value = PixelValue(static_cast<uint32_t>(input->info()->quantization_info().offset));
+            }
+            _border_handler.configure(input, _pooling_layer_kernel.border_size(), border_mode, zero_value);
+            break;
+        }
+        case DataLayout::NHWC:
+            break;
+        default:
+            ARM_COMPUTE_ERROR("Data layout not supported");
     }
-    _border_handler.configure(input, _pooling_layer_kernel.border_size(), border_mode, zero_value);
 }
 
 Status NEPoolingLayer::validate(const ITensorInfo *input, const ITensorInfo *output, const PoolingLayerInfo &pool_info)
@@ -60,9 +74,20 @@ Status NEPoolingLayer::validate(const ITensorInfo *input, const ITensorInfo *out
 
 void NEPoolingLayer::run()
 {
-    // Fill border
-    NEScheduler::get().schedule(&_border_handler, Window::DimY);
+    switch(_data_layout)
+    {
+        case DataLayout::NCHW:
+            // Fill border
+            NEScheduler::get().schedule(&_border_handler, Window::DimY);
 
-    // Run pooling layer
-    NEScheduler::get().schedule(&_pooling_layer_kernel, _is_global_pooling_layer ? Window::DimZ : Window::DimY);
+            // Run pooling layer
+            NEScheduler::get().schedule(&_pooling_layer_kernel, _is_global_pooling_layer ? Window::DimZ : Window::DimY);
+            break;
+        case DataLayout::NHWC:
+            // Run pooling layer
+            NEScheduler::get().schedule(&_pooling_layer_kernel, Window::DimX);
+            break;
+        default:
+            ARM_COMPUTE_ERROR("Data layout not supported");
+    }
 }
\ No newline at end of file