COMPMID-3576: Nightly failure: NEON/PoolingLayer/Float/FP16/MaxUnpooling S10

Extend NEPoolingLayer max pooling to extract indices for FP16

Signed-off-by: Sheri Zhang <sheri.zhang@arm.com>
Change-Id: I5a7c754be353e4c2c5d0ab3794e9427408d0c4fa
Reviewed-on: https://review.mlplatform.org/c/ml/ComputeLibrary/+/3580
Reviewed-by: Sang-Hoon Park <sang-hoon.park@arm.com>
Reviewed-by: Michele Di Giorgio <michele.digiorgio@arm.com>
Tested-by: Arm Jenkins <bsgcomp@arm.com>
Comments-Addressed: Arm Jenkins <bsgcomp@arm.com>

7 files changed, 262 insertions, 140 deletions

diff --git a/arm_compute/core/NEON/kernels/NEMaxUnpoolingLayerKernel.h b/arm_compute/core/NEON/kernels/NEMaxUnpoolingLayerKernel.h
index 9c9b1945f6..f3ea049a87 100644
--- a/arm_compute/core/NEON/kernels/NEMaxUnpoolingLayerKernel.h
+++ b/arm_compute/core/NEON/kernels/NEMaxUnpoolingLayerKernel.h
@@ -55,7 +55,10 @@ public:
      * @note Output shape must be equal to the shape of the original input to pool.
      *
      * @param[in]  input     Source tensor. Data types supported: QASYMM8/QASYMM8_SIGNED/F16/F32.
-     * @param[out] indices   The indices of the maximal values. Data type supported: U32.
+     * @param[in]  indices   Tensor containing the offset to store the input elements in the output tensor.
+     *                       @ref NEPoolingLayerKernel with indices should precede this function in order to
+     *                       properly reconstruct the output tensor.
+     *                       The tensor shape of this tensor has to be equal to the input tensor shape. Data type supported: U32.
      * @param[out] output    Destination tensor. Data types supported: Same as @p input.
      * @param[in]  pool_info Contains pooling operation information described in @ref PoolingLayerInfo.
      */
diff --git a/arm_compute/core/NEON/kernels/NEPoolingLayerKernel.h b/arm_compute/core/NEON/kernels/NEPoolingLayerKernel.h
index 15d63c746b..2be25080cd 100644
--- a/arm_compute/core/NEON/kernels/NEPoolingLayerKernel.h
+++ b/arm_compute/core/NEON/kernels/NEPoolingLayerKernel.h
@@ -91,12 +91,6 @@ private:
      * @param[in] window_input Input region on which to execute the kernel.
      * @param[in] window       Output region on which to execute the kernel.
      */
-    void pooling2_f32_nchw_maxpool_indices(const Window &window_input, const Window &window);
-    /** Function to perform 2x2 pooling and compute the pooling indices. The indices can be used for max unpool.
-     *
-     * @param[in] window_input Input region on which to execute the kernel.
-     * @param[in] window       Output region on which to execute the kernel.
-     */
     void pooling2_f32_nhwc_maxpool_indices(const Window &window_input, const Window &window);
     /** Function to perform MxN pooling for 32-bit floating point values.
      *
@@ -138,6 +132,19 @@ private:
      * @param[in] exclude_padding Flag to specify exclusion of padding from the operation.
      */
     void pooling2_f16_nchw(const Window &window_input, const Window &window, PoolingType pooling_type, bool exclude_padding = false);
+    /** Function to perform 2x2 pooling and compute the pooling indices for FP32/FP16. The indices can be used for max unpool.
+     *
+     * @param[in] window_input Input region on which to execute the kernel.
+     * @param[in] window       Output region on which to execute the kernel.
+     */
+    template <typename T>
+    void pooling2_nchw_maxpool_indices(const Window &window_input, const Window &window);
+    /** Function to perform 2x2 pooling and compute the pooling indices. The indices can be used for max unpool.
+     *
+     * @param[in] window_input Input region on which to execute the kernel.
+     * @param[in] window       Output region on which to execute the kernel.
+     */
+    void pooling2_f16_nhwc_maxpool_indices(const Window &window_input, const Window &window);
     /** Function to perform 3x3 pooling.
      *
      * @param[in] window_input    Input region on which to execute the kernel.
diff --git a/arm_compute/runtime/NEON/functions/NEMaxUnpoolingLayer.h b/arm_compute/runtime/NEON/functions/NEMaxUnpoolingLayer.h
index 468bc70466..f13b4bd9e2 100644
--- a/arm_compute/runtime/NEON/functions/NEMaxUnpoolingLayer.h
+++ b/arm_compute/runtime/NEON/functions/NEMaxUnpoolingLayer.h
@@ -45,7 +45,7 @@ public:
     NEMaxUnpoolingLayer();
     /** Set the input and output tensors.
      *
-     * @note F16 is supported for pool sizes 2 and 3 only
+     * @note Only supported pool size 2
      *
      * @param[in, out] input     Source tensor. (Written to only when padding != 0) Data types supported: QASYMM8/QASYMM8_SIGNED/F16/F32.
      * @param[out]     output    Destination tensor. Data types supported: Same as @p input.
@@ -55,7 +55,7 @@ public:
     void configure(ITensor *input, ITensor *indices, ITensor *output, const PoolingLayerInfo &pool_info);
     /** Static function to check if given info will lead to a valid configuration of @ref NEMaxUnpoolingLayer
      *
-     * @note F16 is supported for pool sizes 2 and 3 only
+     * @note Only supported pool size 2
      *
      * @param[in] input     Source tensor. (Written to only when padding != 0) Data types supported: QASYMM8/QASYMM8_SIGNED/F16/F32.
      * @param[in] indices   The indices of the maximal values. Data type supported: U32.
diff --git a/src/core/NEON/kernels/NEMaxUnpoolingLayerKernel.cpp b/src/core/NEON/kernels/NEMaxUnpoolingLayerKernel.cpp
index 53d70db9e2..821bf53817 100644
--- a/src/core/NEON/kernels/NEMaxUnpoolingLayerKernel.cpp
+++ b/src/core/NEON/kernels/NEMaxUnpoolingLayerKernel.cpp
@@ -43,6 +43,7 @@ Status validate_arguments(const ITensorInfo *input, const ITensorInfo *output, c
     ARM_COMPUTE_RETURN_ERROR_ON_CPU_F16_UNSUPPORTED(input);
     ARM_COMPUTE_RETURN_ERROR_ON_DATA_TYPE_CHANNEL_NOT_IN(input, 1, DataType::QASYMM8, DataType::QASYMM8_SIGNED, DataType::F16, DataType::F32);
     ARM_COMPUTE_RETURN_ERROR_ON_DATA_TYPE_CHANNEL_NOT_IN(indices, 1, DataType::U32);
+    ARM_COMPUTE_RETURN_ERROR_ON_MISMATCHING_SHAPES(input, indices);
 
     int                 pool_stride_x   = 0;
     int                 pool_stride_y   = 0;
diff --git a/src/core/NEON/kernels/NEPoolingLayerKernel.cpp b/src/core/NEON/kernels/NEPoolingLayerKernel.cpp
index e8069d32c9..2bbc307d19 100644
--- a/src/core/NEON/kernels/NEPoolingLayerKernel.cpp
+++ b/src/core/NEON/kernels/NEPoolingLayerKernel.cpp
@@ -725,6 +725,40 @@ void NEPoolingLayerKernel::configure(const ITensor *input, ITensor *output, cons
     INEKernel::configure(win_config.second);
 }
 
+template <typename T = float>
+inline uint32_t offset_no_padding(uint32_t padded_offset, const Coordinates &id, const ITensorInfo &info, int pool_stride_x, int pool_stride_y)
+{
+    const int pad_left    = info.padding().left;
+    const int pad_right   = info.padding().right;
+    const int pad_top     = info.padding().top;
+    const int pad_bottom  = info.padding().bottom;
+    const int in_stride_y = static_cast<int>(info.strides_in_bytes().y());
+    const int in_stride_w = static_cast<int>(info.strides_in_bytes()[3]);
+    const int pad_horiz   = pad_left + pad_right;
+    const int pad_vert    = pad_top + pad_bottom;
+
+    if(info.data_layout() == DataLayout::NCHW)
+    {
+        const uint32_t offset_base = padded_offset
+                                     - sizeof(T) * pad_horiz * id.y() * pool_stride_y                                            /* subtract padding elems per row */
+                                     - pad_top * sizeof(T)                                                                       /* top padding */
+                                     - sizeof(T) * pad_horiz * info.tensor_shape()[1] * id.z() - pad_vert * in_stride_y * id.z() /* for each Z plane there are height*pad_right padding elems */
+                                     - in_stride_w * id[3];
+
+        return offset_base;
+    }
+    else
+    {
+        const uint32_t offset_base = padded_offset
+                                     - sizeof(T) * pad_horiz * id.y() * pool_stride_x                          // subtract padding elems per row
+                                     - pad_top * sizeof(T)                                                     // top padding
+                                     - sizeof(T) * pad_horiz * info.tensor_shape()[1] * id.z() * pool_stride_y // for each Z plane there are width*pad_right padding elems
+                                     - in_stride_w * id[3];
+
+        return offset_base;
+    }
+}
+
 template <typename T>
 void NEPoolingLayerKernel::pooling2_q8_nchw(const Window &window_input, const Window &window, PoolingType pooling_type, bool exclude_padding)
 {
@@ -924,63 +958,130 @@ void NEPoolingLayerKernel::pooling3_f16_nchw(const Window &window_input, const W
 #endif /* __ARM_FEATURE_FP16_VECTOR_ARITHMETIC */
 }
 
+#ifdef __ARM_FEATURE_FP16_VECTOR_ARITHMETIC
+template <typename T>
+inline typename std::enable_if<std::is_same<T, float16_t>::value, float32x2_t>::type
+f16_to_f32(float16x4_t input)
+{
+    float32x2_t output = { static_cast<float>(vget_lane_f16(input, 0)), static_cast<float>(vget_lane_f16(input, 1)) };
+    return output;
+}
+#endif /* __ARM_FEATURE_FP16_VECTOR_ARITHMETIC */
+
+template <typename T>
+inline typename std::enable_if<std::is_same<T, float>::value, float32x2_t>::type
+f16_to_f32(float32x2_t input)
+{
+    return input;
+}
+
+template <typename T = float>
+void NEPoolingLayerKernel::pooling2_nchw_maxpool_indices(const Window &window_input, const Window &window)
+{
+    Iterator  input(_input, window_input);
+    Iterator  output(_output, window);
+    Iterator  indices(_indices, window);
+    const int pool_pad_top  = _pool_info.pad_stride_info.pad_top();
+    const int pool_pad_left = _pool_info.pad_stride_info.pad_left();
+    int       pool_stride_x = 0;
+    int       pool_stride_y = 0;
+    std::tie(pool_stride_x, pool_stride_y) = _pool_info.pad_stride_info.stride();
+    const uint8_t *const input_top_ptr    = _input->ptr_to_element(Coordinates(-static_cast<int>(pool_pad_left), -static_cast<int>(pool_pad_top)));
+    const uint8_t *const input_bottom_ptr = _input->ptr_to_element(Coordinates(-static_cast<int>(pool_pad_left), -static_cast<int>(pool_pad_top) + 1));
+    const int            pad_left         = _input->info()->padding().left;
+    const int            pad_right        = _input->info()->padding().right;
+    const int            in_stride_y      = static_cast<int>(_input->info()->strides_in_bytes().y());
+
+    execute_window_loop(window, [&](const Coordinates & id)
+    {
+        auto        top_data        = wrapper::vload(reinterpret_cast<const T *>(input_top_ptr + input.offset()));
+        auto        bottom_data     = wrapper::vload(reinterpret_cast<const T *>(input_bottom_ptr + input.offset()));
+        float32x2_t top_data_f32    = f16_to_f32<T>(top_data);
+        float32x2_t bottom_data_f32 = f16_to_f32<T>(bottom_data);
+
+        // Calculate max data, compare top first, then bottom, to make sue the first max is recorded.
+        const float32x2_t max_data_top         = vpmax_f32(top_data_f32, top_data_f32);
+        const float32x2_t max_data_bottom      = vpmax_f32(bottom_data_f32, bottom_data_f32);
+        const float32x2_t max_data             = vmax_f32(max_data_top, max_data_bottom);
+        *(reinterpret_cast<T *>(output.ptr())) = static_cast<T>(vget_lane_f32(max_data, 0));
+
+        // Calculate max data indice, which will be used in max unpool.
+        const uint32_t   offset_base              = offset_no_padding<T>(input.offset(), id, *_input->info(), pool_stride_x, pool_stride_y);
+        const uint32_t   offset_top               = (uint32_t)(offset_base / sizeof(T));
+        const uint32_t   offset_bottom            = offset_top + in_stride_y / sizeof(T) - pad_right - pad_left;
+        const uint32x2_t voffset_top              = { offset_top, offset_top + 1u };
+        const uint32x2_t voffset_bottom           = { offset_bottom, offset_bottom + 1u };
+        const uint32x2_t tmp_indices_top          = vbsl_u32(vcge_f32(top_data_f32, vrev64_f32(top_data_f32)), voffset_top, vrev64_u32(voffset_top));
+        const uint32x2_t tmp_indices_bottom       = vbsl_u32(vcge_f32(bottom_data_f32, vrev64_f32(bottom_data_f32)), voffset_bottom, vrev64_u32(voffset_bottom));
+        *(reinterpret_cast<int *>(indices.ptr())) = vget_lane_u32(vbsl_u32(vcge_f32(max_data_top, max_data_bottom), tmp_indices_top, tmp_indices_bottom), 0);
+    },
+    input, output, indices);
+}
+
 void NEPoolingLayerKernel::pooling2_f16_nchw(const Window &window_input, const Window &window, PoolingType pooling_type, bool exclude_padding)
 {
     ARM_COMPUTE_UNUSED(pooling_type);
     ARM_COMPUTE_UNUSED(exclude_padding);
 #ifdef __ARM_FEATURE_FP16_VECTOR_ARITHMETIC
-    Iterator      input(_input, window_input);
-    Iterator      output(_output, window);
-    constexpr int pool_size       = 2;
-    const int     pool_pad_right  = _pool_info.pad_stride_info.pad_right();
-    const int     pool_pad_top    = _pool_info.pad_stride_info.pad_top();
-    const int     pool_pad_left   = _pool_info.pad_stride_info.pad_left();
-    const int     pool_pad_bottom = _pool_info.pad_stride_info.pad_bottom();
-    int           pool_stride_x, pool_stride_y = 0;
-    std::tie(pool_stride_x, pool_stride_y)     = _pool_info.pad_stride_info.stride();
-    const int upper_bound_w = _input->info()->dimension(0) + (exclude_padding ? 0 : pool_pad_right);
-    const int upper_bound_h = _input->info()->dimension(1) + (exclude_padding ? 0 : pool_pad_bottom);
-
-    const unsigned char *const input_top_ptr    = _input->ptr_to_element(Coordinates(-static_cast<int>(pool_pad_left), -static_cast<int>(pool_pad_top)));
-    const unsigned char *const input_bottom_ptr = _input->ptr_to_element(Coordinates(-static_cast<int>(pool_pad_left), -static_cast<int>(pool_pad_top) + 1));
-
-    execute_window_loop(window, [&](const Coordinates & id)
+    if(pooling_type == PoolingType::MAX && _indices)
     {
-        float16x4_t top_data    = vld1_f16(reinterpret_cast<const float16_t *>(input_top_ptr + input.offset()));
-        float16x4_t bottom_data = vld1_f16(reinterpret_cast<const float16_t *>(input_bottom_ptr + input.offset()));
-        float16x4_t res         = {};
+        pooling2_nchw_maxpool_indices<float16_t>(window_input, window);
+    }
+    else
+    {
+        Iterator      input(_input, window_input);
+        Iterator      output(_output, window);
+        constexpr int pool_size       = 2;
+        const int     pool_pad_right  = _pool_info.pad_stride_info.pad_right();
+        const int     pool_pad_top    = _pool_info.pad_stride_info.pad_top();
+        const int     pool_pad_left   = _pool_info.pad_stride_info.pad_left();
+        const int     pool_pad_bottom = _pool_info.pad_stride_info.pad_bottom();
+        int           pool_stride_x, pool_stride_y = 0;
+        std::tie(pool_stride_x, pool_stride_y)     = _pool_info.pad_stride_info.stride();
+        const int upper_bound_w = _input->info()->dimension(0) + (exclude_padding ? 0 : pool_pad_right);
+        const int upper_bound_h = _input->info()->dimension(1) + (exclude_padding ? 0 : pool_pad_bottom);
 
-        // Get power of 2 in case of l2 pooling
-        if(pooling_type == PoolingType::L2)
-        {
-            top_data    = vmul_f16(top_data, top_data);
-            bottom_data = vmul_f16(bottom_data, bottom_data);
-        }
+        const unsigned char *const input_top_ptr    = _input->ptr_to_element(Coordinates(-static_cast<int>(pool_pad_left), -static_cast<int>(pool_pad_top)));
+        const unsigned char *const input_bottom_ptr = _input->ptr_to_element(Coordinates(-static_cast<int>(pool_pad_left), -static_cast<int>(pool_pad_top) + 1));
 
-        if(pooling_type != PoolingType::MAX)
+        execute_window_loop(window, [&](const Coordinates & id)
         {
-            const float       scale   = calculate_avg_scale(exclude_padding, DataLayout::NCHW, id, pool_size, pool_size, upper_bound_w, upper_bound_h, pool_pad_left, pool_pad_top, pool_stride_x, pool_stride_y);
-            const float16x4_t scale_v = vdup_n_f16(scale);
+            float16x4_t top_data    = vld1_f16(reinterpret_cast<const float16_t *>(input_top_ptr + input.offset()));
+            float16x4_t bottom_data = vld1_f16(reinterpret_cast<const float16_t *>(input_bottom_ptr + input.offset()));
+            float16x4_t res         = {};
 
-            const float16x4_t sum_data = vadd_f16(top_data, bottom_data);
-            res                        = vmul_f16(vpadd_f16(sum_data, sum_data), scale_v);
-        }
-        else
-        {
-            const float16x4_t max_data = vmax_f16(top_data, bottom_data);
-            res                        = vpmax_f16(max_data, max_data);
-        }
+            // Get power of 2 in case of l2 pooling
+            if(pooling_type == PoolingType::L2)
+            {
+                top_data    = vmul_f16(top_data, top_data);
+                bottom_data = vmul_f16(bottom_data, bottom_data);
+            }
 
-        // Calculate square-root in case of l2 pooling
-        if(pooling_type == PoolingType::L2)
-        {
-            res = vinv_f16(vinvsqrt_f16(res));
-        }
+            if(pooling_type != PoolingType::MAX)
+            {
+                const float       scale   = calculate_avg_scale(exclude_padding, DataLayout::NCHW, id, pool_size, pool_size, upper_bound_w, upper_bound_h, pool_pad_left, pool_pad_top, pool_stride_x, pool_stride_y);
+                const float16x4_t scale_v = vdup_n_f16(scale);
 
-        // Store result
-        *(reinterpret_cast<float16_t *>(output.ptr())) = vget_lane_f16(res, 0);
-    },
-    input, output);
+                const float16x4_t sum_data = vadd_f16(top_data, bottom_data);
+                res                        = vmul_f16(vpadd_f16(sum_data, sum_data), scale_v);
+            }
+            else
+            {
+                const float16x4_t max_data = vmax_f16(top_data, bottom_data);
+                res                        = vpmax_f16(max_data, max_data);
+            }
+
+            // Calculate square-root in case of l2 pooling
+            if(pooling_type == PoolingType::L2)
+            {
+                res = vinv_f16(vinvsqrt_f16(res));
+            }
+
+            // Store result
+            *(reinterpret_cast<float16_t *>(output.ptr())) = vget_lane_f16(res, 0);
+        },
+        input, output);
+    }
 #else  /* __ARM_FEATURE_FP16_VECTOR_ARITHMETIC */
     ARM_COMPUTE_UNUSED(window_input);
     ARM_COMPUTE_UNUSED(window);
@@ -1267,11 +1368,95 @@ void NEPoolingLayerKernel::poolingMxN_f16_nchw(const Window &window_input, const
 #endif /* __ARM_FEATURE_FP16_VECTOR_ARITHMETIC */
 }
 
+#ifdef __ARM_FEATURE_FP16_VECTOR_ARITHMETIC
+void NEPoolingLayerKernel::pooling2_f16_nhwc_maxpool_indices(const Window &window_input, const Window &window)
+{
+    Iterator input(_input, window_input);
+    Iterator output(_output, window);
+    Iterator indices(_indices, window);
+
+    const int pool_pad_top  = _pool_info.pad_stride_info.pad_top();
+    const int pool_pad_left = _pool_info.pad_stride_info.pad_left();
+
+    int pool_stride_x = 0;
+    int pool_stride_y = 0;
+    std::tie(pool_stride_x, pool_stride_y) = _pool_info.pad_stride_info.stride();
+
+    const int pad_right   = _input->info()->padding().right;
+    const int in_stride_y = static_cast<int>(_input->info()->strides_in_bytes().y());
+    const int in_stride_z = static_cast<int>(_input->info()->strides_in_bytes().z());
+
+    execute_window_loop(window, [&](const Coordinates & id)
+    {
+        const int idx_width    = id.y() * pool_stride_x;
+        const int idx_height   = id.z() * pool_stride_y;
+        const int pool_limit_y = pool_pad_top - idx_height;
+        const int pool_limit_x = pool_pad_left - idx_width;
+
+        const int pool_start_y = std::max(0, window_input.z().start() + pool_limit_y);
+        const int pool_start_x = std::max(0, window_input.y().start() + pool_limit_x);
+        const int in_x0_offset = (pool_start_x - pool_pad_left) * static_cast<int>(_input->info()->strides_in_bytes().y()) + (pool_start_y - pool_pad_top) * static_cast<int>
+                                 (_input->info()->strides_in_bytes().z());
+        const int in_x1_offset = (pool_start_x + 1 - pool_pad_left) * static_cast<int>(_input->info()->strides_in_bytes().y()) + (pool_start_y - pool_pad_top) * static_cast<int>
+                                 (_input->info()->strides_in_bytes().z());
+
+        const int in_x2_offset = (pool_start_x - pool_pad_left) * static_cast<int>(_input->info()->strides_in_bytes().y()) + (pool_start_y + 1 - pool_pad_top) * static_cast<int>
+                                 (_input->info()->strides_in_bytes().z());
+
+        const int in_x3_offset = (pool_start_x + 1 - pool_pad_left) * static_cast<int>(_input->info()->strides_in_bytes().y()) + (pool_start_y + 1 - pool_pad_top) * static_cast<int>
+                                 (_input->info()->strides_in_bytes().z());
+
+        const auto  in_x0_ptr = reinterpret_cast<const float16_t *>(input.ptr() + in_x0_offset);
+        const auto  in_x1_ptr = reinterpret_cast<const float16_t *>(input.ptr() + in_x1_offset);
+        const auto  in_x2_ptr = reinterpret_cast<const float16_t *>(input.ptr() + in_x2_offset);
+        const auto  in_x3_ptr = reinterpret_cast<const float16_t *>(input.ptr() + in_x3_offset);
+        const auto  v_x0      = vld1q_f16(in_x0_ptr);
+        const auto  v_x1      = vld1q_f16(in_x1_ptr);
+        const auto  v_x2      = vld1q_f16(in_x2_ptr);
+        const auto  v_x3      = vld1q_f16(in_x3_ptr);
+        float16x8_t vres      = vmaxq_f16(vmaxq_f16(v_x2, v_x3), vmaxq_f16(v_x0, v_x1));
+        // Store result
+        vst1q_f16(reinterpret_cast<float16_t *>(output.ptr()), vres);
+
+        const uint32_t   offset_base    = offset_no_padding<float16_t>(input.offset(), id, *_input->info(), pool_stride_x, pool_stride_y);
+        const uint32_t   offset_x0      = (uint32_t)offset_base / sizeof(float16_t);
+        const uint32_t   offset_x1      = (uint32_t)offset_x0 + in_stride_y / sizeof(float16_t) - pad_right;
+        const uint32_t   offset_x2      = (uint32_t)offset_x0 + in_stride_z / sizeof(float16_t) - pad_right * _input->info()->tensor_shape()[1];
+        const uint32_t   offset_x3      = (uint32_t)offset_x2 + in_stride_y / sizeof(float16_t) - pad_right;
+        const uint32x4_t voffset_x0_0   = { offset_x0, offset_x0 + 1, offset_x0 + 2, offset_x0 + 3 };
+        const uint32x4_t voffset_x0_1   = { offset_x0 + 4, offset_x0 + 5, offset_x0 + 6, offset_x0 + 7 };
+        const uint16x8_t voffset_x0     = vcombine_u16(vmovn_u32(voffset_x0_0), vmovn_u32(voffset_x0_1));
+        const uint32x4_t voffset_x1_0   = { offset_x1, offset_x1 + 1, offset_x1 + 2, offset_x1 + 3 };
+        const uint32x4_t voffset_x1_1   = { offset_x1 + 4, offset_x1 + 5, offset_x1 + 6, offset_x1 + 7 };
+        const uint16x8_t voffset_x1     = vcombine_u16(vmovn_u32(voffset_x1_0), vmovn_u32(voffset_x1_1));
+        const uint32x4_t voffset_x2_0   = { offset_x2, offset_x2 + 1, offset_x2 + 2, offset_x2 + 3 };
+        const uint32x4_t voffset_x2_1   = { offset_x2 + 4, offset_x2 + 5, offset_x2 + 6, offset_x2 + 7 };
+        const uint16x8_t voffset_x2     = vcombine_u16(vmovn_u32(voffset_x2_0), vmovn_u32(voffset_x2_1));
+        const uint32x4_t voffset_x3_0   = { offset_x3, offset_x3 + 1, offset_x3 + 2, offset_x3 + 3 };
+        const uint32x4_t voffset_x3_1   = { offset_x3 + 4, offset_x3 + 5, offset_x3 + 6, offset_x3 + 7 };
+        const uint16x8_t voffset_x3     = vcombine_u16(vmovn_u32(voffset_x3_0), vmovn_u32(voffset_x3_1));
+        const uint16x8_t tmp_indices0   = vbslq_u16(vcgeq_f16(v_x0, v_x1), voffset_x0, voffset_x1);
+        const uint16x8_t tmp_indices1   = vbslq_u16(vcgeq_f16(v_x2, v_x3), voffset_x2, voffset_x3);
+        const uint16x8_t tmp_indices2   = vbslq_u16(vcgeq_f16(vmaxq_f16(v_x0, v_x1), vmaxq_f16(v_x2, v_x3)), tmp_indices0, tmp_indices1);
+        const uint32x4_t tmp_indeces3_0 = vmovl_u16(vget_low_u16(tmp_indices2));
+        const uint32x4_t tmp_indeces3_1 = vmovl_u16(vget_high_u16(tmp_indices2));
+        // Store indicies
+        vst1q_u32(reinterpret_cast<uint32_t *>(indices.ptr()), tmp_indeces3_0);
+        vst1q_u32(reinterpret_cast<uint32_t *>(indices.ptr() + 16), tmp_indeces3_1);
+    },
+    input, output, indices);
+}
+#endif /* __ARM_FEATURE_FP16_VECTOR_ARITHMETIC */
+
 void NEPoolingLayerKernel::poolingMxN_f16_nhwc(const Window &window_input, const Window &window, PoolingType pooling_type, bool exclude_padding)
 {
     ARM_COMPUTE_UNUSED(pooling_type);
     ARM_COMPUTE_UNUSED(exclude_padding);
 #ifdef __ARM_FEATURE_FP16_VECTOR_ARITHMETIC
+    if(_pool_info.pool_size == Size2D(2, 2) && pooling_type == PoolingType::MAX && _indices)
+    {
+        pooling2_f16_nhwc_maxpool_indices(window_input, window);
+    }
     Iterator input(_input, window_input);
     Iterator output(_output, window);
 
@@ -1488,88 +1673,12 @@ void NEPoolingLayerKernel::poolingMxN_f32_nchw(const Window &window_input, const
     input, output);
 }
 
-inline uint32_t offset_no_padding(uint32_t padded_offset, const Coordinates &id, const ITensorInfo &info, int pool_stride_x, int pool_stride_y)
-{
-    const int pad_left    = info.padding().left;
-    const int pad_right   = info.padding().right;
-    const int pad_top     = info.padding().top;
-    const int pad_bottom  = info.padding().bottom;
-    const int in_stride_y = static_cast<int>(info.strides_in_bytes().y());
-    const int in_stride_w = static_cast<int>(info.strides_in_bytes()[3]);
-    const int pad_horiz   = pad_left + pad_right;
-    const int pad_vert    = pad_top + pad_bottom;
-
-    if(info.data_layout() == DataLayout::NCHW)
-    {
-        const uint32_t offset_base = padded_offset
-                                     - sizeof(float) * pad_horiz * id.y() * pool_stride_y                                            /* subtract padding elems per row */
-                                     - pad_top * sizeof(float)                                                                       /* top padding */
-                                     - sizeof(float) * pad_horiz * info.tensor_shape()[1] * id.z() - pad_vert * in_stride_y * id.z() /* for each Z plane there are height*pad_right padding elems */
-                                     - in_stride_w * id[3];
-
-        return offset_base;
-    }
-    else
-    {
-        const uint32_t offset_base = padded_offset
-                                     - sizeof(float) * pad_horiz * id.y() * pool_stride_x                          // subtract padding elems per row
-                                     - pad_top * sizeof(float)                                                     // top padding
-                                     - sizeof(float) * pad_horiz * info.tensor_shape()[1] * id.z() * pool_stride_y // for each Z plane there are width*pad_right padding elems
-                                     - in_stride_w * id[3];
-
-        return offset_base;
-    }
-}
-
-void NEPoolingLayerKernel::pooling2_f32_nchw_maxpool_indices(const Window &window_input, const Window &window)
-{
-    Iterator  input(_input, window_input);
-    Iterator  output(_output, window);
-    Iterator  indices(_indices, window);
-    int       final_index   = 0;
-    const int pool_pad_top  = _pool_info.pad_stride_info.pad_top();
-    const int pool_pad_left = _pool_info.pad_stride_info.pad_left();
-    int       pool_stride_x = 0;
-    int       pool_stride_y = 0;
-    std::tie(pool_stride_x, pool_stride_y) = _pool_info.pad_stride_info.stride();
-    const uint8_t *const input_top_ptr    = _input->ptr_to_element(Coordinates(-static_cast<int>(pool_pad_left), -static_cast<int>(pool_pad_top)));
-    const uint8_t *const input_bottom_ptr = _input->ptr_to_element(Coordinates(-static_cast<int>(pool_pad_left), -static_cast<int>(pool_pad_top) + 1));
-    const int            pad_left         = _input->info()->padding().left;
-    const int            pad_right        = _input->info()->padding().right;
-    const int            in_stride_y      = static_cast<int>(_input->info()->strides_in_bytes().y());
-    execute_window_loop(window, [&](const Coordinates & id)
-    {
-        const auto        input_offset_top    = input_top_ptr + input.offset();
-        const auto        input_offset_bottom = input_bottom_ptr + input.offset();
-        const auto        in_top_ptr          = reinterpret_cast<const float *>(input_offset_top);
-        const auto        in_bottom_ptr       = reinterpret_cast<const float *>(input_offset_bottom);
-        float32x2_t       top_data            = vld1_f32(in_top_ptr);
-        float32x2_t       bottom_data         = vld1_f32(in_bottom_ptr);
-        float32x2_t       res                 = {};
-        float             final_res           = 0;
-        const float32x2_t max_data            = vmax_f32(top_data, bottom_data);
-        res                                   = vpmax_f32(max_data, max_data);
-        final_res                             = vget_lane_f32(res, 0);
-        // Store result
-        *(reinterpret_cast<float *>(output.ptr())) = final_res;
-        const uint32_t   offset_base               = offset_no_padding(input.offset(), id, *_input->info(), pool_stride_x, pool_stride_y);
-        const uint32_t   offset_top                = (uint32_t)(offset_base / sizeof(float));
-        const uint32_t   offset_bottom             = offset_top + in_stride_y / sizeof(float) - pad_right - pad_left;
-        const uint32x2_t voffset_top               = { offset_top, offset_top + 1u };
-        const uint32x2_t voffset_bottom            = { offset_bottom, offset_bottom + 1u };
-        const uint32x2_t tmp_indices               = vbsl_u32(vcgt_f32(top_data, bottom_data), voffset_top, voffset_bottom);
-        final_index                                = vget_lane_u32(vbsl_u32(vcgt_f32(max_data, vrev64_f32(max_data)), tmp_indices, vrev64_u32(tmp_indices)), 0);
-        *(reinterpret_cast<int *>(indices.ptr()))  = final_index;
-    },
-    input, output, indices);
-}
-
 void NEPoolingLayerKernel::pooling2_f32_nchw(const Window &window_input, const Window &window, PoolingType pooling_type,
                                              bool exclude_padding)
 {
     if(pooling_type == PoolingType::MAX && _indices)
     {
-        pooling2_f32_nchw_maxpool_indices(window_input, window);
+        pooling2_nchw_maxpool_indices<float>(window_input, window);
     }
     else
     {
@@ -1943,12 +2052,11 @@ void NEPoolingLayerKernel::pooling2_f32_nhwc_maxpool_indices(const Window &windo
         const uint32x4_t voffset_x1   = { offset_x1, offset_x1 + 1, offset_x1 + 2, offset_x1 + 3 };
         const uint32x4_t voffset_x2   = { offset_x2, offset_x2 + 1, offset_x2 + 2, offset_x2 + 3 };
         const uint32x4_t voffset_x3   = { offset_x3, offset_x3 + 1, offset_x3 + 2, offset_x3 + 3 };
-        const uint32x4_t tmp_indices0 = vbslq_u32(vcgtq_f32(v_x0, v_x1), voffset_x0, voffset_x1);
-        const uint32x4_t tmp_indices1 = vbslq_u32(vcgtq_f32(v_x2, v_x3), voffset_x2, voffset_x3);
-        const uint32x4_t tmp_indices2 = vbslq_u32(vcgtq_f32(vmaxq_f32(v_x0, v_x1), vmaxq_f32(v_x2, v_x3)), tmp_indices0, tmp_indices1);
-
+        const uint32x4_t tmp_indices0 = vbslq_u32(vcgeq_f32(v_x0, v_x1), voffset_x0, voffset_x1);
+        const uint32x4_t tmp_indices1 = vbslq_u32(vcgeq_f32(v_x2, v_x3), voffset_x2, voffset_x3);
+        const uint32x4_t tmp_indices2 = vbslq_u32(vcgeq_f32(vmaxq_f32(v_x0, v_x1), vmaxq_f32(v_x2, v_x3)), tmp_indices0, tmp_indices1);
+        // Store indices
         vst1q_u32(reinterpret_cast<uint32_t *>(indices.ptr()), tmp_indices2);
-
     },
     input, output, indices);
 }
diff --git a/tests/validation/NEON/MaxUnpoolingLayer.cpp b/tests/validation/NEON/MaxUnpoolingLayer.cpp
index e44021d858..a33ec2885e 100644
--- a/tests/validation/NEON/MaxUnpoolingLayer.cpp
+++ b/tests/validation/NEON/MaxUnpoolingLayer.cpp
@@ -63,11 +63,11 @@ FIXTURE_DATA_TEST_CASE(MaxUnpooling, NEMaxUnpoolingLayerFixture<float>, framewor
 TEST_SUITE_END() // FP32
 #ifdef __ARM_FEATURE_FP16_VECTOR_ARITHMETIC
 TEST_SUITE(FP16)
-FIXTURE_DATA_TEST_CASE(MaxUnpooling, NEMaxUnpoolingLayerFixture<float>, framework::DatasetMode::PRECOMMIT, combine(combine(datasets::SmallShapes(), combine(PoolingLayerIndicesDatasetFPSmall,
-                                                                                                                   framework::dataset::make("DataType", DataType::F16))),
-                                                                                                                   framework::dataset::make("DataLayout", { DataLayout::NCHW, DataLayout::NHWC })
+FIXTURE_DATA_TEST_CASE(MaxUnpooling, NEMaxUnpoolingLayerFixture<half>, framework::DatasetMode::PRECOMMIT, combine(combine(datasets::SmallShapes(), combine(PoolingLayerIndicesDatasetFPSmall,
+                                                                                                                  framework::dataset::make("DataType", DataType::F16))),
+                                                                                                                  framework::dataset::make("DataLayout", { DataLayout::NCHW, DataLayout::NHWC })
 
-                                                                                                                  ))
+                                                                                                                 ))
 {
     // Validate output
     validate(Accessor(_target), _reference);
diff --git a/tests/validation/reference/MaxUnpoolingLayer.cpp b/tests/validation/reference/MaxUnpoolingLayer.cpp
index 134f39726a..880018e99f 100644
--- a/tests/validation/reference/MaxUnpoolingLayer.cpp
+++ b/tests/validation/reference/MaxUnpoolingLayer.cpp
@@ -99,6 +99,9 @@ SimpleTensor<T> max_unpooling_layer(const SimpleTensor<T> &src, const PoolingLay
 template SimpleTensor<float> max_unpooling_layer(const SimpleTensor<float> &src, const PoolingLayerInfo &info,
                                                  const QuantizationInfo &output_qinfo, SimpleTensor<uint32_t> &indices,
                                                  TensorShape output_shape, DataLayout data_layout);
+template SimpleTensor<half> max_unpooling_layer(const SimpleTensor<half> &src, const PoolingLayerInfo &info,
+                                                const QuantizationInfo &output_qinfo, SimpleTensor<uint32_t> &indices,
+                                                TensorShape output_shape, DataLayout data_layout);
 
 } // namespace reference
 } // namespace validation