1 /*
2 * Copyright (c) 2019-2020 Arm Limited.
3 *
4 * SPDX-License-Identifier: MIT
5 *
6 * Permission is hereby granted, free of charge, to any person obtaining a copy
7 * of this software and associated documentation files (the "Software"), to
8 * deal in the Software without restriction, including without limitation the
9 * rights to use, copy, modify, merge, publish, distribute, sublicense, and/or
10 * sell copies of the Software, and to permit persons to whom the Software is
11 * furnished to do so, subject to the following conditions:
12 *
13 * The above copyright notice and this permission notice shall be included in all
14 * copies or substantial portions of the Software.
15 *
16 * THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
17 * IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
18 * FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
19 * AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
20 * LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
21 * OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
22 * SOFTWARE.
23 */
24 #include "src/core/NEON/kernels/NEGatherKernel.h"
25
26 #include "arm_compute/core/Coordinates.h"
27 #include "arm_compute/core/Error.h"
28 #include "arm_compute/core/Helpers.h"
29 #include "arm_compute/core/IAccessWindow.h"
30 #include "arm_compute/core/TensorInfo.h"
31 #include "arm_compute/core/Validate.h"
32 #include "arm_compute/core/Window.h"
33 #include "arm_compute/core/utils/misc/ShapeCalculator.h"
34 #include "src/core/CPP/Validate.h"
35 #include "src/core/helpers/AutoConfiguration.h"
36 #include "src/core/helpers/WindowHelpers.h"
37
38 namespace arm_compute
39 {
40 namespace
41 {
42 /** Validate the indices
43 *
44 * Validate that indices are not negative
45 *
46 * @param[in] indices Indices tensor info.
47 */
48 template <typename U>
validate_indices(const ITensor * indices)49 void validate_indices(const ITensor *indices)
50 {
51 for(size_t i = 0; i < indices->info()->tensor_shape()[0]; ++i)
52 {
53 ARM_COMPUTE_ERROR_ON(*(reinterpret_cast<U *>(indices->ptr_to_element(Coordinates(i)))) < 0);
54 }
55 }
56
validate_arguments(const ITensorInfo * input,const ITensorInfo * indices,const ITensorInfo * output,int axis)57 Status validate_arguments(const ITensorInfo *input, const ITensorInfo *indices, const ITensorInfo *output, int axis)
58 {
59 ARM_COMPUTE_RETURN_ERROR_ON_NULLPTR(input, indices, output);
60 ARM_COMPUTE_RETURN_ERROR_ON(indices->num_dimensions() > 1);
61 ARM_COMPUTE_RETURN_ERROR_ON(input->num_dimensions() > 4);
62
63 if(axis < 0)
64 {
65 axis += input->num_dimensions();
66 }
67
68 ARM_COMPUTE_RETURN_ERROR_ON(0 > axis || axis >= static_cast<int32_t>(input->num_dimensions()));
69 ARM_COMPUTE_RETURN_ERROR_ON(input->data_type() == DataType::UNKNOWN);
70
71 if(output->total_size() != 0)
72 {
73 ARM_COMPUTE_RETURN_ERROR_ON_MISMATCHING_DATA_TYPES(input, output);
74 ARM_COMPUTE_RETURN_ERROR_ON_MISMATCHING_QUANTIZATION_INFO(input, output);
75 TensorShape output_shape = arm_compute::misc::shape_calculator::compute_gather_shape(input->tensor_shape(), indices->tensor_shape(), axis);
76 ARM_COMPUTE_RETURN_ERROR_ON(output_shape.total_size() != output->tensor_shape().total_size());
77 }
78
79 ARM_COMPUTE_RETURN_ERROR_ON_DATA_TYPE_CHANNEL_NOT_IN(indices, 1, DataType::U32, DataType::S32);
80
81 return Status{};
82 }
83 } // namespace
84
NEGatherKernel()85 NEGatherKernel::NEGatherKernel()
86 : _input{}, _indices{}, _axis{}, _output{}, _func{}
87 {
88 }
89
90 template <typename U>
gather_0_axis(const Window & window,const ThreadInfo & info)91 inline void NEGatherKernel::gather_0_axis(const Window &window, const ThreadInfo &info)
92 {
93 ARM_COMPUTE_UNUSED(info);
94
95 // Validate that the indices are not negative
96 validate_indices<U>(_indices);
97
98 Iterator output_it(_output, window);
99 execute_window_loop(window, [&](const Coordinates & id)
100 {
101 Coordinates gather_id(id);
102
103 auto new_index = *(reinterpret_cast<U *>(_indices->ptr_to_element(Coordinates(id[0]))));
104 gather_id.set(0, new_index);
105
106 std::copy_n(_input->ptr_to_element(gather_id), _output->info()->element_size(), output_it.ptr());
107 },
108 output_it);
109 }
110
111 template <typename U>
gather_n_axis(const Window & window,const ThreadInfo & info)112 void NEGatherKernel::gather_n_axis(const Window &window, const ThreadInfo &info)
113 {
114 ARM_COMPUTE_UNUSED(info);
115
116 // Validate that the indices are not negative
117 validate_indices<U>(_indices);
118
119 Window output_window{ window };
120 output_window.set(Window::DimX, Window::Dimension(0, 1, 1));
121
122 Iterator output_it(_output, output_window);
123 execute_window_loop(output_window, [&](const Coordinates & id)
124 {
125 Coordinates gather_id(id);
126
127 auto new_index = *(reinterpret_cast<U *>(_indices->ptr_to_element(Coordinates(id[_axis]))));
128 gather_id.set(_axis, new_index);
129
130 std::copy_n(_input->ptr_to_element(gather_id), _input->info()->dimension(0) * _output->info()->element_size(), output_it.ptr());
131 },
132 output_it);
133 }
134
configure(const ITensor * input,const ITensor * indices,ITensor * output,int axis)135 void NEGatherKernel::configure(const ITensor *input, const ITensor *indices, ITensor *output, int axis)
136 {
137 ARM_COMPUTE_ERROR_ON_NULLPTR(input, output, indices);
138 ARM_COMPUTE_ERROR_THROW_ON(validate_arguments(input->info(), indices->info(), output->info(), axis));
139
140 _input = input;
141 _indices = indices;
142 _output = output;
143 _axis = axis;
144
145 if(_axis < 0)
146 {
147 _axis += input->info()->num_dimensions();
148 }
149 ARM_COMPUTE_ERROR_ON(0 > _axis || _axis >= static_cast<int32_t>(input->info()->num_dimensions()));
150
151 if(0 == _axis)
152 {
153 switch(_indices->info()->data_type())
154 {
155 case DataType::U32:
156 _func = &NEGatherKernel::gather_0_axis<uint32_t>;
157 break;
158 case DataType::S32:
159 _func = &NEGatherKernel::gather_0_axis<int32_t>;
160 break;
161 default:
162 ARM_COMPUTE_ERROR("Not supported");
163 break;
164 }
165 }
166 else
167 {
168 switch(_indices->info()->data_type())
169 {
170 case DataType::U32:
171 _func = &NEGatherKernel::gather_n_axis<uint32_t>;
172 break;
173 case DataType::S32:
174 _func = &NEGatherKernel::gather_n_axis<int32_t>;
175 break;
176 default:
177 ARM_COMPUTE_ERROR("Not supported");
178 break;
179 }
180 }
181 // Output auto initialization if not yet initialized
182 TensorShape output_shape = arm_compute::misc::shape_calculator::compute_gather_shape(input->info()->tensor_shape(), indices->info()->tensor_shape(), _axis);
183 auto_init_if_empty(*output->info(), input->info()->clone()->set_tensor_shape(output_shape));
184
185 // Create window
186 Window win = calculate_max_window(*output->info(), Steps());
187 output->info()->set_valid_region(ValidRegion(Coordinates(), output->info()->tensor_shape()));
188
189 INEKernel::configure(win);
190 }
191
validate(const ITensorInfo * input,const ITensorInfo * indices,const ITensorInfo * output,int axis)192 Status NEGatherKernel::validate(const ITensorInfo *input, const ITensorInfo *indices, const ITensorInfo *output, int axis)
193 {
194 ARM_COMPUTE_RETURN_ON_ERROR(validate_arguments(input, indices, output, axis));
195 return Status{};
196 }
197
run(const Window & window,const ThreadInfo & info)198 void NEGatherKernel::run(const Window &window, const ThreadInfo &info)
199 {
200 ARM_COMPUTE_UNUSED(info);
201 ARM_COMPUTE_ERROR_ON_UNCONFIGURED_KERNEL(this);
202 ARM_COMPUTE_ERROR_ON(_func == nullptr);
203
204 (this->*_func)(window, info);
205 }
206
207 } // namespace arm_compute
208