diff options
Diffstat (limited to 'src/core/NEON/kernels/NEReorderKernel.cpp')
| -rw-r--r-- | src/core/NEON/kernels/NEReorderKernel.cpp | 223 |
1 files changed, 223 insertions, 0 deletions
diff --git a/src/core/NEON/kernels/NEReorderKernel.cpp b/src/core/NEON/kernels/NEReorderKernel.cpp new file mode 100644 index 0000000000..c09745604f --- /dev/null +++ b/src/core/NEON/kernels/NEReorderKernel.cpp @@ -0,0 +1,223 @@ +/* + * Copyright (c) 2023 Arm Limited. + * + * SPDX-License-Identifier: MIT + * + * Permission is hereby granted, free of charge, to any person obtaining a copy + * of this software and associated documentation files (the "Software"), to + * deal in the Software without restriction, including without limitation the + * rights to use, copy, modify, merge, publish, distribute, sublicense, and/or + * sell copies of the Software, and to permit persons to whom the Software is + * furnished to do so, subject to the following conditions: + * + * The above copyright notice and this permission notice shall be included in all + * copies or substantial portions of the Software. + * + * THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR + * IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, + * FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE + * AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER + * LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, + * OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE + * SOFTWARE. + */ + +#include "src/core/NEON/kernels/NEReorderKernel.h" +#include "src/common/utils/Log.h" +#include "src/core/NEON/kernels/arm_gemm/transform.hpp" +#include "arm_compute/core/Helpers.h" +#include "arm_compute/core/Validate.h" + +namespace arm_compute +{ + +void NEReorderKernel::run(const Window &window, const ThreadInfo &info) +{ + ARM_COMPUTE_UNUSED(info); + ARM_COMPUTE_ERROR_ON_UNCONFIGURED_KERNEL(this); + ARM_COMPUTE_ERROR_ON_INVALID_SUBWINDOW(INEKernel::window(), window); + switch(_input->info()->data_type()) + { + case DataType::F32: + { + const int ksize_rows_elements = _xmax * _ksize; + const int jump_rows = ksize_rows_elements * window.x().start(); + const int k_start = window.x().start() * _ksize; + const int k_end = std::min(window.x().end() * _ksize, _kmax); + const int stride = _kmax; + if(k_start < k_end) + { + + switch(_output_wf) + { + case WeightFormat::OHWIo4: + { + arm_gemm::Transform<4, 1, true, arm_gemm::VLType::None>(reinterpret_cast<float *>(_output->buffer()) + jump_rows, reinterpret_cast<float *>(_input->buffer()), stride, k_start, k_end, 0, _xmax); + break; + } +#if defined(ARM_COMPUTE_ENABLE_SVE) + case WeightFormat::OHWIo8: + { + arm_gemm::Transform<1, 1, true, arm_gemm::VLType::SVE>(reinterpret_cast<float *>(_output->buffer()) + jump_rows, reinterpret_cast<float *>(_input->buffer()), stride, k_start, k_end, 0, _xmax); + break; + } +#endif /* ARM_COMPUTE_ENABLE_SVE */ + default: + { + ARM_COMPUTE_ERROR("Unsupported data type!"); + break; + } + } + } + break; + } + default: + ARM_COMPUTE_ERROR("Unsupported data type!"); + } +} + +NEReorderKernel::NEReorderKernel() + : _input(nullptr), _output(nullptr), _ksize(0), _kmax(0), _xmax(0), _input_wf(WeightFormat::ANY), _output_wf(WeightFormat::ANY) +{ +} + +void NEReorderKernel::configure(const ITensor *input, ITensor *output, arm_compute::WeightFormat input_wf, arm_compute::WeightFormat output_wf) +{ + ARM_COMPUTE_LOG_PARAMS(input, output, input_wf, output_wf); + ARM_COMPUTE_ERROR_ON_NULLPTR(input, output); + ARM_COMPUTE_ERROR_THROW_ON(validate(input->info(), output->info(), input_wf, output_wf)); + + // Set variables + _input = input; + _output = output; + _input_wf = input_wf; + _output_wf = output_wf; + + // Setting parameters for transform + auto dims = input->info()->num_dimensions(); + switch(dims) + { + case 2: + { + _xmax = input->info()->dimension(0); // Number of columns in input matrix + _kmax = input->info()->dimension(1); // Number of rows in input matrix + break; + } + case 4: + { + _xmax = input->info()->dimension(2); // Number of columns in input matrix + _kmax = input->info()->dimension(3); // Number of rows in input matrix + break; + } + default: + { + ARM_COMPUTE_ERROR("Only 2 or 4 dimensions supported."); + } + } + + // Configure kernel window + // Window size is set by rows / _ksize + Window win; + int window_size = 0; + switch(_output_wf) + { +#if defined(ARM_COMPUTE_ENABLE_SVE) + case WeightFormat::OHWIo8: + { + _ksize = 8; + window_size = _kmax / _ksize; + break; + } +#endif /* ARM_COMPUTE_ENABLE_SVE */ + case WeightFormat::OHWIo4: + { + _ksize = 4; + window_size = _kmax / _ksize; + break; + } + default: + { + ARM_COMPUTE_ERROR("Unsupported weight format."); + break; + } + } + if(_kmax % _ksize != 0) + { + window_size += 1; + } + + win.set(Window::DimX, Window::Dimension(0, window_size, 1)); + + INEKernel::configure(win); +} + +Status NEReorderKernel::validate(const ITensorInfo *input, const ITensorInfo *output, arm_compute::WeightFormat input_wf, arm_compute::WeightFormat output_wf) +{ + ARM_COMPUTE_RETURN_ERROR_ON_NULLPTR(input, output); + ARM_COMPUTE_RETURN_ERROR_ON(input->data_type() == DataType::UNKNOWN); + if(output->tensor_shape().total_size() != 0) + { + ARM_COMPUTE_RETURN_ERROR_ON_MISMATCHING_DATA_TYPES(input, output); + ARM_COMPUTE_RETURN_ERROR_ON_MISMATCHING_QUANTIZATION_INFO(input, output); + // Only input WeightFormat OHWI supported + ARM_COMPUTE_RETURN_ERROR_ON(input_wf != arm_compute::WeightFormat::OHWI); + int input_x_dim; + int input_k_dim; + int output_x_dim; + int output_k_dim; + auto dims = output->num_dimensions(); + switch(dims) + { + case 2: + { + input_x_dim = input->dimension(0); // Number of columns in input matrix + input_k_dim = input->dimension(1); // Number of rows in input matrix + output_x_dim = output->dimension(0); // Number of columns in output matrix + output_k_dim = output->dimension(1); // Number of rows in output matrix + break; + } + case 4: + { + input_x_dim = input->dimension(2); // Number of columns in input matrix + input_k_dim = input->dimension(3); // Number of rows in input matrix + output_x_dim = output->dimension(2); // Number of columns in output matrix + output_k_dim = output->dimension(3); // Number of rows in output matrix + break; + } + default: + { + ARM_COMPUTE_RETURN_ERROR_MSG("Only 2 or 4 dimensions supported."); + } + } + + int ksize; + switch(output_wf) + { + case WeightFormat::OHWIo8: + { + ksize = 8; + break; + } + case WeightFormat::OHWIo4: + { + ksize = 4; + break; + } + default: + { + ARM_COMPUTE_RETURN_ERROR_MSG("Unsupported weight format."); + break; + } + } + + // output k_dim needs to be same as input but multiple of ksize + int32_t rnd_up_input_kdim = arm_compute::ceil_to_multiple<int32_t, int32_t>(input_k_dim, ksize); + ARM_COMPUTE_RETURN_ERROR_ON(rnd_up_input_kdim != output_k_dim); + // output x_dim needs to be same as input + ARM_COMPUTE_RETURN_ERROR_ON(input_x_dim != output_x_dim); + + } + return Status{}; +} + +} // namespace arm_compute |