xref: /aosp_15_r20/external/ComputeLibrary/src/core/NEON/kernels/NEBatchToSpaceLayerKernel.cpp (revision c217d954acce2dbc11938adb493fc0abd69584f3)
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/NEBatchToSpaceLayerKernel.h"
25 
26 #include "arm_compute/core/Helpers.h"
27 #include "arm_compute/core/ITensor.h"
28 #include "arm_compute/core/Types.h"
29 #include "arm_compute/core/Validate.h"
30 #include "arm_compute/core/utils/misc/ShapeCalculator.h"
31 #include "src/core/helpers/AutoConfiguration.h"
32 #include "src/core/helpers/WindowHelpers.h"
33 
34 using namespace arm_compute::misc::shape_calculator;
35 
36 namespace arm_compute
37 {
38 namespace
39 {
validate_arguments(const ITensorInfo * input,const ITensorInfo * block_info,const ITensorInfo * output)40 Status validate_arguments(const ITensorInfo *input, const ITensorInfo *block_info, const ITensorInfo *output)
41 {
42     ARM_COMPUTE_RETURN_ERROR_ON_NULLPTR(input, block_info, output);
43     ARM_COMPUTE_RETURN_ERROR_ON_DATA_TYPE_CHANNEL_NOT_IN(block_info, 1, DataType::S32);
44     ARM_COMPUTE_RETURN_ERROR_ON(input->num_dimensions() > 4);
45     ARM_COMPUTE_RETURN_ERROR_ON(input->data_type() == DataType::UNKNOWN);
46 
47     // Validate output if initialized
48     if(output->total_size() != 0)
49     {
50         ARM_COMPUTE_RETURN_ERROR_ON(output->num_dimensions() > 4);
51         ARM_COMPUTE_RETURN_ERROR_ON_MISMATCHING_DATA_TYPES(input, output);
52     }
53 
54     return Status{};
55 }
validate_arguments_static(const ITensorInfo * input,const int block_shape_x,const int block_shape_y,const ITensorInfo * output)56 Status validate_arguments_static(const ITensorInfo *input, const int block_shape_x, const int block_shape_y, const ITensorInfo *output)
57 {
58     ARM_COMPUTE_RETURN_ERROR_ON_NULLPTR(input, output);
59     ARM_COMPUTE_RETURN_ERROR_ON(input->num_dimensions() > 4);
60     ARM_COMPUTE_RETURN_ERROR_ON(block_shape_x <= 0);
61     ARM_COMPUTE_RETURN_ERROR_ON(block_shape_y <= 0);
62 
63     const DataLayout data_layout = input->data_layout();
64     const int        idx_batch   = get_data_layout_dimension_index(data_layout, DataLayoutDimension::BATCHES);
65     ARM_COMPUTE_RETURN_ERROR_ON(input->tensor_shape()[idx_batch] % (block_shape_x * block_shape_y) != 0);
66     // Validate output if initialized
67     if(output->total_size() != 0)
68     {
69         const int idx_width   = get_data_layout_dimension_index(data_layout, DataLayoutDimension::WIDTH);
70         const int idx_height  = get_data_layout_dimension_index(data_layout, DataLayoutDimension::HEIGHT);
71         const int idx_channel = get_data_layout_dimension_index(data_layout, DataLayoutDimension::CHANNEL);
72         ARM_COMPUTE_RETURN_ERROR_ON(output->tensor_shape()[idx_width] != (block_shape_x * input->tensor_shape()[idx_width]));
73         ARM_COMPUTE_RETURN_ERROR_ON(output->tensor_shape()[idx_height] != (block_shape_y * input->tensor_shape()[idx_height]));
74         ARM_COMPUTE_RETURN_ERROR_ON(output->tensor_shape()[idx_channel] != input->tensor_shape()[idx_channel]);
75         ARM_COMPUTE_RETURN_ERROR_ON(output->num_dimensions() > 4);
76         ARM_COMPUTE_RETURN_ERROR_ON_MISMATCHING_DATA_TYPES(input, output);
77     }
78 
79     return Status{};
80 }
81 } // namespace
82 
NEBatchToSpaceLayerKernel()83 NEBatchToSpaceLayerKernel::NEBatchToSpaceLayerKernel()
84     : _input(nullptr), _block_shape(nullptr), _output(nullptr), _data_layout(DataLayout::UNKNOWN), _block_shape_x(), _block_shape_y()
85 {
86 }
87 
configure(const ITensor * input,const ITensor * block_shape,ITensor * output)88 void NEBatchToSpaceLayerKernel::configure(const ITensor *input, const ITensor *block_shape, ITensor *output)
89 {
90     ARM_COMPUTE_ERROR_ON_NULLPTR(input, output);
91     ARM_COMPUTE_ERROR_THROW_ON(validate_arguments(input->info(), block_shape->info(), output->info()));
92 
93     _input       = input;
94     _block_shape = block_shape;
95     _output      = output;
96     _data_layout = input->info()->data_layout();
97 
98     // Configure kernel window
99     Window win = calculate_max_window(*input->info(), Steps());
100     ICPPKernel::configure(win);
101 }
102 
configure(const ITensor * input,const int32_t block_shape_x,const int32_t block_shape_y,ITensor * output)103 void NEBatchToSpaceLayerKernel::configure(const ITensor *input, const int32_t block_shape_x, const int32_t block_shape_y, ITensor *output)
104 {
105     ARM_COMPUTE_ERROR_ON_NULLPTR(input, output);
106     TensorShape output_shape = compute_batch_to_space_shape(input->info(), block_shape_x, block_shape_y);
107     // Output auto inizialitation if not yet initialized
108     auto_init_if_empty(*output->info(), input->info()->clone()->set_tensor_shape(output_shape));
109 
110     // Perform validation step
111     ARM_COMPUTE_ERROR_THROW_ON(validate_arguments_static(input->info(), block_shape_x, block_shape_y, output->info()));
112 
113     _input         = input;
114     _output        = output;
115     _block_shape_x = block_shape_x;
116     _block_shape_y = block_shape_y;
117     _data_layout   = input->info()->data_layout();
118 
119     // Configure kernel window
120     Window win = calculate_max_window(*input->info(), Steps());
121     ICPPKernel::configure(win);
122 }
123 
validate(const ITensorInfo * input,const ITensorInfo * block_shape,const ITensorInfo * output)124 Status NEBatchToSpaceLayerKernel::validate(const ITensorInfo *input, const ITensorInfo *block_shape, const ITensorInfo *output)
125 {
126     ARM_COMPUTE_RETURN_ERROR_ON_NULLPTR(input, block_shape, output);
127     ARM_COMPUTE_RETURN_ON_ERROR(validate_arguments(input, block_shape, output));
128     return Status{};
129 }
130 
validate(const ITensorInfo * input,const int32_t block_shape_x,const int32_t block_shape_y,const ITensorInfo * output)131 Status NEBatchToSpaceLayerKernel::validate(const ITensorInfo *input, const int32_t block_shape_x, const int32_t block_shape_y, const ITensorInfo *output)
132 {
133     ARM_COMPUTE_RETURN_ERROR_ON_NULLPTR(input, output);
134     ARM_COMPUTE_RETURN_ON_ERROR(validate_arguments_static(input, block_shape_x, block_shape_y, output));
135     return Status{};
136 }
137 
run(const Window & window,const ThreadInfo & info)138 void NEBatchToSpaceLayerKernel::run(const Window &window, const ThreadInfo &info)
139 {
140     ARM_COMPUTE_UNUSED(info);
141     ARM_COMPUTE_ERROR_ON_UNCONFIGURED_KERNEL(this);
142     ARM_COMPUTE_ERROR_ON_INVALID_SUBWINDOW(ICPPKernel::window(), window);
143 
144     if(_block_shape != nullptr)
145     {
146         // Retrieve the block shapes dynamically
147         _block_shape_x = *(reinterpret_cast<const int *>(_block_shape->ptr_to_element(0)));
148         _block_shape_y = *(reinterpret_cast<const int *>(_block_shape->ptr_to_element(1)));
149     }
150 
151     const int batch_size   = _input->info()->dimension(3);
152     const int r            = (batch_size / (_block_shape_x * _block_shape_y));
153     const int element_size = _input->info()->element_size();
154 
155     Window slice_in  = window.first_slice_window_3D();
156     Window slice_out = window.first_slice_window_4D();
157 
158     // The slice_out slice does not move
159     slice_out.set(Window::DimX, Window::Dimension(0, 0, 0));
160     slice_out.set(Window::DimY, Window::Dimension(0, 0, 0));
161     slice_out.set(Window::DimZ, Window::Dimension(0, 0, 0));
162     slice_out.set(3, Window::Dimension(0, 0, 0));
163 
164     int batch_id = 0;
165     // Main loop for NCHW and NHWC
166     if(_data_layout == DataLayout::NCHW)
167     {
168         do
169         {
170             Iterator in(_input, slice_in);
171             execute_window_loop(slice_in, [&](const Coordinates & id)
172             {
173 
174                 const int x = id.x();
175                 const int y = id.y();
176                 const int z = id.z();
177 
178                 const int   w     = batch_id % r;
179                 const int   out_x = x * _block_shape_x + (batch_id / r) % _block_shape_x;
180                 const int   out_y = y * _block_shape_y + (batch_id / r) / _block_shape_x;
181                 Coordinates output_coords{ out_x, out_y, z, w };
182                 memcpy(_output->ptr_to_element(output_coords), in.ptr(), element_size);
183             },
184             in);
185             ++batch_id;
186         }
187         while(window.slide_window_slice_3D(slice_in));
188     }
189     else
190     {
191         do
192         {
193             Iterator in(_input, slice_in);
194             execute_window_loop(slice_in, [&](const Coordinates & id)
195             {
196 
197                 const int z = id.x();
198                 const int x = id.y();
199                 const int y = id.z();
200 
201                 const int   w     = batch_id % r;
202                 const int   out_x = x * _block_shape_x + (batch_id / r) % _block_shape_x;
203                 const int   out_y = y * _block_shape_y + (batch_id / r) / _block_shape_x;
204                 Coordinates output_coords{ z, out_x, out_y, w };
205                 memcpy(_output->ptr_to_element(output_coords), in.ptr(), element_size);
206             },
207             in);
208             ++batch_id;
209         }
210         while(window.slide_window_slice_3D(slice_in));
211     }
212 }
213 } // namespace arm_compute
214