commit b3a371bc429d2ba45e56baaf239d8200c2662a74
author Jenkins <bsgcomp@arm.com> Wed May 23 11:36:53 2018 +0100
committer Anthony Barbier <anthony.barbier@arm.com> Wed May 23 14:55:11 2018 +0100
tree 554525e415c303d64a08722a755397852ebbb8e4
parent 67c8c91522e5be8156b77f57e63c0253535c902a

arm_compute v18.05

src/runtime/NEON/functions/NEPoolingLayer.cpp

diff --git a/src/runtime/NEON/functions/NEPoolingLayer.cpp b/src/runtime/NEON/functions/NEPoolingLayer.cpp
index bc0b6f8..cbfd684 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 @@
     // 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 @@
 
 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