xref: /aosp_15_r20/external/ComputeLibrary/src/cpu/kernels/elementwise_binary/generic/sve/impl.cpp (revision c217d954acce2dbc11938adb493fc0abd69584f3) 
1 /*
2  * Copyright (c) 2021-2022 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 
25 #include "src/cpu/kernels/elementwise_binary/generic/sve/impl.h"
26 #include "src/core/NEON/SVEMath.h"
27 #include <arm_sve.h>
28 
29 namespace arm_compute
30 {
31 namespace cpu
32 {
33 using namespace arm_compute::wrapper;
34 
35 template <typename ScalarType>
elementwise_arithmetic_op(const ITensor * in1,const ITensor * in2,ITensor * out,ArithmeticOperation op,const Window & window)36 void elementwise_arithmetic_op(const ITensor *in1, const ITensor *in2, ITensor *out, ArithmeticOperation op, const Window &window)
37 {
38     using VectorType = typename sve_vector<ScalarType>::type;
39 
40     const auto all_true_pg = svptrue<ScalarType>();
41 
42     // Create input windows
43     Window input1_win = window.broadcast_if_dimension_le_one(in1->info()->tensor_shape());
44     Window input2_win = window.broadcast_if_dimension_le_one(in2->info()->tensor_shape());
45 
46     // Clear X Dimension on execution window as we handle manually
47     Window win = window;
48     win.set(Window::DimX, Window::Dimension(0, 1, 1));
49 
50     const auto window_start_x        = static_cast<int>(window.x().start());
51     const auto window_end_x          = static_cast<int>(window.x().end());
52     const bool is_broadcast_across_x = in1->info()->tensor_shape().x() != in2->info()->tensor_shape().x();
53 
54     if(is_broadcast_across_x)
55     {
56         const bool     is_broadcast_input_2 = input2_win.x().step() == 0;
57         Window         broadcast_win        = is_broadcast_input_2 ? input2_win : input1_win;
58         Window         non_broadcast_win    = !is_broadcast_input_2 ? input2_win : input1_win;
59         const ITensor *broadcast_tensor     = is_broadcast_input_2 ? in2 : in1;
60         const ITensor *non_broadcast_tensor = !is_broadcast_input_2 ? in2 : in1;
61 
62         // Clear X Dimension on execution window as we handle manually
63         non_broadcast_win.set(Window::DimX, Window::Dimension(0, 1, 1));
64 
65         Iterator broadcast_input(broadcast_tensor, broadcast_win);
66         Iterator non_broadcast_input(non_broadcast_tensor, non_broadcast_win);
67         Iterator output(out, win);
68 
69         execute_window_loop(win, [&](const Coordinates &)
70         {
71             auto             output_ptr              = reinterpret_cast<ScalarType *>(output.ptr());
72             const auto       non_broadcast_input_ptr = reinterpret_cast<const ScalarType *>(non_broadcast_input.ptr());
73             const ScalarType broadcast_value         = *reinterpret_cast<const ScalarType *>(broadcast_input.ptr());
74             const auto       broadcast_vector        = svdup_n(broadcast_value);
75 
76             int x = window_start_x;
77 
78             svbool_t pg = svwhilelt<ScalarType>(x, window_end_x);
79             do
80             {
81                 const auto non_broadcast_vector = svld1(pg, non_broadcast_input_ptr + x);
82                 VectorType res{};
83 
84                 if(is_broadcast_input_2)
85                 {
86                     res = elementwise_arithmetic_op<typename sve_vector<ScalarType>::type>(pg, non_broadcast_vector, broadcast_vector, op);
87                 }
88                 else
89                 {
90                     res = elementwise_arithmetic_op<typename sve_vector<ScalarType>::type>(pg, broadcast_vector, non_broadcast_vector, op);
91                 }
92                 svst1(pg, output_ptr + x, res);
93 
94                 x += svcnt<ScalarType>();
95                 pg = svwhilelt<ScalarType>(x, window_end_x);
96             }
97             while(svptest_any(all_true_pg, pg));
98         },
99         broadcast_input, non_broadcast_input, output);
100     }
101     else
102     {
103         // Clear X Dimension on execution window as we handle manually
104         input1_win.set(Window::DimX, Window::Dimension(0, 1, 1));
105         input2_win.set(Window::DimX, Window::Dimension(0, 1, 1));
106 
107         Iterator input1(in1, input1_win);
108         Iterator input2(in2, input2_win);
109         Iterator output(out, win);
110 
111         execute_window_loop(win, [&](const Coordinates &)
112         {
113             auto       output_ptr = reinterpret_cast<ScalarType *>(output.ptr());
114             const auto input1_ptr = reinterpret_cast<const ScalarType *>(input1.ptr());
115             const auto input2_ptr = reinterpret_cast<const ScalarType *>(input2.ptr());
116 
117             int x = window_start_x;
118 
119             svbool_t pg = svwhilelt<ScalarType>(x, window_end_x);
120             do
121             {
122                 const auto in1 = svld1(pg, input1_ptr + x);
123                 const auto in2 = svld1(pg, input2_ptr + x);
124                 const auto res = elementwise_arithmetic_op<typename sve_vector<ScalarType>::type>(pg, in1, in2, op);
125                 svst1(pg, output_ptr + x, res);
126 
127                 x += svcnt<ScalarType>();
128                 pg = svwhilelt<ScalarType>(x, window_end_x);
129             }
130             while(svptest_any(all_true_pg, pg));
131         },
132         input1, input2, output);
133     }
134 }
135 template void elementwise_arithmetic_op<float32_t>(const ITensor *in1, const ITensor *in2, ITensor *out, const ArithmeticOperation op, const Window &window);
136 template void elementwise_arithmetic_op<float16_t>(const ITensor *in1, const ITensor *in2, ITensor *out, const ArithmeticOperation op, const Window &window);
137 template void elementwise_arithmetic_op<int16_t>(const ITensor *in1, const ITensor *in2, ITensor *out, const ArithmeticOperation op, const Window &window);
138 template void elementwise_arithmetic_op<int32_t>(const ITensor *in1, const ITensor *in2, ITensor *out, const ArithmeticOperation op, const Window &window);
139 
140 template <typename InputScalarType, typename OutputScalarType>
elementwise_comparison_op(const ITensor * in1,const ITensor * in2,ITensor * out,ComparisonOperation op,const Window & window)141 void elementwise_comparison_op(const ITensor *in1, const ITensor *in2, ITensor *out, ComparisonOperation op, const Window &window)
142 {
143     static_assert(sizeof(InputScalarType) >= sizeof(OutputScalarType), "input data type's width should be equal to or greater than output data type's width");
144 
145     using OutputVectorType = typename sve_vector<OutputScalarType>::type;
146     const auto all_true_pg = svptrue<InputScalarType>();
147 
148     // Create input windows
149     Window input1_win = window.broadcast_if_dimension_le_one(in1->info()->tensor_shape());
150     Window input2_win = window.broadcast_if_dimension_le_one(in2->info()->tensor_shape());
151 
152     // Clear X Dimension on execution window as we handle manually
153     Window win = window;
154     win.set(Window::DimX, Window::Dimension(0, 1, 1));
155 
156     const auto window_start_x        = static_cast<int>(window.x().start());
157     const auto window_end_x          = static_cast<int>(window.x().end());
158     const bool is_broadcast_across_x = in1->info()->tensor_shape().x() != in2->info()->tensor_shape().x();
159 
160     if(is_broadcast_across_x)
161     {
162         const bool     is_broadcast_input_2 = input2_win.x().step() == 0;
163         Window         broadcast_win        = is_broadcast_input_2 ? input2_win : input1_win;
164         Window         non_broadcast_win    = !is_broadcast_input_2 ? input2_win : input1_win;
165         const ITensor *broadcast_tensor     = is_broadcast_input_2 ? in2 : in1;
166         const ITensor *non_broadcast_tensor = !is_broadcast_input_2 ? in2 : in1;
167 
168         // Clear X Dimension on execution window as we handle manually
169         non_broadcast_win.set(Window::DimX, Window::Dimension(0, 1, 1));
170 
171         Iterator broadcast_input(broadcast_tensor, broadcast_win);
172         Iterator non_broadcast_input(non_broadcast_tensor, non_broadcast_win);
173         Iterator output(out, win);
174 
175         execute_window_loop(win, [&](const Coordinates &)
176         {
177             auto                  output_ptr              = reinterpret_cast<OutputScalarType *>(output.ptr());
178             const auto            non_broadcast_input_ptr = reinterpret_cast<const InputScalarType *>(non_broadcast_input.ptr());
179             const InputScalarType broadcast_value         = *reinterpret_cast<const InputScalarType *>(broadcast_input.ptr());
180             const auto            broadcast_vector        = svdup_n(broadcast_value);
181 
182             int x = window_start_x;
183 
184             svbool_t pg = svwhilelt<InputScalarType>(x, window_end_x);
185             do
186             {
187                 const auto       non_broadcast_vector = svld1(pg, non_broadcast_input_ptr + x);
188                 const svbool_t   output_pg            = narrow_to_byte_predicate<sizeof(InputScalarType)>(pg);
189                 OutputVectorType res{};
190                 if(is_broadcast_input_2)
191                 {
192                     res = elementwise_comparison_op<typename sve_vector<InputScalarType>::type, typename sve_vector<OutputScalarType>::type>(pg, non_broadcast_vector, broadcast_vector, op);
193                 }
194                 else
195                 {
196                     res = elementwise_comparison_op<typename sve_vector<InputScalarType>::type, typename sve_vector<OutputScalarType>::type>(pg, broadcast_vector, non_broadcast_vector, op);
197                 }
198                 svst1(output_pg, output_ptr + x, res);
199 
200                 x += svcnt<InputScalarType>();
201                 pg = svwhilelt<InputScalarType>(x, window_end_x);
202             }
203             while(svptest_any(all_true_pg, pg));
204         },
205         broadcast_input, non_broadcast_input, output);
206     }
207     else
208     {
209         // Clear X Dimension on execution window as we handle manually
210         input1_win.set(Window::DimX, Window::Dimension(0, 1, 1));
211         input2_win.set(Window::DimX, Window::Dimension(0, 1, 1));
212 
213         Iterator input1(in1, input1_win);
214         Iterator input2(in2, input2_win);
215         Iterator output(out, win);
216 
217         execute_window_loop(win, [&](const Coordinates &)
218         {
219             auto       output_ptr = reinterpret_cast<OutputScalarType *>(output.ptr());
220             const auto input1_ptr = reinterpret_cast<const InputScalarType *>(input1.ptr());
221             const auto input2_ptr = reinterpret_cast<const InputScalarType *>(input2.ptr());
222 
223             int x = window_start_x;
224 
225             svbool_t pg = svwhilelt<InputScalarType>(x, window_end_x);
226             do
227             {
228                 const auto     in1       = svld1(pg, input1_ptr + x);
229                 const auto     in2       = svld1(pg, input2_ptr + x);
230                 const auto     res       = elementwise_comparison_op<typename sve_vector<InputScalarType>::type, typename sve_vector<OutputScalarType>::type>(pg, in1, in2, op);
231                 const svbool_t output_pg = narrow_to_byte_predicate<sizeof(InputScalarType)>(pg);
232                 svst1(output_pg, output_ptr + x, res);
233 
234                 x += svcnt<InputScalarType>();
235                 pg = svwhilelt<InputScalarType>(x, window_end_x);
236             }
237             while(svptest_any(all_true_pg, pg));
238         },
239         input1, input2, output);
240     }
241 }
242 
243 template void elementwise_comparison_op<float32_t>(const ITensor *in1, const ITensor *in2, ITensor *out, const ComparisonOperation op, const Window &window);
244 template void elementwise_comparison_op<float16_t>(const ITensor *in1, const ITensor *in2, ITensor *out, const ComparisonOperation op, const Window &window);
245 template void elementwise_comparison_op<uint8_t>(const ITensor *in1, const ITensor *in2, ITensor *out, const ComparisonOperation op, const Window &window);
246 template void elementwise_comparison_op<int16_t>(const ITensor *in1, const ITensor *in2, ITensor *out, const ComparisonOperation op, const Window &window);
247 template void elementwise_comparison_op<int32_t>(const ITensor *in1, const ITensor *in2, ITensor *out, const ComparisonOperation op, const Window &window);
248 
249 template <>
elementwise_pow(svbool_t & pg,const svint32_t & a,const svint32_t & b)250 svint32_t elementwise_pow<svint32_t>(svbool_t &pg, const svint32_t &a, const svint32_t &b)
251 {
252     return svcvt_s32_z(pg, svpow_z(pg, svcvt_f32_z(pg, a), svcvt_f32_z(pg, b)));
253 }
254 
255 template <>
elementwise_div(svbool_t & pg,const svint32_t & a,const svint32_t & b)256 svint32_t elementwise_div<svint32_t>(svbool_t &pg, const svint32_t &a, const svint32_t &b)
257 {
258     return svcvt_s32_z(pg, svdiv_z(pg, svcvt_f32_z(pg, a), svcvt_f32_z(pg, b)));
259 }
260 
261 template <>
elementwise_div(svbool_t & pg,const svint16_t & a,const svint16_t & b)262 svint16_t elementwise_div<svint16_t>(svbool_t &pg, const svint16_t &a, const svint16_t &b)
263 {
264     ARM_COMPUTE_UNUSED(pg, a, b);
265     ARM_COMPUTE_ERROR("Not supported");
266 }
267 
268 } // namespace cpu
269 } // namespace arm_compute
270