COMPMID-784: Productise Winograd.

a) Added support for kernel size 5.
b) Templatised data type for transforms and batched gemms kernels.

Change-Id: Idb83dda7a5eec19e015888ab31902bd791913297
Reviewed-on: https://eu-gerrit-1.euhpc.arm.com/120540
Reviewed-by: Anthony Barbier <anthony.barbier@arm.com>
Tested-by: Jenkins <bsgcomp@arm.com>

2 files changed, 203 insertions, 111 deletions

diff --git a/src/core/NEON/kernels/NEWinogradLayerKernel.cpp b/src/core/NEON/kernels/NEWinogradLayerKernel.cpp
index b0a36ff46a..b2e44f8e09 100644
--- a/src/core/NEON/kernels/NEWinogradLayerKernel.cpp
+++ b/src/core/NEON/kernels/NEWinogradLayerKernel.cpp
@@ -32,25 +32,25 @@
 namespace arm_compute
 {
 //Batched Gemms
-template <int OutputTileRows, int OutputTileCols, int KernelRows, int KernelCols>
-NEWinogradLayerKernel<OutputTileRows, OutputTileCols, KernelRows, KernelCols>::NEWinogradLayerKernel()
+template <typename TIn, typename TOut, int OutputTileRows, int OutputTileCols, int KernelRows, int KernelCols>
+NEWinogradLayerBatchedGEMMKernel<TIn, TOut, OutputTileRows, OutputTileCols, KernelRows, KernelCols>::NEWinogradLayerBatchedGEMMKernel()
     : _gemms()
 {
 }
 
-template <int OutputTileRows, int OutputTileCols, int KernelRows, int KernelCols>
-void NEWinogradLayerKernel<OutputTileRows, OutputTileCols, KernelRows, KernelCols>::configure(
+template <typename TIn, typename TOut, int OutputTileRows, int OutputTileCols, int KernelRows, int KernelCols>
+void NEWinogradLayerBatchedGEMMKernel<TIn, TOut, OutputTileRows, OutputTileCols, KernelRows, KernelCols>::configure(
     const unsigned int n_gemms,
     const int M, const int K, const int N,
-    const int          a_matrix_stride,
-    const int          a_row_stride,
-    const int          b_matrix_stride,
-    const int          b_row_stride,
-    const int          c_matrix_stride,
-    const int          c_row_stride,
-    const float *const a_ptr,
-    const float *const b_ptr,
-    float *const       c_ptr)
+    const int        a_matrix_stride,
+    const int        a_row_stride,
+    const int        b_matrix_stride,
+    const int        b_row_stride,
+    const int        c_matrix_stride,
+    const int        c_row_stride,
+    const TIn *const a_ptr,
+    const TIn *const b_ptr,
+    TOut *const      c_ptr)
 {
     _gemms = support::cpp14::make_unique<MultiGEMM>(n_gemms, M, K, N, a_matrix_stride, a_row_stride, b_matrix_stride, b_row_stride, c_matrix_stride, c_row_stride, a_ptr, b_ptr, c_ptr);
     Window win;
@@ -59,8 +59,8 @@ void NEWinogradLayerKernel<OutputTileRows, OutputTileCols, KernelRows, KernelCol
     INEKernel::configure(win);
 }
 
-template <int OutputTileRows, int OutputTileCols, int KernelRows, int KernelCols>
-void NEWinogradLayerKernel<OutputTileRows, OutputTileCols, KernelRows, KernelCols>::run(const Window &window, const ThreadInfo &info)
+template <typename TIn, typename TOut, int OutputTileRows, int OutputTileCols, int KernelRows, int KernelCols>
+void NEWinogradLayerBatchedGEMMKernel<TIn, TOut, OutputTileRows, OutputTileCols, KernelRows, KernelCols>::run(const Window &window, const ThreadInfo &info)
 {
     ARM_COMPUTE_UNUSED(info);
     ARM_COMPUTE_ERROR_ON_UNCONFIGURED_KERNEL(this);
@@ -69,36 +69,66 @@ void NEWinogradLayerKernel<OutputTileRows, OutputTileCols, KernelRows, KernelCol
     _gemms->run(first_gemm, last_gemm);
 }
 
-template class NEWinogradLayerKernel<2, 2, 3, 3>;
+template <typename TIn, typename TOut, int OutputTileRows, int OutputTileCols, int KernelRows, int KernelCols>
+unsigned int NEWinogradLayerBatchedGEMMKernel<TIn, TOut, OutputTileRows, OutputTileCols, KernelRows, KernelCols>::get_number_gemms() const
+{
+    return WinogradBase::N_GEMMS;
+}
+
+template <typename TIn, typename TOut, int OutputTileRows, int OutputTileCols, int KernelRows, int KernelCols>
+int NEWinogradLayerBatchedGEMMKernel<TIn, TOut, OutputTileRows, OutputTileCols, KernelRows, KernelCols>::get_output_tile_rows() const
+{
+    return _output_tile_rows;
+}
+
+template <typename TIn, typename TOut, int OutputTileRows, int OutputTileCols, int KernelRows, int KernelCols>
+int NEWinogradLayerBatchedGEMMKernel<TIn, TOut, OutputTileRows, OutputTileCols, KernelRows, KernelCols>::get_output_tile_cols() const
+{
+    return _output_tile_cols;
+}
+
+template <typename TIn, typename TOut, int OutputTileRows, int OutputTileCols, int KernelRows, int KernelCols>
+int NEWinogradLayerBatchedGEMMKernel<TIn, TOut, OutputTileRows, OutputTileCols, KernelRows, KernelCols>::get_number_blocks() const
+{
+    return WinogradConv::N_BLOCK;
+}
+
+template class NEWinogradLayerBatchedGEMMKernel<float, float, 2, 2, 3, 3>;
+template class NEWinogradLayerBatchedGEMMKernel<float, float, 2, 2, 5, 5>;
 
 // Weights transform
 
-template <int OutputTileRows, int OutputTileCols, int KernelRows, int KernelCols>
-unsigned int NEWinogradLayerTransformWeightsKernel<OutputTileRows, OutputTileCols, KernelRows, KernelCols>::get_weight_storage_size(int n_output_channels, int n_input_channels)
+template <typename T, int OutputTileRows, int OutputTileCols, int KernelRows, int KernelCols>
+unsigned int NEWinogradLayerTransformWeightsKernel<T, OutputTileRows, OutputTileCols, KernelRows, KernelCols>::get_weight_storage_size(int n_output_channels, int n_input_channels) const
 {
     const KernelShape shape(n_output_channels, KernelRows, KernelCols, n_input_channels);
     return static_cast<unsigned int>(
-               // WinogradConv returns the size in bytes, we divide by `sizeof(float)` to
-               // express that in units of float.
-               WinogradConv::get_kernel_storage_size(shape) / sizeof(float));
+               // WinogradConv returns the size in bytes, we divide by `sizeof(T)` to express that in units of T
+               WinogradConv::get_kernel_storage_size(shape) / sizeof(T));
 }
 
-template <int OutputTileRows, int OutputTileCols, int KernelRows, int KernelCols>
-NEWinogradLayerTransformWeightsKernel<OutputTileRows, OutputTileCols, KernelRows, KernelCols>::NEWinogradLayerTransformWeightsKernel()
+template <typename T, int OutputTileRows, int OutputTileCols, int KernelRows, int KernelCols>
+NEWinogradLayerTransformWeightsKernel<T, OutputTileRows, OutputTileCols, KernelRows, KernelCols>::NEWinogradLayerTransformWeightsKernel()
     : _transform()
 {
 }
 
-template <int OutputTileRows, int OutputTileCols, int KernelRows, int KernelCols>
-void NEWinogradLayerTransformWeightsKernel<OutputTileRows, OutputTileCols, KernelRows, KernelCols>::configure(
+template <typename T, int OutputTileRows, int OutputTileCols, int KernelRows, int KernelCols>
+int NEWinogradLayerTransformWeightsKernel<T, OutputTileRows, OutputTileCols, KernelRows, KernelCols>::get_matrix_stride(const KernelShape &kernel_shape) const
+{
+    return WinogradConv::get_kernel_matrix_stride(kernel_shape);
+}
+
+template <typename T, int OutputTileRows, int OutputTileCols, int KernelRows, int KernelCols>
+void NEWinogradLayerTransformWeightsKernel<T, OutputTileRows, OutputTileCols, KernelRows, KernelCols>::configure(
     const ITensor *weights_hwio,
-    float *const   output,
+    T *const       output,
     const int      matrix_stride,     /** Stride across matrices in the output. */
     const int      n_output_channels, /** Number of filters. */
     const int      n_input_channels)  /** Number of channels in each filter. */
 {
     const int matrix_row_stride = roundup(n_output_channels, WinogradConv::N_BLOCK);
-    _transform                  = support::cpp14::make_unique<WeightsTransform>(reinterpret_cast<float *>(weights_hwio->buffer()), output, matrix_stride, matrix_row_stride, n_output_channels,
+    _transform                  = support::cpp14::make_unique<WeightsTransform>(reinterpret_cast<T *>(weights_hwio->buffer()), output, matrix_stride, matrix_row_stride, n_output_channels,
                                                                                 n_input_channels);
     Window win;
     auto   win_last = _transform->get_window();
@@ -106,8 +136,8 @@ void NEWinogradLayerTransformWeightsKernel<OutputTileRows, OutputTileCols, Kerne
     INEKernel::configure(win);
 }
 
-template <int OutputTileRows, int OutputTileCols, int KernelRows, int KernelCols>
-void NEWinogradLayerTransformWeightsKernel<OutputTileRows, OutputTileCols, KernelRows, KernelCols>::run(const Window &window, const ThreadInfo &info)
+template <typename T, int OutputTileRows, int OutputTileCols, int KernelRows, int KernelCols>
+void NEWinogradLayerTransformWeightsKernel<T, OutputTileRows, OutputTileCols, KernelRows, KernelCols>::run(const Window &window, const ThreadInfo &info)
 {
     ARM_COMPUTE_UNUSED(info);
     ARM_COMPUTE_ERROR_ON_UNCONFIGURED_KERNEL(this);
@@ -116,50 +146,57 @@ void NEWinogradLayerTransformWeightsKernel<OutputTileRows, OutputTileCols, Kerne
     _transform->run(fst, lst);
 }
 
-template <int OutputTileRows, int OutputTileCols, int KernelRows, int KernelCols>
-bool NEWinogradLayerTransformWeightsKernel<OutputTileRows, OutputTileCols, KernelRows, KernelCols>::is_parallelisable() const
+template <typename T, int OutputTileRows, int OutputTileCols, int KernelRows, int KernelCols>
+bool NEWinogradLayerTransformWeightsKernel<T, OutputTileRows, OutputTileCols, KernelRows, KernelCols>::is_parallelisable() const
 {
     return false;
 }
 
-template class NEWinogradLayerTransformWeightsKernel<2, 2, 3, 3>;
+template class NEWinogradLayerTransformWeightsKernel<float, 2, 2, 3, 3>;
+template class NEWinogradLayerTransformWeightsKernel<float, 2, 2, 5, 5>;
 
 // Input transform
 
-template <int OutputTileRows, int OutputTileCols, int KernelRows, int KernelCols>
-unsigned int NEWinogradLayerTransformInputKernel<OutputTileRows, OutputTileCols, KernelRows, KernelCols>::get_input_storage_size(
+template <typename T, int OutputTileRows, int OutputTileCols, int KernelRows, int KernelCols>
+unsigned int NEWinogradLayerTransformInputKernel<T, OutputTileRows, OutputTileCols, KernelRows, KernelCols>::get_input_storage_size(
     int  n_batches,   /** Number of batches in the input tensor. */
     int  n_channels,  /** Number of feature maps in the input tensor. */
     int  n_rows,      /** Number of rows in each feature map. */
     int  n_cols,      /** Number of columns in each feature map. */
     bool same_padding /** Use "SAME" padding, otherwise use "VALID". */
-)
+) const
 {
     // Construct shapes for the input and kernel tensors.
     const Tensor4DShape input_shape(n_batches, n_rows, n_cols, n_channels);
     const KernelShape   kern_shape(1, KernelRows, KernelCols, n_channels);
     const PaddingType   padding = (same_padding) ? PADDING_SAME : PADDING_VALID;
     // Return the size, converted into units of TIn
-    return static_cast<unsigned int>(
-               WinogradConv::get_input_storage_size(kern_shape, input_shape, padding) / sizeof(float));
+    return static_cast<unsigned int>(WinogradConv::get_input_storage_size(kern_shape, input_shape, padding) / sizeof(T));
+}
+
+template <typename T, int OutputTileRows, int OutputTileCols, int KernelRows, int KernelCols>
+int NEWinogradLayerTransformInputKernel<T, OutputTileRows, OutputTileCols, KernelRows, KernelCols>::get_matrix_stride(
+    const KernelShape &kernel_shape, const Tensor4DShape &input_shape, const PaddingType padding_type) const
+{
+    return WinogradConv::get_input_matrix_stride(kernel_shape, input_shape, padding_type);
 }
 
-template <int OutputTileRows, int OutputTileCols, int KernelRows, int KernelCols>
-NEWinogradLayerTransformInputKernel<OutputTileRows, OutputTileCols, KernelRows, KernelCols>::NEWinogradLayerTransformInputKernel()
+template <typename T, int OutputTileRows, int OutputTileCols, int KernelRows, int KernelCols>
+NEWinogradLayerTransformInputKernel<T, OutputTileRows, OutputTileCols, KernelRows, KernelCols>::NEWinogradLayerTransformInputKernel()
     : _transform()
 {
 }
 
-template <int OutputTileRows, int OutputTileCols, int KernelRows, int KernelCols>
-void NEWinogradLayerTransformInputKernel<OutputTileRows, OutputTileCols, KernelRows, KernelCols>::configure(
-    const float *const input,         /** Input tensor data */
-    const int          n_batches,     /** Number of batches in input tensor. */
-    const int          n_rows,        /** Number of rows in input tensor. */
-    const int          n_cols,        /** Number of columns in input tensor. */
-    const int          n_channels,    /** Number of channels in input tensor. */
-    const PaddingType  padding,       /** Padding type. */
-    float *const       output,        /** Base of output matrices. */
-    const int          matrix_stride) /** Stride between output matrices. */
+template <typename T, int OutputTileRows, int OutputTileCols, int KernelRows, int KernelCols>
+void NEWinogradLayerTransformInputKernel<T, OutputTileRows, OutputTileCols, KernelRows, KernelCols>::configure(
+    const T *const    input,         /** Input tensor data */
+    const int         n_batches,     /** Number of batches in input tensor. */
+    const int         n_rows,        /** Number of rows in input tensor. */
+    const int         n_cols,        /** Number of columns in input tensor. */
+    const int         n_channels,    /** Number of channels in input tensor. */
+    const PaddingType padding,       /** Padding type. */
+    T *const          output,        /** Base of output matrices. */
+    const int         matrix_stride) /** Stride between output matrices. */
 {
     //  _input_matrix_row_stride(n_input_channels),
     _transform = support::cpp14::make_unique<InputTransform>(input, n_batches, n_rows, n_cols, n_channels, padding, output, matrix_stride, n_channels);
@@ -169,8 +206,8 @@ void NEWinogradLayerTransformInputKernel<OutputTileRows, OutputTileCols, KernelR
     INEKernel::configure(win);
 }
 
-template <int OutputTileRows, int OutputTileCols, int KernelRows, int KernelCols>
-void NEWinogradLayerTransformInputKernel<OutputTileRows, OutputTileCols, KernelRows, KernelCols>::run(const Window &window, const ThreadInfo &info)
+template <typename T, int OutputTileRows, int OutputTileCols, int KernelRows, int KernelCols>
+void NEWinogradLayerTransformInputKernel<T, OutputTileRows, OutputTileCols, KernelRows, KernelCols>::run(const Window &window, const ThreadInfo &info)
 {
     ARM_COMPUTE_UNUSED(info);
     ARM_COMPUTE_ERROR_ON_UNCONFIGURED_KERNEL(this);
@@ -179,24 +216,25 @@ void NEWinogradLayerTransformInputKernel<OutputTileRows, OutputTileCols, KernelR
     _transform->run(fst, lst);
 }
 
-template <int OutputTileRows, int OutputTileCols, int KernelRows, int KernelCols>
-bool NEWinogradLayerTransformInputKernel<OutputTileRows, OutputTileCols, KernelRows, KernelCols>::is_parallelisable() const
+template <typename T, int OutputTileRows, int OutputTileCols, int KernelRows, int KernelCols>
+bool NEWinogradLayerTransformInputKernel<T, OutputTileRows, OutputTileCols, KernelRows, KernelCols>::is_parallelisable() const
 {
     return false;
 }
 
-template class NEWinogradLayerTransformInputKernel<2, 2, 3, 3>;
+template class NEWinogradLayerTransformInputKernel<float, 2, 2, 3, 3>;
+template class NEWinogradLayerTransformInputKernel<float, 2, 2, 5, 5>;
 
 // Output transform
 
-template <int OutputTileRows, int OutputTileCols, int KernelRows, int KernelCols>
-unsigned int NEWinogradLayerTransformOutputKernel<OutputTileRows, OutputTileCols, KernelRows, KernelCols>::get_output_storage_size(
+template <typename T, int OutputTileRows, int OutputTileCols, int KernelRows, int KernelCols>
+unsigned int NEWinogradLayerTransformOutputKernel<T, OutputTileRows, OutputTileCols, KernelRows, KernelCols>::get_output_storage_size(
     int  n_batches,         /** Number of batches in the output tensor. */
     int  n_rows,            /** Number of rows in each feature map of the input tensor. */
     int  n_cols,            /** Number of columns in each feature map of the input tensor. */
     int  n_output_channels, /** Number of feature maps in the output tensor. */
     bool same_padding       /** Use "SAME" padding, otherwise use "VALID". */
-)
+) const
 {
     // Construct shapes for the input and kernel tensors.
     const Tensor4DShape input_shape(n_batches, n_rows, n_cols, 1);
@@ -205,25 +243,38 @@ unsigned int NEWinogradLayerTransformOutputKernel<OutputTileRows, OutputTileCols
 
     // Return the size, converted into units of TOut
     return static_cast<unsigned int>(
-               WinogradConv::get_output_storage_size(kern_shape, input_shape, padding) / sizeof(float));
+               WinogradConv::get_output_storage_size(kern_shape, input_shape, padding) / sizeof(T));
 }
 
-template <int OutputTileRows, int OutputTileCols, int KernelRows, int KernelCols>
-NEWinogradLayerTransformOutputKernel<OutputTileRows, OutputTileCols, KernelRows, KernelCols>::NEWinogradLayerTransformOutputKernel()
+template <typename T, int OutputTileRows, int OutputTileCols, int KernelRows, int KernelCols>
+NEWinogradLayerTransformOutputKernel<T, OutputTileRows, OutputTileCols, KernelRows, KernelCols>::NEWinogradLayerTransformOutputKernel()
     : _biases(nullptr), _output_workspace(nullptr), _matrix_stride(0), _matrix_row_stride(0), _output(nullptr), _n_batches(0), _n_rows(0), _n_cols(0), _n_channels(0)
 {
 }
 
-template <int OutputTileRows, int OutputTileCols, int KernelRows, int KernelCols>
-void NEWinogradLayerTransformOutputKernel<OutputTileRows, OutputTileCols, KernelRows, KernelCols>::configure(
-    const ITensor     *biases,
-    const float *const output_workingspace,
-    const int          matrix_stride,
-    float *const       output,
-    const int          n_batches,
-    const int          n_rows,
-    const int          n_cols,
-    const int          n_channels)
+template <typename T, int OutputTileRows, int OutputTileCols, int KernelRows, int KernelCols>
+int NEWinogradLayerTransformOutputKernel<T, OutputTileRows, OutputTileCols, KernelRows, KernelCols>::get_matrix_stride(
+    const KernelShape &kernel_shape, const Tensor4DShape &input_shape, const PaddingType padding_type) const
+{
+    return WinogradConv::get_output_matrix_stride(kernel_shape, input_shape, padding_type);
+}
+template <typename T, int OutputTileRows, int OutputTileCols, int KernelRows, int KernelCols>
+Tensor4DShape NEWinogradLayerTransformOutputKernel<T, OutputTileRows, OutputTileCols, KernelRows, KernelCols>::get_output_shape(
+    const KernelShape &kernel_shape, const Tensor4DShape &in_shape, const PaddingType padding) const
+{
+    return WinogradConv::get_output_shape(kernel_shape, in_shape, padding);
+}
+
+template <typename T, int OutputTileRows, int OutputTileCols, int KernelRows, int KernelCols>
+void NEWinogradLayerTransformOutputKernel<T, OutputTileRows, OutputTileCols, KernelRows, KernelCols>::configure(
+    const ITensor *biases,
+    const T *const output_workingspace,
+    const int      matrix_stride,
+    T *const       output,
+    const int      n_batches,
+    const int      n_rows,
+    const int      n_cols,
+    const int      n_channels)
 {
     _biases            = biases;
     _output_workspace  = output_workingspace;
@@ -243,8 +294,8 @@ void NEWinogradLayerTransformOutputKernel<OutputTileRows, OutputTileCols, Kernel
     INEKernel::configure(win);
 }
 
-template <int OutputTileRows, int OutputTileCols, int KernelRows, int KernelCols>
-void NEWinogradLayerTransformOutputKernel<OutputTileRows, OutputTileCols, KernelRows, KernelCols>::run(const Window &window, const ThreadInfo &info)
+template <typename T, int OutputTileRows, int OutputTileCols, int KernelRows, int KernelCols>
+void NEWinogradLayerTransformOutputKernel<T, OutputTileRows, OutputTileCols, KernelRows, KernelCols>::run(const Window &window, const ThreadInfo &info)
 {
     ARM_COMPUTE_UNUSED(info);
     ARM_COMPUTE_ERROR_ON_UNCONFIGURED_KERNEL(this);
@@ -253,7 +304,7 @@ void NEWinogradLayerTransformOutputKernel<OutputTileRows, OutputTileCols, Kernel
     ARM_COMPUTE_ERROR_ON_NULLPTR(_output);
 
     OutputTransform output_transform(_output_workspace, _matrix_stride, _matrix_row_stride,
-                                     reinterpret_cast<float *>(_biases->buffer()), _output,
+                                     reinterpret_cast<T *>(_biases->buffer()), _output,
                                      _n_batches, _n_rows, _n_cols, _n_channels);
 
     // The code below cannot be moved to configure because biases hasn't been allocated at that point
@@ -262,12 +313,13 @@ void NEWinogradLayerTransformOutputKernel<OutputTileRows, OutputTileCols, Kernel
     output_transform.run(fst, lst);
 }
 
-template <int OutputTileRows, int OutputTileCols, int KernelRows, int KernelCols>
-bool NEWinogradLayerTransformOutputKernel<OutputTileRows, OutputTileCols, KernelRows, KernelCols>::is_parallelisable() const
+template <typename T, int OutputTileRows, int OutputTileCols, int KernelRows, int KernelCols>
+bool NEWinogradLayerTransformOutputKernel<T, OutputTileRows, OutputTileCols, KernelRows, KernelCols>::is_parallelisable() const
 {
     return false;
 }
 
-template class NEWinogradLayerTransformOutputKernel<2, 2, 3, 3>;
+template class NEWinogradLayerTransformOutputKernel<float, 2, 2, 3, 3>;
+template class NEWinogradLayerTransformOutputKernel<float, 2, 2, 5, 5>;
 
 } // namespace arm_compute
diff --git a/src/runtime/NEON/functions/NEWinogradLayer.cpp b/src/runtime/NEON/functions/NEWinogradLayer.cpp
index 215f1bfddf..e343583b36 100644
--- a/src/runtime/NEON/functions/NEWinogradLayer.cpp
+++ b/src/runtime/NEON/functions/NEWinogradLayer.cpp
@@ -28,6 +28,8 @@
 #include "arm_compute/runtime/NEON/NEScheduler.h"
 #include "support/ToolchainSupport.h"
 
+#include "arm_compute/core/NEON/kernels/NEWinogradLayerKernel.h"
+
 #include "arm_compute/core/NEON/kernels/convolution/winograd/winograd_gemm.hpp"
 
 namespace
@@ -45,8 +47,9 @@ inline Tensor4DShape internal_get_input_shape(const arm_compute::ITensor *input)
 namespace arm_compute
 {
 NEWinogradLayer::NEWinogradLayer(std::shared_ptr<IMemoryManager> memory_manager)
-    : _memory_group(std::move(memory_manager)), _winograd_kernel(), _transform_input_kernel(), _transform_output_kernel(), _transform_weights_kernel(), _permute_input(), _permute_weights(),
-      _permute_output(), _input_workspace(), _output_workspace(), _kernel_storage(), _input_nhwc(), _output_nhwc(), _weights_hwio(), _input(), _weights(), _output(), _reshaped_kernel(false)
+    : _memory_group(std::move(memory_manager)), _batched_gemm_kernel(nullptr), _transform_input_kernel(nullptr), _transform_output_kernel(nullptr), _transform_weights_kernel(nullptr), _permute_input(),
+      _permute_weights(), _permute_output(), _input_workspace(), _output_workspace(), _kernel_storage(), _input_nhwc(), _output_nhwc(), _weights_hwio(), _input(), _weights(), _output(),
+      _reshaped_kernel(false)
 {
 } /* arm_compute */
 
@@ -54,7 +57,7 @@ void NEWinogradLayer::configure(const ITensor *input, const ITensor *weights, co
 {
     ARM_COMPUTE_ERROR_ON_DATA_TYPE_CHANNEL_NOT_IN(input, 1, DataType::F32);
     ARM_COMPUTE_ERROR_ON_MISMATCHING_DATA_TYPES(input, weights, biases);
-    ARM_COMPUTE_ERROR_ON_MSG(weights->info()->dimension(1) != 3 || weights->info()->dimension(0) != 3, "Only 3x3 kernels are supported");
+    ARM_COMPUTE_ERROR_ON_MSG(weights->info()->dimension(0) != 3 && weights->info()->dimension(0) != 5, "Only 3 and 5 kernels are supported");
     ARM_COMPUTE_ERROR_ON(weights->info()->num_dimensions() > 4);
 
     if(biases != nullptr)
@@ -67,6 +70,36 @@ void NEWinogradLayer::configure(const ITensor *input, const ITensor *weights, co
     _input   = input;
     _output  = output;
 
+    std::unique_ptr<INEWinogradLayerBatchedGEMMKernel<float, float>> batched_gemm_kernel;
+    std::unique_ptr<INEWinogradLayerTransformInputKernel<float>>   transform_input_kernel;
+    std::unique_ptr<INEWinogradLayerTransformWeightsKernel<float>> transform_weights_kernel;
+    std::unique_ptr<INEWinogradLayerTransformOutputKernel<float>>  transform_output_kernel;
+
+    switch(weights->info()->dimension(0))
+    {
+        case 3:
+        {
+            batched_gemm_kernel      = support::cpp14::make_unique<NEWinogradLayerBatchedGEMMKernel<float, float, 2, 2, 3, 3>>();
+            transform_input_kernel   = support::cpp14::make_unique<NEWinogradLayerTransformInputKernel<float, 2, 2, 3, 3>>();
+            transform_weights_kernel = support::cpp14::make_unique<NEWinogradLayerTransformWeightsKernel<float, 2, 2, 3, 3>>();
+            transform_output_kernel  = support::cpp14::make_unique<NEWinogradLayerTransformOutputKernel<float, 2, 2, 3, 3>>();
+            break;
+        }
+        case 5:
+        {
+            batched_gemm_kernel      = support::cpp14::make_unique<NEWinogradLayerBatchedGEMMKernel<float, float, 2, 2, 5, 5>>();
+            transform_input_kernel   = support::cpp14::make_unique<NEWinogradLayerTransformInputKernel<float, 2, 2, 5, 5>>();
+            transform_weights_kernel = support::cpp14::make_unique<NEWinogradLayerTransformWeightsKernel<float, 2, 2, 5, 5>>();
+            transform_output_kernel  = support::cpp14::make_unique<NEWinogradLayerTransformOutputKernel<float, 2, 2, 5, 5>>();
+            break;
+        }
+        default:
+        {
+            ARM_COMPUTE_ERROR("Not supported.");
+            break;
+        }
+    }
+
     const PaddingType use_padding_type = (conv_info.pad_left() != 0u) ? PADDING_SAME : PADDING_VALID;
     const bool        use_same_padding = use_padding_type == PADDING_SAME;
 
@@ -84,22 +117,19 @@ void NEWinogradLayer::configure(const ITensor *input, const ITensor *weights, co
     const size_t        data_type_size = input->info()->element_size();
     // Get the memory required to instantiate a new Winograd operator.
     constexpr size_t storage_alignment   = 64;
-    const size_t     kernel_storage_size = NEWinogradLayerTransformWeightsKernel<2, 2, 3, 3>::get_weight_storage_size(out_channels, in_channels) * data_type_size;
+    const size_t     kernel_storage_size = transform_weights_kernel->get_weight_storage_size(out_channels, in_channels) * data_type_size;
     _kernel_storage.allocator()->init(TensorInfo(TensorShape{ (kernel_storage_size + storage_alignment - 1) }, 1, DataType::U8));
     _memory_group.manage(&_kernel_storage);
     _memory_group.manage(&_input_nhwc);
     _kernel_storage.allocator()->allocate();
     // Input storage
-
-    using IT                        = NEWinogradLayerTransformInputKernel<2, 2, 3, 3>;
-    const size_t input_storage_size = IT::get_input_storage_size(in_shape.n_batches, in_shape.n_channels, in_shape.n_rows, in_shape.n_cols, use_same_padding) * data_type_size;
+    const size_t input_storage_size = transform_input_kernel->get_input_storage_size(in_shape.n_batches, in_shape.n_channels, in_shape.n_rows, in_shape.n_cols, use_same_padding) * data_type_size;
     _input_workspace.allocator()->init(TensorInfo(TensorShape{ (input_storage_size + storage_alignment - 1) }, 1, DataType::U8));
     _memory_group.manage(&_input_workspace);
     _input_workspace.allocator()->allocate();
 
     // Output storage
-    using OT                         = NEWinogradLayerTransformOutputKernel<2, 2, 3, 3>;
-    const size_t output_storage_size = OT::get_output_storage_size(in_shape.n_batches, in_shape.n_rows, in_shape.n_cols, out_channels, use_same_padding) * data_type_size;
+    const size_t output_storage_size = transform_output_kernel->get_output_storage_size(in_shape.n_batches, in_shape.n_rows, in_shape.n_cols, out_channels, use_same_padding) * data_type_size;
     _output_workspace.allocator()->init(TensorInfo(TensorShape{ (output_storage_size + storage_alignment - 1) }, 1, DataType::U8));
     _memory_group.manage(&_output_workspace);
     _output_workspace.allocator()->allocate();
@@ -119,47 +149,57 @@ void NEWinogradLayer::configure(const ITensor *input, const ITensor *weights, co
     _permute_input.configure(input, &_input_nhwc, PermutationVector(2U, 0U, 1U));
     _input_nhwc.allocator()->allocate();
 
-    using T                          = winograd::WinogradGEMM<2, 2, 3, 3>::Convolution<float, float>;
     const int         weights_width  = weights->info()->dimension(0);
     const int         weights_height = weights->info()->dimension(1);
     const KernelShape kernel_shape({ out_channels, weights_height, weights_width, in_channels });
 
     // Configure the InputTransform
-    const int input_matrix_stride = T::get_input_matrix_stride(kernel_shape, in_shape, use_padding_type);
-    _transform_input_kernel.configure(reinterpret_cast<float *>(_input_nhwc.buffer()), in_shape.n_batches, in_shape.n_rows, in_shape.n_cols, in_shape.n_channels, use_padding_type,
+    const int input_matrix_stride = transform_input_kernel->get_matrix_stride(kernel_shape, in_shape, use_padding_type);
+    transform_input_kernel->configure(reinterpret_cast<float *>(_input_nhwc.buffer()), in_shape.n_batches, in_shape.n_rows, in_shape.n_cols, in_shape.n_channels, use_padding_type,
                                       reinterpret_cast<float *>(_input_workspace.buffer()), input_matrix_stride);
 
     // Configure WeightsTransform
-    const int kernel_matrix_stride = T::get_kernel_matrix_stride(kernel_shape);
-    _transform_weights_kernel.configure(&_weights_hwio, reinterpret_cast<float *>(_kernel_storage.buffer()), kernel_matrix_stride, out_channels, in_channels);
+    const int kernel_matrix_stride = transform_weights_kernel->get_matrix_stride(kernel_shape);
+    transform_weights_kernel->configure(&_weights_hwio, reinterpret_cast<float *>(_kernel_storage.buffer()), kernel_matrix_stride, out_channels, in_channels);
 
     // Configure OutputTransform
     //The biases tensor has not been allocated at this point in time, the output transform will add the biases to the final result in the run() method
-    const int  output_matrix_stride = T::get_output_matrix_stride(kernel_shape, in_shape, use_padding_type);
-    const auto output_shape(T::get_output_shape(kernel_shape, in_shape, use_padding_type));
+    const int  output_matrix_stride = transform_output_kernel->get_matrix_stride(kernel_shape, in_shape, use_padding_type);
+    const auto output_shape(transform_output_kernel->get_output_shape(kernel_shape, in_shape, use_padding_type));
 
-    _transform_output_kernel.configure(biases, reinterpret_cast<float *>(_output_workspace.buffer()),
+    transform_output_kernel->configure(biases, reinterpret_cast<float *>(_output_workspace.buffer()),
                                        output_matrix_stride, reinterpret_cast<float *>(_output_nhwc.buffer()),
                                        in_shape.n_batches, output_shape.n_rows, output_shape.n_cols, out_channels);
 
     // Configure Batched GEMMs
-    const int tile_rows                = iceildiv(output_shape.n_rows, NEWinogradLayerKernel<2, 2, 3, 3>::_output_tile_rows);
-    const int tile_cols                = iceildiv(output_shape.n_cols, NEWinogradLayerKernel<2, 2, 3, 3>::_output_tile_cols);
-    const int m                        = in_shape.n_batches * tile_rows * tile_cols;
-    const int k                        = in_shape.n_channels;
-    const int n                        = out_channels;
-    const int input_matrix_row_stride  = in_shape.n_channels;
-    const int kernel_matrix_row_stride = roundup(out_channels, NEWinogradLayerKernel<2, 2, 3, 3>::WinogradConv::N_BLOCK);
-    const int output_matrix_row_stride = kernel_matrix_row_stride;
-
-    _winograd_kernel.configure(NEWinogradLayerKernel<2, 2, 3, 3>::WinogradBase::N_GEMMS, m, k, n,
-                               input_matrix_stride, input_matrix_row_stride,
-                               kernel_matrix_stride, kernel_matrix_row_stride,
-                               output_matrix_stride, output_matrix_row_stride,
-                               reinterpret_cast<float *>(_input_workspace.buffer()), reinterpret_cast<float *>(_kernel_storage.buffer()), reinterpret_cast<float *>(_output_workspace.buffer()));
+    const int      output_tile_rows         = batched_gemm_kernel->get_output_tile_rows();
+    const int      output_tile_cols         = batched_gemm_kernel->get_output_tile_cols();
+    const int      n_block                  = batched_gemm_kernel->get_number_blocks();
+    const int      tile_rows                = iceildiv(output_shape.n_rows, output_tile_rows);
+    const int      tile_cols                = iceildiv(output_shape.n_cols, output_tile_cols);
+    const int      m                        = in_shape.n_batches * tile_rows * tile_cols;
+    const int      k                        = in_shape.n_channels;
+    const int      n                        = out_channels;
+    const int      input_matrix_row_stride  = in_shape.n_channels;
+    const int      kernel_matrix_row_stride = roundup(out_channels, n_block);
+    const int      output_matrix_row_stride = kernel_matrix_row_stride;
+    const unsigned n_gemms                  = batched_gemm_kernel->get_number_gemms();
+
+    batched_gemm_kernel->configure(n_gemms, m, k, n,
+                                   input_matrix_stride, input_matrix_row_stride,
+                                   kernel_matrix_stride, kernel_matrix_row_stride,
+                                   output_matrix_stride, output_matrix_row_stride,
+                                   reinterpret_cast<float *>(_input_workspace.buffer()),
+                                   reinterpret_cast<float *>(_kernel_storage.buffer()),
+                                   reinterpret_cast<float *>(_output_workspace.buffer()));
 
     // Reorder the convoluted output to ACL's ordering NCHW
     _permute_output.configure(&_output_nhwc, _output, PermutationVector(1U, 2U, 0U));
+
+    _transform_input_kernel   = std::move(transform_input_kernel);
+    _transform_weights_kernel = std::move(transform_weights_kernel);
+    _transform_output_kernel  = std::move(transform_output_kernel);
+    _batched_gemm_kernel      = std::move(batched_gemm_kernel);
 }
 
 void NEWinogradLayer::run()
@@ -169,19 +209,19 @@ void NEWinogradLayer::run()
     {
         _reshaped_kernel = true;
         _permute_weights.run();
-        NEScheduler::get().schedule(&_transform_weights_kernel, Window::DimX);
+        NEScheduler::get().schedule(_transform_weights_kernel.get(), Window::DimX);
     }
     //Bring channels to the front as Winograd code expects the tensor to be in the format NHWC
     _permute_input.run();
 
     // Transform input tensor to the winograd domain
-    NEScheduler::get().schedule(&_transform_input_kernel, Window::DimX);
+    NEScheduler::get().schedule(_transform_input_kernel.get(), Window::DimX);
 
     //Run 16 GEMMs in multiple threads, each kernel runs one or more GEMMs
-    NEScheduler::get().schedule(&_winograd_kernel, Window::DimX);
+    NEScheduler::get().schedule(_batched_gemm_kernel.get(), Window::DimX);
 
     // Transform output tensor to the spatial domain
-    NEScheduler::get().schedule(&_transform_output_kernel, Window::DimX);
+    NEScheduler::get().schedule(_transform_output_kernel.get(), Window::DimX);
 
     // Reorder the convoluted output to ACL's ordering NCHW
     _permute_output.run();