From 894659a98e76d84bf209da27d8ecb6d9ed05b13d Mon Sep 17 00:00:00 2001 From: Pablo Marquez Tello Date: Fri, 13 May 2022 12:20:16 +0100 Subject: Add support for 2d and 3d indices for axis 1 * Resolves COMPMID-5055 Change-Id: I2d14de29d3ec913d20c971bc8bbc9ad71e2d998f Signed-off-by: Pablo Marquez Tello Reviewed-on: https://review.mlplatform.org/c/ml/ComputeLibrary/+/7547 Reviewed-by: SiCong Li Reviewed-by: Michalis Spyrou Comments-Addressed: Arm Jenkins Tested-by: Arm Jenkins Benchmark: Arm Jenkins --- src/core/NEON/kernels/NEGatherKernel.cpp | 114 ++++++++++++++++++++++++------- 1 file changed, 88 insertions(+), 26 deletions(-) (limited to 'src/core/NEON/kernels/NEGatherKernel.cpp') diff --git a/src/core/NEON/kernels/NEGatherKernel.cpp b/src/core/NEON/kernels/NEGatherKernel.cpp index 55c4525dca..085ab7cb18 100644 --- a/src/core/NEON/kernels/NEGatherKernel.cpp +++ b/src/core/NEON/kernels/NEGatherKernel.cpp @@ -44,19 +44,23 @@ namespace * * @param[in] indices Indices tensor info. */ + template void validate_indices(const ITensor *indices) { - for(size_t i = 0; i < indices->info()->tensor_shape()[0]; ++i) + Window window; + window.use_tensor_dimensions(indices->info()->tensor_shape()); + execute_window_loop(window, [&](const Coordinates & id) { - ARM_COMPUTE_ERROR_ON(*(reinterpret_cast(indices->ptr_to_element(Coordinates(i)))) < 0); - } + const auto i = *(reinterpret_cast(indices->ptr_to_element(id))); + ARM_COMPUTE_UNUSED(i); + ARM_COMPUTE_ERROR_ON(i < 0); + }); } Status validate_arguments(const ITensorInfo *input, const ITensorInfo *indices, const ITensorInfo *output, int axis) { ARM_COMPUTE_RETURN_ERROR_ON_NULLPTR(input, indices, output); - ARM_COMPUTE_RETURN_ERROR_ON(indices->num_dimensions() > 1); ARM_COMPUTE_RETURN_ERROR_ON(input->num_dimensions() > 4); if(axis < 0) @@ -65,6 +69,7 @@ Status validate_arguments(const ITensorInfo *input, const ITensorInfo *indices, } ARM_COMPUTE_RETURN_ERROR_ON(0 > axis || axis >= static_cast(input->num_dimensions())); + ARM_COMPUTE_RETURN_ERROR_ON(axis != 1 && indices->num_dimensions() > 1); ARM_COMPUTE_RETURN_ERROR_ON(input->data_type() == DataType::UNKNOWN); if(output->total_size() != 0) @@ -86,6 +91,37 @@ NEGatherKernel::NEGatherKernel() { } +template +inline void NEGatherKernel::gather_multiindices_1_axis(const Window &window, const ThreadInfo &info) +{ + ARM_COMPUTE_UNUSED(info); + ARM_COMPUTE_ERROR_ON(_indices->info()->num_dimensions() < 2 || _indices->info()->num_dimensions() > 3); + validate_indices(_indices); + Window win = window; + win.set(Window::DimX, Window::Dimension(0, 1, 1)); + execute_window_loop(win, [&](const Coordinates & id) + { + auto *dst_ptr = _output->ptr_to_element(id); + Coordinates index_offset; + for(uint32_t k = 0; k < _indices->info()->num_dimensions(); ++k) + { + index_offset.set(k, id[k + 1]); + } + const uint32_t row = *(reinterpret_cast(_indices->ptr_to_element(index_offset))); + Coordinates src_offset; + // Set up input coords to read the row specified by the current index + src_offset.set(0, 0); + src_offset.set(1, row); + for(uint32_t j = 2; j < _input->info()->num_dimensions(); ++j) + { + src_offset.set(j, id[1 + _indices->info()->num_dimensions() + (j - 2)]); + } + const auto in_ptr_row = _input->ptr_to_element(src_offset); + // Copy a row from input to output + memcpy(dst_ptr, in_ptr_row, _input->info()->tensor_shape()[0] * _input->info()->element_size()); + }); +} + template inline void NEGatherKernel::gather_0_axis(const Window &window, const ThreadInfo &info) { @@ -147,38 +183,64 @@ void NEGatherKernel::configure(const ITensor *input, const ITensor *indices, ITe } ARM_COMPUTE_ERROR_ON(0 > _axis || _axis >= static_cast(input->info()->num_dimensions())); - if(0 == _axis) + if(indices->info()->num_dimensions() == 1u) { - switch(_indices->info()->data_type()) + if(_axis == 0) { - case DataType::U32: - _func = &NEGatherKernel::gather_0_axis; - break; - case DataType::S32: - _func = &NEGatherKernel::gather_0_axis; - break; - default: - ARM_COMPUTE_ERROR("Not supported"); - break; + switch(_indices->info()->data_type()) + { + case DataType::U32: + _func = &NEGatherKernel::gather_0_axis; + break; + case DataType::S32: + _func = &NEGatherKernel::gather_0_axis; + break; + default: + ARM_COMPUTE_ERROR("Not supported"); + break; + } + } + else + { + switch(_indices->info()->data_type()) + { + case DataType::U32: + _func = &NEGatherKernel::gather_n_axis; + break; + case DataType::S32: + _func = &NEGatherKernel::gather_n_axis; + break; + default: + ARM_COMPUTE_ERROR("Not supported"); + break; + } } } else { - switch(_indices->info()->data_type()) + if(_axis == 1) + { + switch(_indices->info()->data_type()) + { + case DataType::U32: + _func = &NEGatherKernel::gather_multiindices_1_axis; + break; + case DataType::S32: + _func = &NEGatherKernel::gather_multiindices_1_axis; + break; + default: + ARM_COMPUTE_ERROR("Not supported"); + break; + } + } + else { - case DataType::U32: - _func = &NEGatherKernel::gather_n_axis; - break; - case DataType::S32: - _func = &NEGatherKernel::gather_n_axis; - break; - default: - ARM_COMPUTE_ERROR("Not supported"); - break; + ARM_COMPUTE_ERROR("Not supported"); } } + // Output auto initialization if not yet initialized - TensorShape output_shape = arm_compute::misc::shape_calculator::compute_gather_shape(input->info()->tensor_shape(), indices->info()->tensor_shape(), _axis); + const TensorShape output_shape = arm_compute::misc::shape_calculator::compute_gather_shape(input->info()->tensor_shape(), indices->info()->tensor_shape(), _axis); auto_init_if_empty(*output->info(), input->info()->clone()->set_tensor_shape(output_shape)); // Create window -- cgit v1.2.1