diff options
Diffstat (limited to 'arm_compute/core/NEON/kernels/NEWinogradLayerKernel.h')
| -rw-r--r-- | arm_compute/core/NEON/kernels/NEWinogradLayerKernel.h | 53 |
1 files changed, 41 insertions, 12 deletions
diff --git a/arm_compute/core/NEON/kernels/NEWinogradLayerKernel.h b/arm_compute/core/NEON/kernels/NEWinogradLayerKernel.h index 73b7e8d2b7..95261929ca 100644 --- a/arm_compute/core/NEON/kernels/NEWinogradLayerKernel.h +++ b/arm_compute/core/NEON/kernels/NEWinogradLayerKernel.h @@ -1,5 +1,5 @@ /* - * Copyright (c) 2017, 2018 ARM Limited. + * Copyright (c) 2017-2018 ARM Limited. * * SPDX-License-Identifier: MIT * @@ -25,6 +25,7 @@ #define __ARM_COMPUTE_NEGEMMWINOGRADLAYERKERNEL_H__ #include "arm_compute/core/NEON/INEKernel.h" +#include "arm_compute/core/NEON/kernels/winograd/convolution.hpp" #include "arm_compute/core/NEON/kernels/winograd/tensor.hpp" namespace arm_compute @@ -36,11 +37,25 @@ class Winograd3x3F32 final { public: friend class NEWinogradLayerKernel; - Winograd3x3F32(const KernelShape &kernel_shape, const Tensor4DShape input_shape, const PaddingType padding_type, void *kernel_storage); + Winograd3x3F32( + const int n_batches, /** Number of batches in the input and output tensors. */ + const int n_input_channels, /** Number of feature maps in a batch of the input tensor. */ + const int n_input_rows, /** Number of rows in a feature map of the input tensor. */ + const int n_input_cols, /** Number of columns in a feature map of the input tensor. */ + const int n_output_channels, /** Number of feature maps in the output tensor. */ + const bool same_padding, /** Use "SAME" padding, otherwise use "VALID". */ + const float *const weights, /** Pointer to weight tensor in spatial domain. Must be ordered as "Height x Rows x Input Feature Maps x Output Feature Maps. */ + float *const weights_storage, /** Pointer to storage for weight tensor in the Winograd domain. Must be at least the size returned by `get_weight_storage_size`. */ + const float *const input, /** Pointer to NHWC ordered input tensor, in the spatial domain. */ + float *const winograd_input, /** Pointer to working space for the input tensor in the Winograd domain. Must be at least the size returned by `get_input_storage_size`. */ + float *const output, /** Pointer to NHWC ordered output tensor, in the spatial domain. */ + float *const winograd_output /** Pointer to working space for the output tensor in the Winograd domain. Must be at least the size returned by `get_output_storage_size`. */ + ); + ~Winograd3x3F32(); - void transform_weights(const void *const kernel, void *transform_working_space); - void reshape_input(const Tensor4DShape &input_shape, const PaddingType padding_type, const void *const input, void *working_space); - void reshape_output(const Tensor4DShape &input_shape, const PaddingType padding_type, void *const output); + void transform_weights(); + void transform_input(); + void transform_output(); private: class Private; @@ -75,15 +90,29 @@ public: /* Get the memory required to instantiate a new Winograd operator. */ - static size_t get_kernel_storage_size(const KernelShape &shape); + static size_t get_weight_storage_size( + const int n_output_channels, /** Number of output feature maps. */ + const int n_input_channels /** Number of input feature maps. */ + ); - /* Get the memory required to apply a Winograd operator to some input. - */ - static size_t get_working_space_size(const Tensor4DShape &input_shape, const KernelShape &k_shape, const PaddingType padding); + static unsigned int get_input_storage_size( + const int n_batches, /** Number of batches in the input tensor. */ + const int n_channels, /** Number of feature maps in the input tensor. */ + const int n_rows, /** Number of rows in each feature map. */ + const int n_cols, /** Number of columns in each feature map. */ + const bool same_padding /** Use "SAME" padding, otherwise use "VALID". */ + ); - /* Get the memory required to transform the kernel. - */ - static size_t get_kernel_transform_working_size(const KernelShape &shape); + /** Determine how much memory (in units of TOut) to allocate for the + * (Winograd domain) output. + */ + static unsigned int get_output_storage_size( + const int n_batches, /** Number of batches in the output tensor. */ + const int n_rows, /** Number of rows in each feature map of the input tensor. */ + const int n_cols, /** Number of columns in each feature map of the input tensor. */ + const int n_output_channels, /** Number of feature maps in the output tensor. */ + const bool same_padding /** Use "SAME" padding, otherwise use "VALID". */ + ); protected: Winograd3x3F32 *_convolver; |