Eigen-Contrib  5.0.1
 
Loading...
Searching...
No Matches
TensorReduction.h
1// This file is part of Eigen, a lightweight C++ template library
2// for linear algebra.
3//
4// Copyright (C) 2014 Benoit Steiner <benoit.steiner.goog@gmail.com>
5// Copyright (C) 2016 Mehdi Goli, Codeplay Software Ltd <eigen@codeplay.com>
6//
7// This Source Code Form is subject to the terms of the Mozilla
8// Public License v. 2.0. If a copy of the MPL was not distributed
9// with this file, You can obtain one at http://mozilla.org/MPL/2.0/.
10// SPDX-License-Identifier: MPL-2.0
11
12#ifndef EIGEN_TENSOR_TENSOR_REDUCTION_H
13#define EIGEN_TENSOR_TENSOR_REDUCTION_H
14
15// clang is incompatible with the CUDA syntax wrt making a kernel a class friend,
16// so we'll use a macro to make clang happy.
17#ifndef KERNEL_FRIEND
18#if defined(__clang__) && (defined(__CUDA__) || defined(__HIP__))
19#define KERNEL_FRIEND friend __global__ EIGEN_HIP_LAUNCH_BOUNDS_1024
20#else
21#define KERNEL_FRIEND friend
22#endif
23#endif
24
25// IWYU pragma: private
26#include "./InternalHeaderCheck.h"
27
28namespace Eigen {
29
30namespace internal {
31template <typename Op, typename Dims, typename XprType, template <class> class MakePointer_>
32struct traits<TensorReductionOp<Op, Dims, XprType, MakePointer_> > : traits<XprType> {
33 typedef traits<XprType> XprTraits;
34 typedef typename XprTraits::Scalar Scalar;
35 typedef typename XprTraits::StorageKind StorageKind;
36 typedef typename XprTraits::Index Index;
37 static constexpr int NumDimensions = XprTraits::NumDimensions - array_size<Dims>::value;
38 static constexpr int Layout = XprTraits::Layout;
39 typedef typename XprTraits::PointerType PointerType;
40
41 template <class T>
42 struct MakePointer {
43 typedef typename MakePointer_<T>::Type Type;
44 };
45};
46
47template <typename Op, typename Dims, typename XprType, template <class> class MakePointer_>
48struct eval<TensorReductionOp<Op, Dims, XprType, MakePointer_>, Eigen::Dense> {
49 typedef const TensorReductionOp<Op, Dims, XprType, MakePointer_>& type;
50};
51
52template <typename OutputDims>
53struct DimInitializer {
54 template <typename InputDims, typename ReducedDims>
55 EIGEN_DEVICE_FUNC static void run(const InputDims& input_dims,
56 const array<bool, internal::array_size<InputDims>::value>& reduced,
57 OutputDims* output_dims, ReducedDims* reduced_dims) {
58 const int NumInputDims = internal::array_size<InputDims>::value;
59 int outputIndex = 0;
60 int reduceIndex = 0;
61 for (int i = 0; i < NumInputDims; ++i) {
62 if (reduced[i]) {
63 (*reduced_dims)[reduceIndex] = input_dims[i];
64 ++reduceIndex;
65 } else {
66 (*output_dims)[outputIndex] = input_dims[i];
67 ++outputIndex;
68 }
69 }
70 }
71};
72
73template <>
74struct DimInitializer<Sizes<> > {
75 template <typename InputDims, typename Index, size_t Rank>
76 EIGEN_DEVICE_FUNC static void run(const InputDims& input_dims, const array<bool, Rank>&, Sizes<>*,
77 array<Index, Rank>* reduced_dims) {
78 const int NumInputDims = internal::array_size<InputDims>::value;
79 for (int i = 0; i < NumInputDims; ++i) {
80 (*reduced_dims)[i] = input_dims[i];
81 }
82 }
83};
84
85template <typename ReducedDims, int NumTensorDims, int Layout>
86struct are_inner_most_dims {
87 static constexpr bool value = false;
88};
89template <typename ReducedDims, int NumTensorDims, int Layout>
90struct preserve_inner_most_dims {
91 static constexpr bool value = false;
92};
93
94template <typename ReducedDims, int NumTensorDims>
95struct are_inner_most_dims<ReducedDims, NumTensorDims, ColMajor> {
96 static constexpr bool tmp1 = indices_statically_known_to_increase<ReducedDims>();
97 static constexpr bool tmp2 = index_statically_eq<ReducedDims>(0, 0);
98 static constexpr bool tmp3 =
99 index_statically_eq<ReducedDims>(array_size<ReducedDims>::value - 1, array_size<ReducedDims>::value - 1);
100 static constexpr bool value = tmp1 & tmp2 & tmp3;
101};
102template <typename ReducedDims, int NumTensorDims>
103struct are_inner_most_dims<ReducedDims, NumTensorDims, RowMajor> {
104 static constexpr bool tmp1 = indices_statically_known_to_increase<ReducedDims>();
105 static constexpr bool tmp2 = index_statically_eq<ReducedDims>(0, NumTensorDims - array_size<ReducedDims>::value);
106 static constexpr bool tmp3 = index_statically_eq<ReducedDims>(array_size<ReducedDims>::value - 1, NumTensorDims - 1);
107 static constexpr bool value = tmp1 & tmp2 & tmp3;
108};
109template <typename ReducedDims, int NumTensorDims>
110struct preserve_inner_most_dims<ReducedDims, NumTensorDims, ColMajor> {
111 static constexpr bool tmp1 = indices_statically_known_to_increase<ReducedDims>();
112 static constexpr bool tmp2 = index_statically_gt<ReducedDims>(0, 0);
113 static constexpr bool value = tmp1 & tmp2;
114};
115template <typename ReducedDims, int NumTensorDims>
116struct preserve_inner_most_dims<ReducedDims, NumTensorDims, RowMajor> {
117 static constexpr bool tmp1 = indices_statically_known_to_increase<ReducedDims>();
118 static constexpr bool tmp2 = index_statically_lt<ReducedDims>(array_size<ReducedDims>::value - 1, NumTensorDims - 1);
119 static constexpr bool value = tmp1 & tmp2;
120};
121
122template <int DimIndex, typename Self, typename Op>
123struct GenericDimReducer {
124 static EIGEN_DEVICE_FUNC EIGEN_STRONG_INLINE void reduce(const Self& self, typename Self::Index firstIndex,
125 Op& reducer, typename Self::CoeffReturnType* accum) {
126 EIGEN_STATIC_ASSERT((DimIndex > 0), YOU_MADE_A_PROGRAMMING_MISTAKE);
127 for (int j = 0; j < self.m_reducedDims[DimIndex]; ++j) {
128 const typename Self::Index input = firstIndex + j * self.m_reducedStrides[DimIndex];
129 GenericDimReducer<DimIndex - 1, Self, Op>::reduce(self, input, reducer, accum);
130 }
131 }
132};
133template <typename Self, typename Op>
134struct GenericDimReducer<0, Self, Op> {
135 static EIGEN_DEVICE_FUNC EIGEN_STRONG_INLINE void reduce(const Self& self, typename Self::Index firstIndex,
136 Op& reducer, typename Self::CoeffReturnType* accum) {
137 for (int j = 0; j < self.m_reducedDims[0]; ++j) {
138 const typename Self::Index input = firstIndex + j * self.m_reducedStrides[0];
139 reducer.reduce(self.m_impl.coeff(input), accum);
140 }
141 }
142};
143template <typename Self, typename Op>
144struct GenericDimReducer<-1, Self, Op> {
145 static EIGEN_DEVICE_FUNC EIGEN_STRONG_INLINE void reduce(const Self& self, typename Self::Index index, Op& reducer,
146 typename Self::CoeffReturnType* accum) {
147 reducer.reduce(self.m_impl.coeff(index), accum);
148 }
149};
150
151template <typename Self, typename Op,
152 bool Vectorizable = (Self::InputPacketAccess && Self::ReducerTraits::PacketAccess),
153 bool UseTreeReduction = (!Self::ReducerTraits::IsStateful && !Self::ReducerTraits::IsExactlyAssociative &&
154 // GPU threads can quickly run out of stack space
155 // for moderately sized inputs.
156 !Self::RunningOnGPU)>
157struct InnerMostDimReducer {
158 static EIGEN_DEVICE_FUNC EIGEN_STRONG_INLINE typename Self::CoeffReturnType reduce(
159 const Self& self, typename Self::Index firstIndex, typename Self::Index numValuesToReduce, Op& reducer) {
160 using Index = typename Self::Index;
161 typename Self::CoeffReturnType accum0 = reducer.initialize();
162 Index j = 0;
163 // The accumulators take interleaved operands and are merged by feeding one back through
164 // reduce(), so this needs a pure combine that is also associative and commutative;
165 // reducer_can_reorder_accumulators marks the reducers that guarantee both.
166 EIGEN_IF_CONSTEXPR (reducer_can_reorder_accumulators<Op>::value) {
167 if (numValuesToReduce >= 8) {
168 typename Self::CoeffReturnType accum1 = reducer.initialize(), accum2 = reducer.initialize();
169 typename Self::CoeffReturnType accum3 = reducer.initialize(), accum4 = reducer.initialize();
170 typename Self::CoeffReturnType accum5 = reducer.initialize(), accum6 = reducer.initialize();
171 typename Self::CoeffReturnType accum7 = reducer.initialize();
172 const Index unrolledEnd = numValuesToReduce - numValuesToReduce % 8;
173 for (; j < unrolledEnd; j += 8) {
174 reducer.reduce(self.m_impl.coeff(firstIndex + j + 0), &accum0);
175 reducer.reduce(self.m_impl.coeff(firstIndex + j + 1), &accum1);
176 reducer.reduce(self.m_impl.coeff(firstIndex + j + 2), &accum2);
177 reducer.reduce(self.m_impl.coeff(firstIndex + j + 3), &accum3);
178 reducer.reduce(self.m_impl.coeff(firstIndex + j + 4), &accum4);
179 reducer.reduce(self.m_impl.coeff(firstIndex + j + 5), &accum5);
180 reducer.reduce(self.m_impl.coeff(firstIndex + j + 6), &accum6);
181 reducer.reduce(self.m_impl.coeff(firstIndex + j + 7), &accum7);
182 }
183 reducer.reduce(accum1, &accum0);
184 reducer.reduce(accum2, &accum0);
185 reducer.reduce(accum3, &accum0);
186 reducer.reduce(accum4, &accum0);
187 reducer.reduce(accum5, &accum0);
188 reducer.reduce(accum6, &accum0);
189 reducer.reduce(accum7, &accum0);
190 }
191 }
192 for (; j < numValuesToReduce; ++j) {
193 reducer.reduce(self.m_impl.coeff(firstIndex + j), &accum0);
194 }
195 return reducer.finalize(accum0);
196 }
197};
198
199template <typename Self, typename Op>
200struct InnerMostDimReducer<Self, Op, true, false> {
201 static EIGEN_DEVICE_FUNC EIGEN_STRONG_INLINE typename Self::CoeffReturnType reduce(
202 const Self& self, typename Self::Index firstIndex, typename Self::Index numValuesToReduce, Op& reducer0) {
203 using Index = typename Self::Index;
204 constexpr Index packetSize = internal::unpacket_traits<typename Self::PacketReturnType>::size;
205 Index start = 0;
206 typename Self::PacketReturnType paccum0 = reducer0.template initializePacket<typename Self::PacketReturnType>();
207 EIGEN_IF_CONSTEXPR (!Self::ReducerTraits::IsStateful) {
208 if (numValuesToReduce >= 4 * packetSize) {
209 const Index VectorizedSize4 = (numValuesToReduce / (4 * packetSize)) * (4 * packetSize);
210 typename Self::PacketReturnType paccum1 = reducer0.template initializePacket<typename Self::PacketReturnType>();
211 typename Self::PacketReturnType paccum2 = reducer0.template initializePacket<typename Self::PacketReturnType>();
212 typename Self::PacketReturnType paccum3 = reducer0.template initializePacket<typename Self::PacketReturnType>();
213 const Index offset0 = firstIndex;
214 const Index offset1 = firstIndex + packetSize;
215 const Index offset2 = firstIndex + 2 * packetSize;
216 const Index offset3 = firstIndex + 3 * packetSize;
217 for (Index j = 0; j < VectorizedSize4; j += 4 * packetSize) {
218 reducer0.reducePacket(self.m_impl.template packet<Unaligned>(offset0 + j), &paccum0);
219 reducer0.reducePacket(self.m_impl.template packet<Unaligned>(offset1 + j), &paccum1);
220 reducer0.reducePacket(self.m_impl.template packet<Unaligned>(offset2 + j), &paccum2);
221 reducer0.reducePacket(self.m_impl.template packet<Unaligned>(offset3 + j), &paccum3);
222 }
223 reducer0.reducePacket(paccum1, &paccum0);
224 reducer0.reducePacket(paccum2, &paccum0);
225 reducer0.reducePacket(paccum3, &paccum0);
226 start = VectorizedSize4;
227 }
228 }
229 if (start <= (numValuesToReduce - packetSize)) {
230 const Index VectorizedSize = (numValuesToReduce / packetSize) * packetSize;
231 for (Index j = start; j < VectorizedSize; j += packetSize) {
232 reducer0.reducePacket(self.m_impl.template packet<Unaligned>(firstIndex + j), &paccum0);
233 }
234 start = VectorizedSize;
235 }
236 typename Self::CoeffReturnType accum = reducer0.initialize();
237 for (Index j = start; j < numValuesToReduce; ++j) {
238 reducer0.reduce(self.m_impl.coeff(firstIndex + j), &accum);
239 }
240 return reducer0.finalizeBoth(accum, paccum0);
241 }
242};
243
244#if !defined(EIGEN_HIPCC)
245
246// The following implements tree-based reduction, which improves the accuracy
247// of sum and mean reductions, since each of the n inputs only participates in
248// O(log n) additions.
249template <typename T>
250EIGEN_DEVICE_FUNC inline Index LeafSize() {
251 return 1024;
252}
253template <>
254EIGEN_DEVICE_FUNC inline Index LeafSize<half>() {
255 return 200;
256}
257template <>
258EIGEN_DEVICE_FUNC inline Index LeafSize<bfloat16>() {
259 return 128;
260}
261
262template <typename Self, typename Op>
263struct InnerMostDimReducer<Self, Op, false, true> {
264 static EIGEN_DEVICE_FUNC EIGEN_STRONG_INLINE typename Self::CoeffReturnType reduce(
265 const Self& self, typename Self::Index firstIndex, typename Self::Index numValuesToReduce, Op& reducer) {
266 const Index kLeafSize = LeafSize<typename Self::CoeffReturnType>();
267 typename Self::CoeffReturnType accum = reducer.initialize();
268 if (numValuesToReduce > kLeafSize) {
269 const typename Self::Index half = numValuesToReduce / 2;
270 // Recursively reduce the two halves.
271 reducer.reduce(reduce(self, firstIndex, half, reducer), &accum);
272 reducer.reduce(reduce(self, firstIndex + half, numValuesToReduce - half, reducer), &accum);
273 return reducer.finalize(accum);
274 } else {
275 return InnerMostDimReducer<Self, Op, false, false>::reduce(self, firstIndex, numValuesToReduce, reducer);
276 }
277 }
278};
279
280template <typename Self, typename Op>
281struct InnerMostDimReducer<Self, Op, true, true> {
282 static EIGEN_DEVICE_FUNC EIGEN_STRONG_INLINE typename Self::CoeffReturnType reduce(
283 const Self& self, typename Self::Index firstIndex, typename Self::Index numValuesToReduce, Op& reducer) {
284 const Index kLeafSize = LeafSize<typename Self::CoeffReturnType>();
285 const typename Self::Index packetSize = internal::unpacket_traits<typename Self::PacketReturnType>::size;
286 typename Self::CoeffReturnType accum = reducer.initialize();
287 if (numValuesToReduce > packetSize * kLeafSize) {
288 // Make sure the split point is aligned on a packet boundary.
289 const typename Self::Index split =
290 packetSize *
291 numext::div_ceil(firstIndex + numext::div_ceil(numValuesToReduce, typename Self::Index(2)), packetSize);
292 const typename Self::Index num_left = numext::mini(split - firstIndex, numValuesToReduce);
293 reducer.reduce(reduce(self, firstIndex, num_left, reducer), &accum);
294 if (num_left < numValuesToReduce) {
295 reducer.reduce(reduce(self, split, numValuesToReduce - num_left, reducer), &accum);
296 }
297 return reducer.finalize(accum);
298 } else {
299 return InnerMostDimReducer<Self, Op, true, false>::reduce(self, firstIndex, numValuesToReduce, reducer);
300 }
301 }
302};
303#endif
304
305template <int DimIndex, typename Self, typename Op,
306 bool vectorizable = (Self::InputPacketAccess && Self::ReducerTraits::PacketAccess)>
307struct InnerMostDimPreserver {
308 static EIGEN_DEVICE_FUNC EIGEN_STRONG_INLINE void reduce(const Self&, typename Self::Index, Op&,
309 typename Self::PacketReturnType*) {
310 eigen_assert(false && "should never be called");
311 }
312};
313
314template <int DimIndex, typename Self, typename Op>
315struct InnerMostDimPreserver<DimIndex, Self, Op, true> {
316 static EIGEN_DEVICE_FUNC EIGEN_STRONG_INLINE void reduce(const Self& self, typename Self::Index firstIndex,
317 Op& reducer, typename Self::PacketReturnType* accum) {
318 EIGEN_STATIC_ASSERT((DimIndex > 0), YOU_MADE_A_PROGRAMMING_MISTAKE);
319 for (typename Self::Index j = 0; j < self.m_reducedDims[DimIndex]; ++j) {
320 const typename Self::Index input = firstIndex + j * self.m_reducedStrides[DimIndex];
321 InnerMostDimPreserver<DimIndex - 1, Self, Op>::reduce(self, input, reducer, accum);
322 }
323 }
324};
325
326template <typename Self, typename Op>
327struct InnerMostDimPreserver<0, Self, Op, true> {
328 static EIGEN_DEVICE_FUNC EIGEN_STRONG_INLINE void reduce(const Self& self, typename Self::Index firstIndex,
329 Op& reducer0, typename Self::PacketReturnType* accum0) {
330 using Index = typename Self::Index;
331 const Index stride = self.m_reducedStrides[0];
332 const Index size = self.m_reducedDims[0];
333 EIGEN_IF_CONSTEXPR (!Self::ReducerTraits::IsStateful) {
334 if (size >= 16) {
335 const Index unrolled_size4 = (size / 4) * 4;
336 typename Self::PacketReturnType accum1 = reducer0.template initializePacket<typename Self::PacketReturnType>();
337 typename Self::PacketReturnType accum2 = reducer0.template initializePacket<typename Self::PacketReturnType>();
338 typename Self::PacketReturnType accum3 = reducer0.template initializePacket<typename Self::PacketReturnType>();
339 for (Index j = 0; j < unrolled_size4; j += 4) {
340 const Index input0 = firstIndex + j * stride;
341 reducer0.reducePacket(self.m_impl.template packet<Unaligned>(input0), accum0);
342 const Index input1 = firstIndex + (j + 1) * stride;
343 reducer0.reducePacket(self.m_impl.template packet<Unaligned>(input1), &accum1);
344 const Index input2 = firstIndex + (j + 2) * stride;
345 reducer0.reducePacket(self.m_impl.template packet<Unaligned>(input2), &accum2);
346 const Index input3 = firstIndex + (j + 3) * stride;
347 reducer0.reducePacket(self.m_impl.template packet<Unaligned>(input3), &accum3);
348 }
349 reducer0.reducePacket(accum1, accum0);
350 reducer0.reducePacket(accum2, accum0);
351 reducer0.reducePacket(accum3, accum0);
352 for (Index j = unrolled_size4; j < size; ++j) {
353 Index input = firstIndex + j * stride;
354 reducer0.reducePacket(self.m_impl.template packet<Unaligned>(input), accum0);
355 }
356 return;
357 }
358 }
359 for (Index j = 0; j < size; ++j) {
360 Index input = firstIndex + j * stride;
361 reducer0.reducePacket(self.m_impl.template packet<Unaligned>(input), accum0);
362 }
363 }
364};
365template <typename Self, typename Op>
366struct InnerMostDimPreserver<-1, Self, Op, true> {
367 static EIGEN_DEVICE_FUNC EIGEN_STRONG_INLINE void reduce(const Self&, typename Self::Index, Op&,
368 typename Self::PacketReturnType*) {
369 eigen_assert(false && "should never be called");
370 }
371};
372
373// Default full reducer
374template <typename Self, typename Op, typename Device,
375 bool Vectorizable = (Self::InputPacketAccess && Self::ReducerTraits::PacketAccess)>
376struct FullReducer {
377 static constexpr bool HasOptimizedImplementation = false;
378
379 static EIGEN_DEVICE_FUNC void run(const Self& self, Op& reducer, const Device&,
380 typename Self::EvaluatorPointerType output) {
381 const typename Self::Index num_coeffs = array_prod(self.m_impl.dimensions());
382 *output = InnerMostDimReducer<Self, Op, Vectorizable>::reduce(self, 0, num_coeffs, reducer);
383 }
384};
385
386#ifdef EIGEN_USE_THREADS
387// Multithreaded full reducer
388template <typename Self, typename Op, bool Vectorizable>
389struct FullReducer<Self, Op, ThreadPoolDevice, Vectorizable> {
390 static constexpr bool HasOptimizedImplementation = !Self::ReducerTraits::IsStateful;
391 static constexpr Index PacketSize = unpacket_traits<typename Self::PacketReturnType>::size;
392
393 // launch one reducer per thread and accumulate the result.
394 static void run(const Self& self, Op& reducer, const ThreadPoolDevice& device,
395 typename Self::CoeffReturnType* output) {
396 typedef typename Self::Index Index;
397 const Index num_coeffs = array_prod(self.m_impl.dimensions());
398 if (num_coeffs == 0) {
399 *output = reducer.finalize(reducer.initialize());
400 return;
401 }
402 const TensorOpCost cost = self.m_impl.costPerCoeff(Vectorizable) +
403 TensorOpCost(0, 0, internal::functor_traits<Op>::Cost, Vectorizable, PacketSize);
404 const Index num_threads = TensorCostModel<ThreadPoolDevice>::numThreads(num_coeffs, cost, device.numThreads());
405 if (num_threads == 1) {
406 *output = InnerMostDimReducer<Self, Op, Vectorizable>::reduce(self, 0, num_coeffs, reducer);
407 return;
408 }
409 const Index blocksize = num_coeffs / num_threads;
410 const Index numblocks = blocksize > 0 ? num_coeffs / blocksize : 0;
411 eigen_assert(num_coeffs >= numblocks * blocksize);
412
413 Barrier barrier(internal::convert_index<unsigned int>(numblocks));
414 MaxSizeVector<typename Self::CoeffReturnType> shards(numblocks, reducer.initialize());
415 for (Index i = 0; i < numblocks; ++i) {
416 auto run_shard = [i, blocksize, &self, &barrier, &shards, &reducer]() {
417 shards[i] = InnerMostDimReducer<Self, Op, Vectorizable>::reduce(self, i * blocksize, blocksize, reducer);
418 barrier.Notify();
419 };
420 device.enqueue(std::move(run_shard));
421 }
422 typename Self::CoeffReturnType finalShard;
423 if (numblocks * blocksize < num_coeffs) {
424 finalShard = InnerMostDimReducer<Self, Op, Vectorizable>::reduce(self, numblocks * blocksize,
425 num_coeffs - numblocks * blocksize, reducer);
426 } else {
427 finalShard = reducer.initialize();
428 }
429 barrier.Wait();
430
431 for (Index i = 0; i < numblocks; ++i) {
432 reducer.reduce(shards[i], &finalShard);
433 }
434 *output = reducer.finalize(finalShard);
435 }
436};
437
438#endif
439
440// Default inner reducer
441template <typename Self, typename Op, typename Device>
442struct InnerReducer {
443 static constexpr bool HasOptimizedImplementation = false;
444
445 EIGEN_DEVICE_FUNC static bool run(const Self&, Op&, const Device&, typename Self::CoeffReturnType*,
446 typename Self::Index, typename Self::Index) {
447 eigen_assert(false && "Not implemented");
448 return true;
449 }
450};
451
452// Default outer reducer
453template <typename Self, typename Op, typename Device>
454struct OuterReducer {
455 static constexpr bool HasOptimizedImplementation = false;
456
457 EIGEN_DEVICE_FUNC static bool run(const Self&, Op&, const Device&, typename Self::CoeffReturnType*,
458 typename Self::Index, typename Self::Index) {
459 eigen_assert(false && "Not implemented");
460 return true;
461 }
462};
463
464#ifdef EIGEN_USE_SYCL
465// Default Generic reducer
466template <typename Self, typename Op, typename Device>
467struct GenericReducer {
468 static constexpr bool HasOptimizedImplementation = false;
469
470 EIGEN_DEVICE_FUNC static bool run(const Self&, Op&, const Device&, typename Self::CoeffReturnType*,
471 typename Self::Index, typename Self::Index) {
472 eigen_assert(false && "Not implemented");
473 return true;
474 }
475};
476#endif
477
478#if defined(EIGEN_USE_GPU) && (defined(EIGEN_GPUCC))
479template <int B, int N, typename S, typename R, typename I_>
480__global__ EIGEN_HIP_LAUNCH_BOUNDS_1024 void FullReductionKernel(R, const S, I_, typename S::CoeffReturnType*,
481 unsigned int*);
482
483#if defined(EIGEN_GPUCC)
484// The scratch parameter is half* as in the definitions (TensorReductionGpu.h): declared as
485// packet_traits<half>::type*, it is Packet4h2* in the device pass, which makes the name denote two templates there.
486template <typename S, typename R, typename I_>
487__global__ EIGEN_HIP_LAUNCH_BOUNDS_1024 void ReductionInitFullReduxKernelHalfFloat(R, const S, I_, half*);
488template <int B, int N, typename S, typename R, typename I_>
489__global__ EIGEN_HIP_LAUNCH_BOUNDS_1024 void FullReductionKernelHalfFloat(R, const S, I_, half*, half*);
490template <int NPT, typename S, typename R, typename I_>
491__global__ EIGEN_HIP_LAUNCH_BOUNDS_1024 void InnerReductionKernelHalfFloat(R, const S, I_, I_, half*);
492
493#endif
494
495template <int NPT, typename S, typename R, typename I_>
496__global__ EIGEN_HIP_LAUNCH_BOUNDS_1024 void InnerReductionKernel(R, const S, I_, I_, typename S::CoeffReturnType*);
497
498template <int NPT, typename S, typename R, typename I_>
499__global__ EIGEN_HIP_LAUNCH_BOUNDS_1024 void OuterReductionKernel(R, const S, I_, I_, typename S::CoeffReturnType*);
500#endif
501
510template <typename Op, typename CoeffReturnType>
511struct ReductionReturnType {
512#if defined(EIGEN_USE_SYCL)
513 typedef std::remove_const_t<decltype(std::declval<Op>().initialize())> type;
514#else
515 typedef std::remove_const_t<CoeffReturnType> type;
516#endif
517};
518
519} // end namespace internal
520
527template <typename Op, typename Dims, typename XprType, template <class> class MakePointer_>
528class TensorReductionOp : public TensorBase<TensorReductionOp<Op, Dims, XprType, MakePointer_>, ReadOnlyAccessors> {
529 public:
530 typedef typename Eigen::internal::traits<TensorReductionOp>::Scalar Scalar;
531 typedef typename Eigen::NumTraits<Scalar>::Real RealScalar;
532 typedef std::remove_const_t<typename XprType::CoeffReturnType> CoeffReturnType;
533 typedef typename Eigen::internal::ref_selector<TensorReductionOp>::type Nested;
534 typedef typename Eigen::internal::traits<TensorReductionOp>::StorageKind StorageKind;
535 typedef typename Eigen::internal::traits<TensorReductionOp>::Index Index;
536
537 EIGEN_DEVICE_FUNC EIGEN_STRONG_INLINE TensorReductionOp(const XprType& expr, const Dims& dims)
538 : m_expr(expr), m_dims(dims) {}
539 EIGEN_DEVICE_FUNC EIGEN_STRONG_INLINE TensorReductionOp(const XprType& expr, const Dims& dims, const Op& reducer)
540 : m_expr(expr), m_dims(dims), m_reducer(reducer) {}
541
542 EIGEN_DEVICE_FUNC EIGEN_STRONG_INLINE const XprType& expression() const { return m_expr; }
543 EIGEN_DEVICE_FUNC EIGEN_STRONG_INLINE const Dims& dims() const { return m_dims; }
544 EIGEN_DEVICE_FUNC EIGEN_STRONG_INLINE const Op& reducer() const { return m_reducer; }
545
546 // Rank 0 guarantees one coefficient, so only rank-0 reductions convert directly to a scalar.
547 template <int NumDims = internal::traits<TensorReductionOp>::NumDimensions, EIGEN_SFINAE_ENABLE_IF(NumDims == 0)>
548 EIGEN_STRONG_INLINE operator CoeffReturnType() const {
549 TensorEvaluator<const TensorReductionOp, DefaultDevice> evaluator(*this, DefaultDevice());
550 evaluator.evalSubExprsIfNeeded(nullptr);
551 const CoeffReturnType result = evaluator.coeff(0);
552 evaluator.cleanup();
553 return result;
554 }
555
556#if !defined(EIGEN_PARSED_BY_DOXYGEN)
557 // Exact-match friends keep mixed-scalar arithmetic from being ambiguous between the tensor
558 // operators and built-in arithmetic on the scalar conversion above. is_scalar_operand keeps
559 // reduction-with-reduction arithmetic on the tensor-tensor path, which the conversion to
560 // Scalar would otherwise make ambiguous as well.
561#define EIGEN_TENSOR_REDUCTION_SCALAR_BINOP(op, name) \
562 template <typename T, EIGEN_SFINAE_ENABLE_IF((internal::is_scalar_operand<T, Scalar>::value))> \
563 EIGEN_DEVICE_FUNC EIGEN_STRONG_INLINE friend const TensorCwiseUnaryOp< \
564 internal::bind1st_op<internal::scalar_##name##_op<Scalar> >, const TensorReductionOp> \
565 op(const T& lhs, const TensorReductionOp& rhs) { \
566 return rhs.unaryExpr(internal::bind1st_op<internal::scalar_##name##_op<Scalar> >(static_cast<Scalar>(lhs))); \
567 } \
568 template <typename T, EIGEN_SFINAE_ENABLE_IF((internal::is_scalar_operand<T, Scalar>::value))> \
569 EIGEN_DEVICE_FUNC EIGEN_STRONG_INLINE friend const TensorCwiseUnaryOp< \
570 internal::bind2nd_op<internal::scalar_##name##_op<Scalar> >, const TensorReductionOp> \
571 op(const TensorReductionOp& lhs, const T& rhs) { \
572 return lhs.unaryExpr(internal::bind2nd_op<internal::scalar_##name##_op<Scalar> >(static_cast<Scalar>(rhs))); \
573 }
574
575 EIGEN_TENSOR_REDUCTION_SCALAR_BINOP(operator+, sum)
576 EIGEN_TENSOR_REDUCTION_SCALAR_BINOP(operator-, difference)
577 EIGEN_TENSOR_REDUCTION_SCALAR_BINOP(operator*, product)
578 EIGEN_TENSOR_REDUCTION_SCALAR_BINOP(operator/, quotient)
579#undef EIGEN_TENSOR_REDUCTION_SCALAR_BINOP
580#endif
581
582 protected:
583 typename XprType::Nested m_expr;
584 const Dims m_dims;
585 const Op m_reducer;
586};
587
588template <typename ArgType, typename Device>
589struct TensorReductionEvaluatorBase;
590
591// Eval as rvalue
592template <typename Op, typename Dims, typename ArgType, template <class> class MakePointer_, typename Device>
593struct TensorReductionEvaluatorBase<const TensorReductionOp<Op, Dims, ArgType, MakePointer_>, Device> {
594 typedef internal::reducer_traits<Op, Device> ReducerTraits;
595 typedef Dims ReducedDims;
596 typedef TensorReductionOp<Op, Dims, ArgType, MakePointer_> XprType;
597 typedef typename XprType::Index Index;
598 typedef ArgType ChildType;
599 typedef typename TensorEvaluator<ArgType, Device>::Dimensions InputDimensions;
600 static constexpr int NumInputDims = internal::array_size<InputDimensions>::value;
601 static constexpr int NumReducedDims = internal::array_size<Dims>::value;
602 static constexpr int NumOutputDims = NumInputDims - NumReducedDims;
603 typedef std::conditional_t<NumOutputDims == 0, Sizes<>, DSizes<Index, NumOutputDims> > Dimensions;
604 typedef typename XprType::Scalar Scalar;
605 typedef TensorReductionEvaluatorBase<const TensorReductionOp<Op, Dims, ArgType, MakePointer_>, Device> Self;
606 static constexpr bool InputPacketAccess = TensorEvaluator<ArgType, Device>::PacketAccess;
607 typedef typename internal::ReductionReturnType<Op, typename XprType::CoeffReturnType>::type CoeffReturnType;
608 typedef typename PacketType<CoeffReturnType, Device>::type PacketReturnType;
609 static constexpr Index PacketSize = PacketType<CoeffReturnType, Device>::size;
610
611 typedef typename Eigen::internal::traits<XprType>::PointerType TensorPointerType;
612 typedef StorageMemory<CoeffReturnType, Device> Storage;
613 typedef typename Storage::Type EvaluatorPointerType;
614
615 // Subset of strides of the input tensor for the non-reduced dimensions.
616 // Indexed by output dimensions.
617 static constexpr int NumPreservedStrides = max_n_1<NumOutputDims>::size;
618
619 // For full reductions
620#if defined(EIGEN_USE_GPU) && (defined(EIGEN_GPUCC))
621 static constexpr bool RunningOnGPU = std::is_same<Device, Eigen::GpuDevice>::value;
622 static constexpr bool RunningOnSycl = false;
623#elif defined(EIGEN_USE_SYCL)
624 static constexpr bool RunningOnSycl = std::is_same<internal::remove_all_t<Device>, Eigen::SyclDevice>::value;
625 static constexpr bool RunningOnGPU = false;
626#else
627 static constexpr bool RunningOnGPU = false;
628 static constexpr bool RunningOnSycl = false;
629#endif
630
631 static constexpr int Layout = TensorEvaluator<ArgType, Device>::Layout;
632 enum {
633 IsAligned = false,
634 PacketAccess = Self::InputPacketAccess && ReducerTraits::PacketAccess,
635 BlockAccess = false,
636 PreferBlockAccess = true,
637 CoordAccess = false, // to be implemented
638 RawAccess = false
639 };
640
641 typedef std::remove_const_t<Scalar> ScalarNoConst;
642
643 //===- Tensor block evaluation strategy (see TensorBlock.h) -------------===//
644 typedef internal::TensorBlockNotImplemented TensorBlock;
645 //===--------------------------------------------------------------------===//
646
647 static constexpr bool ReducingInnerMostDims = internal::are_inner_most_dims<Dims, NumInputDims, Layout>::value;
648 static constexpr bool PreservingInnerMostDims = internal::preserve_inner_most_dims<Dims, NumInputDims, Layout>::value;
649 static constexpr bool RunningFullReduction = (NumOutputDims == 0);
650
651 EIGEN_STRONG_INLINE TensorReductionEvaluatorBase(const XprType& op, const Device& device)
652 : m_impl(op.expression(), device), m_reducer(op.reducer()), m_result(nullptr), m_device(device) {
653 EIGEN_STATIC_ASSERT((NumInputDims >= NumReducedDims), YOU_MADE_A_PROGRAMMING_MISTAKE);
654 EIGEN_STATIC_ASSERT((!ReducingInnerMostDims | !PreservingInnerMostDims | (NumReducedDims == NumInputDims)),
655 YOU_MADE_A_PROGRAMMING_MISTAKE);
656
657 // Build the bitmap indicating if an input dimension is reduced or not.
658 for (int i = 0; i < NumInputDims; ++i) {
659 m_reduced[i] = false;
660 }
661 for (int i = 0; i < NumReducedDims; ++i) {
662 eigen_assert(op.dims()[i] >= 0);
663 eigen_assert(op.dims()[i] < NumInputDims);
664 m_reduced[op.dims()[i]] = true;
665 }
666
667 const typename TensorEvaluator<ArgType, Device>::Dimensions& input_dims = m_impl.dimensions();
668 internal::DimInitializer<Dimensions>::run(input_dims, m_reduced, &m_dimensions, &m_reducedDims);
669
670 // Precompute output strides.
671 EIGEN_IF_CONSTEXPR (NumOutputDims > 0) {
672 EIGEN_IF_CONSTEXPR (static_cast<int>(Layout) == static_cast<int>(ColMajor)) {
673 m_outputStrides[0] = 1;
674 for (int i = 1; i < NumOutputDims; ++i) {
675 m_outputStrides[i] = m_outputStrides[i - 1] * m_dimensions[i - 1];
676 m_fastOutputStrides[i] = internal::TensorIntDivisor<Index>(m_outputStrides[i]);
677 }
678 } else {
679 m_outputStrides[static_cast<size_t>(NumOutputDims - 1)] = 1;
680 for (int i = NumOutputDims - 2; i >= 0; --i) {
681 m_outputStrides[i] = m_outputStrides[i + 1] * m_dimensions[i + 1];
682 m_fastOutputStrides[i] = internal::TensorIntDivisor<Index>(m_outputStrides[i]);
683 }
684 }
685 }
686
687 // Precompute input strides.
688 EIGEN_IF_CONSTEXPR (NumInputDims > 0) {
689 array<Index, NumInputDims> input_strides;
690 EIGEN_IF_CONSTEXPR (static_cast<int>(Layout) == static_cast<int>(ColMajor)) {
691 input_strides[0] = 1;
692 for (int i = 1; i < NumInputDims; ++i) {
693 input_strides[i] = input_strides[i - 1] * input_dims[i - 1];
694 }
695 } else {
696 input_strides.back() = 1;
697 for (int i = NumInputDims - 2; i >= 0; --i) {
698 input_strides[i] = input_strides[i + 1] * input_dims[i + 1];
699 }
700 }
701
702 int outputIndex = 0;
703 int reduceIndex = 0;
704 for (int i = 0; i < NumInputDims; ++i) {
705 if (m_reduced[i]) {
706 m_reducedStrides[reduceIndex] = input_strides[i];
707 ++reduceIndex;
708 } else {
709 m_preservedStrides[outputIndex] = input_strides[i];
710 m_output_to_input_dim_map[outputIndex] = i;
711 ++outputIndex;
712 }
713 }
714 }
715
716 // Special case for full reductions
717 EIGEN_IF_CONSTEXPR (NumOutputDims == 0) {
718 m_preservedStrides[0] = internal::array_prod(input_dims);
719 }
720
721 m_numValuesToReduce = NumOutputDims == 0 ? internal::array_prod(input_dims)
722 : (static_cast<int>(Layout) == static_cast<int>(ColMajor))
723 ? m_preservedStrides[0]
724 : m_preservedStrides[static_cast<size_t>(NumOutputDims - 1)];
725
726 // Runtime mirror of the static `ReducingInnerMostDims` predicate, set when
727 // the reduce dims aren't statically known (e.g. a plain std::array).
728 m_reducingInnerMostDims = (NumReducedDims > 0);
729 for (int i = 0; i < NumReducedDims && m_reducingInnerMostDims; ++i) {
730 const int axis = (static_cast<int>(Layout) == static_cast<int>(ColMajor)) ? i : NumInputDims - 1 - i;
731 if (!m_reduced[axis]) m_reducingInnerMostDims = false;
732 }
733 }
734
735 EIGEN_DEVICE_FUNC EIGEN_STRONG_INLINE const Dimensions& dimensions() const { return m_dimensions; }
736
737 EIGEN_STRONG_INLINE bool evalSubExprsIfNeededCommon(EvaluatorPointerType data) {
738 // Use the FullReducer if possible.
739 EIGEN_IF_CONSTEXPR (RunningFullReduction) {
740 if (RunningOnSycl || (internal::FullReducer<Self, Op, Device>::HasOptimizedImplementation &&
741 ((RunningOnGPU && (m_device.majorDeviceVersion() >= 3)) || !RunningOnGPU))) {
742 bool need_assign = false;
743 if (!data) {
744 m_result = static_cast<EvaluatorPointerType>(
745 m_device.get((CoeffReturnType*)m_device.allocate_temp(sizeof(CoeffReturnType))));
746 data = m_result;
747 need_assign = true;
748 }
749 Op reducer(m_reducer);
750 internal::FullReducer<Self, Op, Device>::run(*this, reducer, m_device, data);
751 return need_assign;
752 }
753 }
754
755 // Attempt to use an optimized reduction.
756 EIGEN_IF_CONSTEXPR (RunningOnGPU || RunningOnSycl) {
757 if ((RunningOnGPU && (m_device.majorDeviceVersion() >= 3)) || (RunningOnSycl)) {
758 bool reducing_inner_dims = true;
759 for (int i = 0; i < NumReducedDims; ++i) {
760 EIGEN_IF_CONSTEXPR (static_cast<int>(Layout) == static_cast<int>(ColMajor)) {
761 reducing_inner_dims &= m_reduced[i];
762 } else {
763 reducing_inner_dims &= m_reduced[NumInputDims - 1 - i];
764 }
765 }
766 EIGEN_IF_CONSTEXPR ((internal::InnerReducer<Self, Op, Device>::HasOptimizedImplementation)) {
767 if (reducing_inner_dims || ReducingInnerMostDims) {
768 const Index num_values_to_reduce = internal::array_prod(m_reducedDims);
769 const Index num_coeffs_to_preserve = static_cast<Index>(internal::array_prod(m_dimensions));
770 if (!data) {
771 if ((num_coeffs_to_preserve < 1024 && num_values_to_reduce > num_coeffs_to_preserve &&
772 num_values_to_reduce > 128) ||
773 (RunningOnSycl)) {
774 data = static_cast<EvaluatorPointerType>(m_device.get(
775 (CoeffReturnType*)m_device.allocate_temp(sizeof(CoeffReturnType) * num_coeffs_to_preserve)));
776 m_result = data;
777 } else {
778 return true;
779 }
780 }
781 Op reducer(m_reducer);
782 // For SYCL, this always returns false.
783 if (internal::InnerReducer<Self, Op, Device>::run(*this, reducer, m_device, data, num_values_to_reduce,
784 num_coeffs_to_preserve)) {
785 if (m_result) {
786 m_device.deallocate_temp(m_result);
787 m_result = nullptr;
788 }
789 return true;
790 } else {
791 return (m_result != nullptr);
792 }
793 }
794 }
795
796 bool preserving_inner_dims = true;
797 for (int i = 0; i < NumReducedDims; ++i) {
798 EIGEN_IF_CONSTEXPR (static_cast<int>(Layout) == static_cast<int>(ColMajor)) {
799 preserving_inner_dims &= m_reduced[NumInputDims - 1 - i];
800 } else {
801 preserving_inner_dims &= m_reduced[i];
802 }
803 }
804 EIGEN_IF_CONSTEXPR ((internal::OuterReducer<Self, Op, Device>::HasOptimizedImplementation)) {
805 if (preserving_inner_dims) {
806 const Index num_values_to_reduce = internal::array_prod(m_reducedDims);
807 const Index num_coeffs_to_preserve = static_cast<Index>(internal::array_prod(m_dimensions));
808 if (!data) {
809 if ((num_coeffs_to_preserve < 1024 && num_values_to_reduce > num_coeffs_to_preserve &&
810 num_values_to_reduce > 32) ||
811 (RunningOnSycl)) {
812 data = static_cast<EvaluatorPointerType>(m_device.get(
813 (CoeffReturnType*)m_device.allocate_temp(sizeof(CoeffReturnType) * num_coeffs_to_preserve)));
814 m_result = data;
815 } else {
816 return true;
817 }
818 }
819 Op reducer(m_reducer);
820 // For SYCL, this always returns false.
821 if (internal::OuterReducer<Self, Op, Device>::run(*this, reducer, m_device, data, num_values_to_reduce,
822 num_coeffs_to_preserve)) {
823 if (m_result) {
824 m_device.deallocate_temp(m_result);
825 m_result = nullptr;
826 }
827 return true;
828 } else {
829 return (m_result != nullptr);
830 }
831 }
832 }
833#if defined(EIGEN_USE_SYCL)
834 // If there is no Optimised version for SYCL, the reduction expression
835 // must break into two subexpression and use the SYCL generic Reducer on the device.
836 EIGEN_IF_CONSTEXPR (RunningOnSycl) {
837 const Index num_values_to_reduce = internal::array_prod(m_reducedDims);
838 const Index num_coeffs_to_preserve = static_cast<Index>(internal::array_prod(m_dimensions));
839 if (!data) {
840 data = static_cast<EvaluatorPointerType>(m_device.get(
841 (CoeffReturnType*)m_device.allocate_temp(sizeof(CoeffReturnType) * num_coeffs_to_preserve)));
842 m_result = data;
843 }
844 Op reducer(m_reducer);
845 internal::GenericReducer<Self, Op, Device>::run(*this, reducer, m_device, data, num_values_to_reduce,
846 num_coeffs_to_preserve);
847 return (m_result != nullptr);
848 }
849#endif
850 }
851 }
852 return true;
853 }
854
855#ifdef EIGEN_USE_THREADS
856 template <typename EvalSubExprsCallback>
857 EIGEN_STRONG_INLINE void evalSubExprsIfNeededAsync(EvaluatorPointerType data, EvalSubExprsCallback done) {
858 m_impl.evalSubExprsIfNeededAsync(nullptr, [this, data, done](bool) { done(evalSubExprsIfNeededCommon(data)); });
859 }
860#endif
861
862 EIGEN_STRONG_INLINE bool evalSubExprsIfNeeded(EvaluatorPointerType data) {
863 m_impl.evalSubExprsIfNeeded(nullptr);
864 return evalSubExprsIfNeededCommon(data);
865 }
866
867 EIGEN_STRONG_INLINE void cleanup() {
868 m_impl.cleanup();
869 if (m_result) {
870 m_device.deallocate_temp(m_result);
871 m_result = nullptr;
872 }
873 }
874
875 EIGEN_DEVICE_FUNC EIGEN_STRONG_INLINE CoeffReturnType coeff(Index index) const {
876 EIGEN_IF_CONSTEXPR (RunningFullReduction || RunningOnGPU) {
877 if (m_result) {
878 return *(m_result + index);
879 }
880 }
881 Op reducer(m_reducer);
882 EIGEN_IF_CONSTEXPR (ReducingInnerMostDims || RunningFullReduction) {
883 const Index num_values_to_reduce = (static_cast<int>(Layout) == static_cast<int>(ColMajor))
884 ? m_preservedStrides[0]
885 : m_preservedStrides[NumPreservedStrides - 1];
886 return internal::InnerMostDimReducer<Self, Op>::reduce(*this, firstInput(index), num_values_to_reduce, reducer);
887 } else if (m_reducingInnerMostDims) {
888 return internal::InnerMostDimReducer<Self, Op>::reduce(*this, index * m_numValuesToReduce, m_numValuesToReduce,
889 reducer);
890 } else {
891 typename Self::CoeffReturnType accum = reducer.initialize();
892 internal::GenericDimReducer<NumReducedDims - 1, Self, Op>::reduce(*this, firstInput(index), reducer, &accum);
893 return reducer.finalize(accum);
894 }
895 }
896
897 // TODO(bsteiner): provide a more efficient implementation.
898 template <int LoadMode>
899 EIGEN_DEVICE_FUNC EIGEN_STRONG_INLINE PacketReturnType packet(Index index) const {
900 eigen_assert(index + PacketSize - 1 < Index(internal::array_prod(dimensions())));
901
902 EIGEN_IF_CONSTEXPR (RunningOnGPU) {
903 if (m_result) {
904 return internal::pload<PacketReturnType>(m_result + index);
905 }
906 }
907
908 EIGEN_ALIGN_TO_BOUNDARY(internal::unpacket_traits<PacketReturnType>::alignment)
909 std::remove_const_t<CoeffReturnType> values[PacketSize];
910 // Materializing reduction results before storing them avoids a GCC ICE in emit_move_insn. Fixed in GCC 6.5, 7.4,
911 // and 8.1. See issue #1647 and https://gcc.gnu.org/bugzilla/show_bug.cgi?id=85496.
912 EIGEN_IF_CONSTEXPR (ReducingInnerMostDims) {
913 const Index num_values_to_reduce = (static_cast<int>(Layout) == static_cast<int>(ColMajor))
914 ? m_preservedStrides[0]
915 : m_preservedStrides[NumPreservedStrides - 1];
916 const Index firstIndex = firstInput(index);
917 for (Index i = 0; i < PacketSize; ++i) {
918 Op reducer(m_reducer);
919 const CoeffReturnType value = internal::InnerMostDimReducer<Self, Op>::reduce(
920 *this, firstIndex + i * num_values_to_reduce, num_values_to_reduce, reducer);
921 values[i] = value;
922 }
923 } else EIGEN_IF_CONSTEXPR (PreservingInnerMostDims) {
924 const Index firstIndex = firstInput(index);
925 constexpr int innermost_dim = (static_cast<int>(Layout) == static_cast<int>(ColMajor)) ? 0 : NumOutputDims - 1;
926 // TBD: extend this to the n innermost dimensions that we preserve.
927 if (((firstIndex % m_dimensions[innermost_dim]) + PacketSize - 1) < m_dimensions[innermost_dim]) {
928 Op reducer(m_reducer);
929 typename Self::PacketReturnType accum = reducer.template initializePacket<typename Self::PacketReturnType>();
930 internal::InnerMostDimPreserver<NumReducedDims - 1, Self, Op>::reduce(*this, firstIndex, reducer, &accum);
931 return reducer.finalizePacket(accum);
932 } else {
933 for (int i = 0; i < PacketSize; ++i) {
934 values[i] = coeff(index + i);
935 }
936 }
937 } else if (m_reducingInnerMostDims) {
938 // Mirror the static `ReducingInnerMostDims` packet path so we don't fall
939 // back to PS coeff() calls that route through the scalar GenericDimReducer.
940 const Index firstIndex = index * m_numValuesToReduce;
941 for (Index i = 0; i < PacketSize; ++i) {
942 Op reducer(m_reducer);
943 const CoeffReturnType value = internal::InnerMostDimReducer<Self, Op>::reduce(
944 *this, firstIndex + i * m_numValuesToReduce, m_numValuesToReduce, reducer);
945 values[i] = value;
946 }
947 } else {
948 for (int i = 0; i < PacketSize; ++i) {
949 values[i] = coeff(index + i);
950 }
951 }
952 PacketReturnType rslt = internal::pload<PacketReturnType>(values);
953 return rslt;
954 }
955
956 // Must be called after evalSubExprsIfNeeded().
957 EIGEN_DEVICE_FUNC EIGEN_STRONG_INLINE TensorOpCost costPerCoeff(bool vectorized) const {
958 EIGEN_IF_CONSTEXPR (RunningFullReduction) {
959 if (m_result) {
960 return TensorOpCost(sizeof(CoeffReturnType), 0, 0, vectorized, PacketSize);
961 }
962 }
963 const Index num_values_to_reduce = internal::array_prod(m_reducedDims);
964 const double compute_cost = num_values_to_reduce * internal::functor_traits<Op>::Cost;
965 return m_impl.costPerCoeff(vectorized) * num_values_to_reduce +
966 TensorOpCost(0, 0, compute_cost, vectorized, PacketSize);
967 }
968
969 EIGEN_DEVICE_FUNC EvaluatorPointerType data() const { return m_result; }
970 EIGEN_DEVICE_FUNC const TensorEvaluator<ArgType, Device>& impl() const { return m_impl; }
971 EIGEN_DEVICE_FUNC const Device& device() const { return m_device; }
972
973 private:
974 template <int, typename, typename>
975 friend struct internal::GenericDimReducer;
976 template <typename, typename, bool, bool>
977 friend struct internal::InnerMostDimReducer;
978 template <int, typename, typename, bool>
979 friend struct internal::InnerMostDimPreserver;
980 template <typename S, typename O, typename D, bool V>
981 friend struct internal::FullReducer;
982#if defined(EIGEN_USE_GPU) && (defined(EIGEN_GPUCC))
983 template <int B, int N, typename S, typename R, typename I_>
984 KERNEL_FRIEND void internal::FullReductionKernel(R, const S, I_, typename S::CoeffReturnType*, unsigned int*);
985#if defined(EIGEN_GPUCC)
986 template <typename S, typename R, typename I_>
987 KERNEL_FRIEND void internal::ReductionInitFullReduxKernelHalfFloat(R, const S, I_, half*);
988 template <int B, int N, typename S, typename R, typename I_>
989 KERNEL_FRIEND void internal::FullReductionKernelHalfFloat(R, const S, I_, half*, half*);
990 template <int NPT, typename S, typename R, typename I_>
991 KERNEL_FRIEND void internal::InnerReductionKernelHalfFloat(R, const S, I_, I_, half*);
992#endif
993 template <int NPT, typename S, typename R, typename I_>
994 KERNEL_FRIEND void internal::InnerReductionKernel(R, const S, I_, I_, typename S::CoeffReturnType*);
995
996 template <int NPT, typename S, typename R, typename I_>
997 KERNEL_FRIEND void internal::OuterReductionKernel(R, const S, I_, I_, typename S::CoeffReturnType*);
998#endif
999
1000#if defined(EIGEN_USE_SYCL)
1001 template <typename Evaluator_, typename Op__>
1002 friend class TensorSycl::internal::GenericNondeterministicReducer;
1003 // SYCL needs the generic reducer when the reduction is neither inner, outer, nor full.
1004 template <typename, typename, typename>
1005 friend struct internal::GenericReducer;
1006#endif
1007
1008 template <typename S, typename O, typename D>
1009 friend struct internal::InnerReducer;
1010
1011 // Returns the Index in the input tensor of the first value that needs to be
1012 // used to compute the reduction at output index "index".
1013 EIGEN_DEVICE_FUNC EIGEN_STRONG_INLINE Index firstInput(Index index) const {
1014 EIGEN_IF_CONSTEXPR (ReducingInnerMostDims) {
1015 EIGEN_IF_CONSTEXPR (static_cast<int>(Layout) == static_cast<int>(ColMajor)) {
1016 return index * m_preservedStrides[0];
1017 } else {
1018 return index * m_preservedStrides[NumPreservedStrides - 1];
1019 }
1020 }
1021 // TBD: optimize the case where we preserve the innermost dimensions.
1022 Index startInput = 0;
1023 EIGEN_IF_CONSTEXPR (static_cast<int>(Layout) == static_cast<int>(ColMajor)) {
1024 for (int i = NumOutputDims - 1; i > 0; --i) {
1025 // This is index_i in the output tensor.
1026 const Index idx = index / m_outputStrides[i];
1027 startInput += idx * m_preservedStrides[i];
1028 index -= idx * m_outputStrides[i];
1029 }
1030 EIGEN_IF_CONSTEXPR (PreservingInnerMostDims) {
1031 eigen_assert(m_preservedStrides[0] == 1);
1032 startInput += index;
1033 } else {
1034 startInput += index * m_preservedStrides[0];
1035 }
1036 } else {
1037 for (int i = 0; i < NumOutputDims - 1; ++i) {
1038 // This is index_i in the output tensor.
1039 const Index idx = index / m_outputStrides[i];
1040 startInput += idx * m_preservedStrides[i];
1041 index -= idx * m_outputStrides[i];
1042 }
1043 EIGEN_IF_CONSTEXPR (PreservingInnerMostDims) {
1044 eigen_assert(m_preservedStrides[NumPreservedStrides - 1] == 1);
1045 startInput += index;
1046 } else {
1047 startInput += index * m_preservedStrides[NumPreservedStrides - 1];
1048 }
1049 }
1050 return startInput;
1051 }
1052
1053 // Bitmap indicating if an input dimension is reduced or not.
1054 array<bool, NumInputDims> m_reduced;
1055 // Dimensions of the output of the operation.
1056 Dimensions m_dimensions;
1057 // Precomputed strides for the output tensor.
1058 // Avoid zero-sized arrays, since element access fails to compile on GPU.
1059 array<Index, (std::max)(NumOutputDims, 1)> m_outputStrides;
1060 array<internal::TensorIntDivisor<Index>, (std::max)(NumOutputDims, 1)> m_fastOutputStrides;
1061 array<Index, (std::max)(NumPreservedStrides, 1)> m_preservedStrides;
1062 // Map from output to input dimension index.
1063 array<Index, (std::max)(NumOutputDims, 1)> m_output_to_input_dim_map;
1064 // How many values go into each reduction
1065 Index m_numValuesToReduce;
1066
1067 // Runtime mirror of `ReducingInnerMostDims` (set when Dims is non-static).
1068 bool m_reducingInnerMostDims;
1069
1070 // Subset of strides of the input tensor for the reduced dimensions.
1071 // Indexed by reduced dimensions.
1072 array<Index, NumReducedDims> m_reducedStrides;
1073 // Size of the input dimensions that are reduced.
1074 // Indexed by reduced dimensions.
1075 array<Index, NumReducedDims> m_reducedDims;
1076
1077 // Evaluator for the input expression.
1078 TensorEvaluator<ArgType, Device> m_impl;
1079
1080 // Operation to apply for computing the reduction.
1081 Op m_reducer;
1082
1083 EvaluatorPointerType m_result;
1084
1085 const Device EIGEN_DEVICE_REF m_device;
1086};
1087
1088template <typename Op, typename Dims, typename ArgType, template <class> class MakePointer_, typename Device>
1089struct TensorEvaluator<const TensorReductionOp<Op, Dims, ArgType, MakePointer_>, Device>
1090 : public TensorReductionEvaluatorBase<const TensorReductionOp<Op, Dims, ArgType, MakePointer_>, Device> {
1091 typedef TensorReductionEvaluatorBase<const TensorReductionOp<Op, Dims, ArgType, MakePointer_>, Device> Base;
1092 EIGEN_STRONG_INLINE TensorEvaluator(const typename Base::XprType& op, const Device& device) : Base(op, device) {}
1093};
1094
1095template <typename Op, typename Dims, typename ArgType, template <class> class MakePointer_>
1096struct TensorEvaluator<const TensorReductionOp<Op, Dims, ArgType, MakePointer_>, Eigen::SyclDevice>
1097 : public TensorReductionEvaluatorBase<const TensorReductionOp<Op, Dims, ArgType, MakePointer_>, Eigen::SyclDevice> {
1098 typedef TensorReductionEvaluatorBase<const TensorReductionOp<Op, Dims, ArgType, MakePointer_>, Eigen::SyclDevice>
1099 Base;
1100 EIGEN_STRONG_INLINE TensorEvaluator(const typename Base::XprType& op, const Eigen::SyclDevice& device)
1101 : Base(op, device) {}
1102 // The base coeff function uses a recursive method that is not standard layout and cannot be used in
1103 // SYCL kernels, so it must be overridden.
1104 EIGEN_DEVICE_FUNC EIGEN_STRONG_INLINE typename Base::CoeffReturnType coeff(typename Base::Index index) const {
1105 return *(this->data() + index);
1106 }
1107 // The base packet function uses a recursive method that is not standard layout and cannot be used in
1108 // SYCL kernels, so it must be overridden.
1109 template <int LoadMode>
1110 EIGEN_DEVICE_FUNC EIGEN_STRONG_INLINE typename Base::PacketReturnType packet(typename Base::Index index) const {
1111 return internal::pload<typename Base::PacketReturnType>(this->data() + index);
1112 }
1113};
1114
1115} // end namespace Eigen
1116
1117#endif // EIGEN_TENSOR_TENSOR_REDUCTION_H
Eigen::TensorBase
The tensor base class.
Definition TensorForwardDeclarations.h:69
Eigen::TensorReductionOp
Tensor reduction class.
Definition TensorReduction.h:528
Eigen::ColMajor
ColMajor
Eigen::RowMajor
RowMajor
Eigen
Namespace containing all symbols from the Eigen library.
Eigen::TensorEvaluator
The tensor evaluator class.
Definition TensorEvaluator.h:47
Eigen::internal::ReductionReturnType
Definition TensorReduction.h:511