COMPMID-403: Add 7x7 NEON Pooling support.

Change-Id: I2f1e808884f215b9cf79e1f2015ef901e66b3e5f
Reviewed-on: http://mpd-gerrit.cambridge.arm.com/78146
Reviewed-by: Georgios Pinitas <georgios.pinitas@arm.com>
Tested-by: Kaizen <jeremy.johnson+kaizengerrit@arm.com>

1 files changed, 111 insertions, 14 deletions

diff --git a/src/core/NEON/kernels/NEPoolingLayerKernel.cpp b/src/core/NEON/kernels/NEPoolingLayerKernel.cpp
index 30b67b64b9..8991e9b9ee 100644
--- a/src/core/NEON/kernels/NEPoolingLayerKernel.cpp
+++ b/src/core/NEON/kernels/NEPoolingLayerKernel.cpp
@@ -37,6 +37,7 @@
 #include <algorithm>
 #include <arm_neon.h>
 #include <limits>
+#include <set>
 #include <string>
 #include <tuple>
 
@@ -93,11 +94,15 @@ void NEPoolingLayerKernel::configure(const ITensor *input, ITensor *output, cons
     std::tie(pool_pad_x, pool_pad_y)       = pad_stride_info.pad();
     std::tie(pool_stride_x, pool_stride_y) = pad_stride_info.stride();
 
+    static const std::set<int> supported_pool_sizes = { 2, 3, 7 };
+    ARM_COMPUTE_UNUSED(supported_pool_sizes);
+
     ARM_COMPUTE_ERROR_ON_DATA_TYPE_CHANNEL_NOT_IN(input, 1, DataType::QS8, DataType::F32);
     ARM_COMPUTE_ERROR_ON_DATA_TYPE_CHANNEL_NOT_IN(output, 1, DataType::QS8, DataType::F32);
     ARM_COMPUTE_ERROR_ON_MISMATCHING_DATA_TYPES(input, output);
     ARM_COMPUTE_ERROR_ON_MISMATCHING_FIXED_POINT(input, output);
-    ARM_COMPUTE_ERROR_ON(2 != pool_size && 3 != pool_size);
+    ARM_COMPUTE_ERROR_ON(supported_pool_sizes.find(pool_size) == supported_pool_sizes.end());
+    ARM_COMPUTE_ERROR_ON(7 == pool_size && input->info()->data_type() != DataType::F32);
     ARM_COMPUTE_ERROR_ON(pool_pad_x >= pool_size || pool_pad_y >= pool_size);
     ARM_COMPUTE_ERROR_ON(input->info()->data_type() == DataType::QS8 && pool_type == PoolingType::AVG && input->info()->fixed_point_position() > 6);
     ARM_COMPUTE_ERROR_ON(input->info()->data_type() == DataType::QS8 && pool_stride_x > 2);
@@ -118,12 +123,35 @@ void NEPoolingLayerKernel::configure(const ITensor *input, ITensor *output, cons
     switch(input->info()->data_type())
     {
         case DataType::QS8:
-            num_elems_read_per_iteration      = 16;
-            num_elems_processed_per_iteration = (pool_size == 2) ? 8 : 7;
-            num_elems_horizontal_window       = 8;
+            num_elems_read_per_iteration = 16;
+            switch(pool_size)
+            {
+                case 2:
+                    num_elems_processed_per_iteration = 8;
+                    break;
+                case 3:
+                    num_elems_processed_per_iteration = 7;
+                    break;
+                default:
+                    ARM_COMPUTE_ERROR("Pooling size not supported");
+            }
+            num_elems_horizontal_window = 8;
             break;
         case DataType::F32:
-            num_elems_read_per_iteration      = (pool_size == 2) ? 2 : 4; // We use vload4 for pooling3
+            switch(pool_size)
+            {
+                case 2:
+                    num_elems_read_per_iteration = 2;
+                    break;
+                case 3:
+                    num_elems_read_per_iteration = 4; // We use vload4 for pooling3
+                    break;
+                case 7:
+                    num_elems_read_per_iteration = 8; // We use vload8 for pooling7
+                    break;
+                default:
+                    ARM_COMPUTE_ERROR("Pooling size not supported");
+            }
             num_elems_processed_per_iteration = 1;
             num_elems_horizontal_window       = 1;
             break;
@@ -169,6 +197,9 @@ void NEPoolingLayerKernel::configure(const ITensor *input, ITensor *output, cons
                 _func = (PoolingType::AVG == pool_type) ? &NEPoolingLayerKernel::pooling3_f32<PoolingType::AVG> : &NEPoolingLayerKernel::pooling3_f32<PoolingType::MAX>;
             }
             break;
+        case 7:
+            _func = (PoolingType::AVG == pool_type) ? &NEPoolingLayerKernel::pooling7_f32<PoolingType::AVG> : &NEPoolingLayerKernel::pooling7_f32<PoolingType::MAX>;
+            break;
         default:
             ARM_COMPUTE_ERROR("Unsupported pooling size");
             break;
@@ -234,15 +265,18 @@ void NEPoolingLayerKernel::pooling2_f32(const Window &window_input, const Window
     Iterator input(_input, window_input);
     Iterator output(_output, window);
 
-    constexpr int pool_size = 2;
-    int           pool_pad_x, pool_pad_y, pool_stride_x, pool_stride_y = 0;
+    constexpr int pool_size     = 2;
+    int           pool_pad_x    = 0;
+    int           pool_pad_y    = 0;
+    int           pool_stride_x = 0;
+    int           pool_stride_y = 0;
     std::tie(pool_pad_x, pool_pad_y)       = _pool_info.pad_stride_info().pad();
     std::tie(pool_stride_x, pool_stride_y) = _pool_info.pad_stride_info().stride();
     const int upper_bound_w = _input->info()->dimension(0) + pool_pad_x;
     const int upper_bound_h = _input->info()->dimension(1) + pool_pad_y;
 
-    const unsigned char *const input_top_ptr    = _input->ptr_to_element(Coordinates(-static_cast<int>(pool_pad_x), -static_cast<int>(pool_pad_y)));
-    const unsigned char *const input_bottom_ptr = _input->ptr_to_element(Coordinates(-static_cast<int>(pool_pad_x), -static_cast<int>(pool_pad_y) + 1));
+    const uint8_t *const input_top_ptr    = _input->ptr_to_element(Coordinates(-static_cast<int>(pool_pad_x), -static_cast<int>(pool_pad_y)));
+    const uint8_t *const input_bottom_ptr = _input->ptr_to_element(Coordinates(-static_cast<int>(pool_pad_x), -static_cast<int>(pool_pad_y) + 1));
 
     execute_window_loop(window, [&](const Coordinates & id)
     {
@@ -348,16 +382,19 @@ void NEPoolingLayerKernel::pooling3_f32(const Window &window_input, const Window
     Iterator input(_input, window_input);
     Iterator output(_output, window);
 
-    constexpr const int pool_size = 3;
-    int                 pool_pad_x, pool_pad_y, pool_stride_x, pool_stride_y = 0;
+    constexpr const int pool_size     = 3;
+    int                 pool_pad_x    = 0;
+    int                 pool_pad_y    = 0;
+    int                 pool_stride_x = 0;
+    int                 pool_stride_y = 0;
     std::tie(pool_pad_x, pool_pad_y)       = _pool_info.pad_stride_info().pad();
     std::tie(pool_stride_x, pool_stride_y) = _pool_info.pad_stride_info().stride();
     const int upper_bound_w = _input->info()->dimension(0) + pool_pad_x;
     const int upper_bound_h = _input->info()->dimension(1) + pool_pad_y;
 
-    const unsigned char *const input_top_ptr    = _input->ptr_to_element(Coordinates(-static_cast<int>(pool_pad_x), -static_cast<int>(pool_pad_y)));
-    const unsigned char *const input_middle_ptr = _input->ptr_to_element(Coordinates(-static_cast<int>(pool_pad_x), -static_cast<int>(pool_pad_y) + 1));
-    const unsigned char *const input_bottom_ptr = _input->ptr_to_element(Coordinates(-static_cast<int>(pool_pad_x), -static_cast<int>(pool_pad_y) + 2));
+    const uint8_t *const input_top_ptr    = _input->ptr_to_element(Coordinates(-static_cast<int>(pool_pad_x), -static_cast<int>(pool_pad_y)));
+    const uint8_t *const input_middle_ptr = _input->ptr_to_element(Coordinates(-static_cast<int>(pool_pad_x), -static_cast<int>(pool_pad_y) + 1));
+    const uint8_t *const input_bottom_ptr = _input->ptr_to_element(Coordinates(-static_cast<int>(pool_pad_x), -static_cast<int>(pool_pad_y) + 2));
 
     execute_window_loop(window, [&](const Coordinates & id)
     {
@@ -387,6 +424,66 @@ void NEPoolingLayerKernel::pooling3_f32(const Window &window_input, const Window
     input, output);
 }
 
+template <PoolingType pooling_type>
+void NEPoolingLayerKernel::pooling7_f32(const Window &window_input, const Window &window)
+{
+    Iterator input(_input, window_input);
+    Iterator output(_output, window);
+
+    constexpr const int pool_size     = 7;
+    int                 pool_pad_x    = 0;
+    int                 pool_pad_y    = 0;
+    int                 pool_stride_x = 0;
+    int                 pool_stride_y = 0;
+    std::tie(pool_pad_x, pool_pad_y)       = _pool_info.pad_stride_info().pad();
+    std::tie(pool_stride_x, pool_stride_y) = _pool_info.pad_stride_info().stride();
+    const int upper_bound_w = _input->info()->dimension(0) + pool_pad_x;
+    const int upper_bound_h = _input->info()->dimension(1) + pool_pad_y;
+
+    std::array<const uint8_t *, pool_size> input_ptrs{ {} };
+    for(int i = 0; i < pool_size; ++i)
+    {
+        input_ptrs[i] = _input->ptr_to_element(Coordinates(-static_cast<int>(pool_pad_x), -static_cast<int>(pool_pad_y) + i));
+    }
+
+    execute_window_loop(window, [&](const Coordinates & id)
+    {
+        float32x2_t res = {};
+        if(pooling_type == PoolingType::AVG)
+        {
+            // Calculate scale
+            float             scale   = calculate_avg_scale(id, pool_size, upper_bound_w, upper_bound_h, pool_pad_x, pool_pad_y, pool_stride_x, pool_stride_y);
+            const float32x2_t scale_v = vdup_n_f32(scale);
+
+            // Perform pooling
+            float32x4x2_t data     = vld2q_f32(reinterpret_cast<const float *>(input_ptrs[0] + input.offset()));
+            float32x4_t   sum_data = vaddq_f32(data.val[0], vsetq_lane_f32(0.f, data.val[1], 3));
+            for(int i = 1; i < pool_size; ++i)
+            {
+                data     = vld2q_f32(reinterpret_cast<const float *>(input_ptrs[i] + input.offset()));
+                sum_data = vaddq_f32(sum_data, data.val[0]);
+                sum_data = vaddq_f32(sum_data, vsetq_lane_f32(0.f, data.val[1], 3));
+            }
+            res = vpadd_f32(vget_high_f32(sum_data), vget_low_f32(sum_data));
+            res = vmul_f32(vpadd_f32(res, res), scale_v);
+        }
+        else
+        {
+            float32x4x2_t max_data = vld2q_f32(reinterpret_cast<const float *>(input_ptrs[0] + input.offset()));
+            for(int i = 1; i < pool_size; ++i)
+            {
+                const float32x4x2_t data = vld2q_f32(reinterpret_cast<const float *>(input_ptrs[i] + input.offset()));
+                max_data                 = vmax2q_f32(max_data, data);
+            }
+            res = vpmax_f32(vget_high_f32(vsetq_lane_f32(-std::numeric_limits<float>::max(), max_data.val[1], 3)), vget_low_f32(max_data.val[1]));
+            res = vpmax_f32(res, vpmax_f32(vget_high_f32(max_data.val[0]), vget_low_f32(max_data.val[0])));
+            res = vpmax_f32(res, res);
+        }
+        *(reinterpret_cast<float *>(output.ptr())) = vget_lane_f32(res, 0);
+    },
+    input, output);
+}
+
 void NEPoolingLayerKernel::run(const Window &window)
 {
     ARM_COMPUTE_ERROR_ON_UNCONFIGURED_KERNEL(this);