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
36 template<
typename Op,
typename Dims,
typename XprType,
template <
class>
class MakePointer_ >
44 typedef typename XprType::Nested
Nested;
46 static const int Layout = XprTraits::Layout;
56 template<
typename Op,
typename Dims,
typename XprType,
template <
class>
class MakePointer_>
62 template<
typename Op,
typename Dims,
typename XprType,
template <
class>
class MakePointer_>
71 static void run(
const InputDims& input_dims,
73 OutputDims* output_dims, ReducedDims* reduced_dims) {
77 for (
int i = 0;
i < NumInputDims; ++
i) {
79 (*reduced_dims)[reduceIndex] = input_dims[
i];
82 (*output_dims)[outputIndex] = input_dims[
i];
94 for (
int i = 0;
i < NumInputDims; ++
i) {
95 (*reduced_dims)[
i] = input_dims[
i];
101 template <
typename ReducedDims,
int NumTensorDims,
int Layout>
105 template <
typename ReducedDims,
int NumTensorDims,
int Layout>
110 #if EIGEN_HAS_CONSTEXPR && EIGEN_HAS_VARIADIC_TEMPLATES
111 template <
typename ReducedDims,
int NumTensorDims>
113 static const bool tmp1 = indices_statically_known_to_increase<ReducedDims>();
114 static const bool tmp2 = index_statically_eq<ReducedDims>(0, 0);
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>();
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>();
137 static const bool value = tmp1 & tmp2;
142 template <
int DimIndex,
typename Self,
typename Op>
146 for (
int j = 0;
j <
self.m_reducedDims[DimIndex]; ++
j) {
147 const typename Self::Index input = firstIndex +
j *
self.m_reducedStrides[DimIndex];
152 template <
typename Self,
typename Op>
155 for (
int j = 0;
j <
self.m_reducedDims[0]; ++
j) {
156 const typename Self::Index input = firstIndex +
j *
self.m_reducedStrides[0];
157 reducer.reduce(
self.m_impl.coeff(input), accum);
161 template <
typename Self,
typename Op>
164 reducer.reduce(
self.m_impl.coeff(index), accum);
168 template <
typename Self,
typename Op,
bool Vectorizable = (Self::InputPacketAccess && Self::ReducerTraits::PacketAccess),
169 bool UseTreeReduction = (!Self::ReducerTraits::IsStateful &&
170 !Self::ReducerTraits::IsExactlyAssociative)>
173 typename Self::CoeffReturnType accum = reducer.initialize();
175 reducer.reduce(
self.m_impl.coeff(firstIndex +
j), &accum);
177 return reducer.finalize(accum);
181 template <
typename Self,
typename Op>
185 const typename Self::Index VectorizedSize = (numValuesToReduce / packetSize) * packetSize;
186 typename Self::PacketReturnType paccum = reducer.template initializePacket<typename Self::PacketReturnType>();
187 for (
typename Self::Index j = 0;
j < VectorizedSize;
j += packetSize) {
188 reducer.reducePacket(
self.m_impl.template packet<Unaligned>(firstIndex +
j), &paccum);
190 typename Self::CoeffReturnType accum = reducer.initialize();
191 for (
typename Self::Index j = VectorizedSize;
j < numValuesToReduce; ++
j) {
192 reducer.reduce(
self.m_impl.coeff(firstIndex +
j), &accum);
194 return reducer.finalizeBoth(accum, paccum);
198 #if !defined(EIGEN_HIPCC)
201 template <
typename Self,
typename Op>
205 typename Self::Index numValuesToReduce, Op& reducer) {
206 typename Self::CoeffReturnType accum = reducer.initialize();
209 reducer.reduce(
reduce(
self, firstIndex,
half, reducer), &accum);
211 reduce(
self, firstIndex +
half, numValuesToReduce -
half, reducer),
215 reducer.reduce(
self.m_impl.coeff(firstIndex +
j), &accum);
218 return reducer.finalize(accum);
222 template <
typename Self,
typename Op>
226 typename Self::Index numValuesToReduce, Op& reducer) {
229 typename Self::CoeffReturnType accum = reducer.initialize();
230 if (numValuesToReduce > packetSize *
kLeafSize) {
238 reducer.reduce(
reduce(
self, firstIndex, num_left, reducer), &accum);
239 if (num_left < numValuesToReduce) {
241 reduce(
self,
split, numValuesToReduce - num_left, reducer), &accum);
243 return reducer.finalize(accum);
246 (numValuesToReduce / (2*packetSize)) * 2*packetSize;
248 (numValuesToReduce / packetSize) * packetSize;
249 typename Self::PacketReturnType paccum =
250 reducer.template initializePacket<typename Self::PacketReturnType>();
251 typename Self::PacketReturnType paccum2 =
252 reducer.template initializePacket<typename Self::PacketReturnType>();
253 for (
typename Self::Index j = 0;
j < UnrollSize;
j += packetSize * 2) {
254 reducer.reducePacket(
255 self.m_impl.template packet<Unaligned>(firstIndex +
j), &paccum);
256 reducer.reducePacket(
257 self.m_impl.template packet<Unaligned>(firstIndex +
j + packetSize),
260 for (
typename Self::Index j = UnrollSize;
j < VectorizedSize;
j+= packetSize) {
261 reducer.reducePacket(
self.m_impl.template packet<Unaligned>(
262 firstIndex +
j), &paccum);
264 reducer.reducePacket(paccum2, &paccum);
265 for (
typename Self::Index j = VectorizedSize;
j < numValuesToReduce;
267 reducer.reduce(
self.m_impl.coeff(firstIndex +
j), &accum);
269 return reducer.finalizeBoth(accum, paccum);
275 template <
int DimIndex,
typename Self,
typename Op,
bool vectorizable = (Self::InputPacketAccess && Self::ReducerTraits::PacketAccess)>
282 template <
int DimIndex,
typename Self,
typename Op>
286 for (
typename Self::Index j = 0;
j <
self.m_reducedDims[DimIndex]; ++
j) {
287 const typename Self::Index input = firstIndex +
j *
self.m_reducedStrides[DimIndex];
293 template <
typename Self,
typename Op>
296 for (
typename Self::Index j = 0;
j <
self.m_reducedDims[0]; ++
j) {
297 const typename Self::Index input = firstIndex +
j *
self.m_reducedStrides[0];
298 reducer.reducePacket(
self.m_impl.template packet<Unaligned>(input), accum);
302 template <
typename Self,
typename Op>
310 template <
typename Self,
typename Op,
typename Device,
bool Vectorizable = (Self::InputPacketAccess && Self::ReducerTraits::PacketAccess)>
314 static EIGEN_DEVICE_FUNC void run(
const Self&
self, Op& reducer,
const Device&,
typename Self::EvaluatorPointerType output) {
321 #ifdef EIGEN_USE_THREADS
323 template <
typename Self,
typename Op,
324 bool Vectorizable = (Self::InputPacketAccess && Self::ReducerTraits::PacketAccess)>
325 struct FullReducerShard {
327 typename Self::Index numValuesToReduce, Op& reducer,
328 typename Self::CoeffReturnType* output) {
330 self, firstIndex, numValuesToReduce, reducer);
335 template <
typename Self,
typename Op,
bool Vectorizable>
336 struct FullReducer<Self, Op, ThreadPoolDevice, Vectorizable> {
337 static const bool HasOptimizedImplementation = !Self::ReducerTraits::IsStateful;
338 static const Index PacketSize =
342 static void run(
const Self&
self, Op& reducer,
const ThreadPoolDevice& device,
343 typename Self::CoeffReturnType* output) {
346 if (num_coeffs == 0) {
347 *output = reducer.finalize(reducer.initialize());
350 const TensorOpCost cost =
351 self.m_impl.costPerCoeff(Vectorizable) +
355 num_coeffs, cost, device.numThreads());
356 if (num_threads == 1) {
361 const Index blocksize =
362 std::floor<Index>(
static_cast<float>(num_coeffs) / num_threads);
363 const Index numblocks = blocksize > 0 ? num_coeffs / blocksize : 0;
366 Barrier barrier(internal::convert_index<unsigned int>(numblocks));
367 MaxSizeVector<typename Self::CoeffReturnType> shards(numblocks, reducer.initialize());
368 for (
Index i = 0;
i < numblocks; ++
i) {
370 self,
i * blocksize, blocksize, reducer,
373 typename Self::CoeffReturnType finalShard;
374 if (numblocks * blocksize < num_coeffs) {
376 self, numblocks * blocksize, num_coeffs - numblocks * blocksize,
379 finalShard = reducer.initialize();
383 for (
Index i = 0;
i < numblocks; ++
i) {
384 reducer.reduce(shards[
i], &finalShard);
386 *output = reducer.finalize(finalShard);
394 template <
typename Self,
typename Op,
typename Device>
405 template <
typename Self,
typename Op,
typename Device>
415 #ifdef EIGEN_USE_SYCL
417 template <
typename Self,
typename Op,
typename Device>
418 struct GenericReducer {
419 static const bool HasOptimizedImplementation =
false;
428 #if defined(EIGEN_USE_GPU) && (defined(EIGEN_GPUCC))
429 template <
int B,
int N,
typename S,
typename R,
typename I_>
433 #if defined(EIGEN_HAS_GPU_FP16)
434 template <
typename S,
typename R,
typename I_>
436 template <
int B,
int N,
typename S,
typename R,
typename I_>
438 template <
int NPT,
typename S,
typename R,
typename I_>
443 template <
int NPT,
typename S,
typename R,
typename I_>
446 template <
int NPT,
typename S,
typename R,
typename I_>
458 template <
typename Op,
typename CoeffReturnType>
460 #if defined(EIGEN_USE_SYCL)
470 template <
typename Op,
typename Dims,
typename XprType,
template <
class>
class MakePointer_>
500 template<
typename ArgType,
typename Device>
504 template<
typename Op,
typename Dims,
typename ArgType,
template <
class>
class MakePointer_,
typename Device>
515 static const int NumOutputDims = NumInputDims - NumReducedDims;
534 PacketAccess = Self::InputPacketAccess && ReducerTraits::PacketAccess,
536 PreferBlockAccess =
true,
550 static const bool RunningFullReduction = (NumOutputDims==0);
553 : m_impl(op.expression(), device), m_reducer(op.reducer()), m_result(
NULL), m_device(device)
556 EIGEN_STATIC_ASSERT((!ReducingInnerMostDims | !PreservingInnerMostDims | (NumReducedDims == NumInputDims)),
557 YOU_MADE_A_PROGRAMMING_MISTAKE);
560 for (
int i = 0;
i < NumInputDims; ++
i) {
561 m_reduced[
i] =
false;
563 for (
int i = 0;
i < NumReducedDims; ++
i) {
566 m_reduced[op.
dims()[
i]] =
true;
573 if (NumOutputDims > 0) {
574 if (
static_cast<int>(Layout) ==
static_cast<int>(
ColMajor)) {
575 m_outputStrides[0] = 1;
576 for (
int i = 1;
i < NumOutputDims; ++
i) {
577 m_outputStrides[
i] = m_outputStrides[
i - 1] * m_dimensions[
i - 1];
581 m_outputStrides[NumOutputDims - 1] = 1;
582 for (
int i = NumOutputDims - 2;
i >= 0; --
i) {
583 m_outputStrides[
i] = m_outputStrides[
i + 1] * m_dimensions[
i + 1];
590 if (NumInputDims > 0) {
592 if (
static_cast<int>(Layout) ==
static_cast<int>(
ColMajor)) {
593 input_strides[0] = 1;
594 for (
int i = 1;
i < NumInputDims; ++
i) {
595 input_strides[
i] = input_strides[
i-1] * input_dims[
i-1];
598 input_strides.
back() = 1;
599 for (
int i = NumInputDims - 2;
i >= 0; --
i) {
600 input_strides[
i] = input_strides[
i + 1] * input_dims[
i + 1];
606 for (
int i = 0;
i < NumInputDims; ++
i) {
608 m_reducedStrides[reduceIndex] = input_strides[
i];
611 m_preservedStrides[outputIndex] = input_strides[
i];
612 m_output_to_input_dim_map[outputIndex] =
i;
619 if (NumOutputDims == 0) {
623 m_numValuesToReduce =
626 : (
static_cast<int>(Layout) ==
static_cast<int>(
ColMajor))
627 ? m_preservedStrides[0]
628 : m_preservedStrides[NumOutputDims - 1];
636 if ((RunningFullReduction && RunningOnSycl) ||(RunningFullReduction &&
638 ((RunningOnGPU && (m_device.majorDeviceVersion() >= 3)) ||
640 bool need_assign =
false;
646 Op reducer(m_reducer);
652 else if ((RunningOnGPU && (m_device.majorDeviceVersion() >= 3)) || (RunningOnSycl)) {
653 bool reducing_inner_dims =
true;
654 for (
int i = 0;
i < NumReducedDims; ++
i) {
655 if (
static_cast<int>(Layout) ==
static_cast<int>(
ColMajor)) {
656 reducing_inner_dims &= m_reduced[
i];
658 reducing_inner_dims &= m_reduced[NumInputDims - 1 -
i];
662 (reducing_inner_dims || ReducingInnerMostDims)) {
666 if ((num_coeffs_to_preserve < 1024 && num_values_to_reduce > num_coeffs_to_preserve && num_values_to_reduce > 128) || (RunningOnSycl)) {
674 Op reducer(m_reducer);
678 m_device.deallocate_temp(m_result);
683 return (m_result !=
NULL);
687 bool preserving_inner_dims =
true;
688 for (
int i = 0;
i < NumReducedDims; ++
i) {
689 if (
static_cast<int>(Layout) ==
static_cast<int>(
ColMajor)) {
690 preserving_inner_dims &= m_reduced[NumInputDims - 1 -
i];
692 preserving_inner_dims &= m_reduced[
i];
696 preserving_inner_dims) {
700 if ((num_coeffs_to_preserve < 1024 && num_values_to_reduce > num_coeffs_to_preserve && num_values_to_reduce > 32) || (RunningOnSycl)) {
708 Op reducer(m_reducer);
712 m_device.deallocate_temp(m_result);
717 return (m_result !=
NULL);
720 #if defined(EIGEN_USE_SYCL)
730 Op reducer(m_reducer);
732 return (m_result !=
NULL);
739 #ifdef EIGEN_USE_THREADS
740 template <
typename EvalSubExprsCallback>
743 evalSubExprsIfNeededAsync(EvaluatorPointerType
data,
744 EvalSubExprsCallback done) {
745 m_impl.evalSubExprsIfNeededAsync(
NULL, [
this,
data, done](
bool) {
746 done(evalSubExprsIfNeededCommon(
data));
753 m_impl.evalSubExprsIfNeeded(
NULL);
754 return evalSubExprsIfNeededCommon(
data);
760 m_device.deallocate_temp(m_result);
767 if (( RunningFullReduction || RunningOnGPU) && m_result ) {
768 return *(m_result + index);
770 Op reducer(m_reducer);
771 if (ReducingInnerMostDims || RunningFullReduction) {
772 const Index num_values_to_reduce =
773 (
static_cast<int>(Layout) ==
static_cast<int>(
ColMajor)) ? m_preservedStrides[0] : m_preservedStrides[NumPreservedStrides - 1];
775 num_values_to_reduce, reducer);
779 return reducer.finalize(accum);
784 template<
int LoadMode>
790 if (RunningOnGPU && m_result) {
791 return internal::pload<PacketReturnType>(m_result + index);
795 if (ReducingInnerMostDims) {
796 const Index num_values_to_reduce =
797 (
static_cast<int>(Layout) ==
static_cast<int>(
ColMajor)) ? m_preservedStrides[0] : m_preservedStrides[NumPreservedStrides - 1];
798 const Index firstIndex = firstInput(index);
799 for (
Index i = 0;
i < PacketSize; ++
i) {
800 Op reducer(m_reducer);
802 num_values_to_reduce, reducer);
804 }
else if (PreservingInnerMostDims) {
805 const Index firstIndex = firstInput(index);
806 const int innermost_dim = (
static_cast<int>(Layout) ==
static_cast<int>(
ColMajor)) ? 0 : NumOutputDims - 1;
808 if (((firstIndex % m_dimensions[innermost_dim]) + PacketSize - 1) < m_dimensions[innermost_dim]) {
809 Op reducer(m_reducer);
810 typename Self::PacketReturnType accum = reducer.template initializePacket<typename Self::PacketReturnType>();
812 return reducer.finalizePacket(accum);
814 for (
int i = 0;
i < PacketSize; ++
i) {
819 for (
int i = 0;
i < PacketSize; ++
i) {
829 if (RunningFullReduction && m_result) {
834 return m_impl.costPerCoeff(vectorized) * num_values_to_reduce +
835 TensorOpCost(0, 0, compute_cost, vectorized, PacketSize);
842 #ifdef EIGEN_USE_SYCL
855 #ifdef EIGEN_USE_THREADS
856 template <
typename S,
typename O,
bool V>
friend struct internal::FullReducerShard;
858 #if defined(EIGEN_USE_GPU) && (defined(EIGEN_GPUCC))
859 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*);
860 #if defined(EIGEN_HAS_GPU_FP16)
863 template <
int NPT,
typename S,
typename R,
typename I_>
KERNEL_FRIEND void internal::InnerReductionKernelHalfFloat(
R,
const S, I_, I_,
half*);
865 template <
int NPT,
typename S,
typename R,
typename I_>
KERNEL_FRIEND void internal::InnerReductionKernel(
R,
const S, I_, I_,
typename S::CoeffReturnType*);
867 template <
int NPT,
typename S,
typename R,
typename I_>
KERNEL_FRIEND void internal::OuterReductionKernel(
R,
const S, I_, I_,
typename S::CoeffReturnType*);
870 #if defined(EIGEN_USE_SYCL)
873 template <
typename,
typename,
typename>
friend struct internal::GenericReducer;
879 struct BlockIteratorState {
888 if (ReducingInnerMostDims) {
889 if (
static_cast<int>(Layout) ==
static_cast<int>(
ColMajor)) {
890 return index * m_preservedStrides[0];
892 return index * m_preservedStrides[NumPreservedStrides - 1];
896 Index startInput = 0;
897 if (
static_cast<int>(Layout) ==
static_cast<int>(
ColMajor)) {
898 for (
int i = NumOutputDims - 1;
i > 0; --
i) {
900 const Index idx = index / m_outputStrides[
i];
901 startInput += idx * m_preservedStrides[
i];
902 index -= idx * m_outputStrides[
i];
904 if (PreservingInnerMostDims) {
908 startInput += index * m_preservedStrides[0];
911 for (
int i = 0;
i < NumOutputDims - 1; ++
i) {
913 const Index idx = index / m_outputStrides[
i];
914 startInput += idx * m_preservedStrides[
i];
915 index -= idx * m_outputStrides[
i];
917 if (PreservingInnerMostDims) {
918 eigen_assert(m_preservedStrides[NumPreservedStrides - 1] == 1);
921 startInput += index * m_preservedStrides[NumPreservedStrides - 1];
954 #if defined(EIGEN_USE_GPU) && (defined(EIGEN_GPUCC))
956 static const bool RunningOnSycl =
false;
957 #elif defined(EIGEN_USE_SYCL)
959 static const bool RunningOnGPU =
false;
961 static const bool RunningOnGPU =
false;
962 static const bool RunningOnSycl =
false;
969 template<
typename Op,
typename Dims,
typename ArgType,
template <
class>
class MakePointer_,
typename Device>
977 template<
typename Op,
typename Dims,
typename ArgType,
template <
class>
class MakePointer_>
986 return *(this->
data() + index);
990 template<
int LoadMode>
992 return internal::pload<typename Base::PacketReturnType>(this->
data() + index);
998 #endif // EIGEN_CXX11_TENSOR_TENSOR_REDUCTION_H