xref: /aosp_15_r20/external/eigen/unsupported/test/cxx11_tensor_reduction.cpp (revision bf2c37156dfe67e5dfebd6d394bad8b2ab5804d4)
1 // This file is part of Eigen, a lightweight C++ template library
2 // for linear algebra.
3 //
4 // Copyright (C) 2014 Benoit Steiner <[email protected]>
5 //
6 // This Source Code Form is subject to the terms of the Mozilla
7 // Public License v. 2.0. If a copy of the MPL was not distributed
8 // with this file, You can obtain one at http://mozilla.org/MPL/2.0/.
9 
10 #include "main.h"
11 #include <limits>
12 #include <numeric>
13 #include <Eigen/CXX11/Tensor>
14 
15 using Eigen::Tensor;
16 
17 template <int DataLayout>
test_trivial_reductions()18 static void test_trivial_reductions() {
19   {
20     Tensor<float, 0, DataLayout> tensor;
21     tensor.setRandom();
22     array<ptrdiff_t, 0> reduction_axis;
23 
24     Tensor<float, 0, DataLayout> result = tensor.sum(reduction_axis);
25     VERIFY_IS_EQUAL(result(), tensor());
26   }
27 
28   {
29     Tensor<float, 1, DataLayout> tensor(7);
30     tensor.setRandom();
31     array<ptrdiff_t, 0> reduction_axis;
32 
33     Tensor<float, 1, DataLayout> result = tensor.sum(reduction_axis);
34     VERIFY_IS_EQUAL(result.dimension(0), 7);
35     for (int i = 0; i < 7; ++i) {
36       VERIFY_IS_EQUAL(result(i), tensor(i));
37     }
38   }
39 
40   {
41     Tensor<float, 2, DataLayout> tensor(2, 3);
42     tensor.setRandom();
43     array<ptrdiff_t, 0> reduction_axis;
44 
45     Tensor<float, 2, DataLayout> result = tensor.sum(reduction_axis);
46     VERIFY_IS_EQUAL(result.dimension(0), 2);
47     VERIFY_IS_EQUAL(result.dimension(1), 3);
48     for (int i = 0; i < 2; ++i) {
49       for (int j = 0; j < 3; ++j) {
50         VERIFY_IS_EQUAL(result(i, j), tensor(i, j));
51       }
52     }
53   }
54 }
55 
56 template <typename Scalar,int DataLayout>
test_simple_reductions()57 static void test_simple_reductions() {
58   Tensor<Scalar, 4, DataLayout> tensor(2, 3, 5, 7);
59   tensor.setRandom();
60   // Add a little offset so that the product reductions won't be close to zero.
61   tensor += tensor.constant(Scalar(0.5f));
62   array<ptrdiff_t, 2> reduction_axis2;
63   reduction_axis2[0] = 1;
64   reduction_axis2[1] = 3;
65 
66   Tensor<Scalar, 2, DataLayout> result = tensor.sum(reduction_axis2);
67   VERIFY_IS_EQUAL(result.dimension(0), 2);
68   VERIFY_IS_EQUAL(result.dimension(1), 5);
69   for (int i = 0; i < 2; ++i) {
70     for (int j = 0; j < 5; ++j) {
71       Scalar sum = Scalar(0.0f);
72       for (int k = 0; k < 3; ++k) {
73         for (int l = 0; l < 7; ++l) {
74           sum += tensor(i, k, j, l);
75         }
76       }
77       VERIFY_IS_APPROX(result(i, j), sum);
78     }
79   }
80 
81   {
82     Tensor<Scalar, 0, DataLayout> sum1 = tensor.sum();
83     VERIFY_IS_EQUAL(sum1.rank(), 0);
84 
85     array<ptrdiff_t, 4> reduction_axis4;
86     reduction_axis4[0] = 0;
87     reduction_axis4[1] = 1;
88     reduction_axis4[2] = 2;
89     reduction_axis4[3] = 3;
90     Tensor<Scalar, 0, DataLayout> sum2 = tensor.sum(reduction_axis4);
91     VERIFY_IS_EQUAL(sum2.rank(), 0);
92 
93     VERIFY_IS_APPROX(sum1(), sum2());
94   }
95 
96   reduction_axis2[0] = 0;
97   reduction_axis2[1] = 2;
98   result = tensor.prod(reduction_axis2);
99   VERIFY_IS_EQUAL(result.dimension(0), 3);
100   VERIFY_IS_EQUAL(result.dimension(1), 7);
101   for (int i = 0; i < 3; ++i) {
102     for (int j = 0; j < 7; ++j) {
103       Scalar prod = Scalar(1.0f);
104       for (int k = 0; k < 2; ++k) {
105         for (int l = 0; l < 5; ++l) {
106           prod *= tensor(k, i, l, j);
107         }
108       }
109       VERIFY_IS_APPROX(result(i, j), prod);
110     }
111   }
112 
113   {
114     Tensor<Scalar, 0, DataLayout> prod1 = tensor.prod();
115     VERIFY_IS_EQUAL(prod1.rank(), 0);
116 
117     array<ptrdiff_t, 4> reduction_axis4;
118     reduction_axis4[0] = 0;
119     reduction_axis4[1] = 1;
120     reduction_axis4[2] = 2;
121     reduction_axis4[3] = 3;
122     Tensor<Scalar, 0, DataLayout> prod2 = tensor.prod(reduction_axis4);
123     VERIFY_IS_EQUAL(prod2.rank(), 0);
124 
125     VERIFY_IS_APPROX(prod1(), prod2());
126   }
127 
128   reduction_axis2[0] = 0;
129   reduction_axis2[1] = 2;
130   result = tensor.maximum(reduction_axis2);
131   VERIFY_IS_EQUAL(result.dimension(0), 3);
132   VERIFY_IS_EQUAL(result.dimension(1), 7);
133   for (int i = 0; i < 3; ++i) {
134     for (int j = 0; j < 7; ++j) {
135       Scalar max_val = std::numeric_limits<Scalar>::lowest();
136       for (int k = 0; k < 2; ++k) {
137         for (int l = 0; l < 5; ++l) {
138           max_val = (std::max)(max_val, tensor(k, i, l, j));
139         }
140       }
141       VERIFY_IS_APPROX(result(i, j), max_val);
142     }
143   }
144 
145   {
146     Tensor<Scalar, 0, DataLayout> max1 = tensor.maximum();
147     VERIFY_IS_EQUAL(max1.rank(), 0);
148 
149     array<ptrdiff_t, 4> reduction_axis4;
150     reduction_axis4[0] = 0;
151     reduction_axis4[1] = 1;
152     reduction_axis4[2] = 2;
153     reduction_axis4[3] = 3;
154     Tensor<Scalar, 0, DataLayout> max2 = tensor.maximum(reduction_axis4);
155     VERIFY_IS_EQUAL(max2.rank(), 0);
156 
157     VERIFY_IS_APPROX(max1(), max2());
158   }
159 
160   reduction_axis2[0] = 0;
161   reduction_axis2[1] = 1;
162   result = tensor.minimum(reduction_axis2);
163   VERIFY_IS_EQUAL(result.dimension(0), 5);
164   VERIFY_IS_EQUAL(result.dimension(1), 7);
165   for (int i = 0; i < 5; ++i) {
166     for (int j = 0; j < 7; ++j) {
167       Scalar min_val = (std::numeric_limits<Scalar>::max)();
168       for (int k = 0; k < 2; ++k) {
169         for (int l = 0; l < 3; ++l) {
170           min_val = (std::min)(min_val, tensor(k, l, i, j));
171         }
172       }
173       VERIFY_IS_APPROX(result(i, j), min_val);
174     }
175   }
176 
177   {
178     Tensor<Scalar, 0, DataLayout> min1 = tensor.minimum();
179     VERIFY_IS_EQUAL(min1.rank(), 0);
180 
181     array<ptrdiff_t, 4> reduction_axis4;
182     reduction_axis4[0] = 0;
183     reduction_axis4[1] = 1;
184     reduction_axis4[2] = 2;
185     reduction_axis4[3] = 3;
186     Tensor<Scalar, 0, DataLayout> min2 = tensor.minimum(reduction_axis4);
187     VERIFY_IS_EQUAL(min2.rank(), 0);
188 
189     VERIFY_IS_APPROX(min1(), min2());
190   }
191 
192   reduction_axis2[0] = 0;
193   reduction_axis2[1] = 1;
194   result = tensor.mean(reduction_axis2);
195   VERIFY_IS_EQUAL(result.dimension(0), 5);
196   VERIFY_IS_EQUAL(result.dimension(1), 7);
197   for (int i = 0; i < 5; ++i) {
198     for (int j = 0; j < 7; ++j) {
199       Scalar sum = Scalar(0.0f);
200       int count = 0;
201       for (int k = 0; k < 2; ++k) {
202         for (int l = 0; l < 3; ++l) {
203           sum += tensor(k, l, i, j);
204           ++count;
205         }
206       }
207       VERIFY_IS_APPROX(result(i, j), sum / Scalar(count));
208     }
209   }
210 
211   {
212     Tensor<Scalar, 0, DataLayout> mean1 = tensor.mean();
213     VERIFY_IS_EQUAL(mean1.rank(), 0);
214 
215     array<ptrdiff_t, 4> reduction_axis4;
216     reduction_axis4[0] = 0;
217     reduction_axis4[1] = 1;
218     reduction_axis4[2] = 2;
219     reduction_axis4[3] = 3;
220     Tensor<Scalar, 0, DataLayout> mean2 = tensor.mean(reduction_axis4);
221     VERIFY_IS_EQUAL(mean2.rank(), 0);
222 
223     VERIFY_IS_APPROX(mean1(), mean2());
224   }
225 
226   {
227     Tensor<int, 1> ints(10);
228     std::iota(ints.data(), ints.data() + ints.dimension(0), 0);
229 
230     TensorFixedSize<bool, Sizes<> > all_;
231     all_ = ints.all();
232     VERIFY(!all_());
233     all_ = (ints >= ints.constant(0)).all();
234     VERIFY(all_());
235 
236     TensorFixedSize<bool, Sizes<> > any;
237     any = (ints > ints.constant(10)).any();
238     VERIFY(!any());
239     any = (ints < ints.constant(1)).any();
240     VERIFY(any());
241   }
242 }
243 
244 
245 template <int DataLayout>
test_reductions_in_expr()246 static void test_reductions_in_expr() {
247   Tensor<float, 4, DataLayout> tensor(2, 3, 5, 7);
248   tensor.setRandom();
249   array<ptrdiff_t, 2> reduction_axis2;
250   reduction_axis2[0] = 1;
251   reduction_axis2[1] = 3;
252 
253   Tensor<float, 2, DataLayout> result(2, 5);
254   result = result.constant(1.0f) - tensor.sum(reduction_axis2);
255   VERIFY_IS_EQUAL(result.dimension(0), 2);
256   VERIFY_IS_EQUAL(result.dimension(1), 5);
257   for (int i = 0; i < 2; ++i) {
258     for (int j = 0; j < 5; ++j) {
259       float sum = 0.0f;
260       for (int k = 0; k < 3; ++k) {
261         for (int l = 0; l < 7; ++l) {
262           sum += tensor(i, k, j, l);
263         }
264       }
265       VERIFY_IS_APPROX(result(i, j), 1.0f - sum);
266     }
267   }
268 }
269 
270 
271 template <int DataLayout>
test_full_reductions()272 static void test_full_reductions() {
273   Tensor<float, 2, DataLayout> tensor(2, 3);
274   tensor.setRandom();
275   array<ptrdiff_t, 2> reduction_axis;
276   reduction_axis[0] = 0;
277   reduction_axis[1] = 1;
278 
279   Tensor<float, 0, DataLayout> result = tensor.sum(reduction_axis);
280   VERIFY_IS_EQUAL(result.rank(), 0);
281 
282   float sum = 0.0f;
283   for (int i = 0; i < 2; ++i) {
284     for (int j = 0; j < 3; ++j) {
285       sum += tensor(i, j);
286     }
287   }
288   VERIFY_IS_APPROX(result(0), sum);
289 
290   result = tensor.square().sum(reduction_axis).sqrt();
291   VERIFY_IS_EQUAL(result.rank(), 0);
292 
293   sum = 0.0f;
294   for (int i = 0; i < 2; ++i) {
295     for (int j = 0; j < 3; ++j) {
296       sum += tensor(i, j) * tensor(i, j);
297     }
298   }
299   VERIFY_IS_APPROX(result(), sqrtf(sum));
300 }
301 
302 struct UserReducer {
303   static const bool PacketAccess = false;
UserReducerUserReducer304   UserReducer(float offset) : offset_(offset) {}
reduceUserReducer305   void reduce(const float val, float* accum) { *accum += val * val; }
initializeUserReducer306   float initialize() const { return 0; }
finalizeUserReducer307   float finalize(const float accum) const { return 1.0f / (accum + offset_); }
308 
309  private:
310   const float offset_;
311 };
312 
313 template <int DataLayout>
test_user_defined_reductions()314 static void test_user_defined_reductions() {
315   Tensor<float, 2, DataLayout> tensor(5, 7);
316   tensor.setRandom();
317   array<ptrdiff_t, 1> reduction_axis;
318   reduction_axis[0] = 1;
319 
320   UserReducer reducer(10.0f);
321   Tensor<float, 1, DataLayout> result = tensor.reduce(reduction_axis, reducer);
322   VERIFY_IS_EQUAL(result.dimension(0), 5);
323   for (int i = 0; i < 5; ++i) {
324     float expected = 10.0f;
325     for (int j = 0; j < 7; ++j) {
326       expected += tensor(i, j) * tensor(i, j);
327     }
328     expected = 1.0f / expected;
329     VERIFY_IS_APPROX(result(i), expected);
330   }
331 }
332 
333 template <int DataLayout>
test_tensor_maps()334 static void test_tensor_maps() {
335   int inputs[2 * 3 * 5 * 7];
336   TensorMap<Tensor<int, 4, DataLayout> > tensor_map(inputs, 2, 3, 5, 7);
337   TensorMap<Tensor<const int, 4, DataLayout> > tensor_map_const(inputs, 2, 3, 5,
338                                                                 7);
339   const TensorMap<Tensor<const int, 4, DataLayout> > tensor_map_const_const(
340       inputs, 2, 3, 5, 7);
341 
342   tensor_map.setRandom();
343   array<ptrdiff_t, 2> reduction_axis;
344   reduction_axis[0] = 1;
345   reduction_axis[1] = 3;
346 
347   Tensor<int, 2, DataLayout> result = tensor_map.sum(reduction_axis);
348   Tensor<int, 2, DataLayout> result2 = tensor_map_const.sum(reduction_axis);
349   Tensor<int, 2, DataLayout> result3 =
350       tensor_map_const_const.sum(reduction_axis);
351 
352   for (int i = 0; i < 2; ++i) {
353     for (int j = 0; j < 5; ++j) {
354       int sum = 0;
355       for (int k = 0; k < 3; ++k) {
356         for (int l = 0; l < 7; ++l) {
357           sum += tensor_map(i, k, j, l);
358         }
359       }
360       VERIFY_IS_EQUAL(result(i, j), sum);
361       VERIFY_IS_EQUAL(result2(i, j), sum);
362       VERIFY_IS_EQUAL(result3(i, j), sum);
363     }
364   }
365 }
366 
367 template <int DataLayout>
test_static_dims()368 static void test_static_dims() {
369   Tensor<float, 4, DataLayout> in(72, 53, 97, 113);
370   Tensor<float, 2, DataLayout> out(72, 97);
371   in.setRandom();
372 
373 #if !EIGEN_HAS_CONSTEXPR
374   array<int, 2> reduction_axis;
375   reduction_axis[0] = 1;
376   reduction_axis[1] = 3;
377 #else
378   Eigen::IndexList<Eigen::type2index<1>, Eigen::type2index<3> > reduction_axis;
379 #endif
380 
381   out = in.maximum(reduction_axis);
382 
383   for (int i = 0; i < 72; ++i) {
384     for (int j = 0; j < 97; ++j) {
385       float expected = -1e10f;
386       for (int k = 0; k < 53; ++k) {
387         for (int l = 0; l < 113; ++l) {
388           expected = (std::max)(expected, in(i, k, j, l));
389         }
390       }
391       VERIFY_IS_EQUAL(out(i, j), expected);
392     }
393   }
394 }
395 
396 template <int DataLayout>
test_innermost_last_dims()397 static void test_innermost_last_dims() {
398   Tensor<float, 4, DataLayout> in(72, 53, 97, 113);
399   Tensor<float, 2, DataLayout> out(97, 113);
400   in.setRandom();
401 
402 // Reduce on the innermost dimensions.
403 #if !EIGEN_HAS_CONSTEXPR
404   array<int, 2> reduction_axis;
405   reduction_axis[0] = 0;
406   reduction_axis[1] = 1;
407 #else
408   // This triggers the use of packets for ColMajor.
409   Eigen::IndexList<Eigen::type2index<0>, Eigen::type2index<1> > reduction_axis;
410 #endif
411 
412   out = in.maximum(reduction_axis);
413 
414   for (int i = 0; i < 97; ++i) {
415     for (int j = 0; j < 113; ++j) {
416       float expected = -1e10f;
417       for (int k = 0; k < 53; ++k) {
418         for (int l = 0; l < 72; ++l) {
419           expected = (std::max)(expected, in(l, k, i, j));
420         }
421       }
422       VERIFY_IS_EQUAL(out(i, j), expected);
423     }
424   }
425 }
426 
427 template <int DataLayout>
test_innermost_first_dims()428 static void test_innermost_first_dims() {
429   Tensor<float, 4, DataLayout> in(72, 53, 97, 113);
430   Tensor<float, 2, DataLayout> out(72, 53);
431   in.setRandom();
432 
433 // Reduce on the innermost dimensions.
434 #if !EIGEN_HAS_CONSTEXPR
435   array<int, 2> reduction_axis;
436   reduction_axis[0] = 2;
437   reduction_axis[1] = 3;
438 #else
439   // This triggers the use of packets for RowMajor.
440   Eigen::IndexList<Eigen::type2index<2>, Eigen::type2index<3>> reduction_axis;
441 #endif
442 
443   out = in.maximum(reduction_axis);
444 
445   for (int i = 0; i < 72; ++i) {
446     for (int j = 0; j < 53; ++j) {
447       float expected = -1e10f;
448       for (int k = 0; k < 97; ++k) {
449         for (int l = 0; l < 113; ++l) {
450           expected = (std::max)(expected, in(i, j, k, l));
451         }
452       }
453       VERIFY_IS_EQUAL(out(i, j), expected);
454     }
455   }
456 }
457 
458 template <int DataLayout>
test_reduce_middle_dims()459 static void test_reduce_middle_dims() {
460   Tensor<float, 4, DataLayout> in(72, 53, 97, 113);
461   Tensor<float, 2, DataLayout> out(72, 53);
462   in.setRandom();
463 
464 // Reduce on the innermost dimensions.
465 #if !EIGEN_HAS_CONSTEXPR
466   array<int, 2> reduction_axis;
467   reduction_axis[0] = 1;
468   reduction_axis[1] = 2;
469 #else
470   // This triggers the use of packets for RowMajor.
471   Eigen::IndexList<Eigen::type2index<1>, Eigen::type2index<2>> reduction_axis;
472 #endif
473 
474   out = in.maximum(reduction_axis);
475 
476   for (int i = 0; i < 72; ++i) {
477     for (int j = 0; j < 113; ++j) {
478       float expected = -1e10f;
479       for (int k = 0; k < 53; ++k) {
480         for (int l = 0; l < 97; ++l) {
481           expected = (std::max)(expected, in(i, k, l, j));
482         }
483       }
484       VERIFY_IS_EQUAL(out(i, j), expected);
485     }
486   }
487 }
488 
test_sum_accuracy()489 static void test_sum_accuracy() {
490   Tensor<float, 3> tensor(101, 101, 101);
491   for (float prescribed_mean : {1.0f, 10.0f, 100.0f, 1000.0f, 10000.0f}) {
492     tensor.setRandom();
493     tensor += tensor.constant(prescribed_mean);
494 
495     Tensor<float, 0> sum = tensor.sum();
496     double expected_sum = 0.0;
497     for (int i = 0; i < 101; ++i) {
498       for (int j = 0; j < 101; ++j) {
499         for (int k = 0; k < 101; ++k) {
500           expected_sum += static_cast<double>(tensor(i, j, k));
501         }
502       }
503     }
504     VERIFY_IS_APPROX(sum(), static_cast<float>(expected_sum));
505   }
506 }
507 
EIGEN_DECLARE_TEST(cxx11_tensor_reduction)508 EIGEN_DECLARE_TEST(cxx11_tensor_reduction) {
509   CALL_SUBTEST(test_trivial_reductions<ColMajor>());
510   CALL_SUBTEST(test_trivial_reductions<RowMajor>());
511   CALL_SUBTEST(( test_simple_reductions<float,ColMajor>() ));
512   CALL_SUBTEST(( test_simple_reductions<float,RowMajor>() ));
513   CALL_SUBTEST(( test_simple_reductions<Eigen::half,ColMajor>() ));
514   CALL_SUBTEST(( test_simple_reductions<Eigen::bfloat16,ColMajor>() ));
515   CALL_SUBTEST(test_reductions_in_expr<ColMajor>());
516   CALL_SUBTEST(test_reductions_in_expr<RowMajor>());
517   CALL_SUBTEST(test_full_reductions<ColMajor>());
518   CALL_SUBTEST(test_full_reductions<RowMajor>());
519   CALL_SUBTEST(test_user_defined_reductions<ColMajor>());
520   CALL_SUBTEST(test_user_defined_reductions<RowMajor>());
521   CALL_SUBTEST(test_tensor_maps<ColMajor>());
522   CALL_SUBTEST(test_tensor_maps<RowMajor>());
523   CALL_SUBTEST(test_static_dims<ColMajor>());
524   CALL_SUBTEST(test_static_dims<RowMajor>());
525   CALL_SUBTEST(test_innermost_last_dims<ColMajor>());
526   CALL_SUBTEST(test_innermost_last_dims<RowMajor>());
527   CALL_SUBTEST(test_innermost_first_dims<ColMajor>());
528   CALL_SUBTEST(test_innermost_first_dims<RowMajor>());
529   CALL_SUBTEST(test_reduce_middle_dims<ColMajor>());
530   CALL_SUBTEST(test_reduce_middle_dims<RowMajor>());
531   CALL_SUBTEST(test_sum_accuracy());
532 }
533