COMPMID-432 - Extended Convolution Layer to support rectangular kernels

Change-Id: I99be1efede4de6dd63ce103fb11196c413757621
Reviewed-on: http://mpd-gerrit.cambridge.arm.com/79252
Tested-by: Kaizen <jeremy.johnson+kaizengerrit@arm.com>
Reviewed-by: Moritz Pflanzer <moritz.pflanzer@arm.com>

1 files changed, 18 insertions, 13 deletions

diff --git a/src/core/NEON/kernels/NEIm2ColKernel.cpp b/src/core/NEON/kernels/NEIm2ColKernel.cpp
index 875c08ed42..99daa2e5e7 100644
--- a/src/core/NEON/kernels/NEIm2ColKernel.cpp
+++ b/src/core/NEON/kernels/NEIm2ColKernel.cpp
@@ -27,6 +27,7 @@
 #include "arm_compute/core/FixedPoint.h"
 #include "arm_compute/core/Helpers.h"
 #include "arm_compute/core/ITensor.h"
+#include "arm_compute/core/Size2D.h"
 #include "arm_compute/core/TensorInfo.h"
 #include "arm_compute/core/Types.h"
 #include "arm_compute/core/Validate.h"
@@ -47,7 +48,8 @@ inline void linearize_volume(const uint8_t *const in_ptr,
                              bool                 has_bias,
                              int                  top_left_x,
                              int                  top_left_y,
-                             int                  kernel_size,
+                             int                  kernel_width,
+                             int                  kernel_height,
                              int                  kernel_depth,
                              int                  input_w,
                              int                  input_h,
@@ -56,9 +58,9 @@ inline void linearize_volume(const uint8_t *const in_ptr,
                              int                  input_stride_z,
                              int                  fixed_point_position)
 {
-    const int kernel_size2 = kernel_size * kernel_size;
-    const int x_e          = top_left_x + kernel_size;
-    const int y_e          = top_left_y + kernel_size;
+    const int kernel_size2 = kernel_width * kernel_height;
+    const int x_e          = top_left_x + kernel_width;
+    const int y_e          = top_left_y + kernel_height;
 
     // Linearize volume
     int d = 0;
@@ -109,8 +111,8 @@ inline void linearize_volume(const uint8_t *const in_ptr,
             if((y < 0 || y >= input_h) && has_pads)
             {
                 // All the values will be zeros
-                memset(out_ptr, 0, kernel_size * sizeof(T));
-                out_ptr += kernel_size;
+                memset(out_ptr, 0, kernel_width * sizeof(T));
+                out_ptr += kernel_width;
             }
             else
             {
@@ -199,7 +201,8 @@ void NEIm2ColKernel::run_generic(const Window &window)
                                       _has_bias,
                                       top_left_x,
                                       top_left_y,
-                                      static_cast<int>(_kernel_size),
+                                      static_cast<int>(_kernel_width),
+                                      static_cast<int>(_kernel_height),
                                       kernel_depth,
                                       input_w,
                                       input_h,
@@ -260,22 +263,24 @@ void NEIm2ColKernel::run_reduced(const Window &window)
 }
 
 NEIm2ColKernel::NEIm2ColKernel()
-    : _func(), _input(nullptr), _output(nullptr), _convolved_dims(), _conv_info(), _kernel_size(0), _has_bias(false)
+    : _func(), _input(nullptr), _output(nullptr), _convolved_dims(), _conv_info(), _kernel_width(0), _kernel_height(0), _has_bias(false)
 {
 }
 
-void NEIm2ColKernel::configure(const ITensor *input, ITensor *output, std::pair<unsigned int, unsigned int> convolved_dims, const PadStrideInfo &conv_info, bool has_bias)
+void NEIm2ColKernel::configure(const ITensor *input, ITensor *output, const Size2D &kernel_dims, const PadStrideInfo &conv_info, bool has_bias)
 {
     ARM_COMPUTE_ERROR_ON_DATA_TYPE_CHANNEL_NOT_IN(input, 1, DataType::F16, DataType::F32, DataType::QS8);
-    ARM_COMPUTE_ERROR_ON_DATA_TYPE_CHANNEL_NOT_IN(output, 1, DataType::F16, DataType::F32, DataType::QS8);
     ARM_COMPUTE_ERROR_ON_MISMATCHING_DATA_TYPES(input, output);
 
     _input          = input;
     _output         = output;
-    _convolved_dims = convolved_dims;
     _conv_info      = conv_info;
-    _kernel_size    = std::sqrt((output->info()->dimension(0) - (has_bias ? 1 : 0)) / input->info()->dimension(2));
-    _has_bias       = has_bias;
+    _kernel_width   = kernel_dims.width;
+    _kernel_height  = kernel_dims.height,
+    _convolved_dims = scaled_dimensions(input->info()->dimension(0), input->info()->dimension(1),
+                                        _kernel_width, _kernel_height,
+                                        _conv_info);
+    _has_bias = has_bias;
 
     unsigned int pad_x, pad_y, stride_x, stride_y = 0;
     std::tie(pad_x, pad_y)       = conv_info.pad();