11 #ifndef EIGEN_CXX11_TENSOR_TENSOR_REDUCTION_H
12 #define EIGEN_CXX11_TENSOR_TENSOR_REDUCTION_H
17 #if defined(__clang__) && (defined(__CUDA__) || defined(__HIP__))
18 #define KERNEL_FRIEND friend __global__ EIGEN_HIP_LAUNCH_BOUNDS_1024
20 #define KERNEL_FRIEND friend
24 #include "./InternalHeaderCheck.h"
37 template<
typename Op,
typename Dims,
typename XprType,
template <
class>
class MakePointer_ >
38 struct traits<TensorReductionOp<Op, Dims, XprType, MakePointer_> >
41 typedef traits<XprType> XprTraits;
42 typedef typename XprTraits::Scalar Scalar;
43 typedef typename XprTraits::StorageKind StorageKind;
44 typedef typename XprTraits::Index
Index;
45 typedef typename XprType::Nested Nested;
46 static constexpr
int NumDimensions = XprTraits::NumDimensions - array_size<Dims>::value;
47 static constexpr
int Layout = XprTraits::Layout;
48 typedef typename XprTraits::PointerType PointerType;
50 template <
class T>
struct MakePointer {
52 typedef MakePointer_<T> MakePointerT;
53 typedef typename MakePointerT::Type Type;
57 template<
typename Op,
typename Dims,
typename XprType,
template <
class>
class MakePointer_>
58 struct eval<TensorReductionOp<Op, Dims, XprType, MakePointer_>,
Eigen::Dense>
60 typedef const TensorReductionOp<Op, Dims, XprType, MakePointer_>& type;
63 template<
typename Op,
typename Dims,
typename XprType,
template <
class>
class MakePointer_>
64 struct nested<TensorReductionOp<Op, Dims, XprType, MakePointer_>, 1, typename eval<TensorReductionOp<Op, Dims, XprType, MakePointer_> >::type>
66 typedef TensorReductionOp<Op, Dims, XprType, MakePointer_> type;
70 template <
typename OutputDims>
struct DimInitializer {
71 template <
typename InputDims,
typename ReducedDims> EIGEN_DEVICE_FUNC
72 static void run(
const InputDims& input_dims,
73 const array<
bool, internal::array_size<InputDims>::value>& reduced,
74 OutputDims* output_dims, ReducedDims* reduced_dims) {
75 const int NumInputDims = internal::array_size<InputDims>::value;
78 for (
int i = 0; i < NumInputDims; ++i) {
80 (*reduced_dims)[reduceIndex] = input_dims[i];
83 (*output_dims)[outputIndex] = input_dims[i];
90 template <>
struct DimInitializer<Sizes<> > {
91 template <
typename InputDims,
typename Index,
size_t Rank> EIGEN_DEVICE_FUNC
92 static void run(
const InputDims& input_dims,
const array<bool, Rank>&,
93 Sizes<>*, array<Index, Rank>* reduced_dims) {
94 const int NumInputDims = internal::array_size<InputDims>::value;
95 for (
int i = 0; i < NumInputDims; ++i) {
96 (*reduced_dims)[i] = input_dims[i];
102 template <
typename ReducedDims,
int NumTensorDims,
int Layout>
103 struct are_inner_most_dims {
104 static const bool value =
false;
106 template <
typename ReducedDims,
int NumTensorDims,
int Layout>
107 struct preserve_inner_most_dims {
108 static const bool value =
false;
111 template <
typename ReducedDims,
int NumTensorDims>
112 struct are_inner_most_dims<ReducedDims, NumTensorDims,
ColMajor>{
113 static const bool tmp1 = indices_statically_known_to_increase<ReducedDims>();
114 static const bool tmp2 = index_statically_eq<ReducedDims>(0, 0);
115 static const bool tmp3 = index_statically_eq<ReducedDims>(array_size<ReducedDims>::value-1, array_size<ReducedDims>::value-1);
116 static const bool value = tmp1 & tmp2 & tmp3;
118 template <
typename ReducedDims,
int NumTensorDims>
119 struct are_inner_most_dims<ReducedDims, NumTensorDims,
RowMajor>{
120 static const bool tmp1 = indices_statically_known_to_increase<ReducedDims>();
121 static const bool tmp2 = index_statically_eq<ReducedDims>(0, NumTensorDims - array_size<ReducedDims>::value);
122 static const bool tmp3 = index_statically_eq<ReducedDims>(array_size<ReducedDims>::value - 1, NumTensorDims - 1);
123 static const bool value = tmp1 & tmp2 & tmp3;
126 template <
typename ReducedDims,
int NumTensorDims>
127 struct preserve_inner_most_dims<ReducedDims, NumTensorDims,
ColMajor>{
128 static const bool tmp1 = indices_statically_known_to_increase<ReducedDims>();
129 static const bool tmp2 = index_statically_gt<ReducedDims>(0, 0);
130 static const bool value = tmp1 & tmp2;
133 template <
typename ReducedDims,
int NumTensorDims>
134 struct preserve_inner_most_dims<ReducedDims, NumTensorDims,
RowMajor>{
135 static const bool tmp1 = indices_statically_known_to_increase<ReducedDims>();
136 static const bool tmp2 = index_statically_lt<ReducedDims>(array_size<ReducedDims>::value - 1, NumTensorDims - 1);
137 static const bool value = tmp1 & tmp2;
141 template <
int DimIndex,
typename Self,
typename Op>
142 struct GenericDimReducer {
143 static EIGEN_DEVICE_FUNC EIGEN_STRONG_INLINE
void reduce(
const Self&
self,
typename Self::Index firstIndex, Op& reducer,
typename Self::CoeffReturnType* accum) {
144 EIGEN_STATIC_ASSERT((DimIndex > 0), YOU_MADE_A_PROGRAMMING_MISTAKE);
145 for (
int j = 0; j <
self.m_reducedDims[DimIndex]; ++j) {
146 const typename Self::Index input = firstIndex + j *
self.m_reducedStrides[DimIndex];
147 GenericDimReducer<DimIndex-1, Self, Op>::reduce(
self, input, reducer, accum);
151 template <
typename Self,
typename Op>
152 struct GenericDimReducer<0, Self, Op> {
153 static EIGEN_DEVICE_FUNC EIGEN_STRONG_INLINE
void reduce(
const Self&
self,
typename Self::Index firstIndex, Op& reducer,
typename Self::CoeffReturnType* accum) {
154 for (
int j = 0; j <
self.m_reducedDims[0]; ++j) {
155 const typename Self::Index input = firstIndex + j *
self.m_reducedStrides[0];
156 reducer.reduce(
self.m_impl.coeff(input), accum);
160 template <
typename Self,
typename Op>
161 struct GenericDimReducer<-1, Self, Op> {
162 static EIGEN_DEVICE_FUNC EIGEN_STRONG_INLINE
void reduce(
const Self&
self,
typename Self::Index index, Op& reducer,
typename Self::CoeffReturnType* accum) {
163 reducer.reduce(
self.m_impl.coeff(index), accum);
167 template <
typename Self,
typename Op,
bool Vectorizable = (Self::InputPacketAccess && Self::ReducerTraits::PacketAccess),
168 bool UseTreeReduction = (!Self::ReducerTraits::IsStateful &&
169 !Self::ReducerTraits::IsExactlyAssociative &&
174 struct InnerMostDimReducer {
175 static EIGEN_DEVICE_FUNC EIGEN_STRONG_INLINE
typename Self::CoeffReturnType reduce(
const Self&
self,
typename Self::Index firstIndex,
typename Self::Index numValuesToReduce, Op& reducer) {
176 typename Self::CoeffReturnType accum = reducer.initialize();
177 for (
typename Self::Index j = 0; j < numValuesToReduce; ++j) {
178 reducer.reduce(
self.m_impl.coeff(firstIndex + j), &accum);
180 return reducer.finalize(accum);
184 template <
typename Self,
typename Op>
185 struct InnerMostDimReducer<Self, Op, true, false> {
186 static EIGEN_DEVICE_FUNC EIGEN_STRONG_INLINE
typename Self::CoeffReturnType reduce(
const Self&
self,
typename Self::Index firstIndex,
typename Self::Index numValuesToReduce, Op& reducer0) {
187 using Index =
typename Self::Index;
188 constexpr
Index packetSize = internal::unpacket_traits<typename Self::PacketReturnType>::size;
190 typename Self::PacketReturnType paccum0 = reducer0.template initializePacket<typename Self::PacketReturnType>();
191 if (!Self::ReducerTraits::IsStateful && numValuesToReduce >= 4*packetSize) {
192 const Index VectorizedSize4 = (numValuesToReduce / (4*packetSize)) * (4*packetSize);
193 typename Self::PacketReturnType paccum1 = reducer0.template initializePacket<typename Self::PacketReturnType>();
194 typename Self::PacketReturnType paccum2 = reducer0.template initializePacket<typename Self::PacketReturnType>();
195 typename Self::PacketReturnType paccum3 = reducer0.template initializePacket<typename Self::PacketReturnType>();
196 const Index offset0 = firstIndex;
197 const Index offset1 = firstIndex + packetSize;
198 const Index offset2 = firstIndex + 2*packetSize;
199 const Index offset3 = firstIndex + 3*packetSize;
200 for (
Index j = 0; j < VectorizedSize4; j += 4*packetSize) {
201 reducer0.reducePacket(
self.m_impl.template packet<Unaligned>(offset0 + j), &paccum0);
202 reducer0.reducePacket(
self.m_impl.template packet<Unaligned>(offset1 + j), &paccum1);
203 reducer0.reducePacket(
self.m_impl.template packet<Unaligned>(offset2 + j), &paccum2);
204 reducer0.reducePacket(
self.m_impl.template packet<Unaligned>(offset3 + j), &paccum3);
206 reducer0.reducePacket(paccum1, &paccum0);
207 reducer0.reducePacket(paccum2, &paccum0);
208 reducer0.reducePacket(paccum3, &paccum0);
209 start = VectorizedSize4;
211 if (start <= (numValuesToReduce - packetSize)) {
212 const Index VectorizedSize = (numValuesToReduce / packetSize) * packetSize;
213 for (
Index j = start; j < VectorizedSize; j += packetSize) {
214 reducer0.reducePacket(
self.m_impl.template packet<Unaligned>(firstIndex + j), &paccum0);
216 start = VectorizedSize;
218 typename Self::CoeffReturnType accum = reducer0.initialize();
219 for (
Index j = start; j < numValuesToReduce; ++j) {
220 reducer0.reduce(
self.m_impl.coeff(firstIndex + j), &accum);
222 return reducer0.finalizeBoth(accum, paccum0);
227 #if !defined(EIGEN_HIPCC)
232 template <
typename T>
233 EIGEN_DEVICE_FUNC
inline Index LeafSize() {
return 1024; }
235 EIGEN_DEVICE_FUNC
inline Index LeafSize<half>() {
return 200; }
237 EIGEN_DEVICE_FUNC
inline Index LeafSize<bfloat16>() {
return 128; }
239 template <
typename Self,
typename Op>
240 struct InnerMostDimReducer<Self, Op, false, true> {
241 static EIGEN_DEVICE_FUNC EIGEN_STRONG_INLINE
typename Self::CoeffReturnType
242 reduce(
const Self&
self,
typename Self::Index firstIndex,
243 typename Self::Index numValuesToReduce, Op& reducer) {
244 const Index kLeafSize = LeafSize<typename Self::CoeffReturnType>();
245 typename Self::CoeffReturnType accum = reducer.initialize();
246 if (numValuesToReduce > kLeafSize) {
247 const typename Self::Index half = numValuesToReduce / 2;
249 reducer.reduce(reduce(
self, firstIndex, half, reducer), &accum);
251 reduce(
self, firstIndex + half, numValuesToReduce - half, reducer),
253 return reducer.finalize(accum);
255 return InnerMostDimReducer<Self, Op, false, false>::reduce(
self, firstIndex, numValuesToReduce, reducer);
260 template <
typename Self,
typename Op>
261 struct InnerMostDimReducer<Self, Op, true, true> {
262 static EIGEN_DEVICE_FUNC EIGEN_STRONG_INLINE
typename Self::CoeffReturnType
263 reduce(
const Self&
self,
typename Self::Index firstIndex,
264 typename Self::Index numValuesToReduce, Op& reducer) {
265 const Index kLeafSize = LeafSize<typename Self::CoeffReturnType>();
266 const typename Self::Index packetSize =
267 internal::unpacket_traits<typename Self::PacketReturnType>::size;
268 typename Self::CoeffReturnType accum = reducer.initialize();
269 if (numValuesToReduce > packetSize * kLeafSize) {
271 const typename Self::Index split =
273 divup(firstIndex + divup(numValuesToReduce,
typename Self::Index(2)),
275 const typename Self::Index num_left =
276 numext::mini(split - firstIndex, numValuesToReduce);
277 reducer.reduce(reduce(
self, firstIndex, num_left, reducer), &accum);
278 if (num_left < numValuesToReduce) {
280 reduce(
self, split, numValuesToReduce - num_left, reducer), &accum);
282 return reducer.finalize(accum);
284 return InnerMostDimReducer<Self, Op, true, false>::reduce(
self, firstIndex, numValuesToReduce, reducer);
290 template <
int DimIndex,
typename Self,
typename Op,
bool vectorizable = (Self::InputPacketAccess && Self::ReducerTraits::PacketAccess)>
291 struct InnerMostDimPreserver {
292 static EIGEN_DEVICE_FUNC EIGEN_STRONG_INLINE
void reduce(
const Self&,
typename Self::Index, Op&,
typename Self::PacketReturnType*) {
293 eigen_assert(
false &&
"should never be called");
297 template <
int DimIndex,
typename Self,
typename Op>
298 struct InnerMostDimPreserver<DimIndex, Self, Op, true> {
299 static EIGEN_DEVICE_FUNC EIGEN_STRONG_INLINE
void reduce(
const Self&
self,
typename Self::Index firstIndex, Op& reducer,
typename Self::PacketReturnType* accum) {
300 EIGEN_STATIC_ASSERT((DimIndex > 0), YOU_MADE_A_PROGRAMMING_MISTAKE);
301 for (
typename Self::Index j = 0; j <
self.m_reducedDims[DimIndex]; ++j) {
302 const typename Self::Index input = firstIndex + j *
self.m_reducedStrides[DimIndex];
303 InnerMostDimPreserver<DimIndex-1, Self, Op>::reduce(
self, input, reducer, accum);
308 template <
typename Self,
typename Op>
309 struct InnerMostDimPreserver<0, Self, Op, true> {
310 static EIGEN_DEVICE_FUNC EIGEN_STRONG_INLINE
void reduce(
const Self&
self,
typename Self::Index firstIndex, Op& reducer0,
typename Self::PacketReturnType* accum0) {
311 using Index =
typename Self::Index;
312 const Index stride =
self.m_reducedStrides[0];
313 const Index size =
self.m_reducedDims[0];
314 if (!Self::ReducerTraits::IsStateful && size >= 16) {
315 const Index unrolled_size4 = (size / 4) * 4;
316 typename Self::PacketReturnType accum1 = reducer0.template initializePacket<typename Self::PacketReturnType>();
317 typename Self::PacketReturnType accum2 = reducer0.template initializePacket<typename Self::PacketReturnType>();
318 typename Self::PacketReturnType accum3 = reducer0.template initializePacket<typename Self::PacketReturnType>();
319 for (
Index j = 0; j < unrolled_size4; j += 4) {
320 const Index input0 = firstIndex + j * stride;
321 reducer0.reducePacket(
self.m_impl.template packet<Unaligned>(input0), accum0);
322 const Index input1 = firstIndex + (j+1) * stride;
323 reducer0.reducePacket(
self.m_impl.template packet<Unaligned>(input1), &accum1);
324 const Index input2 = firstIndex + (j+2) * stride;
325 reducer0.reducePacket(
self.m_impl.template packet<Unaligned>(input2), &accum2);
326 const Index input3 = firstIndex + (j+3) * stride;
327 reducer0.reducePacket(
self.m_impl.template packet<Unaligned>(input3), &accum3);
329 reducer0.reducePacket(accum1, accum0);
330 reducer0.reducePacket(accum2, accum0);
331 reducer0.reducePacket(accum3, accum0);
332 for (
Index j = unrolled_size4; j < size; ++j) {
333 Index input = firstIndex + j * stride;
334 reducer0.reducePacket(
self.m_impl.template packet<Unaligned>(input), accum0);
337 for (
Index j = 0; j < size; ++j) {
338 Index input = firstIndex + j * stride;
339 reducer0.reducePacket(
self.m_impl.template packet<Unaligned>(input), accum0);
344 template <
typename Self,
typename Op>
345 struct InnerMostDimPreserver<-1, Self, Op, true> {
346 static EIGEN_DEVICE_FUNC EIGEN_STRONG_INLINE
void reduce(
const Self&,
typename Self::Index, Op&,
typename Self::PacketReturnType*) {
347 eigen_assert(
false &&
"should never be called");
352 template <
typename Self,
typename Op,
typename Device,
bool Vectorizable = (Self::InputPacketAccess && Self::ReducerTraits::PacketAccess)>
354 static constexpr
bool HasOptimizedImplementation =
false;
356 static EIGEN_DEVICE_FUNC
void run(
const Self&
self, Op& reducer,
const Device&,
typename Self::EvaluatorPointerType output) {
357 const typename Self::Index num_coeffs = array_prod(
self.m_impl.dimensions());
358 *output = InnerMostDimReducer<Self, Op, Vectorizable>::reduce(
self, 0, num_coeffs, reducer);
363 #ifdef EIGEN_USE_THREADS
365 template <
typename Self,
typename Op,
366 bool Vectorizable = (Self::InputPacketAccess && Self::ReducerTraits::PacketAccess)>
367 struct FullReducerShard {
368 static EIGEN_DEVICE_FUNC EIGEN_STRONG_INLINE
void run(
const Self&
self,
typename Self::Index firstIndex,
369 typename Self::Index numValuesToReduce, Op& reducer,
370 typename Self::CoeffReturnType* output) {
371 *output = InnerMostDimReducer<Self, Op, Vectorizable>::reduce(
372 self, firstIndex, numValuesToReduce, reducer);
377 template <
typename Self,
typename Op,
bool Vectorizable>
378 struct FullReducer<Self, Op, ThreadPoolDevice, Vectorizable> {
379 static constexpr
bool HasOptimizedImplementation = !Self::ReducerTraits::IsStateful;
380 static constexpr
Index PacketSize =
381 unpacket_traits<typename Self::PacketReturnType>::size;
384 static void run(
const Self&
self, Op& reducer,
const ThreadPoolDevice& device,
385 typename Self::CoeffReturnType* output) {
386 typedef typename Self::Index
Index;
387 const Index num_coeffs = array_prod(
self.m_impl.dimensions());
388 if (num_coeffs == 0) {
389 *output = reducer.finalize(reducer.initialize());
392 const TensorOpCost cost =
393 self.m_impl.costPerCoeff(Vectorizable) +
394 TensorOpCost(0, 0, internal::functor_traits<Op>::Cost, Vectorizable,
396 const Index num_threads = TensorCostModel<ThreadPoolDevice>::numThreads(
397 num_coeffs, cost, device.numThreads());
398 if (num_threads == 1) {
400 InnerMostDimReducer<Self, Op, Vectorizable>::reduce(
self, 0, num_coeffs, reducer);
403 const Index blocksize = num_coeffs / num_threads;
404 const Index numblocks = blocksize > 0 ? num_coeffs / blocksize : 0;
405 eigen_assert(num_coeffs >= numblocks * blocksize);
407 Barrier barrier(internal::convert_index<unsigned int>(numblocks));
408 MaxSizeVector<typename Self::CoeffReturnType> shards(numblocks, reducer.initialize());
409 for (
Index i = 0; i < numblocks; ++i) {
410 device.enqueue_with_barrier(&barrier, &FullReducerShard<Self, Op, Vectorizable>::run,
411 self, i * blocksize, blocksize, reducer,
414 typename Self::CoeffReturnType finalShard;
415 if (numblocks * blocksize < num_coeffs) {
416 finalShard = InnerMostDimReducer<Self, Op, Vectorizable>::reduce(
417 self, numblocks * blocksize, num_coeffs - numblocks * blocksize,
420 finalShard = reducer.initialize();
424 for (
Index i = 0; i < numblocks; ++i) {
425 reducer.reduce(shards[i], &finalShard);
427 *output = reducer.finalize(finalShard);
435 template <
typename Self,
typename Op,
typename Device>
436 struct InnerReducer {
437 static constexpr
bool HasOptimizedImplementation =
false;
439 EIGEN_DEVICE_FUNC
static bool run(
const Self&, Op&,
const Device&,
typename Self::CoeffReturnType*,
typename Self::Index,
typename Self::Index) {
440 eigen_assert(
false &&
"Not implemented");
446 template <
typename Self,
typename Op,
typename Device>
447 struct OuterReducer {
448 static constexpr
bool HasOptimizedImplementation =
false;
450 EIGEN_DEVICE_FUNC
static bool run(
const Self&, Op&,
const Device&,
typename Self::CoeffReturnType*,
typename Self::Index,
typename Self::Index) {
451 eigen_assert(
false &&
"Not implemented");
456 #ifdef EIGEN_USE_SYCL
458 template <
typename Self,
typename Op,
typename Device>
459 struct GenericReducer {
460 static constexpr
bool HasOptimizedImplementation =
false;
462 EIGEN_DEVICE_FUNC
static bool run(
const Self&, Op&,
const Device&,
typename Self::CoeffReturnType*,
typename Self::Index,
typename Self::Index) {
463 eigen_assert(
false &&
"Not implemented");
469 #if defined(EIGEN_USE_GPU) && (defined(EIGEN_GPUCC))
470 template <
int B,
int N,
typename S,
typename R,
typename I_>
471 __global__ EIGEN_HIP_LAUNCH_BOUNDS_1024
void FullReductionKernel(R,
const S, I_,
typename S::CoeffReturnType*,
unsigned int*);
474 #if defined(EIGEN_HAS_GPU_FP16)
475 template <
typename S,
typename R,
typename I_>
476 __global__ EIGEN_HIP_LAUNCH_BOUNDS_1024
void ReductionInitFullReduxKernelHalfFloat(R,
const S, I_, internal::packet_traits<half>::type*);
477 template <
int B,
int N,
typename S,
typename R,
typename I_>
478 __global__ EIGEN_HIP_LAUNCH_BOUNDS_1024
void FullReductionKernelHalfFloat(R,
const S, I_, half*, internal::packet_traits<half>::type*);
479 template <
int NPT,
typename S,
typename R,
typename I_>
480 __global__ EIGEN_HIP_LAUNCH_BOUNDS_1024
void InnerReductionKernelHalfFloat(R,
const S, I_, I_, half*);
484 template <
int NPT,
typename S,
typename R,
typename I_>
485 __global__ EIGEN_HIP_LAUNCH_BOUNDS_1024
void InnerReductionKernel(R,
const S, I_, I_,
typename S::CoeffReturnType*);
487 template <
int NPT,
typename S,
typename R,
typename I_>
488 __global__ EIGEN_HIP_LAUNCH_BOUNDS_1024
void OuterReductionKernel(R,
const S, I_, I_,
typename S::CoeffReturnType*);
499 template <
typename Op,
typename CoeffReturnType>
500 struct ReductionReturnType {
501 #if defined(EIGEN_USE_SYCL)
502 typedef std::remove_const_t<decltype(std::declval<Op>().initialize())> type;
504 typedef std::remove_const_t<CoeffReturnType> type;
511 template <
typename Op,
typename Dims,
typename XprType,
template <
class>
class MakePointer_>
512 class TensorReductionOp :
public TensorBase<TensorReductionOp<Op, Dims, XprType, MakePointer_>, ReadOnlyAccessors> {
514 typedef typename Eigen::internal::traits<TensorReductionOp>::Scalar Scalar;
516 typedef std::remove_const_t<typename XprType::CoeffReturnType> CoeffReturnType;
517 typedef typename Eigen::internal::nested<TensorReductionOp>::type Nested;
518 typedef typename Eigen::internal::traits<TensorReductionOp>::StorageKind StorageKind;
519 typedef typename Eigen::internal::traits<TensorReductionOp>::Index
Index;
521 EIGEN_DEVICE_FUNC EIGEN_STRONG_INLINE
522 TensorReductionOp(
const XprType& expr,
const Dims& dims) : m_expr(expr), m_dims(dims)
524 EIGEN_DEVICE_FUNC EIGEN_STRONG_INLINE
525 TensorReductionOp(
const XprType& expr,
const Dims& dims,
const Op& reducer) : m_expr(expr), m_dims(dims), m_reducer(reducer)
528 EIGEN_DEVICE_FUNC EIGEN_STRONG_INLINE
529 const XprType& expression()
const {
return m_expr; }
530 EIGEN_DEVICE_FUNC EIGEN_STRONG_INLINE
531 const Dims& dims()
const {
return m_dims; }
532 EIGEN_DEVICE_FUNC EIGEN_STRONG_INLINE
533 const Op& reducer()
const {
return m_reducer; }
536 typename XprType::Nested m_expr;
541 template<
typename ArgType,
typename Device>
542 struct TensorReductionEvaluatorBase;
545 template<
typename Op,
typename Dims,
typename ArgType,
template <
class>
class MakePointer_,
typename Device>
546 struct TensorReductionEvaluatorBase<const TensorReductionOp<Op, Dims, ArgType, MakePointer_>, Device>
548 typedef internal::reducer_traits<Op, Device> ReducerTraits;
549 typedef Dims ReducedDims;
550 typedef TensorReductionOp<Op, Dims, ArgType, MakePointer_> XprType;
551 typedef typename XprType::Index
Index;
552 typedef ArgType ChildType;
553 typedef typename TensorEvaluator<ArgType, Device>::Dimensions InputDimensions;
554 static constexpr
int NumInputDims = internal::array_size<InputDimensions>::value;
555 static constexpr
int NumReducedDims = internal::array_size<Dims>::value;
556 static constexpr
int NumOutputDims = NumInputDims - NumReducedDims;
557 typedef std::conditional_t<NumOutputDims==0, Sizes<>, DSizes<Index, NumOutputDims> > Dimensions;
558 typedef typename XprType::Scalar Scalar;
559 typedef TensorReductionEvaluatorBase<const TensorReductionOp<Op, Dims, ArgType, MakePointer_>, Device> Self;
560 static constexpr
bool InputPacketAccess = TensorEvaluator<ArgType, Device>::PacketAccess;
561 typedef typename internal::ReductionReturnType<Op, typename XprType::CoeffReturnType>::type CoeffReturnType;
562 typedef typename PacketType<CoeffReturnType, Device>::type PacketReturnType;
563 static constexpr
Index PacketSize = PacketType<CoeffReturnType, Device>::size;
565 typedef typename Eigen::internal::traits<XprType>::PointerType TensorPointerType;
566 typedef StorageMemory<CoeffReturnType, Device> Storage;
567 typedef typename Storage::Type EvaluatorPointerType;
571 static constexpr
int NumPreservedStrides = max_n_1<NumOutputDims>::size;
574 #if defined(EIGEN_USE_GPU) && (defined(EIGEN_GPUCC))
575 static constexpr
bool RunningOnGPU = internal::is_same<Device, Eigen::GpuDevice>::value;
576 static constexpr
bool RunningOnSycl =
false;
577 #elif defined(EIGEN_USE_SYCL)
578 static constexpr
bool RunningOnSycl = internal::is_same<internal::remove_all_t<Device>, Eigen::SyclDevice>::value;
579 static constexpr
bool RunningOnGPU =
false;
581 static constexpr
bool RunningOnGPU =
false;
582 static constexpr
bool RunningOnSycl =
false;
585 static constexpr
int Layout = TensorEvaluator<ArgType, Device>::Layout;
588 PacketAccess = Self::InputPacketAccess && ReducerTraits::PacketAccess,
590 PreferBlockAccess =
true,
595 typedef std::remove_const_t<Scalar> ScalarNoConst;
598 typedef internal::TensorBlockNotImplemented TensorBlock;
601 static constexpr
bool ReducingInnerMostDims = internal::are_inner_most_dims<Dims, NumInputDims, Layout>::value;
602 static constexpr
bool PreservingInnerMostDims = internal::preserve_inner_most_dims<Dims, NumInputDims, Layout>::value;
603 static constexpr
bool RunningFullReduction = (NumOutputDims==0);
605 EIGEN_STRONG_INLINE TensorReductionEvaluatorBase(
const XprType& op,
const Device& device)
606 : m_impl(op.expression(), device), m_reducer(op.reducer()), m_result(NULL), m_device(device)
608 EIGEN_STATIC_ASSERT((NumInputDims >= NumReducedDims), YOU_MADE_A_PROGRAMMING_MISTAKE);
609 EIGEN_STATIC_ASSERT((!ReducingInnerMostDims | !PreservingInnerMostDims | (NumReducedDims == NumInputDims)),
610 YOU_MADE_A_PROGRAMMING_MISTAKE);
613 for (
int i = 0; i < NumInputDims; ++i) {
614 m_reduced[i] =
false;
616 for (
int i = 0; i < NumReducedDims; ++i) {
617 eigen_assert(op.dims()[i] >= 0);
618 eigen_assert(op.dims()[i] < NumInputDims);
619 m_reduced[op.dims()[i]] =
true;
622 const typename TensorEvaluator<ArgType, Device>::Dimensions& input_dims = m_impl.dimensions();
623 internal::DimInitializer<Dimensions>::run(input_dims, m_reduced, &m_dimensions, &m_reducedDims);
626 if (NumOutputDims > 0) {
627 if (
static_cast<int>(Layout) ==
static_cast<int>(
ColMajor)) {
628 m_outputStrides[0] = 1;
629 for (
int i = 1; i < NumOutputDims; ++i) {
630 m_outputStrides[i] = m_outputStrides[i - 1] * m_dimensions[i - 1];
631 m_fastOutputStrides[i] = internal::TensorIntDivisor<Index>(m_outputStrides[i]);
634 m_outputStrides[
static_cast<size_t>(NumOutputDims - 1)] = 1;
635 for (
int i = NumOutputDims - 2; i >= 0; --i) {
636 m_outputStrides[i] = m_outputStrides[i + 1] * m_dimensions[i + 1];
637 m_fastOutputStrides[i] = internal::TensorIntDivisor<Index>(m_outputStrides[i]);
643 if (NumInputDims > 0) {
644 array<Index, NumInputDims> input_strides;
645 if (
static_cast<int>(Layout) ==
static_cast<int>(
ColMajor)) {
646 input_strides[0] = 1;
647 for (
int i = 1; i < NumInputDims; ++i) {
648 input_strides[i] = input_strides[i-1] * input_dims[i-1];
651 input_strides.back() = 1;
652 for (
int i = NumInputDims - 2; i >= 0; --i) {
653 input_strides[i] = input_strides[i + 1] * input_dims[i + 1];
659 for (
int i = 0; i < NumInputDims; ++i) {
661 m_reducedStrides[reduceIndex] = input_strides[i];
664 m_preservedStrides[outputIndex] = input_strides[i];
665 m_output_to_input_dim_map[outputIndex] = i;
672 if (NumOutputDims == 0) {
673 m_preservedStrides[0] = internal::array_prod(input_dims);
676 m_numValuesToReduce =
678 ? internal::array_prod(input_dims)
679 : (static_cast<int>(Layout) == static_cast<int>(
ColMajor))
680 ? m_preservedStrides[0]
681 : m_preservedStrides[static_cast<size_t>(NumOutputDims - 1)];
684 EIGEN_DEVICE_FUNC EIGEN_STRONG_INLINE
const Dimensions& dimensions()
const {
return m_dimensions; }
687 bool evalSubExprsIfNeededCommon(EvaluatorPointerType data) {
689 if ((RunningFullReduction && RunningOnSycl) ||(RunningFullReduction &&
690 internal::FullReducer<Self, Op, Device>::HasOptimizedImplementation &&
691 ((RunningOnGPU && (m_device.majorDeviceVersion() >= 3)) ||
693 bool need_assign =
false;
695 m_result =
static_cast<EvaluatorPointerType
>(m_device.get((CoeffReturnType*)m_device.allocate_temp(
sizeof(CoeffReturnType))));
699 Op reducer(m_reducer);
700 internal::FullReducer<Self, Op, Device>::run(*
this, reducer, m_device, data);
705 else if ((RunningOnGPU && (m_device.majorDeviceVersion() >= 3)) || (RunningOnSycl)) {
706 bool reducing_inner_dims =
true;
707 for (
int i = 0; i < NumReducedDims; ++i) {
708 if (
static_cast<int>(Layout) ==
static_cast<int>(
ColMajor)) {
709 reducing_inner_dims &= m_reduced[i];
711 reducing_inner_dims &= m_reduced[NumInputDims - 1 - i];
714 if (internal::InnerReducer<Self, Op, Device>::HasOptimizedImplementation &&
715 (reducing_inner_dims || ReducingInnerMostDims)) {
716 const Index num_values_to_reduce = internal::array_prod(m_reducedDims);
717 const Index num_coeffs_to_preserve = internal::array_prod(m_dimensions);
719 if ((num_coeffs_to_preserve < 1024 && num_values_to_reduce > num_coeffs_to_preserve && num_values_to_reduce > 128) || (RunningOnSycl)) {
720 data =
static_cast<EvaluatorPointerType
>(m_device.get((CoeffReturnType*)m_device.allocate_temp(
sizeof(CoeffReturnType) * num_coeffs_to_preserve)));
727 Op reducer(m_reducer);
729 if (internal::InnerReducer<Self, Op, Device>::run(*
this, reducer, m_device, data, num_values_to_reduce, num_coeffs_to_preserve)) {
731 m_device.deallocate_temp(m_result);
736 return (m_result != NULL);
740 bool preserving_inner_dims =
true;
741 for (
int i = 0; i < NumReducedDims; ++i) {
742 if (
static_cast<int>(Layout) ==
static_cast<int>(
ColMajor)) {
743 preserving_inner_dims &= m_reduced[NumInputDims - 1 - i];
745 preserving_inner_dims &= m_reduced[i];
748 if (internal::OuterReducer<Self, Op, Device>::HasOptimizedImplementation &&
749 preserving_inner_dims) {
750 const Index num_values_to_reduce = internal::array_prod(m_reducedDims);
751 const Index num_coeffs_to_preserve = internal::array_prod(m_dimensions);
753 if ((num_coeffs_to_preserve < 1024 && num_values_to_reduce > num_coeffs_to_preserve && num_values_to_reduce > 32) || (RunningOnSycl)) {
754 data =
static_cast<EvaluatorPointerType
>(m_device.get((CoeffReturnType*)m_device.allocate_temp(
sizeof(CoeffReturnType) * num_coeffs_to_preserve)));
761 Op reducer(m_reducer);
763 if (internal::OuterReducer<Self, Op, Device>::run(*
this, reducer, m_device, data, num_values_to_reduce, num_coeffs_to_preserve)) {
765 m_device.deallocate_temp(m_result);
770 return (m_result != NULL);
773 #if defined(EIGEN_USE_SYCL)
777 const Index num_values_to_reduce = internal::array_prod(m_reducedDims);
778 const Index num_coeffs_to_preserve = internal::array_prod(m_dimensions);
780 data =
static_cast<EvaluatorPointerType
>(m_device.get((CoeffReturnType*)m_device.allocate_temp(
sizeof(CoeffReturnType) * num_coeffs_to_preserve)));
783 Op reducer(m_reducer);
784 internal::GenericReducer<Self, Op, Device>::run(*
this, reducer, m_device, data, num_values_to_reduce, num_coeffs_to_preserve);
785 return (m_result != NULL);
792 #ifdef EIGEN_USE_THREADS
793 template <
typename EvalSubExprsCallback>
796 evalSubExprsIfNeededAsync(EvaluatorPointerType data,
797 EvalSubExprsCallback done) {
798 m_impl.evalSubExprsIfNeededAsync(NULL, [
this, data, done](
bool) {
799 done(evalSubExprsIfNeededCommon(data));
805 bool evalSubExprsIfNeeded(EvaluatorPointerType data) {
806 m_impl.evalSubExprsIfNeeded(NULL);
807 return evalSubExprsIfNeededCommon(data);
810 EIGEN_STRONG_INLINE
void cleanup() {
813 m_device.deallocate_temp(m_result);
818 EIGEN_DEVICE_FUNC EIGEN_STRONG_INLINE CoeffReturnType coeff(
Index index)
const
820 if (( RunningFullReduction || RunningOnGPU) && m_result ) {
821 return *(m_result + index);
823 Op reducer(m_reducer);
824 if (ReducingInnerMostDims || RunningFullReduction) {
825 const Index num_values_to_reduce =
826 (
static_cast<int>(Layout) ==
static_cast<int>(
ColMajor)) ? m_preservedStrides[0] : m_preservedStrides[NumPreservedStrides - 1];
827 return internal::InnerMostDimReducer<Self, Op>::reduce(*
this, firstInput(index),
828 num_values_to_reduce, reducer);
830 typename Self::CoeffReturnType accum = reducer.initialize();
831 internal::GenericDimReducer<NumReducedDims-1, Self, Op>::reduce(*
this, firstInput(index), reducer, &accum);
832 return reducer.finalize(accum);
837 template<
int LoadMode>
838 EIGEN_DEVICE_FUNC EIGEN_STRONG_INLINE PacketReturnType packet(
Index index)
const
840 eigen_assert(index + PacketSize - 1 <
Index(internal::array_prod(dimensions())));
842 if (RunningOnGPU && m_result) {
843 return internal::pload<PacketReturnType>(m_result + index);
846 EIGEN_ALIGN_MAX std::remove_const_t<CoeffReturnType> values[PacketSize];
847 if (ReducingInnerMostDims) {
848 const Index num_values_to_reduce =
849 (
static_cast<int>(Layout) ==
static_cast<int>(
ColMajor)) ? m_preservedStrides[0] : m_preservedStrides[NumPreservedStrides - 1];
850 const Index firstIndex = firstInput(index);
851 for (
Index i = 0; i < PacketSize; ++i) {
852 Op reducer(m_reducer);
853 values[i] = internal::InnerMostDimReducer<Self, Op>::reduce(*
this, firstIndex + i * num_values_to_reduce,
854 num_values_to_reduce, reducer);
856 }
else if (PreservingInnerMostDims) {
857 const Index firstIndex = firstInput(index);
858 const int innermost_dim = (
static_cast<int>(Layout) ==
static_cast<int>(
ColMajor)) ? 0 : NumOutputDims - 1;
860 if (((firstIndex % m_dimensions[innermost_dim]) + PacketSize - 1) < m_dimensions[innermost_dim]) {
861 Op reducer(m_reducer);
862 typename Self::PacketReturnType accum = reducer.template initializePacket<typename Self::PacketReturnType>();
863 internal::InnerMostDimPreserver<NumReducedDims-1, Self, Op>::reduce(*
this, firstIndex, reducer, &accum);
864 return reducer.finalizePacket(accum);
866 for (
int i = 0; i < PacketSize; ++i) {
867 values[i] = coeff(index + i);
871 for (
int i = 0; i < PacketSize; ++i) {
872 values[i] = coeff(index + i);
875 PacketReturnType rslt = internal::pload<PacketReturnType>(values);
880 EIGEN_DEVICE_FUNC EIGEN_STRONG_INLINE TensorOpCost costPerCoeff(
bool vectorized)
const {
881 if (RunningFullReduction && m_result) {
882 return TensorOpCost(
sizeof(CoeffReturnType), 0, 0, vectorized, PacketSize);
884 const Index num_values_to_reduce = internal::array_prod(m_reducedDims);
885 const double compute_cost = num_values_to_reduce * internal::functor_traits<Op>::Cost;
886 return m_impl.costPerCoeff(vectorized) * num_values_to_reduce +
887 TensorOpCost(0, 0, compute_cost, vectorized, PacketSize);
891 EIGEN_DEVICE_FUNC EvaluatorPointerType data()
const {
return m_result; }
892 EIGEN_DEVICE_FUNC
const TensorEvaluator<ArgType, Device>& impl()
const {
return m_impl; }
893 EIGEN_DEVICE_FUNC
const Device& device()
const {
return m_device; }
894 #ifdef EIGEN_USE_SYCL
896 EIGEN_DEVICE_FUNC EIGEN_STRONG_INLINE
void bind(cl::sycl::handler &cgh)
const {
903 template <
int,
typename,
typename>
friend struct internal::GenericDimReducer;
904 template <
typename,
typename,
bool,
bool>
friend struct internal::InnerMostDimReducer;
905 template <
int,
typename,
typename,
bool>
friend struct internal::InnerMostDimPreserver;
906 template <
typename S,
typename O,
typename D,
bool V>
friend struct internal::FullReducer;
907 #ifdef EIGEN_USE_THREADS
908 template <
typename S,
typename O,
bool V>
friend struct internal::FullReducerShard;
910 #if defined(EIGEN_USE_GPU) && (defined(EIGEN_GPUCC))
911 template <
int B,
int N,
typename S,
typename R,
typename I_> KERNEL_FRIEND
void internal::FullReductionKernel(R,
const S, I_,
typename S::CoeffReturnType*,
unsigned int*);
912 #if defined(EIGEN_HAS_GPU_FP16)
913 template <
typename S,
typename R,
typename I_> KERNEL_FRIEND
void internal::ReductionInitFullReduxKernelHalfFloat(R,
const S, I_, internal::packet_traits<Eigen::half>::type*);
914 template <
int B,
int N,
typename S,
typename R,
typename I_> KERNEL_FRIEND
void internal::FullReductionKernelHalfFloat(R,
const S, I_, half*, internal::packet_traits<Eigen::half>::type*);
915 template <
int NPT,
typename S,
typename R,
typename I_> KERNEL_FRIEND
void internal::InnerReductionKernelHalfFloat(R,
const S, I_, I_, half*);
917 template <
int NPT,
typename S,
typename R,
typename I_> KERNEL_FRIEND
void internal::InnerReductionKernel(R,
const S, I_, I_,
typename S::CoeffReturnType*);
919 template <
int NPT,
typename S,
typename R,
typename I_> KERNEL_FRIEND
void internal::OuterReductionKernel(R,
const S, I_, I_,
typename S::CoeffReturnType*);
922 #if defined(EIGEN_USE_SYCL)
923 template <
typename Evaluator_,
typename Op__>
friend class TensorSycl::internal::GenericNondeterministicReducer;
925 template <
typename,
typename,
typename>
friend struct internal::GenericReducer;
929 template <
typename S,
typename O,
typename D>
friend struct internal::InnerReducer;
931 struct BlockIteratorState {
939 EIGEN_DEVICE_FUNC EIGEN_STRONG_INLINE
Index firstInput(
Index index)
const {
940 if (ReducingInnerMostDims) {
941 if (
static_cast<int>(Layout) ==
static_cast<int>(
ColMajor)) {
942 return index * m_preservedStrides[0];
944 return index * m_preservedStrides[NumPreservedStrides - 1];
948 Index startInput = 0;
949 if (
static_cast<int>(Layout) ==
static_cast<int>(
ColMajor)) {
950 for (
int i = NumOutputDims - 1; i > 0; --i) {
952 const Index idx = index / m_outputStrides[i];
953 startInput += idx * m_preservedStrides[i];
954 index -= idx * m_outputStrides[i];
956 if (PreservingInnerMostDims) {
957 eigen_assert(m_preservedStrides[0] == 1);
960 startInput += index * m_preservedStrides[0];
963 for (
int i = 0; i < NumOutputDims - 1; ++i) {
965 const Index idx = index / m_outputStrides[i];
966 startInput += idx * m_preservedStrides[i];
967 index -= idx * m_outputStrides[i];
969 if (PreservingInnerMostDims) {
970 eigen_assert(m_preservedStrides[NumPreservedStrides - 1] == 1);
973 startInput += index * m_preservedStrides[NumPreservedStrides - 1];
980 array<bool, NumInputDims> m_reduced;
982 Dimensions m_dimensions;
984 array<Index, NumOutputDims> m_outputStrides;
985 array<internal::TensorIntDivisor<Index>, NumOutputDims> m_fastOutputStrides;
986 array<Index, NumPreservedStrides> m_preservedStrides;
988 array<Index, NumOutputDims> m_output_to_input_dim_map;
990 Index m_numValuesToReduce;
994 array<Index, NumReducedDims> m_reducedStrides;
997 array<Index, NumReducedDims> m_reducedDims;
1000 TensorEvaluator<ArgType, Device> m_impl;
1005 EvaluatorPointerType m_result;
1007 const Device EIGEN_DEVICE_REF m_device;
1010 template<
typename Op,
typename Dims,
typename ArgType,
template <
class>
class MakePointer_,
typename Device>
1011 struct TensorEvaluator<const TensorReductionOp<Op, Dims, ArgType, MakePointer_>, Device>
1012 :
public TensorReductionEvaluatorBase<const TensorReductionOp<Op, Dims, ArgType, MakePointer_>, Device> {
1013 typedef TensorReductionEvaluatorBase<const TensorReductionOp<Op, Dims, ArgType, MakePointer_>, Device> Base;
1014 EIGEN_STRONG_INLINE TensorEvaluator(
const typename Base::XprType& op,
const Device& device) : Base(op, device){}
1018 template<
typename Op,
typename Dims,
typename ArgType,
template <
class>
class MakePointer_>
1019 struct TensorEvaluator<const TensorReductionOp<Op, Dims, ArgType, MakePointer_>,
Eigen::SyclDevice>
1020 :
public TensorReductionEvaluatorBase<const TensorReductionOp<Op, Dims, ArgType, MakePointer_>, Eigen::SyclDevice> {
1022 typedef TensorReductionEvaluatorBase<const TensorReductionOp<Op, Dims, ArgType, MakePointer_>, Eigen::SyclDevice> Base;
1023 EIGEN_STRONG_INLINE TensorEvaluator(
const typename Base::XprType& op,
const Eigen::SyclDevice& device) : Base(op, device){}
1026 EIGEN_DEVICE_FUNC EIGEN_STRONG_INLINE
typename Base::CoeffReturnType coeff(
typename Base::Index index)
const {
1027 return *(this->data() + index);
1031 template<
int LoadMode>
1032 EIGEN_DEVICE_FUNC EIGEN_STRONG_INLINE
typename Base::PacketReturnType packet(
typename Base::Index index)
const {
1033 return internal::pload<typename Base::PacketReturnType>(this->data() + index);
Namespace containing all symbols from the Eigen library.
EIGEN_DEFAULT_DENSE_INDEX_TYPE Index