COMPMID-1944 Add support for "reflect" padding mode in NEPad

Change-Id: I56c42524497d37d44708648571fa211ac1afbd98
Signed-off-by: Usama Arif <usama.arif@arm.com>
Reviewed-on: https://review.mlplatform.org/c/885
Tested-by: Arm Jenkins <bsgcomp@arm.com>
Comments-Addressed: Arm Jenkins <bsgcomp@arm.com>
Reviewed-by: Pablo Marquez <pablo.tello@arm.com>

4 files changed, 84 insertions, 12 deletions

diff --git a/arm_compute/core/Helpers.h b/arm_compute/core/Helpers.h
index c7c7110ef5..235657a38a 100644
--- a/arm_compute/core/Helpers.h
+++ b/arm_compute/core/Helpers.h
@@ -707,6 +707,15 @@ inline int coords2index(const TensorShape &shape, const Coordinates &coord);
  */
 inline size_t get_data_layout_dimension_index(const DataLayout data_layout, const DataLayoutDimension data_layout_dimension);
 
+/** Get the DataLayoutDimension of a given index and layout.
+ *
+ * @param[in] data_layout The data layout.
+ * @param[in] index       The data layout index.
+ *
+ * @return The dimension which this index is requested for.
+ */
+inline DataLayoutDimension get_index_data_layout_dimension(const DataLayout data_layout, const size_t index);
+
 /** Calculate the normalization dimension index for a given normalization type
  *
  * @param[in] layout Data layout of the input and output tensor
diff --git a/arm_compute/core/Helpers.inl b/arm_compute/core/Helpers.inl
index c0e4ab8d7d..aeb290b23e 100644
--- a/arm_compute/core/Helpers.inl
+++ b/arm_compute/core/Helpers.inl
@@ -1,5 +1,5 @@
 /*
- * Copyright (c) 2016-2018 ARM Limited.
+ * Copyright (c) 2016-2019 ARM Limited.
  *
  * SPDX-License-Identifier: MIT
  *
@@ -351,4 +351,33 @@ inline size_t get_data_layout_dimension_index(const DataLayout data_layout, cons
             break;
     }
 }
+
+inline DataLayoutDimension get_index_data_layout_dimension(const DataLayout data_layout, const size_t index)
+{
+    ARM_COMPUTE_ERROR_ON_MSG(data_layout == DataLayout::UNKNOWN, "Cannot retrieve the dimension index for an unknown layout!");
+
+    /* Return the index based on the data layout
+    * [N C H W]
+    * [3 2 1 0]
+    * [N H W C]
+    */
+    switch(index)
+    {
+        case 0:
+            return (data_layout == DataLayout::NCHW) ? DataLayoutDimension::WIDTH : DataLayoutDimension::CHANNEL;
+            break;
+        case 1:
+            return (data_layout == DataLayout::NCHW) ? DataLayoutDimension::HEIGHT : DataLayoutDimension::WIDTH;
+            break;
+        case 2:
+            return (data_layout == DataLayout::NCHW) ? DataLayoutDimension::CHANNEL : DataLayoutDimension::HEIGHT;
+            break;
+        case 3:
+            return DataLayoutDimension::BATCHES;
+            break;
+        default:
+            ARM_COMPUTE_ERROR("Index value not supported!");
+            break;
+    }
+}
 } // namespace arm_compute
diff --git a/arm_compute/runtime/CL/functions/CLPadLayer.h b/arm_compute/runtime/CL/functions/CLPadLayer.h
index 1ecf82fa7c..0179441af2 100644
--- a/arm_compute/runtime/CL/functions/CLPadLayer.h
+++ b/arm_compute/runtime/CL/functions/CLPadLayer.h
@@ -53,9 +53,12 @@ public:
      * @param[out] output         Output tensor. Data type supported: same as @p input
      * @param[in]  padding        The padding for each spatial dimension of the input tensor. The pair padding[i]
      *                            specifies the front and the end padding in the i-th dimension.
-     * @param[in]  constant_value (Optional) Constant value to be used for the padding
+     * @param[in]  constant_value (Optional) Constant value to be used for the padding.
+     * @param[in]  mode           (Optional) Controls whether the padding should be filled with @p constant_value using CONSTANT,
+     *                            or reflect the input, either including the border values (SYMMETRIC) or not (REFLECT). Only CONSTANT
+     *                              is currently supported.
      */
-    void configure(ICLTensor *input, ICLTensor *output, const PaddingList &padding, PixelValue constant_value = PixelValue());
+    void configure(ICLTensor *input, ICLTensor *output, const PaddingList &padding, PixelValue constant_value = PixelValue(), PaddingMode mode = PaddingMode::CONSTANT);
 
     /**  Static function to check if given info will lead to a valid configuration of @ref CLPadLayer.
      *
@@ -64,8 +67,11 @@ public:
      * @param[in] padding        The padding for each spatial dimension of the input tensor. The pair padding[i]
      *                           specifies the front and the end padding in the i-th dimension.
      * @param[in] constant_value (Optional) Constant value to be used for the padding
+     * @param[in] mode           (Optional) Controls whether the padding should be filled with @p constant_value using CONSTANT,
+     *                            or reflect the input, either including the border values (SYMMETRIC) or not (REFLECT). Only CONSTANT
+     *                              is currently supported.
      */
-    static Status validate(const ITensorInfo *input, const ITensorInfo *output, const PaddingList &padding, PixelValue constant_value = PixelValue());
+    static Status validate(const ITensorInfo *input, const ITensorInfo *output, const PaddingList &padding, PixelValue constant_value = PixelValue(), PaddingMode mode = PaddingMode::CONSTANT);
 
     // Inherited methods overridden:
     void run() override;
diff --git a/arm_compute/runtime/NEON/functions/NEPadLayer.h b/arm_compute/runtime/NEON/functions/NEPadLayer.h
index 3a0863802a..78dbc1f1f9 100644
--- a/arm_compute/runtime/NEON/functions/NEPadLayer.h
+++ b/arm_compute/runtime/NEON/functions/NEPadLayer.h
@@ -25,17 +25,17 @@
 #define __ARM_COMPUTE_NEPADLAYER_H__
 
 #include "arm_compute/runtime/IFunction.h"
+#include "arm_compute/runtime/NEON/functions/NEConcatenateLayer.h"
+#include "arm_compute/runtime/NEON/functions/NEStridedSlice.h"
 #include "arm_compute/runtime/SubTensor.h"
 
 #include "arm_compute/core/NEON/kernels/NECopyKernel.h"
 #include "arm_compute/core/NEON/kernels/NEMemsetKernel.h"
 #include "arm_compute/core/Types.h"
+#include "arm_compute/runtime/Tensor.h"
 
 namespace arm_compute
 {
-// Forward declarations
-class ITensor;
-
 /** Basic function to pad a tensor. This function calls the following NEON kernels:
  *
  *  -# @ref NEMemsetKernel
@@ -53,8 +53,10 @@ public:
      * @param[in]  padding        The padding for each spatial dimension of the input tensor. The pair padding[i]
      *                            specifies the front and the end padding in the i-th dimension.
      * @param[in]  constant_value (Optional) Constant value to be used for the padding
+     * @param[in]  mode           (Optional) Controls whether the padding should be filled with @p constant_value using CONSTANT,
+     *                            or reflect the input, either including the border values (SYMMETRIC) or not (REFLECT).
      */
-    void configure(ITensor *input, ITensor *output, const PaddingList &padding, PixelValue constant_value = PixelValue());
+    void configure(ITensor *input, ITensor *output, const PaddingList &padding, const PixelValue constant_value = PixelValue(), const PaddingMode mode = PaddingMode::CONSTANT);
     /**  Static function to check if given info will lead to a valid configuration of @ref NEPadLayer.
      *
      * @param[in] input          Source tensor info. Data types supported: U8/S8/QASYMM8/U16/S16/F16/U32/S32/F32.
@@ -62,18 +64,44 @@ public:
      * @param[in] padding        The padding for each spatial dimension of the input tensor. The pair padding[i]
      *                           specifies the front and the end padding in the i-th dimension.
      * @param[in] constant_value (Optional) Constant value to be used for the padding
+     * @param[in] mode           (Optional) Controls whether the padding should be filled with @p constant_value using CONSTANT,
+     *                     or reflect the input, either including the border values (SYMMETRIC) or not (REFLECT).
      *
      * @return a status
      */
-    static Status validate(const ITensorInfo *input, const ITensorInfo *output, const PaddingList &padding, PixelValue constant_value = PixelValue());
+    static Status validate(const ITensorInfo *input, const ITensorInfo *output, const PaddingList &padding, const PixelValue constant_value = PixelValue(), const PaddingMode mode = PaddingMode::CONSTANT);
 
     // Inherited methods overridden:
     void run() override;
 
 private:
-    NEMemsetKernel _memset_kernel;
-    NECopyKernel   _copy_kernel;
-    SubTensor      _output_subtensor;
+    /** Configure kernels for when constant padding is used.
+     *
+     * @param[in]  input          Source tensor. Data types supported: U8/S8/QASYMM8/U16/S16/F16/U32/S32/F32.
+     * @param[out] output         Output tensor. Data type supported: same as @p input
+     * @param[in]  padding        The padding for each spatial dimension of the input tensor. The pair padding[i]
+     *                            specifies the front and the end padding in the i-th dimension.
+     * @param[in]  constant_value Constant value to be used for the padding
+     */
+    void configure_constant_mode(ITensor *input, ITensor *output, const PaddingList &padding, const PixelValue constant_value);
+    /** Configure functions for when reflect or symmetric padding is used.
+     *
+     * @param[in]  input  Source tensor. Data types supported: U8/S8/QASYMM8/U16/S16/F16/U32/S32/F32.
+     * @param[out] output Output tensor. Data type supported: same as @p input
+     */
+    void configure_reflect_symmetric_mode(ITensor *input, ITensor *output);
+
+private:
+    NECopyKernel                          _copy_kernel;
+    PaddingMode                           _mode;
+    PaddingList                           _padding;
+    NEMemsetKernel                        _memset_kernel;
+    uint32_t                              _num_dimensions;
+    std::unique_ptr<NEStridedSlice[]>     _slice_functions;
+    std::unique_ptr<NEConcatenateLayer[]> _concat_functions;
+    std::unique_ptr<Tensor[]>             _slice_results;
+    std::unique_ptr<Tensor[]>             _concat_results;
+    SubTensor                             _output_subtensor;
 };
 } // namespace arm_compute
 #endif /*__ARM_COMPUTE_NEPADLAYER_H__ */