diff options
| author | Michalis Spyrou <michalis.spyrou@arm.com> | 2018-04-23 15:17:31 +0100 |
|---|---|---|
| committer | Anthony Barbier <anthony.barbier@arm.com> | 2018-11-02 16:51:37 +0000 |
| commit | e250389ed6d78153a55382fa5b3519c151bfd79f (patch) | |
| tree | 80c63793769ad18fd0406e7f8b40840aed7ac3ce /src/core/NEON/kernels/NEWeightsReshapeKernel.cpp | |
| parent | 79ffadebd8dff7eaecbcfa3a28106736f240f1c5 (diff) | |
| download | ComputeLibrary-e250389ed6d78153a55382fa5b3519c151bfd79f.tar.gz | |
COMPMID-810 Add NHWC data format support for NEON convolution
Change-Id: I2a7b49a12da7f3bc3f04749243b1dc111160de6e
Reviewed-on: https://eu-gerrit-1.euhpc.arm.com/129348
Tested-by: Jenkins <bsgcomp@arm.com>
Reviewed-by: Anthony Barbier <anthony.barbier@arm.com>
Diffstat (limited to 'src/core/NEON/kernels/NEWeightsReshapeKernel.cpp')
| -rw-r--r-- | src/core/NEON/kernels/NEWeightsReshapeKernel.cpp | 27 |
1 files changed, 17 insertions, 10 deletions
diff --git a/src/core/NEON/kernels/NEWeightsReshapeKernel.cpp b/src/core/NEON/kernels/NEWeightsReshapeKernel.cpp index 150140271d..3031a87637 100644 --- a/src/core/NEON/kernels/NEWeightsReshapeKernel.cpp +++ b/src/core/NEON/kernels/NEWeightsReshapeKernel.cpp @@ -34,12 +34,16 @@ using namespace arm_compute; namespace { -template <typename T> +template <typename T, bool is_nhwc> void weights_reshape(const ITensor *input, const ITensor *bias, ITensor *output, const Window &window) { - const unsigned int kernel_size_x = input->info()->dimension(0); - const unsigned int kernel_size_y = input->info()->dimension(1); - const unsigned int kernel_depth = input->info()->dimension(2); + DataLayout data_layout = input->info()->data_layout(); + const int idx_width = get_data_layout_dimension_index(data_layout, DataLayoutDimension::WIDTH); + const int idx_height = get_data_layout_dimension_index(data_layout, DataLayoutDimension::HEIGHT); + const int idx_channel = get_data_layout_dimension_index(data_layout, DataLayoutDimension::CHANNEL); + const unsigned int kernel_size_x = input->info()->dimension(idx_width); + const unsigned int kernel_size_y = input->info()->dimension(idx_height); + const unsigned int kernel_depth = input->info()->dimension(idx_channel); const unsigned int input_stride_x = input->info()->strides_in_bytes().x(); const unsigned int input_stride_y = input->info()->strides_in_bytes().y(); const unsigned int input_stride_z = input->info()->strides_in_bytes().z(); @@ -67,13 +71,13 @@ void weights_reshape(const ITensor *input, const ITensor *bias, ITensor *output, for(unsigned int i = 0; i < kernel_size_x; ++i) { *(reinterpret_cast<T *>(tmp_output_ptr)) = *(reinterpret_cast<const T *>(tmp_input_ptr)); - tmp_input_ptr += input_stride_x; + tmp_input_ptr += is_nhwc ? input_stride_y : input_stride_x; tmp_output_ptr += output_stride_y; } - curr_input_row_ptr += input_stride_y; + curr_input_row_ptr += is_nhwc ? input_stride_z : input_stride_y; tmp_input_ptr = curr_input_row_ptr; } - curr_input_depth_ptr += input_stride_z; + curr_input_depth_ptr += is_nhwc ? input_stride_x : input_stride_z; curr_input_row_ptr = curr_input_depth_ptr; tmp_input_ptr = curr_input_depth_ptr; } @@ -161,21 +165,24 @@ void NEWeightsReshapeKernel::configure(const ITensor *input, const ITensor *bias _bias = bias; _output = output; + const DataLayout data_layout = input->info()->data_layout(); + const bool is_nhwc = data_layout == DataLayout::NHWC; + switch(_input->info()->element_size()) { case 4: { - _func = &weights_reshape<uint32_t>; + _func = is_nhwc ? &weights_reshape<uint32_t, true> : &weights_reshape<uint32_t, false>; break; } case 2: { - _func = &weights_reshape<uint16_t>; + _func = is_nhwc ? &weights_reshape<uint16_t, true> : &weights_reshape<uint16_t, false>; break; } case 1: { - _func = &weights_reshape<uint8_t>; + _func = is_nhwc ? &weights_reshape<uint8_t, true> : &weights_reshape<uint8_t, false>; break; } default: |