diff options
Diffstat (limited to 'arm_compute/core/NEON/kernels/convolution/winograd/transforms/output.hpp')
-rw-r--r-- | arm_compute/core/NEON/kernels/convolution/winograd/transforms/output.hpp | 75 |
1 files changed, 48 insertions, 27 deletions
diff --git a/arm_compute/core/NEON/kernels/convolution/winograd/transforms/output.hpp b/arm_compute/core/NEON/kernels/convolution/winograd/transforms/output.hpp index 401b2816be..700ca76c68 100644 --- a/arm_compute/core/NEON/kernels/convolution/winograd/transforms/output.hpp +++ b/arm_compute/core/NEON/kernels/convolution/winograd/transforms/output.hpp @@ -23,7 +23,7 @@ */ #pragma once -#include "arm_compute/core/NEON/kernels/convolution/winograd/winograd_gemm.hpp" +#include "../winograd_gemm.hpp" namespace winograd { @@ -31,7 +31,13 @@ namespace winograd int kernel_rows, int kernel_cols> template <typename T> void WinogradGEMM<output_tile_rows, output_tile_cols, kernel_rows, kernel_cols>::OutputTransform<T>::execute( - const Tensor4DShape &output_shape, + const int n_batches, + const int output_batch_stride, + const int n_rows, + const int output_row_stride, + const int n_cols, + const int output_col_stride, + const int n_channels, const T* const matrix_base, const int matrix_stride, const int matrix_row_stride, @@ -41,19 +47,16 @@ namespace winograd { // Compute the number of tiles and hence the padding required on the bottom // and right of the image. - const int tile_M = iceildiv(output_shape.n_rows, output_tile_rows); - const int tile_N = iceildiv(output_shape.n_cols, output_tile_cols); - const int pad_bottom = output_tile_rows*tile_M - output_shape.n_rows; - const int pad_right = output_tile_cols*tile_N - output_shape.n_cols; + const int tile_M = iceildiv(n_rows, output_tile_rows); + const int tile_N = iceildiv(n_cols, output_tile_cols); + const int pad_bottom = output_tile_rows*tile_M - n_rows; + const int pad_right = output_tile_cols*tile_N - n_cols; const int matrix_tile_row_stride = tile_N * matrix_row_stride; const int matrix_batch_stride = tile_M * matrix_tile_row_stride; - const int output_col_stride = output_shape.n_channels; - const int output_row_stride = output_shape.n_cols * output_col_stride; - const int output_batch_stride = output_shape.n_rows * output_row_stride; // Perform the output transformation for each batch - for (int batch = 0; batch < output_shape.n_batches; batch++) + for (int batch = 0; batch < n_batches; batch++) { // Get batch offset for input and outputs. const T* const matrix_batch = matrix_base + batch*matrix_batch_stride; @@ -69,7 +72,7 @@ namespace winograd // Process the row process_tile_row( - tile_N, output_shape.n_channels, matrix_tile_row, matrix_stride, + tile_N, n_channels, matrix_tile_row, matrix_stride, matrix_row_stride, biases, outptr_row, output_row_stride, output_col_stride, row_pad_bottom, pad_right @@ -139,12 +142,18 @@ namespace winograd const int n_batches, const int n_rows, const int n_cols, - const int n_channels + const int n_channels, + const int out_batch_stride, + const int out_row_stride, + const int out_col_stride ) : _matrix_base(matrix_base), _biases(biases), _matrix_stride(matrix_stride), _matrix_row_stride(matrix_row_stride), _outptr(output), _n_batches(n_batches), _n_rows(n_rows), _n_cols(n_cols), _n_channels(n_channels), _tile_M(iceildiv(n_rows, output_tile_rows)), - _tile_N(iceildiv(n_cols, output_tile_cols)) + _tile_N(iceildiv(n_cols, output_tile_cols)), + _out_col_stride(out_col_stride ? out_col_stride : n_channels), + _out_row_stride(out_row_stride ? out_row_stride : n_cols * _out_col_stride), + _out_batch_stride(out_batch_stride ? out_batch_stride : n_rows * _out_row_stride) { } @@ -152,10 +161,9 @@ namespace winograd template <typename T> unsigned int WinogradGEMM<otr, otc, kr, kc>::OutputTransform<T>::get_window() const { - // TODO When the output transform supports multithreading, return the total - // number of tile rows (allowing for multiple batches). For now we return 1 - // to indicate that the activations must be transformed as a single block. - return 1; // TODO _tile_M * _n_batches; + // The final window includes the tail, all other windows will be a multiple + // of the window block in size. + return iceildiv(_n_channels, WINDOW_BLOCK); } template <int otr, int otc, int kr, int kc> @@ -164,18 +172,31 @@ namespace winograd const unsigned int start, const unsigned int stop ) { - // TODO When the output transform supports multithreading call execute for a - // portion of the tile rows. - (void) start; - (void) stop; + if (start >= get_window()) + { + return; + } + + // Determine the window of work to perform + const unsigned int start_channel = start * WINDOW_BLOCK; + const unsigned int stop_channel = std::min<const unsigned int>( + stop * WINDOW_BLOCK, _n_channels + ); + const unsigned int n_channels = stop_channel - start_channel; - // For now, just do all of the work. - const Tensor4DShape output_shape = { - _n_batches, _n_rows, _n_cols, _n_channels, NHWC - }; execute( - output_shape, _matrix_base, _matrix_stride, _matrix_row_stride, _biases, - _outptr + _n_batches, + _out_batch_stride, + _n_rows, + _out_row_stride, + _n_cols, + _out_col_stride, + n_channels, + _matrix_base + start_channel, + _matrix_stride, + _matrix_row_stride, + (_biases)?(_biases + start_channel):(nullptr), + _outptr + start_channel ); } } // namespace winograd |