xref: /aosp_15_r20/external/eigen/unsupported/Eigen/CXX11/src/Tensor/TensorFunctors.h (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 #ifndef EIGEN_CXX11_TENSOR_TENSOR_FUNCTORS_H
11 #define EIGEN_CXX11_TENSOR_TENSOR_FUNCTORS_H
12 
13 namespace Eigen {
14 namespace internal {
15 
16 
17 /** \internal
18  * \brief Template functor to compute the modulo between an array and a scalar.
19  */
20 template <typename Scalar>
21 struct scalar_mod_op {
scalar_mod_opscalar_mod_op22   EIGEN_DEVICE_FUNC scalar_mod_op(const Scalar& divisor) : m_divisor(divisor) {}
operatorscalar_mod_op23   EIGEN_DEVICE_FUNC EIGEN_STRONG_INLINE Scalar operator() (const Scalar& a) const { return a % m_divisor; }
24   const Scalar m_divisor;
25 };
26 template <typename Scalar>
27 struct functor_traits<scalar_mod_op<Scalar> >
28 { enum { Cost = scalar_div_cost<Scalar,false>::value, PacketAccess = false }; };
29 
30 
31 /** \internal
32  * \brief Template functor to compute the modulo between 2 arrays.
33  */
34 template <typename Scalar>
35 struct scalar_mod2_op {
36   EIGEN_EMPTY_STRUCT_CTOR(scalar_mod2_op)
37   EIGEN_DEVICE_FUNC EIGEN_STRONG_INLINE Scalar operator() (const Scalar& a, const Scalar& b) const { return a % b; }
38 };
39 template <typename Scalar>
40 struct functor_traits<scalar_mod2_op<Scalar> >
41 { enum { Cost = scalar_div_cost<Scalar,false>::value, PacketAccess = false }; };
42 
43 template <typename Scalar>
44 struct scalar_fmod_op {
45   EIGEN_EMPTY_STRUCT_CTOR(scalar_fmod_op)
46   EIGEN_DEVICE_FUNC EIGEN_STRONG_INLINE Scalar
47   operator()(const Scalar& a, const Scalar& b) const {
48     return numext::fmod(a, b);
49   }
50 };
51 template <typename Scalar>
52 struct functor_traits<scalar_fmod_op<Scalar> > {
53   enum { Cost = 13,  // Reciprocal throughput of FPREM on Haswell.
54          PacketAccess = false };
55 };
56 
57 template<typename Reducer, typename Device>
58 struct reducer_traits {
59   enum {
60     Cost = 1,
61     PacketAccess = false,
62     IsStateful = false,
63     IsExactlyAssociative = true
64   };
65 };
66 
67 // Standard reduction functors
68 template <typename T> struct SumReducer
69 {
70   EIGEN_DEVICE_FUNC EIGEN_STRONG_INLINE void reduce(const T t, T* accum) const {
71     internal::scalar_sum_op<T> sum_op;
72     *accum = sum_op(*accum, t);
73   }
74   template <typename Packet>
75   EIGEN_DEVICE_FUNC EIGEN_STRONG_INLINE void reducePacket(const Packet& p, Packet* accum) const {
76     (*accum) = padd<Packet>(*accum, p);
77   }
78 
79   EIGEN_DEVICE_FUNC EIGEN_STRONG_INLINE T initialize() const {
80     internal::scalar_cast_op<int, T> conv;
81     return conv(0);
82   }
83   template <typename Packet>
84   EIGEN_DEVICE_FUNC EIGEN_STRONG_INLINE Packet initializePacket() const {
85     return pset1<Packet>(initialize());
86   }
87   EIGEN_DEVICE_FUNC EIGEN_STRONG_INLINE T finalize(const T accum) const {
88     return accum;
89   }
90   template <typename Packet>
91   EIGEN_DEVICE_FUNC EIGEN_STRONG_INLINE Packet finalizePacket(const Packet& vaccum) const {
92     return vaccum;
93   }
94   template <typename Packet>
95   EIGEN_DEVICE_FUNC EIGEN_STRONG_INLINE T finalizeBoth(const T saccum, const Packet& vaccum) const {
96     internal::scalar_sum_op<T> sum_op;
97     return sum_op(saccum, predux(vaccum));
98   }
99 };
100 
101 template <typename T, typename Device>
102 struct reducer_traits<SumReducer<T>, Device> {
103   enum {
104     Cost = NumTraits<T>::AddCost,
105     PacketAccess = PacketType<T, Device>::HasAdd,
106     IsStateful = false,
107     IsExactlyAssociative = NumTraits<T>::IsInteger
108   };
109 };
110 
111 template <typename T> struct MeanReducer
112 {
113   EIGEN_DEVICE_FUNC EIGEN_STRONG_INLINE
114   MeanReducer() : scalarCount_(0), packetCount_(0) { }
115 
116   EIGEN_DEVICE_FUNC EIGEN_STRONG_INLINE void reduce(const T t, T* accum) {
117     internal::scalar_sum_op<T> sum_op;
118     *accum = sum_op(*accum, t);
119     scalarCount_++;
120   }
121   template <typename Packet>
122   EIGEN_DEVICE_FUNC EIGEN_STRONG_INLINE void reducePacket(const Packet& p, Packet* accum) {
123     (*accum) = padd<Packet>(*accum, p);
124     packetCount_++;
125   }
126 
127   EIGEN_DEVICE_FUNC EIGEN_STRONG_INLINE T initialize() const {
128     internal::scalar_cast_op<int, T> conv;
129     return conv(0);
130   }
131   template <typename Packet>
132   EIGEN_DEVICE_FUNC EIGEN_STRONG_INLINE Packet initializePacket() const {
133     return pset1<Packet>(initialize());
134   }
135   EIGEN_DEVICE_FUNC EIGEN_STRONG_INLINE T finalize(const T accum) const {
136     internal::scalar_quotient_op<T> quotient_op;
137     return quotient_op(accum, T(scalarCount_));
138   }
139   template <typename Packet>
140   EIGEN_DEVICE_FUNC EIGEN_STRONG_INLINE Packet finalizePacket(const Packet& vaccum) const {
141     return pdiv(vaccum, pset1<Packet>(T(packetCount_)));
142   }
143   template <typename Packet>
144   EIGEN_DEVICE_FUNC EIGEN_STRONG_INLINE T finalizeBoth(const T saccum, const Packet& vaccum) const {
145     internal::scalar_sum_op<T> sum_op;
146     internal::scalar_quotient_op<T> quotient_op;
147     return quotient_op(
148         sum_op(saccum, predux(vaccum)),
149         T(scalarCount_ + packetCount_ * unpacket_traits<Packet>::size));
150   }
151 
152   protected:
153     DenseIndex scalarCount_;
154     DenseIndex packetCount_;
155 };
156 
157 template <typename T, typename Device>
158 struct reducer_traits<MeanReducer<T>, Device> {
159   enum {
160     Cost = NumTraits<T>::AddCost,
161     PacketAccess = PacketType<T, Device>::HasAdd &&
162                    PacketType<T, Device>::HasDiv && !NumTraits<T>::IsInteger,
163     IsStateful = true,
164     IsExactlyAssociative = NumTraits<T>::IsInteger
165   };
166 };
167 
168 
169 template <typename T, bool IsMax = true, bool IsInteger = true>
170 struct MinMaxBottomValue {
171   EIGEN_DEVICE_FUNC static EIGEN_STRONG_INLINE T bottom_value() {
172     return Eigen::NumTraits<T>::lowest();
173   }
174 };
175 template <typename T>
176 struct MinMaxBottomValue<T, true, false> {
177   EIGEN_DEVICE_FUNC static EIGEN_STRONG_INLINE T bottom_value() {
178     return -Eigen::NumTraits<T>::infinity();
179   }
180 };
181 template <typename T>
182 struct MinMaxBottomValue<T, false, true> {
183   EIGEN_DEVICE_FUNC static EIGEN_STRONG_INLINE T bottom_value() {
184     return Eigen::NumTraits<T>::highest();
185   }
186 };
187 template <typename T>
188 struct MinMaxBottomValue<T, false, false> {
189   EIGEN_DEVICE_FUNC static EIGEN_STRONG_INLINE T bottom_value() {
190     return Eigen::NumTraits<T>::infinity();
191   }
192 };
193 
194 
195 template <typename T, int NaNPropagation=PropagateFast> struct MaxReducer
196 {
197   EIGEN_DEVICE_FUNC EIGEN_STRONG_INLINE void reduce(const T t, T* accum) const {
198     scalar_max_op<T, T, NaNPropagation> op;
199     *accum = op(t, *accum);
200   }
201   template <typename Packet>
202   EIGEN_DEVICE_FUNC EIGEN_STRONG_INLINE void reducePacket(const Packet& p, Packet* accum) const {
203     scalar_max_op<T, T, NaNPropagation> op;
204     (*accum) = op.packetOp(*accum, p);
205   }
206   EIGEN_DEVICE_FUNC EIGEN_STRONG_INLINE T initialize() const {
207     return MinMaxBottomValue<T, /*IsMax=*/true, Eigen::NumTraits<T>::IsInteger>::bottom_value();
208   }
209   template <typename Packet>
210   EIGEN_DEVICE_FUNC EIGEN_STRONG_INLINE Packet initializePacket() const {
211     return pset1<Packet>(initialize());
212   }
213   EIGEN_DEVICE_FUNC EIGEN_STRONG_INLINE T finalize(const T accum) const {
214     return accum;
215   }
216   template <typename Packet>
217   EIGEN_DEVICE_FUNC EIGEN_STRONG_INLINE Packet finalizePacket(const Packet& vaccum) const {
218     return vaccum;
219   }
220   template <typename Packet>
221   EIGEN_DEVICE_FUNC EIGEN_STRONG_INLINE T finalizeBoth(const T saccum, const Packet& vaccum) const {
222     scalar_max_op<T, T, NaNPropagation> op;
223     return op(saccum, op.predux(vaccum));
224   }
225 };
226 
227 template <typename T, typename Device, int NaNPropagation>
228     struct reducer_traits<MaxReducer<T, NaNPropagation>, Device> {
229   enum {
230     Cost = NumTraits<T>::AddCost,
231     PacketAccess = PacketType<T, Device>::HasMax,
232     IsStateful = false,
233     IsExactlyAssociative = (NaNPropagation!=PropagateFast)
234   };
235 };
236 
237 template <typename T, int NaNPropagation=PropagateFast> struct MinReducer
238 {
239   EIGEN_DEVICE_FUNC EIGEN_STRONG_INLINE void reduce(const T t, T* accum) const {
240     scalar_min_op<T, T, NaNPropagation> op;
241     *accum = op(t, *accum);
242   }
243   template <typename Packet>
244   EIGEN_DEVICE_FUNC EIGEN_STRONG_INLINE void reducePacket(const Packet& p, Packet* accum) const {
245     scalar_min_op<T, T, NaNPropagation> op;
246     (*accum) = op.packetOp(*accum, p);
247   }
248   EIGEN_DEVICE_FUNC EIGEN_STRONG_INLINE T initialize() const {
249     return MinMaxBottomValue<T, /*IsMax=*/false, Eigen::NumTraits<T>::IsInteger>::bottom_value();
250   }
251   template <typename Packet>
252   EIGEN_DEVICE_FUNC EIGEN_STRONG_INLINE Packet initializePacket() const {
253     return pset1<Packet>(initialize());
254   }
255   EIGEN_DEVICE_FUNC EIGEN_STRONG_INLINE T finalize(const T accum) const {
256     return accum;
257   }
258   template <typename Packet>
259   EIGEN_DEVICE_FUNC EIGEN_STRONG_INLINE Packet finalizePacket(const Packet& vaccum) const {
260     return vaccum;
261   }
262   template <typename Packet>
263   EIGEN_DEVICE_FUNC EIGEN_STRONG_INLINE T finalizeBoth(const T saccum, const Packet& vaccum) const {
264     scalar_min_op<T, T, NaNPropagation> op;
265     return op(saccum, op.predux(vaccum));
266   }
267 };
268 
269 template <typename T, typename Device, int NaNPropagation>
270     struct reducer_traits<MinReducer<T, NaNPropagation>, Device> {
271   enum {
272     Cost = NumTraits<T>::AddCost,
273     PacketAccess = PacketType<T, Device>::HasMin,
274     IsStateful = false,
275     IsExactlyAssociative = (NaNPropagation!=PropagateFast)
276   };
277 };
278 
279 template <typename T> struct ProdReducer
280 {
281   EIGEN_DEVICE_FUNC EIGEN_STRONG_INLINE void reduce(const T t, T* accum) const {
282     internal::scalar_product_op<T> prod_op;
283     (*accum) = prod_op(*accum, t);
284   }
285   template <typename Packet>
286   EIGEN_DEVICE_FUNC EIGEN_STRONG_INLINE void reducePacket(const Packet& p, Packet* accum) const {
287     (*accum) = pmul<Packet>(*accum, p);
288   }
289   EIGEN_DEVICE_FUNC EIGEN_STRONG_INLINE T initialize() const {
290     internal::scalar_cast_op<int, T> conv;
291     return conv(1);
292   }
293   template <typename Packet>
294   EIGEN_DEVICE_FUNC EIGEN_STRONG_INLINE Packet initializePacket() const {
295     return pset1<Packet>(initialize());
296   }
297   EIGEN_DEVICE_FUNC EIGEN_STRONG_INLINE T finalize(const T accum) const {
298     return accum;
299   }
300   template <typename Packet>
301   EIGEN_DEVICE_FUNC EIGEN_STRONG_INLINE Packet finalizePacket(const Packet& vaccum) const {
302     return vaccum;
303   }
304   template <typename Packet>
305   EIGEN_DEVICE_FUNC EIGEN_STRONG_INLINE T finalizeBoth(const T saccum, const Packet& vaccum) const {
306     internal::scalar_product_op<T> prod_op;
307     return prod_op(saccum, predux_mul(vaccum));
308   }
309 };
310 
311 template <typename T, typename Device>
312 struct reducer_traits<ProdReducer<T>, Device> {
313   enum {
314     Cost = NumTraits<T>::MulCost,
315     PacketAccess = PacketType<T, Device>::HasMul,
316     IsStateful = false,
317     IsExactlyAssociative = true
318   };
319 };
320 
321 
322 struct AndReducer
323 {
324   EIGEN_DEVICE_FUNC EIGEN_STRONG_INLINE void reduce(bool t, bool* accum) const {
325     *accum = *accum && t;
326   }
327   EIGEN_DEVICE_FUNC EIGEN_STRONG_INLINE bool initialize() const {
328     return true;
329   }
330   EIGEN_DEVICE_FUNC EIGEN_STRONG_INLINE bool finalize(bool accum) const {
331     return accum;
332   }
333 };
334 
335 template <typename Device>
336 struct reducer_traits<AndReducer, Device> {
337   enum {
338     Cost = 1,
339     PacketAccess = false,
340     IsStateful = false,
341     IsExactlyAssociative = true
342   };
343 };
344 
345 
346 struct OrReducer {
347   EIGEN_DEVICE_FUNC EIGEN_STRONG_INLINE void reduce(bool t, bool* accum) const {
348     *accum = *accum || t;
349   }
350   EIGEN_DEVICE_FUNC EIGEN_STRONG_INLINE bool initialize() const {
351     return false;
352   }
353   EIGEN_DEVICE_FUNC EIGEN_STRONG_INLINE bool finalize(bool accum) const {
354     return accum;
355   }
356 };
357 
358 template <typename Device>
359 struct reducer_traits<OrReducer, Device> {
360   enum {
361     Cost = 1,
362     PacketAccess = false,
363     IsStateful = false,
364     IsExactlyAssociative = true
365   };
366 };
367 
368 // Argmin/Argmax reducers.  Returns the first occurrence if multiple locations
369 // contain the same min/max value.
370 template <typename T> struct ArgMaxTupleReducer
371 {
372   EIGEN_DEVICE_FUNC EIGEN_STRONG_INLINE void reduce(const T t, T* accum) const {
373     if (t.second < accum->second) {
374       return;
375     } else if (t.second > accum->second || accum->first > t.first ) {
376       *accum = t;
377     }
378   }
379   EIGEN_DEVICE_FUNC EIGEN_STRONG_INLINE T initialize() const {
380     return T(0, NumTraits<typename T::second_type>::lowest());
381   }
382   EIGEN_DEVICE_FUNC EIGEN_STRONG_INLINE T finalize(const T& accum) const {
383     return accum;
384   }
385 };
386 
387 template <typename T, typename Device>
388 struct reducer_traits<ArgMaxTupleReducer<T>, Device> {
389   enum {
390     Cost = NumTraits<T>::AddCost,
391     PacketAccess = false,
392     IsStateful = false,
393     IsExactlyAssociative = true
394   };
395 };
396 
397 
398 template <typename T> struct ArgMinTupleReducer
399 {
400   EIGEN_DEVICE_FUNC EIGEN_STRONG_INLINE void reduce(const T& t, T* accum) const {
401     if (t.second > accum->second) {
402       return;
403     } else if (t.second < accum->second || accum->first > t.first) {
404       *accum = t;
405     }
406   }
407   EIGEN_DEVICE_FUNC EIGEN_STRONG_INLINE T initialize() const {
408     return T(0, NumTraits<typename T::second_type>::highest());
409   }
410   EIGEN_DEVICE_FUNC EIGEN_STRONG_INLINE T finalize(const T& accum) const {
411     return accum;
412   }
413 };
414 
415 template <typename T, typename Device>
416 struct reducer_traits<ArgMinTupleReducer<T>, Device> {
417   enum {
418     Cost = NumTraits<T>::AddCost,
419     PacketAccess = false,
420     IsStateful = false,
421     IsExactlyAssociative = true
422   };
423 };
424 
425 
426 template <typename T, typename Index, size_t NumDims>
427 class GaussianGenerator {
428  public:
429   static const bool PacketAccess = false;
430 
431   EIGEN_DEVICE_FUNC GaussianGenerator(const array<T, NumDims>& means,
432                                       const array<T, NumDims>& std_devs)
433       : m_means(means)
434   {
435     EIGEN_UNROLL_LOOP
436     for (size_t i = 0; i < NumDims; ++i) {
437       m_two_sigmas[i] = std_devs[i] * std_devs[i] * 2;
438     }
439   }
440 
441   EIGEN_DEVICE_FUNC T operator()(const array<Index, NumDims>& coordinates) const {
442     T tmp = T(0);
443     EIGEN_UNROLL_LOOP
444     for (size_t i = 0; i < NumDims; ++i) {
445       T offset = coordinates[i] - m_means[i];
446       tmp += offset * offset / m_two_sigmas[i];
447     }
448     return numext::exp(-tmp);
449   }
450 
451  private:
452   array<T, NumDims> m_means;
453   array<T, NumDims> m_two_sigmas;
454 };
455 
456 template <typename T, typename Index, size_t NumDims>
457 struct functor_traits<GaussianGenerator<T, Index, NumDims> > {
458   enum {
459     Cost = NumDims * (2 * NumTraits<T>::AddCost + NumTraits<T>::MulCost +
460                       functor_traits<scalar_quotient_op<T, T> >::Cost) +
461            functor_traits<scalar_exp_op<T> >::Cost,
462     PacketAccess = GaussianGenerator<T, Index, NumDims>::PacketAccess
463   };
464 };
465 
466 template <typename Scalar>
467 struct scalar_clamp_op {
468   EIGEN_DEVICE_FUNC inline scalar_clamp_op(const Scalar& _min, const Scalar& _max) : m_min(_min), m_max(_max) {}
469   EIGEN_DEVICE_FUNC EIGEN_STRONG_INLINE const Scalar
470   operator()(const Scalar& x) const {
471     return numext::mini(numext::maxi(x, m_min), m_max);
472   }
473   template <typename Packet>
474   EIGEN_DEVICE_FUNC EIGEN_STRONG_INLINE const Packet
475   packetOp(const Packet& x) const {
476     return internal::pmin(internal::pmax(x, pset1<Packet>(m_min)), pset1<Packet>(m_max));
477   }
478   const Scalar m_min;
479   const Scalar m_max;
480 };
481 template<typename Scalar>
482 struct functor_traits<scalar_clamp_op<Scalar> >
483 { enum { Cost = 2 * NumTraits<Scalar>::AddCost, PacketAccess = (packet_traits<Scalar>::HasMin && packet_traits<Scalar>::HasMax)}; };
484 
485 } // end namespace internal
486 } // end namespace Eigen
487 
488 #endif // EIGEN_CXX11_TENSOR_TENSOR_FUNCTORS_H
489