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
37 template<
typename Op,
typename Dims,
typename XprType,
template <
class>
class MakePointer_ >
38 struct traits<TensorReductionOp<Op, Dims, XprType, MakePointer_> >
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 {
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<> > {
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) {
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>
239 template <
typename Self,
typename Op>
240 struct InnerMostDimReducer<Self, Op, false, true> {
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> {
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 =
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*) {
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) {
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);
333 Index input = firstIndex +
j * stride;
334 reducer0.reducePacket(
self.m_impl.template packet<Unaligned>(input), accum0);
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*) {
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;
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,
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;
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) {
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) {
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) {
469 #if defined(EIGEN_USE_GPU) && (defined(EIGEN_GPUCC))
470 template <
int B,
int N,
typename S,
typename R,
typename I_>
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_>
479 template <
int NPT,
typename S,
typename R,
typename I_>
484 template <
int NPT,
typename S,
typename R,
typename I_>
487 template <
int NPT,
typename S,
typename R,
typename I_>
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_>
514 typedef typename Eigen::internal::traits<TensorReductionOp>::Scalar
Scalar;
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;
541 template<
typename ArgType,
typename Device>
545 template<
typename Op,
typename Dims,
typename ArgType,
template <
class>
class MakePointer_,
typename Device>
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;
561 typedef typename internal::ReductionReturnType<Op, typename XprType::CoeffReturnType>::type
CoeffReturnType;
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;
588 PacketAccess = Self::InputPacketAccess && ReducerTraits::PacketAccess,
590 PreferBlockAccess =
true,
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);
606 : m_impl(op.expression(), device),
607 m_reducer(op.reducer()),
611 EIGEN_STATIC_ASSERT((!ReducingInnerMostDims | !PreservingInnerMostDims | (NumReducedDims == NumInputDims)),
612 YOU_MADE_A_PROGRAMMING_MISTAKE);
615 for (
int i = 0;
i < NumInputDims; ++
i) {
616 m_reduced[
i] =
false;
618 for (
int i = 0;
i < NumReducedDims; ++
i) {
621 m_reduced[op.
dims()[
i]] =
true;
625 internal::DimInitializer<Dimensions>::run(input_dims, m_reduced, &m_dimensions, &m_reducedDims);
628 if (NumOutputDims > 0) {
629 if (
static_cast<int>(Layout) ==
static_cast<int>(
ColMajor)) {
630 m_outputStrides[0] = 1;
631 for (
int i = 1;
i < NumOutputDims; ++
i) {
632 m_outputStrides[
i] = m_outputStrides[
i - 1] * m_dimensions[
i - 1];
633 m_fastOutputStrides[
i] = internal::TensorIntDivisor<Index>(m_outputStrides[
i]);
636 m_outputStrides[
static_cast<size_t>(NumOutputDims - 1)] = 1;
637 for (
int i = NumOutputDims - 2;
i >= 0; --
i) {
638 m_outputStrides[
i] = m_outputStrides[
i + 1] * m_dimensions[
i + 1];
639 m_fastOutputStrides[
i] = internal::TensorIntDivisor<Index>(m_outputStrides[
i]);
645 if (NumInputDims > 0) {
647 if (
static_cast<int>(Layout) ==
static_cast<int>(
ColMajor)) {
648 input_strides[0] = 1;
649 for (
int i = 1;
i < NumInputDims; ++
i) {
650 input_strides[
i] = input_strides[
i-1] * input_dims[
i-1];
653 input_strides.back() = 1;
654 for (
int i = NumInputDims - 2;
i >= 0; --
i) {
655 input_strides[
i] = input_strides[
i + 1] * input_dims[
i + 1];
661 for (
int i = 0;
i < NumInputDims; ++
i) {
663 m_reducedStrides[reduceIndex] = input_strides[
i];
666 m_preservedStrides[outputIndex] = input_strides[
i];
667 m_output_to_input_dim_map[outputIndex] =
i;
674 if (NumOutputDims == 0) {
678 m_numValuesToReduce =
681 : (
static_cast<int>(Layout) ==
static_cast<int>(
ColMajor))
682 ? m_preservedStrides[0]
683 : m_preservedStrides[
static_cast<size_t>(NumOutputDims - 1)];
691 if ((RunningFullReduction && RunningOnSycl) ||(RunningFullReduction &&
692 internal::FullReducer<Self, Op, Device>::HasOptimizedImplementation &&
693 ((RunningOnGPU && (m_device.majorDeviceVersion() >= 3)) ||
695 bool need_assign =
false;
701 Op reducer(m_reducer);
702 internal::FullReducer<Self, Op, Device>::run(*
this, reducer, m_device,
data);
707 else if ((RunningOnGPU && (m_device.majorDeviceVersion() >= 3)) || (RunningOnSycl)) {
708 bool reducing_inner_dims =
true;
709 for (
int i = 0;
i < NumReducedDims; ++
i) {
710 if (
static_cast<int>(Layout) ==
static_cast<int>(
ColMajor)) {
711 reducing_inner_dims &= m_reduced[
i];
713 reducing_inner_dims &= m_reduced[NumInputDims - 1 -
i];
716 if (internal::InnerReducer<Self, Op, Device>::HasOptimizedImplementation &&
717 (reducing_inner_dims || ReducingInnerMostDims)) {
721 if ((num_coeffs_to_preserve < 1024 && num_values_to_reduce > num_coeffs_to_preserve && num_values_to_reduce > 128) || (RunningOnSycl)) {
729 Op reducer(m_reducer);
731 if (internal::InnerReducer<Self, Op, Device>::run(*
this, reducer, m_device,
data, num_values_to_reduce, num_coeffs_to_preserve)) {
733 m_device.deallocate_temp(m_result);
738 return (m_result != NULL);
742 bool preserving_inner_dims =
true;
743 for (
int i = 0;
i < NumReducedDims; ++
i) {
744 if (
static_cast<int>(Layout) ==
static_cast<int>(
ColMajor)) {
745 preserving_inner_dims &= m_reduced[NumInputDims - 1 -
i];
747 preserving_inner_dims &= m_reduced[
i];
750 if (internal::OuterReducer<Self, Op, Device>::HasOptimizedImplementation &&
751 preserving_inner_dims) {
755 if ((num_coeffs_to_preserve < 1024 && num_values_to_reduce > num_coeffs_to_preserve && num_values_to_reduce > 32) || (RunningOnSycl)) {
763 Op reducer(m_reducer);
765 if (internal::OuterReducer<Self, Op, Device>::run(*
this, reducer, m_device,
data, num_values_to_reduce, num_coeffs_to_preserve)) {
767 m_device.deallocate_temp(m_result);
772 return (m_result != NULL);
775 #if defined(EIGEN_USE_SYCL)
785 Op reducer(m_reducer);
786 internal::GenericReducer<Self, Op, Device>::run(*
this, reducer, m_device,
data, num_values_to_reduce, num_coeffs_to_preserve);
787 return (m_result != NULL);
794 #ifdef EIGEN_USE_THREADS
795 template <
typename EvalSubExprsCallback>
798 evalSubExprsIfNeededAsync(EvaluatorPointerType data,
799 EvalSubExprsCallback done) {
800 m_impl.evalSubExprsIfNeededAsync(NULL, [
this, data, done](
bool) {
801 done(evalSubExprsIfNeededCommon(data));
808 m_impl.evalSubExprsIfNeeded(NULL);
809 return evalSubExprsIfNeededCommon(
data);
815 m_device.deallocate_temp(m_result);
822 if (( RunningFullReduction || RunningOnGPU) && m_result ) {
823 return *(m_result + index);
825 Op reducer(m_reducer);
826 if (ReducingInnerMostDims || RunningFullReduction) {
827 const Index num_values_to_reduce =
828 (
static_cast<int>(Layout) ==
static_cast<int>(
ColMajor)) ? m_preservedStrides[0] : m_preservedStrides[NumPreservedStrides - 1];
829 return internal::InnerMostDimReducer<Self, Op>::reduce(*
this, firstInput(index),
830 num_values_to_reduce, reducer);
833 internal::GenericDimReducer<NumReducedDims-1, Self, Op>::reduce(*
this, firstInput(index), reducer, &accum);
834 return reducer.finalize(accum);
839 template<
int LoadMode>
844 if (RunningOnGPU && m_result) {
845 return internal::pload<PacketReturnType>(m_result + index);
848 EIGEN_ALIGN_MAX std::remove_const_t<CoeffReturnType> values[PacketSize];
849 if (ReducingInnerMostDims) {
850 const Index num_values_to_reduce =
851 (
static_cast<int>(Layout) ==
static_cast<int>(
ColMajor)) ? m_preservedStrides[0] : m_preservedStrides[NumPreservedStrides - 1];
852 const Index firstIndex = firstInput(index);
853 for (
Index i = 0;
i < PacketSize; ++
i) {
854 Op reducer(m_reducer);
855 values[
i] = internal::InnerMostDimReducer<Self, Op>::reduce(*
this, firstIndex +
i * num_values_to_reduce,
856 num_values_to_reduce, reducer);
858 }
else if (PreservingInnerMostDims) {
859 const Index firstIndex = firstInput(index);
860 const int innermost_dim = (
static_cast<int>(Layout) ==
static_cast<int>(
ColMajor)) ? 0 : NumOutputDims - 1;
862 if (((firstIndex % m_dimensions[innermost_dim]) + PacketSize - 1) < m_dimensions[innermost_dim]) {
863 Op reducer(m_reducer);
864 typename Self::PacketReturnType accum = reducer.template initializePacket<typename Self::PacketReturnType>();
865 internal::InnerMostDimPreserver<NumReducedDims-1, Self, Op>::reduce(*
this, firstIndex, reducer, &accum);
866 return reducer.finalizePacket(accum);
868 for (
int i = 0;
i < PacketSize; ++
i) {
869 values[
i] = coeff(index +
i);
873 for (
int i = 0;
i < PacketSize; ++
i) {
874 values[
i] = coeff(index +
i);
883 if (RunningFullReduction && m_result) {
887 const double compute_cost = num_values_to_reduce * internal::functor_traits<Op>::Cost;
888 return m_impl.costPerCoeff(vectorized) * num_values_to_reduce +
889 TensorOpCost(0, 0, compute_cost, vectorized, PacketSize);
898 template <
int,
typename,
typename>
friend struct internal::GenericDimReducer;
899 template <
typename,
typename,
bool,
bool>
friend struct internal::InnerMostDimReducer;
900 template <
int,
typename,
typename,
bool>
friend struct internal::InnerMostDimPreserver;
901 template <
typename S,
typename O,
typename D,
bool V>
friend struct internal::FullReducer;
902 #ifdef EIGEN_USE_THREADS
903 template <
typename S,
typename O,
bool V>
friend struct internal::FullReducerShard;
905 #if defined(EIGEN_USE_GPU) && (defined(EIGEN_GPUCC))
906 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*);
907 #if defined(EIGEN_HAS_GPU_FP16)
908 template <
typename S,
typename R,
typename I_>
KERNEL_FRIEND void internal::ReductionInitFullReduxKernelHalfFloat(R,
const S, I_, internal::packet_traits<Eigen::half>::type*);
909 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*);
910 template <
int NPT,
typename S,
typename R,
typename I_>
KERNEL_FRIEND void internal::InnerReductionKernelHalfFloat(R,
const S, I_, I_,
half*);
912 template <
int NPT,
typename S,
typename R,
typename I_>
KERNEL_FRIEND void internal::InnerReductionKernel(R,
const S, I_, I_,
typename S::CoeffReturnType*);
914 template <
int NPT,
typename S,
typename R,
typename I_>
KERNEL_FRIEND void internal::OuterReductionKernel(R,
const S, I_, I_,
typename S::CoeffReturnType*);
917 #if defined(EIGEN_USE_SYCL)
918 template <
typename Evaluator_,
typename Op__>
friend class TensorSycl::internal::GenericNondeterministicReducer;
920 template <
typename,
typename,
typename>
friend struct internal::GenericReducer;
924 template <
typename S,
typename O,
typename D>
friend struct internal::InnerReducer;
926 struct BlockIteratorState {
935 if (ReducingInnerMostDims) {
936 if (
static_cast<int>(Layout) ==
static_cast<int>(
ColMajor)) {
937 return index * m_preservedStrides[0];
939 return index * m_preservedStrides[NumPreservedStrides - 1];
943 Index startInput = 0;
944 if (
static_cast<int>(Layout) ==
static_cast<int>(
ColMajor)) {
945 for (
int i = NumOutputDims - 1;
i > 0; --
i) {
947 const Index idx = index / m_outputStrides[
i];
948 startInput += idx * m_preservedStrides[
i];
949 index -= idx * m_outputStrides[
i];
951 if (PreservingInnerMostDims) {
955 startInput += index * m_preservedStrides[0];
958 for (
int i = 0;
i < NumOutputDims - 1; ++
i) {
960 const Index idx = index / m_outputStrides[
i];
961 startInput += idx * m_preservedStrides[
i];
962 index -= idx * m_outputStrides[
i];
964 if (PreservingInnerMostDims) {
965 eigen_assert(m_preservedStrides[NumPreservedStrides - 1] == 1);
968 startInput += index * m_preservedStrides[NumPreservedStrides - 1];
1005 template<
typename Op,
typename Dims,
typename ArgType,
template <
class>
class MakePointer_,
typename Device>
1013 template<
typename Op,
typename Dims,
typename ArgType,
template <
class>
class MakePointer_>
1018 EIGEN_STRONG_INLINE
TensorEvaluator(
const typename Base::XprType& op,
const Eigen::SyclDevice& device) :
Base(op, device){}
1022 return *(this->
data() + index);
1026 template<
int LoadMode>
1028 return internal::pload<typename Base::PacketReturnType>(this->
data() + index);
#define EIGEN_DEVICE_FUNC
#define EIGEN_HIP_LAUNCH_BOUNDS_1024
#define EIGEN_STATIC_ASSERT(X, MSG)
static int numThreads(double output_size, const TensorOpCost &cost_per_coeff, int max_threads)
Eigen::internal::traits< TensorReductionOp >::Scalar Scalar
std::remove_const_t< typename XprType::CoeffReturnType > CoeffReturnType
Eigen::internal::nested< TensorReductionOp >::type Nested
Eigen::internal::traits< TensorReductionOp >::StorageKind StorageKind
Eigen::internal::traits< TensorReductionOp >::Index Index
const XprType & expression() const
const Op & reducer() const
TensorReductionOp(const XprType &expr, const Dims &dims, const Op &reducer)
Eigen::NumTraits< Scalar >::Real RealScalar
TensorReductionOp(const XprType &expr, const Dims &dims)
const Dims & dims() const
constexpr auto array_prod(const array< T, N > &arr) -> decltype(array_reduce< product_op, T, N >(arr, static_cast< T >(1)))
Index LeafSize< bfloat16 >()
EIGEN_ALWAYS_INLINE T mini(const T &x, const T &y)
: TensorContractionSycl.h, provides various tensor contraction kernel for SYCL backend
EIGEN_DEFAULT_DENSE_INDEX_TYPE Index
EIGEN_ALWAYS_INLINE T divup(const X x, const Y y)
internal::packet_traits< Scalar >::type type
TensorEvaluator(const typename Base::XprType &op, const Eigen::SyclDevice &device)
TensorReductionEvaluatorBase< const TensorReductionOp< Op, Dims, ArgType, MakePointer_ >, Eigen::SyclDevice > Base
Base::CoeffReturnType coeff(typename Base::Index index) const
Base::PacketReturnType packet(typename Base::Index index) const
TensorEvaluator(const typename Base::XprType &op, const Device &device)
TensorReductionEvaluatorBase< const TensorReductionOp< Op, Dims, ArgType, MakePointer_ >, Device > Base
A cost model used to limit the number of threads used for evaluating tensor expression.
const Dimensions & dimensions() const
EvaluatorPointerType data() const
Index m_numValuesToReduce
TensorEvaluator< ArgType, Device > m_impl
array< Index, NumOutputDims > m_outputStrides
const Dimensions & dimensions() const
array< bool, NumInputDims > m_reduced
const TensorEvaluator< ArgType, Device > & impl() const
TensorReductionEvaluatorBase(const XprType &op, const Device &device)
internal::TensorBlockNotImplemented TensorBlock
array< Index, NumReducedDims > m_reducedStrides
const Device EIGEN_DEVICE_REF m_device
PacketType< CoeffReturnType, Device >::type PacketReturnType
EvaluatorPointerType m_result
PacketReturnType packet(Index index) const
array< Index, NumReducedDims > m_reducedDims
Storage::Type EvaluatorPointerType
std::remove_const_t< Scalar > ScalarNoConst
TensorOpCost costPerCoeff(bool vectorized) const
TensorReductionEvaluatorBase< const TensorReductionOp< Op, Dims, ArgType, MakePointer_ >, Device > Self
Index firstInput(Index index) const
std::conditional_t< NumOutputDims==0, Sizes<>, DSizes< Index, NumOutputDims > > Dimensions
const Device & device() const
CoeffReturnType coeff(Index index) const
internal::ReductionReturnType< Op, typename XprType::CoeffReturnType >::type CoeffReturnType
StorageMemory< CoeffReturnType, Device > Storage
bool evalSubExprsIfNeeded(EvaluatorPointerType data)
bool evalSubExprsIfNeededCommon(EvaluatorPointerType data)
internal::reducer_traits< Op, Device > ReducerTraits
Eigen::internal::traits< XprType >::PointerType TensorPointerType
array< Index, NumOutputDims > m_output_to_input_dim_map
TensorReductionOp< Op, Dims, ArgType, MakePointer_ > XprType
array< Index, NumPreservedStrides > m_preservedStrides
EvaluatorPointerType data() const
array< internal::TensorIntDivisor< Index >, NumOutputDims > m_fastOutputStrides
TensorEvaluator< ArgType, Device >::Dimensions InputDimensions