| // This file is part of Eigen, a lightweight C++ template library |
| // for linear algebra. |
| // |
| // Copyright (C) 2014 Benoit Steiner <benoit.steiner.goog@gmail.com> |
| // |
| // This Source Code Form is subject to the terms of the Mozilla |
| // Public License v. 2.0. If a copy of the MPL was not distributed |
| // with this file, You can obtain one at http://mozilla.org/MPL/2.0/. |
| |
| #ifndef EIGEN_CXX11_TENSOR_TENSOR_REDUCTION_H |
| #define EIGEN_CXX11_TENSOR_TENSOR_REDUCTION_H |
| |
| // Kernel friend declarations may or not need the __global__ annotation |
| // depending on the compiler. |
| #ifndef KERNEL_FRIEND |
| #if defined(__clang__) && defined(__CUDA__) |
| #define KERNEL_FRIEND friend __global__ |
| #else |
| #define KERNEL_FRIEND friend |
| #endif |
| #endif |
| |
| namespace Eigen { |
| |
| /** \class TensorReduction |
| * \ingroup CXX11_Tensor_Module |
| * |
| * \brief Tensor reduction class. |
| * |
| */ |
| |
| namespace internal { |
| template<typename Op, typename Dims, typename XprType> |
| struct traits<TensorReductionOp<Op, Dims, XprType> > |
| : traits<XprType> |
| { |
| typedef typename traits<XprType>::Scalar Scalar; |
| typedef typename traits<XprType>::StorageKind StorageKind; |
| typedef typename traits<XprType>::Index Index; |
| typedef typename XprType::Nested Nested; |
| }; |
| |
| template<typename Op, typename Dims, typename XprType> |
| struct eval<TensorReductionOp<Op, Dims, XprType>, Eigen::Dense> |
| { |
| typedef const TensorReductionOp<Op, Dims, XprType>& type; |
| }; |
| |
| template<typename Op, typename Dims, typename XprType> |
| struct nested<TensorReductionOp<Op, Dims, XprType>, 1, typename eval<TensorReductionOp<Op, Dims, XprType> >::type> |
| { |
| typedef TensorReductionOp<Op, Dims, XprType> type; |
| }; |
| |
| |
| |
| template <typename InputDims, typename OutputDims, typename ReducedDims> EIGEN_DEVICE_FUNC |
| static void partition_dims(const InputDims& input_dims, |
| const array<bool, internal::array_size<InputDims>::value>& reduced, |
| OutputDims* output_dims, ReducedDims* reduced_dims) { |
| const int NumInputDims = internal::array_size<InputDims>::value; |
| int outputIndex = 0; |
| int reduceIndex = 0; |
| for (int i = 0; i < NumInputDims; ++i) { |
| if (OutputDims::count == 0 || reduced[i]) { |
| (*reduced_dims)[reduceIndex] = input_dims[i]; |
| ++reduceIndex; |
| } else { |
| (*output_dims)[outputIndex] = input_dims[i]; |
| ++outputIndex; |
| } |
| } |
| } |
| |
| |
| |
| template <typename ReducedDims, int NumTensorDims, int Layout> |
| struct are_inner_most_dims { |
| static const bool value = false; |
| }; |
| template <typename ReducedDims, int NumTensorDims, int Layout> |
| struct preserve_inner_most_dims { |
| static const bool value = false; |
| }; |
| |
| #if defined(EIGEN_HAS_CONSTEXPR) && defined(EIGEN_HAS_VARIADIC_TEMPLATES) |
| // The use of the tmp1, tmp2, tmp3 intermediate variables is needed for nvcc 7 |
| // to compile the code below. NVidia is working on a fix. |
| template <typename ReducedDims, int NumTensorDims> |
| struct are_inner_most_dims<ReducedDims, NumTensorDims, ColMajor>{ |
| static const bool tmp1 = indices_statically_known_to_increase<ReducedDims>()(); |
| static const bool tmp2 = index_statically_eq<ReducedDims>()(0, 0); |
| static const bool tmp3 = index_statically_eq<ReducedDims>()(array_size<ReducedDims>::value-1, array_size<ReducedDims>::value-1); |
| static const bool value = tmp1 & tmp2 & tmp3; |
| }; |
| template <typename ReducedDims, int NumTensorDims> |
| struct are_inner_most_dims<ReducedDims, NumTensorDims, RowMajor>{ |
| static const bool tmp1 = indices_statically_known_to_increase<ReducedDims>()(); |
| static const bool tmp2 = index_statically_eq<ReducedDims>()(0, NumTensorDims - array_size<ReducedDims>::value); |
| static const bool tmp3 = index_statically_eq<ReducedDims>()(array_size<ReducedDims>::value - 1, NumTensorDims - 1); |
| static const bool value = tmp1 & tmp2 & tmp3; |
| |
| }; |
| template <typename ReducedDims, int NumTensorDims> |
| struct preserve_inner_most_dims<ReducedDims, NumTensorDims, ColMajor>{ |
| static const bool tmp1 = indices_statically_known_to_increase<ReducedDims>()(); |
| static const bool tmp2 = index_statically_gt<ReducedDims>()(0, 0); |
| static const bool value = tmp1 & tmp2; |
| |
| }; |
| template <typename ReducedDims, int NumTensorDims> |
| struct preserve_inner_most_dims<ReducedDims, NumTensorDims, RowMajor>{ |
| static const bool tmp1 = indices_statically_known_to_increase<ReducedDims>()(); |
| static const bool tmp2 = index_statically_lt<ReducedDims>()(array_size<ReducedDims>::value - 1, NumTensorDims - 1); |
| static const bool value = tmp1 & tmp2; |
| }; |
| #endif |
| |
| |
| template <int DimIndex, typename Self, typename Op> |
| struct GenericDimReducer { |
| static EIGEN_DEVICE_FUNC EIGEN_STRONG_INLINE void reduce(const Self& self, typename Self::Index firstIndex, Op& reducer, typename Self::CoeffReturnType* accum) { |
| EIGEN_STATIC_ASSERT(DimIndex >= 0, YOU_MADE_A_PROGRAMMING_MISTAKE); |
| for (int j = 0; j < self.m_reducedDims[DimIndex]; ++j) { |
| const typename Self::Index input = firstIndex + j * self.m_reducedStrides[DimIndex]; |
| GenericDimReducer<DimIndex-1, Self, Op>::reduce(self, input, reducer, accum); |
| } |
| } |
| }; |
| template <typename Self, typename Op> |
| struct GenericDimReducer<-1, Self, Op> { |
| static EIGEN_DEVICE_FUNC EIGEN_STRONG_INLINE void reduce(const Self& self, typename Self::Index firstIndex, Op& reducer, typename Self::CoeffReturnType* accum) { |
| reducer.reduce(self.m_impl.coeff(firstIndex), accum); |
| } |
| }; |
| |
| template <typename Self, typename Op, bool Vectorizable = (Self::InputPacketAccess & Op::PacketAccess)> |
| struct InnerMostDimReducer { |
| static EIGEN_DEVICE_FUNC EIGEN_STRONG_INLINE typename Self::CoeffReturnType reduce(const Self& self, typename Self::Index firstIndex, typename Self::Index numValuesToReduce, Op& reducer) { |
| typename Self::CoeffReturnType accum = reducer.initialize(); |
| for (typename Self::Index j = 0; j < numValuesToReduce; ++j) { |
| reducer.reduce(self.m_impl.coeff(firstIndex + j), &accum); |
| } |
| return reducer.finalize(accum); |
| } |
| }; |
| |
| template <typename Self, typename Op> |
| struct InnerMostDimReducer<Self, Op, true> { |
| static EIGEN_DEVICE_FUNC EIGEN_STRONG_INLINE typename Self::CoeffReturnType reduce(const Self& self, typename Self::Index firstIndex, typename Self::Index numValuesToReduce, Op& reducer) { |
| const int packetSize = internal::unpacket_traits<typename Self::PacketReturnType>::size; |
| const typename Self::Index VectorizedSize = (numValuesToReduce / packetSize) * packetSize; |
| typename Self::PacketReturnType p = reducer.template initializePacket<typename Self::PacketReturnType>(); |
| for (typename Self::Index j = 0; j < VectorizedSize; j += packetSize) { |
| reducer.reducePacket(self.m_impl.template packet<Unaligned>(firstIndex + j), &p); |
| } |
| typename Self::CoeffReturnType accum = reducer.initialize(); |
| for (typename Self::Index j = VectorizedSize; j < numValuesToReduce; ++j) { |
| reducer.reduce(self.m_impl.coeff(firstIndex + j), &accum); |
| } |
| return reducer.finalizeBoth(accum, p); |
| } |
| }; |
| |
| template <int DimIndex, typename Self, typename Op, bool vectorizable = (Self::InputPacketAccess & Op::PacketAccess)> |
| struct InnerMostDimPreserver { |
| static EIGEN_DEVICE_FUNC EIGEN_STRONG_INLINE void reduce(const Self& self, typename Self::Index firstIndex, Op& reducer, typename Self::PacketReturnType* accum) { |
| eigen_assert(false && "should never be called"); |
| } |
| }; |
| |
| template <int DimIndex, typename Self, typename Op> |
| struct InnerMostDimPreserver<DimIndex, Self, Op, true> { |
| static EIGEN_DEVICE_FUNC EIGEN_STRONG_INLINE void reduce(const Self& self, typename Self::Index firstIndex, Op& reducer, typename Self::PacketReturnType* accum) { |
| EIGEN_STATIC_ASSERT(DimIndex >= 0, YOU_MADE_A_PROGRAMMING_MISTAKE); |
| for (typename Self::Index j = 0; j < self.m_reducedDims[DimIndex]; ++j) { |
| const typename Self::Index input = firstIndex + j * self.m_reducedStrides[DimIndex]; |
| InnerMostDimPreserver<DimIndex-1, Self, Op>::reduce(self, input, reducer, accum); |
| } |
| } |
| }; |
| |
| template <typename Self, typename Op> |
| struct InnerMostDimPreserver<-1, Self, Op, true> { |
| static EIGEN_DEVICE_FUNC EIGEN_STRONG_INLINE void reduce(const Self& self, typename Self::Index firstIndex, Op& reducer, typename Self::PacketReturnType* accum) { |
| reducer.reducePacket(self.m_impl.template packet<Unaligned>(firstIndex), accum); |
| } |
| }; |
| |
| // Default full reducer |
| template <typename Self, typename Op, typename Device, |
| bool vectorizable = (Self::InputPacketAccess & Op::PacketAccess)> |
| struct FullReducer { |
| static const bool HasOptimizedImplementation = false; |
| |
| static EIGEN_DEVICE_FUNC void run(const Self& self, Op& reducer, |
| const Device&, |
| typename Self::CoeffReturnType* output) { |
| const typename Self::Index num_coeffs = |
| array_prod(self.m_impl.dimensions()); |
| *output = InnerMostDimReducer<Self, Op, vectorizable>::reduce( |
| self, 0, num_coeffs, reducer); |
| } |
| }; |
| |
| |
| #ifdef EIGEN_USE_THREADS |
| template <typename Self, typename Op, |
| bool vectorizable = (Self::InputPacketAccess & Op::PacketAccess)> |
| struct FullReducerShard { |
| static EIGEN_DEVICE_FUNC EIGEN_STRONG_INLINE void run(const Self& self, typename Self::Index firstIndex, |
| typename Self::Index numValuesToReduce, Op& reducer, |
| typename Self::CoeffReturnType* output) { |
| *output = InnerMostDimReducer<Self, Op, vectorizable>::reduce( |
| self, firstIndex, numValuesToReduce, reducer); |
| } |
| }; |
| |
| // Multithreaded full reducer |
| template <typename Self, typename Op, bool Vectorizable> |
| struct FullReducer<Self, Op, ThreadPoolDevice, Vectorizable> { |
| static const bool HasOptimizedImplementation = !Op::IsStateful; |
| static const int PacketSize = |
| unpacket_traits<typename Self::PacketReturnType>::size; |
| |
| // launch one reducer per thread and accumulate the result. |
| static void run(const Self& self, Op& reducer, const ThreadPoolDevice& device, |
| typename Self::CoeffReturnType* output) { |
| typedef typename Self::Index Index; |
| const Index num_coeffs = array_prod(self.m_impl.dimensions()); |
| if (num_coeffs == 0) { |
| *output = reducer.finalize(reducer.initialize()); |
| return; |
| } |
| const TensorOpCost cost = |
| self.m_impl.costPerCoeff(Vectorizable) + |
| TensorOpCost(0, 0, internal::functor_traits<Op>::Cost, Vectorizable, |
| PacketSize); |
| const int num_threads = TensorCostModel<ThreadPoolDevice>::numThreads( |
| num_coeffs, cost, device.numThreads()); |
| if (num_threads == 1) { |
| *output = |
| InnerMostDimReducer<Self, Op, Vectorizable>::reduce(self, 0, num_coeffs, reducer); |
| return; |
| } |
| const Index blocksize = |
| std::floor<Index>(static_cast<float>(num_coeffs) / num_threads); |
| const Index numblocks = blocksize > 0 ? num_coeffs / blocksize : 0; |
| eigen_assert(num_coeffs >= numblocks * blocksize); |
| |
| Barrier barrier(numblocks); |
| FixedSizeVector<typename Self::CoeffReturnType> shards(numblocks, reducer.initialize()); |
| for (Index i = 0; i < numblocks; ++i) { |
| device.enqueue_with_barrier(&barrier, &FullReducerShard<Self, Op, Vectorizable>::run, |
| self, i * blocksize, blocksize, reducer, |
| &shards[i]); |
| } |
| typename Self::CoeffReturnType finalShard; |
| if (numblocks * blocksize < num_coeffs) { |
| finalShard = InnerMostDimReducer<Self, Op, Vectorizable>::reduce( |
| self, numblocks * blocksize, num_coeffs - numblocks * blocksize, |
| reducer); |
| } else { |
| finalShard = reducer.initialize(); |
| } |
| barrier.Wait(); |
| |
| for (Index i = 0; i < numblocks; ++i) { |
| reducer.reduce(shards[i], &finalShard); |
| } |
| *output = reducer.finalize(finalShard); |
| } |
| }; |
| |
| #endif |
| |
| |
| #if defined(EIGEN_USE_GPU) && defined(__CUDACC__) |
| // Full reducers for GPU, don't vectorize for now |
| |
| // Reducer function that enables multiple cuda thread to safely accumulate at the same |
| // output address. It basically reads the current value of the output variable, and |
| // attempts to update it with the new value. If in the meantime another cuda thread |
| // updated the content of the output address it will try again. |
| template <typename T, typename R> |
| __device__ EIGEN_ALWAYS_INLINE void atomicReduce(T* output, T accum, R& reducer) { |
| #if __CUDA_ARCH__ >= 300 |
| if (sizeof(T) == 4) |
| { |
| unsigned int oldval = *reinterpret_cast<unsigned int*>(output); |
| unsigned int newval = oldval; |
| reducer.reduce(accum, reinterpret_cast<T*>(&newval)); |
| if (newval == oldval) { |
| return; |
| } |
| unsigned int readback; |
| while ((readback = atomicCAS((unsigned int*)output, oldval, newval)) != oldval) { |
| oldval = readback; |
| newval = oldval; |
| reducer.reduce(accum, reinterpret_cast<T*>(&newval)); |
| if (newval == oldval) { |
| return; |
| } |
| } |
| } |
| else if (sizeof(T) == 8) { |
| unsigned long long oldval = *reinterpret_cast<unsigned long long*>(output); |
| unsigned long long newval = oldval; |
| reducer.reduce(accum, reinterpret_cast<T*>(&newval)); |
| if (newval == oldval) { |
| return; |
| } |
| unsigned long long readback; |
| while ((readback = atomicCAS((unsigned long long*)output, oldval, newval)) != oldval) { |
| oldval = readback; |
| newval = oldval; |
| reducer.reduce(accum, reinterpret_cast<T*>(&newval)); |
| if (newval == oldval) { |
| return; |
| } |
| } |
| } |
| else { |
| assert(0 && "Wordsize not supported"); |
| } |
| #else |
| assert(0 && "Shouldn't be called on unsupported device"); |
| #endif |
| } |
| |
| template <typename T> |
| __device__ inline void atomicReduce(T* output, T accum, SumReducer<T>&) { |
| #if __CUDA_ARCH__ >= 300 |
| atomicAdd(output, accum); |
| #else |
| assert(0 && "Shouldn't be called on unsupported device"); |
| #endif |
| } |
| |
| #ifdef EIGEN_HAS_CUDA_FP16 |
| template <template <typename T> class R> |
| __device__ inline void atomicReduce(half2* output, half2 accum, R<half>& reducer) { |
| #if __CUDA_ARCH__ >= 300 |
| unsigned int oldval = *reinterpret_cast<unsigned int*>(output); |
| unsigned int newval = oldval; |
| reducer.reducePacket(accum, reinterpret_cast<half2*>(&newval)); |
| if (newval == oldval) { |
| return; |
| } |
| unsigned int readback; |
| while ((readback = atomicCAS((unsigned int*)output, oldval, newval)) != oldval) { |
| oldval = readback; |
| newval = oldval; |
| reducer.reducePacket(accum, reinterpret_cast<half2*>(&newval)); |
| if (newval == oldval) { |
| return; |
| } |
| } |
| #else |
| assert(0 && "Shouldn't be called on unsupported device"); |
| #endif |
| } |
| #endif |
| |
| |
| template <typename CoeffType, typename Index> |
| __global__ void ReductionInitKernel(const CoeffType val, Index num_preserved_coeffs, CoeffType* output) { |
| const Index thread_id = blockIdx.x * blockDim.x + threadIdx.x; |
| const Index num_threads = blockDim.x * gridDim.x; |
| for (Index i = thread_id; i < num_preserved_coeffs; i += num_threads) { |
| output[i] = val; |
| } |
| } |
| |
| |
| template <int BlockSize, int NumPerThread, typename Self, |
| typename Reducer, typename Index> |
| __global__ void FullReductionKernel(Reducer reducer, const Self input, Index num_coeffs, |
| typename Self::CoeffReturnType* output) { |
| const Index first_index = blockIdx.x * BlockSize * NumPerThread + threadIdx.x; |
| |
| // Initialize the output value if it wasn't initialized by the ReductionInitKernel |
| if (gridDim.x == 1) { |
| if (first_index == 0) { |
| *output = reducer.initialize(); |
| } |
| __syncthreads(); |
| } |
| |
| typename Self::CoeffReturnType accum = reducer.initialize(); |
| |
| // Process up to NumPerThread*BlockSize coefficient per thread while making sure that we don't go past the last coefficient of the tensor. |
| Index max_iter = numext::mini<Index>(num_coeffs - first_index, NumPerThread*BlockSize); |
| #pragma unroll 8 |
| for (Index i = 0; i < max_iter; i += BlockSize) { |
| const Index index = first_index + i; |
| eigen_assert(index < num_coeffs); |
| typename Self::CoeffReturnType val = input.m_impl.coeff(index); |
| reducer.reduce(val, &accum); |
| } |
| |
| #pragma unroll |
| for (int offset = warpSize/2; offset > 0; offset /= 2) { |
| reducer.reduce(__shfl_down(accum, offset), &accum); |
| } |
| |
| if ((threadIdx.x & (warpSize - 1)) == 0) { |
| atomicReduce(output, accum, reducer); |
| } |
| } |
| |
| #ifdef EIGEN_HAS_CUDA_FP16 |
| template <typename Self, typename Reducer, typename Index> |
| __global__ void ReductionInitFullReduxKernelHalfFloat( |
| Reducer reducer, const Self input, Index num_coeffs, half2* scratch) { |
| eigen_assert(threadIdx.x == 1); |
| if (num_coeffs % 2 != 0) { |
| half last = input.m_impl.coeff(num_coeffs - 1); |
| *scratch = __halves2half2(last, reducer.initialize()); |
| } else { |
| *scratch = reducer.template initializePacket<half2>(); |
| } |
| } |
| |
| template <typename Self, typename Reducer, typename Index> |
| __global__ void ReductionInitKernelHalfFloat(Reducer reducer, const Self input, Index num_coeffs, half* output) { |
| const Index thread_id = blockIdx.x * blockDim.x + threadIdx.x; |
| const Index num_threads = blockDim.x * gridDim.x; |
| const Index num_packets = num_coeffs / 2; |
| for (Index i = thread_id; i < num_packets; i += num_threads) { |
| ((half2*)output)[i] = reducer.template initializePacket<half2>(); |
| } |
| |
| if (thread_id == 0 && num_coeffs & 0x1 != 0) { |
| output[num_coeffs-1] = reducer.initialize(); |
| } |
| } |
| |
| template <int BlockSize, int NumPerThread, typename Self, typename Reducer, |
| typename Index> |
| __global__ void FullReductionKernelHalfFloat(Reducer reducer, |
| const Self input, |
| Index num_coeffs, |
| half* output, |
| half2* scratch) { |
| eigen_assert(NumPerThread & 0x1 == 0); |
| |
| const Index first_index = |
| blockIdx.x * BlockSize * NumPerThread + 2 * threadIdx.x; |
| |
| // Initialize the output value if it wasn't initialized by the |
| // ReductionInitKernel |
| if (gridDim.x == 1 && first_index == 0) { |
| if (num_coeffs & 0x1 != 0) { |
| half last = input.m_impl.coeff(num_coeffs - 1); |
| *scratch = __halves2half2(last, reducer.initialize()); |
| } else { |
| *scratch = reducer.template initializePacket<half2>(); |
| } |
| } |
| |
| half2 accum = reducer.template initializePacket<half2>(); |
| const Index max_iter = numext::mini<Index>((num_coeffs - first_index) / 2, |
| NumPerThread * BlockSize / 2); |
| for (Index i = 0; i < max_iter; i += BlockSize) { |
| const Index index = first_index + 2 * i; |
| eigen_assert(index + 1 < num_coeffs); |
| half2 val = input.m_impl.template packet<Unaligned>(index); |
| reducer.reducePacket(val, &accum); |
| } |
| |
| #pragma unroll |
| for (int offset = warpSize / 2; offset > 0; offset /= 2) { |
| reducer.reducePacket(__shfl_down(accum, offset, warpSize), &accum); |
| } |
| |
| if ((threadIdx.x & (warpSize - 1)) == 0) { |
| atomicReduce(scratch, accum, reducer); |
| } |
| |
| __syncthreads(); |
| |
| if (gridDim.x == 1 && first_index == 0) { |
| half tmp = __low2half(*scratch); |
| reducer.reduce(__high2half(*scratch), &tmp); |
| *output = tmp; |
| } |
| } |
| |
| template <typename Op> |
| __global__ void ReductionCleanupKernelHalfFloat(Op& reducer, half* output, |
| half2* scratch) { |
| eigen_assert(threadIdx.x == 1); |
| half tmp = __low2half(*scratch); |
| reducer.reduce(__high2half(*scratch), &tmp); |
| *output = tmp; |
| } |
| |
| #endif |
| |
| template <typename Self, typename Op, typename OutputType, bool PacketAccess> |
| struct FullReductionLauncher { |
| static void run(const Self&, Op&, const GpuDevice&, OutputType*, typename Self::Index) { |
| assert(false && "Should only be called on floats and half floats"); |
| } |
| }; |
| |
| // Launch a full reduction on fp32. Packet access isn't required in this case. |
| template <typename Self, typename Op, bool PacketAccess> |
| struct FullReductionLauncher<Self, Op, float, PacketAccess> { |
| static void run(const Self& self, Op& reducer, const GpuDevice& device, float* output, typename Self::Index num_coeffs) { |
| typedef typename Self::Index Index; |
| typedef typename Self::CoeffReturnType Scalar; |
| const int block_size = 256; |
| const int num_per_thread = 128; |
| const int num_blocks = divup<int>(num_coeffs, block_size * num_per_thread); |
| |
| if (num_blocks > 1) { |
| // We initialize the outputs outside the reduction kernel when we can't |
| // be sure that there won't be a race conditions between multiple |
| // thread blocks. |
| LAUNCH_CUDA_KERNEL((ReductionInitKernel<Scalar, Index>), 1, 32, 0, device, |
| reducer.initialize(), 1, output); |
| } |
| |
| LAUNCH_CUDA_KERNEL( |
| (FullReductionKernel<block_size, num_per_thread, Self, Op, Index>), |
| num_blocks, block_size, 0, device, reducer, self, num_coeffs, output); |
| } |
| }; |
| |
| #ifdef EIGEN_HAS_CUDA_FP16 |
| // Launch a full reduction on fp16. Packet access is required in this case, so make sure we abort if a fp16 optimized reduction is attempted when packet accessors aren't available (note that this would be a programming mistake, in practice we should always fallback to the generic reduction code to handle such cases) |
| template <typename Self, typename Op> |
| struct FullReductionLauncher<Self, Op, Eigen::half, false> { |
| // Leave the function unimplemented to create a linking error. |
| static void run(const Self&, Op&, const GpuDevice&, half*, typename Self::Index); |
| }; |
| |
| |
| template <typename Self, typename Op> |
| struct FullReductionLauncher<Self, Op, Eigen::half, true> { |
| static void run(const Self& self, Op& reducer, const GpuDevice& device, half* output, typename Self::Index num_coeffs) { |
| typedef typename Self::Index Index; |
| |
| const int block_size = 256; |
| const int num_per_thread = 128; |
| const int num_blocks = divup<int>(num_coeffs, block_size * num_per_thread); |
| half2* scratch = static_cast<half2*>(device.scratchpad()); |
| |
| if (num_blocks > 1) { |
| // We initialize the output and the scrathpad outside the reduction kernel when we can't be sure that there |
| // won't be a race conditions between multiple thread blocks. |
| LAUNCH_CUDA_KERNEL((ReductionInitFullReduxKernelHalfFloat<Self, Op, Index>), |
| 1, 1, 0, device, reducer, self, num_coeffs, scratch); |
| } |
| |
| LAUNCH_CUDA_KERNEL((FullReductionKernelHalfFloat<block_size, num_per_thread, Self, Op, Index>), |
| num_blocks, block_size, 0, device, reducer, self, num_coeffs, output, scratch); |
| |
| if (num_blocks > 1) { |
| // TODO(bsteiner): refactor the code to avoid the need for 3 kernel launches. |
| LAUNCH_CUDA_KERNEL((ReductionCleanupKernelHalfFloat<Op>), |
| 1, 1, 0, device, reducer, output, scratch); |
| } |
| } |
| }; |
| #endif |
| |
| template <typename Self, typename Op, bool Vectorizable> |
| struct FullReducer<Self, Op, GpuDevice, Vectorizable> { |
| // Unfortunately nvidia doesn't support well exotic types such as complex, |
| // so reduce the scope of the optimized version of the code to the simple case |
| // of floats and half floats. |
| #ifdef EIGEN_HAS_CUDA_FP16 |
| static const bool HasOptimizedImplementation = |
| !Op::IsStateful && |
| (internal::is_same<typename Self::CoeffReturnType, float>::value || |
| (internal::is_same<typename Self::CoeffReturnType, Eigen::half>::value && |
| reducer_traits<Op, GpuDevice>::PacketAccess)); |
| #else |
| static const bool HasOptimizedImplementation = |
| !Op::IsStateful && |
| internal::is_same<typename Self::CoeffReturnType, float>::value; |
| #endif |
| |
| template <typename OutputType> |
| static void run(const Self& self, Op& reducer, const GpuDevice& device, |
| OutputType* output) { |
| assert(HasOptimizedImplementation && "Should only be called on floats or half floats"); |
| const Index num_coeffs = array_prod(self.m_impl.dimensions()); |
| // Don't crash when we're called with an input tensor of size 0. |
| if (num_coeffs == 0) { |
| return; |
| } |
| |
| FullReductionLauncher<Self, Op, OutputType, reducer_traits<Op, GpuDevice>::PacketAccess>::run(self, reducer, device, output, |
| num_coeffs); |
| } |
| }; |
| |
| |
| template <int NumPerThread, typename Self, |
| typename Reducer, typename Index> |
| __global__ void InnerReductionKernel(Reducer reducer, const Self input, Index num_coeffs_to_reduce, Index num_preserved_coeffs, |
| typename Self::CoeffReturnType* output) { |
| eigen_assert(blockDim.y == 1); |
| eigen_assert(blockDim.z == 1); |
| eigen_assert(gridDim.y == 1); |
| eigen_assert(gridDim.z == 1); |
| |
| const int unroll_times = 16; |
| eigen_assert(NumPerThread % unroll_times == 0); |
| |
| const Index input_col_blocks = divup<Index>(num_coeffs_to_reduce, blockDim.x * NumPerThread); |
| const Index num_input_blocks = input_col_blocks * num_preserved_coeffs; |
| |
| const Index num_threads = blockDim.x * gridDim.x; |
| const Index thread_id = blockIdx.x * blockDim.x + threadIdx.x; |
| // Initialize the output values if they weren't initialized by the ReductionInitKernel |
| if (gridDim.x == 1) { |
| for (Index i = thread_id; i < num_preserved_coeffs; i += num_threads) { |
| output[i] = reducer.initialize(); |
| } |
| // We don't need a __syncthreads() here, because we don't write to output |
| // until after a __syncthreads(). |
| } |
| |
| for (Index i = blockIdx.x; i < num_input_blocks; i += gridDim.x) { |
| const Index row = i / input_col_blocks; |
| |
| if (row < num_preserved_coeffs) { |
| const Index col_block = i % input_col_blocks; |
| const Index col_begin = col_block * blockDim.x * NumPerThread + threadIdx.x; |
| |
| float reduced_val = reducer.initialize(); |
| |
| for (Index j = 0; j < NumPerThread; j += unroll_times) { |
| const Index last_col = col_begin + blockDim.x * (j + unroll_times - 1); |
| if (last_col >= num_coeffs_to_reduce) { |
| for (Index col = col_begin + blockDim.x * j; col < num_coeffs_to_reduce; col +=blockDim.x) { |
| const float val =input.m_impl.coeff(row * num_coeffs_to_reduce + col); |
| reducer.reduce(val, &reduced_val); |
| } |
| break; |
| } else { |
| // Faster version of the loop with no branches after unrolling. |
| #pragma unroll |
| for (int k = 0; k < unroll_times; ++k) { |
| const Index col = col_begin + blockDim.x * (j + k); |
| reducer.reduce(input.m_impl.coeff(row * num_coeffs_to_reduce + col), &reduced_val); |
| } |
| } |
| } |
| |
| __syncthreads(); |
| |
| for (int offset = warpSize/2; offset > 0; offset /= 2) { |
| reducer.reduce(__shfl_down(reduced_val, offset), &reduced_val); |
| } |
| |
| if ((threadIdx.x & (warpSize - 1)) == 0) { |
| atomicReduce(&(output[row]), reduced_val, reducer); |
| } |
| } |
| |
| __syncthreads(); |
| } |
| } |
| |
| template <typename Self, typename Op, typename OutputType, bool PacketAccess> |
| struct InnerReductionLauncher { |
| // Leave the function unimplemented to create a linking error. |
| static EIGEN_DEVICE_FUNC void run(const Self&, Op&, const GpuDevice&, OutputType*, typename Self::Index, typename Self::Index); |
| }; |
| |
| template <typename Self, typename Op, bool PacketAccess> |
| struct InnerReductionLauncher<Self, Op, float, PacketAccess> { |
| static void run(const Self& self, Op& reducer, const GpuDevice& device, float* output, typename Self::Index num_coeffs_to_reduce, typename Self::Index num_preserved_vals) { |
| typedef typename Self::Index Index; |
| |
| const Index num_coeffs = num_coeffs_to_reduce * num_preserved_vals; |
| const int block_size = 256; |
| const int num_per_thread = 128; |
| const int dyn_blocks = divup<int>(num_coeffs, block_size * num_per_thread); |
| const int max_blocks = device.getNumCudaMultiProcessors() * |
| device.maxCudaThreadsPerMultiProcessor() / block_size; |
| const int num_blocks = numext::mini<int>(max_blocks, dyn_blocks); |
| |
| if (num_blocks > 1) { |
| // We initialize the outputs in the reduction kernel itself when we don't have to worry |
| // about race conditions between multiple thread blocks. |
| const int dyn_blocks = divup<int>(num_preserved_vals, 1024); |
| const int max_blocks = device.getNumCudaMultiProcessors() * |
| device.maxCudaThreadsPerMultiProcessor() / 1024; |
| const int num_blocks = numext::mini<int>(max_blocks, dyn_blocks); |
| LAUNCH_CUDA_KERNEL((ReductionInitKernel<float, Index>), |
| num_blocks, 1024, 0, device, reducer.initialize(), |
| num_preserved_vals, output); |
| } |
| |
| LAUNCH_CUDA_KERNEL((InnerReductionKernel<num_per_thread, Self, Op, Index>), |
| num_blocks, block_size, 0, device, reducer, self, num_coeffs_to_reduce, num_preserved_vals, output); |
| } |
| }; |
| |
| |
| template <typename Self, typename Op> |
| struct InnerGpuReducer { |
| // Unfortunately nvidia doesn't support well exotic types such as complex, |
| // so reduce the scope of the optimized version of the code to the simple case |
| // of floats. |
| static const bool HasOptimizedImplementation = !Op::IsStateful && |
| internal::is_same<typename Self::CoeffReturnType, float>::value; |
| |
| |
| template <typename Device, typename OutputType> |
| static EIGEN_DEVICE_FUNC void run(const Self&, Op&, const Device&, OutputType*, typename Self::Index, typename Self::Index) { |
| assert(false && "Should only be called to reduce floats on a gpu device"); |
| } |
| |
| template <typename OutputType> |
| static void run(const Self& self, Op& reducer, const GpuDevice& device, OutputType* output, typename Self::Index num_coeffs_to_reduce, typename Self::Index num_preserved_vals) { |
| assert(HasOptimizedImplementation && "Should only be called on floats or half floats"); |
| const Index num_coeffs = num_coeffs_to_reduce * num_preserved_vals; |
| // Don't crash when we're called with an input tensor of size 0. |
| if (num_coeffs == 0) { |
| return; |
| } |
| InnerReductionLauncher<Self, Op, OutputType, reducer_traits<Op, GpuDevice>::PacketAccess>::run(self, reducer, device, output, num_coeffs_to_reduce, num_preserved_vals); |
| } |
| }; |
| |
| |
| template <int NumPerThread, typename Self, |
| typename Reducer, typename Index> |
| __global__ void OuterReductionKernel(Reducer reducer, const Self input, Index num_coeffs_to_reduce, Index num_preserved_coeffs, |
| typename Self::CoeffReturnType* output) { |
| const Index num_threads = blockDim.x * gridDim.x; |
| const Index thread_id = blockIdx.x * blockDim.x + threadIdx.x; |
| // Initialize the output values if they weren't initialized by the ReductionInitKernel |
| if (gridDim.x == 1) { |
| for (Index i = thread_id; i < num_preserved_coeffs; i += num_threads) { |
| output[i] = reducer.initialize(); |
| } |
| __syncthreads(); |
| } |
| |
| // Do the reduction. |
| const Index max_iter = num_preserved_coeffs * divup<Index>(num_coeffs_to_reduce, NumPerThread); |
| for (Index i = thread_id; i < max_iter; i += num_threads) { |
| const Index input_col = i % num_preserved_coeffs; |
| const Index input_row = (i / num_preserved_coeffs) * NumPerThread; |
| typename Self::CoeffReturnType reduced_val = reducer.initialize(); |
| const Index max_row = numext::mini(input_row + NumPerThread, num_coeffs_to_reduce); |
| for (Index j = input_row; j < max_row; j++) { |
| typename Self::CoeffReturnType val = input.m_impl.coeff(j * num_preserved_coeffs + input_col); |
| reducer.reduce(val, &reduced_val); |
| } |
| atomicReduce(&(output[input_col]), reduced_val, reducer); |
| } |
| } |
| |
| |
| template <typename Self, typename Op> |
| struct OuterGpuReducer { |
| // Unfortunately nvidia doesn't support well exotic types such as complex, |
| // so reduce the scope of the optimized version of the code to the simple case |
| // of floats. |
| static const bool HasOptimizedImplementation = !Op::IsStateful && |
| internal::is_same<typename Self::CoeffReturnType, float>::value; |
| |
| template <typename Device, typename OutputType> |
| static EIGEN_DEVICE_FUNC void run(const Self&, Op&, const Device&, OutputType*, typename Self::Index, typename Self::Index) { |
| assert(false && "Should only be called to reduce floats on a gpu device"); |
| } |
| |
| static void run(const Self& self, Op& reducer, const GpuDevice& device, float* output, typename Self::Index num_coeffs_to_reduce, typename Self::Index num_preserved_vals) { |
| typedef typename Self::Index Index; |
| |
| const Index num_coeffs = num_coeffs_to_reduce * num_preserved_vals; |
| // Don't crash when we're called with an input tensor of size 0. |
| if (num_coeffs == 0) { |
| return; |
| } |
| |
| const int block_size = 256; |
| const int num_per_thread = 16; |
| const int dyn_blocks = divup<int>(num_coeffs, block_size * num_per_thread); |
| const int max_blocks = device.getNumCudaMultiProcessors() * |
| device.maxCudaThreadsPerMultiProcessor() / block_size; |
| const int num_blocks = numext::mini<int>(max_blocks, dyn_blocks); |
| |
| if (num_blocks > 1) { |
| // We initialize the outputs in the reduction kernel itself when we don't have to worry |
| // about race conditions between multiple thread blocks. |
| const int dyn_blocks = divup<int>(num_preserved_vals, 1024); |
| const int max_blocks = device.getNumCudaMultiProcessors() * |
| device.maxCudaThreadsPerMultiProcessor() / 1024; |
| const int num_blocks = numext::mini<int>(max_blocks, dyn_blocks); |
| LAUNCH_CUDA_KERNEL((ReductionInitKernel<float, Index>), |
| num_blocks, 1024, 0, device, reducer.initialize(), |
| num_preserved_vals, output); |
| } |
| |
| LAUNCH_CUDA_KERNEL((OuterReductionKernel<num_per_thread, Self, Op, Index>), |
| num_blocks, block_size, 0, device, reducer, self, num_coeffs_to_reduce, num_preserved_vals, output); |
| } |
| }; |
| |
| #endif |
| |
| |
| template <typename Self, typename Op, |
| bool Vectorizable = (Self::InputPacketAccess & Op::PacketAccess)> |
| class BlockReducer { |
| public: |
| typedef typename Self::Index Index; |
| typedef typename Self::Scalar Scalar; |
| typedef typename Self::CoeffReturnType CoeffReturnType; |
| typedef typename Self::PacketReturnType PacketReturnType; |
| explicit BlockReducer(const Op& reducer) : op_(reducer) { |
| accum_ = op_.initialize(); |
| } |
| void Reduce(Index index, Index num_values_to_reduce, Scalar* data) { |
| for (Index i = 0; i < num_values_to_reduce; ++i) { |
| op_.reduce(data[index + i], &accum_); |
| } |
| } |
| CoeffReturnType Finalize() { |
| return op_.finalize(accum_); |
| } |
| PacketReturnType FinalizePacket() { |
| // TODO(andydavis) This function should not be called for Scalar |
| // reductions: clean this up or add an assert here. |
| return PacketReturnType(); |
| } |
| |
| private: |
| CoeffReturnType accum_; |
| Op op_; |
| }; |
| |
| template <typename Self, typename Op> |
| class BlockReducer<Self, Op, true> { |
| public: |
| typedef typename Self::Index Index; |
| typedef typename Self::Scalar Scalar; |
| typedef typename Self::CoeffReturnType CoeffReturnType; |
| typedef typename Self::PacketReturnType PacketReturnType; |
| static const int PacketSize = internal::unpacket_traits<PacketReturnType>::size; |
| |
| explicit BlockReducer(const Op& reducer) : op_(reducer) { |
| vaccum_ = op_.template initializePacket<PacketReturnType>(); |
| accum_ = op_.initialize(); |
| } |
| void Reduce(Index index, Index num_values_to_reduce, Scalar* data) { |
| const Index vectorized_size = (num_values_to_reduce / PacketSize) * |
| PacketSize; |
| for (Index i = 0; i < vectorized_size; i += PacketSize) { |
| op_.reducePacket(internal::ploadt<PacketReturnType, Unaligned>( |
| &data[index + i]), &vaccum_); |
| } |
| for (Index i = vectorized_size; i < num_values_to_reduce; ++i) { |
| op_.reduce(data[index + i], &accum_); |
| } |
| } |
| CoeffReturnType Finalize() { |
| return op_.finalizeBoth(accum_, vaccum_); |
| } |
| PacketReturnType FinalizePacket() { |
| return op_.finalizePacket(vaccum_); |
| } |
| |
| private: |
| PacketReturnType vaccum_; |
| CoeffReturnType accum_; |
| Op op_; |
| }; |
| |
| } // end namespace internal |
| |
| |
| template <typename Op, typename Dims, typename XprType> |
| class TensorReductionOp : public TensorBase<TensorReductionOp<Op, Dims, XprType>, ReadOnlyAccessors> { |
| public: |
| typedef typename Eigen::internal::traits<TensorReductionOp>::Scalar Scalar; |
| typedef typename Eigen::NumTraits<Scalar>::Real RealScalar; |
| typedef typename internal::remove_const<typename XprType::CoeffReturnType>::type CoeffReturnType; |
| typedef typename Eigen::internal::nested<TensorReductionOp>::type Nested; |
| typedef typename Eigen::internal::traits<TensorReductionOp>::StorageKind StorageKind; |
| typedef typename Eigen::internal::traits<TensorReductionOp>::Index Index; |
| |
| EIGEN_DEVICE_FUNC EIGEN_STRONG_INLINE |
| TensorReductionOp(const XprType& expr, const Dims& dims) : m_expr(expr), m_dims(dims) |
| { } |
| EIGEN_DEVICE_FUNC EIGEN_STRONG_INLINE |
| TensorReductionOp(const XprType& expr, const Dims& dims, const Op& reducer) : m_expr(expr), m_dims(dims), m_reducer(reducer) |
| { } |
| |
| EIGEN_DEVICE_FUNC EIGEN_STRONG_INLINE |
| const XprType& expression() const { return m_expr; } |
| EIGEN_DEVICE_FUNC EIGEN_STRONG_INLINE |
| const Dims& dims() const { return m_dims; } |
| EIGEN_DEVICE_FUNC EIGEN_STRONG_INLINE |
| const Op& reducer() const { return m_reducer; } |
| |
| protected: |
| typename XprType::Nested m_expr; |
| const Dims m_dims; |
| const Op m_reducer; |
| }; |
| |
| |
| // Eval as rvalue |
| template<typename Op, typename Dims, typename ArgType, typename Device> |
| struct TensorEvaluator<const TensorReductionOp<Op, Dims, ArgType>, Device> |
| { |
| typedef TensorReductionOp<Op, Dims, ArgType> XprType; |
| typedef typename XprType::Index Index; |
| typedef typename TensorEvaluator<ArgType, Device>::Dimensions InputDimensions; |
| static const int NumInputDims = internal::array_size<InputDimensions>::value; |
| static const int NumReducedDims = internal::array_size<Dims>::value; |
| EIGEN_STATIC_ASSERT(NumInputDims >= NumReducedDims, YOU_MADE_A_PROGRAMMING_MISTAKE) |
| static const int NumOutputDims = NumInputDims - NumReducedDims; |
| typedef DSizes<Index, NumOutputDims> Dimensions; |
| typedef typename XprType::Scalar Scalar; |
| typedef typename internal::remove_const<Scalar>::type ScalarNonConst; |
| typedef TensorEvaluator<const TensorReductionOp<Op, Dims, ArgType>, Device> Self; |
| static const bool InputPacketAccess = TensorEvaluator<ArgType, Device>::PacketAccess; |
| typedef typename internal::remove_const<typename XprType::CoeffReturnType>::type CoeffReturnType; |
| typedef typename PacketType<CoeffReturnType, Device>::type PacketReturnType; |
| static const int PacketSize = internal::unpacket_traits<PacketReturnType>::size; |
| |
| enum { |
| IsAligned = false, |
| PacketAccess = Self::InputPacketAccess && Op::PacketAccess, |
| BlockAccess = TensorEvaluator<ArgType, Device>::BlockAccess, |
| Layout = TensorEvaluator<ArgType, Device>::Layout, |
| CoordAccess = false, // to be implemented |
| RawAccess = false |
| }; |
| |
| typedef typename internal::TensorBlock<Index, ScalarNonConst, NumOutputDims, |
| Layout> OutputTensorBlock; |
| typedef typename internal::TensorBlock<Index, ScalarNonConst, NumInputDims, |
| Layout> InputTensorBlock; |
| |
| static const bool ReducingInnerMostDims = internal::are_inner_most_dims<Dims, NumInputDims, Layout>::value; |
| static const bool PreservingInnerMostDims = internal::preserve_inner_most_dims<Dims, NumInputDims, Layout>::value; |
| static const bool RunningFullReduction = (NumInputDims==NumReducedDims); |
| |
| EIGEN_DEVICE_FUNC EIGEN_STRONG_INLINE TensorEvaluator(const XprType& op, const Device& device) |
| : m_impl(op.expression(), device), m_reducer(op.reducer()), m_result(NULL), m_device(device) |
| { |
| EIGEN_STATIC_ASSERT((!ReducingInnerMostDims | !PreservingInnerMostDims | (NumReducedDims == NumInputDims)), |
| YOU_MADE_A_PROGRAMMING_MISTAKE); |
| for (int i = 0; i < NumInputDims; ++i) { |
| m_reduced_dim[i] = false; |
| } |
| for (int i = 0; i < NumReducedDims; ++i) { |
| eigen_assert(op.dims()[i] >= 0); |
| eigen_assert(op.dims()[i] < NumInputDims); |
| m_reduced_dim[op.dims()[i]] = true; |
| } |
| |
| const typename TensorEvaluator<ArgType, Device>::Dimensions& input_dims = m_impl.dimensions(); |
| internal::partition_dims(input_dims, m_reduced_dim, &m_dimensions, &m_reducedDims); |
| |
| // Precompute output strides. |
| if (NumOutputDims > 0) { |
| if (static_cast<int>(Layout) == static_cast<int>(ColMajor)) { |
| m_outputStrides[0] = 1; |
| for (int i = 1; i < NumOutputDims; ++i) { |
| m_outputStrides[i] = m_outputStrides[i - 1] * m_dimensions[i - 1]; |
| m_fastOutputStrides[i] = internal::TensorIntDivisor<Index>(m_outputStrides[i]); |
| } |
| } else { |
| m_outputStrides[NumOutputDims - 1] = 1; |
| for (int i = NumOutputDims - 2; i >= 0; --i) { |
| m_outputStrides[i] = m_outputStrides[i + 1] * m_dimensions[i + 1]; |
| m_fastOutputStrides[i] = internal::TensorIntDivisor<Index>(m_outputStrides[i]); |
| } |
| } |
| } |
| |
| // Precompute input strides. |
| if (NumInputDims > 0) { |
| if (static_cast<int>(Layout) == static_cast<int>(ColMajor)) { |
| m_inputStrides[0] = 1; |
| for (int i = 1; i < NumInputDims; ++i) { |
| m_inputStrides[i] = m_inputStrides[i-1] * input_dims[i-1]; |
| } |
| } else { |
| m_inputStrides[NumInputDims - 1] = 1; |
| for (int i = NumInputDims - 2; i >= 0; --i) { |
| m_inputStrides[i] = m_inputStrides[i + 1] * input_dims[i + 1]; |
| } |
| } |
| } |
| |
| int outputIndex = 0; |
| int reduceIndex = 0; |
| for (int i = 0; i < NumInputDims; ++i) { |
| if (m_reduced_dim[i]) { |
| m_reducedStrides[reduceIndex] = m_inputStrides[i]; |
| ++reduceIndex; |
| } else { |
| m_preservedStrides[outputIndex] = m_inputStrides[i]; |
| m_output_to_input_dim_map[outputIndex] = i; |
| ++outputIndex; |
| } |
| } |
| |
| m_numValuesToReduce |
| = NumOutputDims == 0 ? internal::array_prod(input_dims) |
| : (static_cast<int>(Layout) == static_cast<int>(ColMajor)) |
| ? m_preservedStrides[0] : m_preservedStrides[NumOutputDims - 1]; |
| |
| m_block_total_size_max = numext::maxi(static_cast<std::size_t>(1), |
| device.lastLevelCacheSize() / |
| sizeof(Scalar)); |
| } |
| |
| EIGEN_DEVICE_FUNC EIGEN_STRONG_INLINE const Dimensions& dimensions() const { return m_dimensions; } |
| |
| EIGEN_DEVICE_FUNC EIGEN_STRONG_INLINE bool evalSubExprsIfNeeded(CoeffReturnType* data) { |
| m_impl.evalSubExprsIfNeeded(NULL); |
| |
| // Use the FullReducer if possible. |
| if (RunningFullReduction && |
| internal::FullReducer<Self, Op, Device>::HasOptimizedImplementation && |
| ((RunningOnGPU && (m_device.majorDeviceVersion() >= 3)) || |
| !RunningOnGPU)) { |
| bool need_assign = false; |
| if (!data) { |
| m_result = static_cast<CoeffReturnType*>(m_device.allocate(sizeof(CoeffReturnType))); |
| // Make sure the data is aligned. |
| eigen_assert((reinterpret_cast<size_t>(m_result) & 0x3) == 0); |
| data = m_result; |
| need_assign = true; |
| } |
| else if ((reinterpret_cast<size_t>(data) & 0x3) != 0) { |
| // The data isn't aligned, so revert to the slow code path |
| return true; |
| } |
| Op reducer(m_reducer); |
| internal::FullReducer<Self, Op, Device>::run(*this, reducer, m_device, data); |
| return need_assign; |
| } |
| // Attempt to use an optimized reduction. This requires that the data is aligned. |
| #if defined(EIGEN_USE_GPU) && defined(__CUDACC__) |
| else if (RunningOnGPU && ((reinterpret_cast<size_t>(data) & 0x3) == 0) && (m_device.majorDeviceVersion() >= 3)) { |
| bool reducing_inner_dims = true; |
| for (int i = 0; i < NumReducedDims; ++i) { |
| if (static_cast<int>(Layout) == static_cast<int>(ColMajor)) { |
| reducing_inner_dims &= m_reduced_dim[i]; |
| } else { |
| reducing_inner_dims &= m_reduced_dim[NumInputDims - 1 - i]; |
| } |
| } |
| if (internal::InnerGpuReducer<Self, Op>::HasOptimizedImplementation && |
| (reducing_inner_dims || ReducingInnerMostDims)) { |
| const Index num_values_to_reduce = internal::array_prod(m_reducedDims); |
| if (num_values_to_reduce <= 32) { |
| // It's faster to call the usual codepath if we have fewer than warpSize |
| // values to reduce. |
| return true; |
| } |
| const Index num_coeffs_to_preserve = internal::array_prod(m_dimensions); |
| bool need_assign = false; |
| if (!data) { |
| m_result = static_cast<CoeffReturnType*>(m_device.allocate(sizeof(CoeffReturnType) * num_coeffs_to_preserve)); |
| data = m_result; |
| need_assign = true; |
| } |
| Op reducer(m_reducer); |
| internal::InnerGpuReducer<Self, Op>::run(*this, reducer, m_device, data, num_values_to_reduce, num_coeffs_to_preserve); |
| return need_assign; |
| } |
| |
| bool preserving_inner_dims = true; |
| for (int i = 0; i < NumReducedDims; ++i) { |
| if (static_cast<int>(Layout) == static_cast<int>(ColMajor)) { |
| preserving_inner_dims &= m_reduced_dim[NumInputDims - 1 - i]; |
| } else { |
| preserving_inner_dims &= m_reduced_dim[i]; |
| } |
| } |
| if (internal::OuterGpuReducer<Self, Op>::HasOptimizedImplementation && |
| (preserving_inner_dims)) { |
| const Index num_values_to_reduce = internal::array_prod(m_reducedDims); |
| if (num_values_to_reduce <= 32) { |
| // It's faster to call the usual codepath if we have fewer than warpSize |
| // values to reduce. |
| return true; |
| } |
| const Index num_coeffs_to_preserve = internal::array_prod(m_dimensions); |
| bool need_assign = false; |
| if (!data) { |
| m_result = static_cast<CoeffReturnType*>(m_device.allocate(sizeof(CoeffReturnType) * num_coeffs_to_preserve)); |
| data = m_result; |
| need_assign = true; |
| } |
| Op reducer(m_reducer); |
| internal::OuterGpuReducer<Self, Op>::run(*this, reducer, m_device, data, num_values_to_reduce, num_coeffs_to_preserve); |
| return need_assign; |
| } |
| } |
| #endif |
| return true; |
| } |
| |
| EIGEN_DEVICE_FUNC EIGEN_STRONG_INLINE void cleanup() { |
| m_impl.cleanup(); |
| |
| if (m_result) { |
| m_device.deallocate(m_result); |
| } |
| } |
| |
| EIGEN_DEVICE_FUNC EIGEN_STRONG_INLINE CoeffReturnType coeff(Index index) const |
| { |
| if ((RunningFullReduction || RunningOnGPU) && m_result) { |
| return *(m_result + index); |
| } |
| Op reducer(m_reducer); |
| if (ReducingInnerMostDims) { |
| return internal::InnerMostDimReducer<Self, Op>::reduce(*this, firstInput(index), |
| m_numValuesToReduce, reducer); |
| } else { |
| typename Self::CoeffReturnType accum = reducer.initialize(); |
| internal::GenericDimReducer<NumReducedDims-1, Self, Op>::reduce(*this, firstInput(index), reducer, &accum); |
| return reducer.finalize(accum); |
| } |
| } |
| |
| // TODO(bsteiner): provide a more efficient implementation. |
| template<int LoadMode> |
| EIGEN_DEVICE_FUNC EIGEN_STRONG_INLINE PacketReturnType packet(Index index) const |
| { |
| const int packetSize = internal::unpacket_traits<PacketReturnType>::size; |
| EIGEN_STATIC_ASSERT(packetSize > 1, YOU_MADE_A_PROGRAMMING_MISTAKE) |
| eigen_assert(index + packetSize - 1 < dimensions().TotalSize()); |
| |
| if (RunningOnGPU && m_result) { |
| return internal::pload<PacketReturnType>(m_result + index); |
| } |
| |
| EIGEN_ALIGN_DEFAULT typename internal::remove_const<CoeffReturnType>::type values[packetSize]; |
| if (ReducingInnerMostDims) { |
| const Index num_values_to_reduce = m_numValuesToReduce; |
| const Index firstIndex = firstInput(index); |
| for (Index i = 0; i < packetSize; ++i) { |
| Op reducer(m_reducer); |
| values[i] = internal::InnerMostDimReducer<Self, Op>::reduce(*this, firstIndex + i * num_values_to_reduce, |
| num_values_to_reduce, reducer); |
| } |
| } else if (PreservingInnerMostDims) { |
| const Index firstIndex = firstInput(index); |
| const int innermost_dim = (static_cast<int>(Layout) == static_cast<int>(ColMajor)) ? 0 : NumOutputDims - 1; |
| // TBD: extend this the the n innermost dimensions that we preserve. |
| if (((firstIndex % m_dimensions[innermost_dim]) + packetSize - 1) < m_dimensions[innermost_dim]) { |
| Op reducer(m_reducer); |
| typename Self::PacketReturnType accum = reducer.template initializePacket<typename Self::PacketReturnType>(); |
| internal::InnerMostDimPreserver<NumReducedDims-1, Self, Op>::reduce(*this, firstIndex, reducer, &accum); |
| return reducer.finalizePacket(accum); |
| } else { |
| for (int i = 0; i < packetSize; ++i) { |
| values[i] = coeff(index + i); |
| } |
| } |
| } else { |
| for (int i = 0; i < packetSize; ++i) { |
| values[i] = coeff(index + i); |
| } |
| } |
| PacketReturnType rslt = internal::pload<PacketReturnType>(values); |
| return rslt; |
| } |
| |
| EIGEN_DEVICE_FUNC EIGEN_STRONG_INLINE void getResourceRequirements( |
| std::vector<internal::TensorOpResourceRequirements>* resources) const { |
| resources->push_back(internal::TensorOpResourceRequirements( |
| internal::kSkewedInnerDims, m_block_total_size_max)); |
| m_impl.getResourceRequirements(resources); |
| } |
| |
| // Must be called after evalSubExprsIfNeeded(). |
| EIGEN_DEVICE_FUNC EIGEN_STRONG_INLINE TensorOpCost costPerCoeff(bool vectorized) const { |
| if (RunningFullReduction && m_result) { |
| return TensorOpCost(sizeof(CoeffReturnType), 0, 0, vectorized, PacketSize); |
| } else { |
| const Index num_values_to_reduce = internal::array_prod(m_reducedDims); |
| const double compute_cost = num_values_to_reduce * internal::functor_traits<Op>::Cost; |
| return m_impl.costPerCoeff(vectorized) * num_values_to_reduce + |
| TensorOpCost(0, 0, compute_cost, vectorized, PacketSize); |
| } |
| } |
| |
| EIGEN_DEVICE_FUNC EIGEN_DONT_INLINE void block( |
| OutputTensorBlock* output_block) const { |
| // Special case full reductions to avoid input block copy below. |
| if (NumInputDims == NumReducedDims) { |
| eigen_assert(output_block->first_coeff_index() == 0); |
| eigen_assert(output_block->block_sizes().TotalSize() == 1); |
| Op reducer(m_reducer); |
| output_block->data()[0] = internal::InnerMostDimReducer<Self, Op>::reduce( |
| *this, 0, m_numValuesToReduce, reducer); |
| return; |
| } |
| |
| // Calculate input tensor 'slice' required to reduce output block coeffs. |
| DSizes<Index, NumInputDims> input_slice_sizes(m_impl.dimensions()); |
| for (int i = 0; i < NumOutputDims; ++i) { |
| // Clip preserved input dimensions by output block size. |
| input_slice_sizes[m_output_to_input_dim_map[i]] = |
| output_block->block_sizes()[i]; |
| } |
| |
| // Shard input tensor slice into blocks (because it could be large if we |
| // need to reduce along several dimensions to calculate required output |
| // coefficients). |
| const Index max_coeff_count = |
| numext::mini(((m_device.firstLevelCacheSize()) / sizeof(Scalar)), |
| input_slice_sizes.TotalSize()); |
| |
| // Calculate max output shard size needed to keep working set of reducers |
| // in L1, while leaving enough space for reducer overhead and 'PacketSize' |
| // reductions. |
| DSizes<Index, NumInputDims> target_input_block_sizes; |
| CalculateTargetInputBlockShape(max_coeff_count, input_slice_sizes, |
| &target_input_block_sizes); |
| // Calculate indices for first preserved dimension. |
| const Index first_preserved_dim_output_index = |
| static_cast<int>(Layout) == static_cast<int>(ColMajor) ? |
| 0 : NumOutputDims - 1; |
| const Index first_preserved_dim_input_index = m_output_to_input_dim_map[ |
| first_preserved_dim_output_index]; |
| const bool inner_most_dim_preserved = first_preserved_dim_input_index == |
| (static_cast<int>(Layout) == static_cast<int>(ColMajor) ? 0 : |
| NumInputDims - 1) | PreservingInnerMostDims; |
| |
| // Calculate output block inner/outer dimension sizes. |
| const Index output_block_inner_dim_size = output_block->block_sizes()[ |
| first_preserved_dim_output_index]; |
| const Index output_block_outer_dim_size = |
| output_block->block_sizes().TotalSize() / output_block_inner_dim_size; |
| // Calculate shard size for first preserved dimension. |
| const Index output_shard_size = target_input_block_sizes[ |
| first_preserved_dim_input_index]; |
| const Index num_output_shards = |
| (output_block_inner_dim_size + output_shard_size - 1) / |
| output_shard_size; |
| |
| // Initialize 'tensor_slice_offsets' from input coords of output index. |
| DSizes<Index, NumInputDims> tensor_slice_offsets; |
| GetInputCoordsForOutputIndex(output_block->first_coeff_index(), |
| &tensor_slice_offsets); |
| |
| // Store tensor slice offset in first preserved dimension to be used |
| // to update tensor slice extents in loop below. |
| const Index first_preserved_dim_offset_start = tensor_slice_offsets[ |
| first_preserved_dim_input_index]; |
| |
| array<BlockIteratorState, NumOutputDims> block_iter_state; |
| |
| // Initialize state used to iterate through output coefficients |
| // and update 'tensor_slice_offsets' in outer preserved dims. |
| for (int i = 0; i < NumOutputDims - 1; ++i) { |
| const int dim = static_cast<int>(Layout) == static_cast<int>(ColMajor) |
| ? i + 1 : NumOutputDims - i - 2; |
| block_iter_state[i].input_dim = m_output_to_input_dim_map[dim]; |
| block_iter_state[i].output_size = output_block->block_sizes()[dim]; |
| block_iter_state[i].output_count = 0; |
| } |
| |
| // Allocate input block memory. |
| ScalarNonConst* input_block_data = static_cast<ScalarNonConst*>( |
| m_device.allocate(max_coeff_count * sizeof(Scalar))); |
| // Allocate reducer memory. |
| const bool packet_reductions_enabled = (Self::InputPacketAccess & |
| Op::PacketAccess); |
| const Index num_reducers = |
| (inner_most_dim_preserved && packet_reductions_enabled) ? |
| (output_shard_size / PacketSize + output_shard_size % PacketSize + |
| PacketSize) : output_shard_size; |
| typedef internal::BlockReducer<Self, Op> BlockReducer; |
| BlockReducer* reducers = static_cast<BlockReducer*>( |
| m_device.allocate(num_reducers * sizeof(BlockReducer))); |
| |
| InputDimensions input_tensor_dims(m_impl.dimensions()); |
| for (Index output_outer_index = 0; |
| output_outer_index < output_block_outer_dim_size; |
| ++output_outer_index) { |
| for (Index output_shard_index = 0; |
| output_shard_index < num_output_shards; |
| ++output_shard_index) { |
| // Initialize 'tensor_slice_extents' for this output shard. |
| DSizes<Index, NumInputDims> tensor_slice_extents(input_slice_sizes); |
| for (int i = 0; i < NumInputDims; ++i) { |
| if (i == first_preserved_dim_input_index) { |
| // Clip first preserved dim size to output shard size. |
| tensor_slice_extents[i] = numext::mini( |
| output_shard_size, |
| input_slice_sizes[i] - (tensor_slice_offsets[i] - |
| first_preserved_dim_offset_start)); |
| |
| } else if (!m_reduced_dim[i]) { |
| // Clip outer preserved dims to size 1, so that we reduce a |
| // contiguous set of output coefficients. |
| tensor_slice_extents[i] = 1; |
| } |
| } |
| |
| // Intialize output coefficient reducers. |
| for (int i = 0; i < num_reducers; ++i) { |
| new (&reducers[i]) BlockReducer(m_reducer); |
| } |
| |
| typedef internal::TensorSliceBlockMapper< |
| Index, ScalarNonConst, NumInputDims, Layout> TensorSliceBlockMapper; |
| |
| // TODO(andydavis) Consider removing 'input_block_stride_order' if we |
| // find that scattered reads are not worth supporting in |
| // TensorSliceBlockMapper. |
| TensorSliceBlockMapper block_mapper( |
| input_tensor_dims, tensor_slice_offsets, tensor_slice_extents, |
| target_input_block_sizes, DimensionList<Index, NumInputDims>()); |
| |
| const Index num_outputs_to_update = tensor_slice_extents[ |
| first_preserved_dim_input_index]; |
| const Index preserved_dim_vector_reducer_count = |
| (inner_most_dim_preserved && packet_reductions_enabled) ? |
| num_outputs_to_update / PacketSize: 0; |
| const Index preserved_dim_vector_coeff_count = |
| inner_most_dim_preserved ? preserved_dim_vector_reducer_count * |
| PacketSize : 0; |
| const Index preserved_dim_reducer_limit = |
| (inner_most_dim_preserved && packet_reductions_enabled) ? |
| (preserved_dim_vector_reducer_count + |
| num_outputs_to_update % PacketSize) : num_outputs_to_update; |
| |
| const Index total_block_count = block_mapper.total_block_count(); |
| for (Index b = 0; b < total_block_count; ++b) { |
| InputTensorBlock input_block = block_mapper.GetBlockForIndex( |
| b, input_block_data); |
| // Read. |
| m_impl.block(&input_block); |
| |
| Index num_values_to_reduce = 1; |
| for (Index i = 0; i < NumInputDims; ++i) { |
| if (m_reduced_dim[i]) { |
| num_values_to_reduce *= input_block.block_sizes()[i]; |
| } |
| } |
| // Reduce. |
| if (inner_most_dim_preserved) { |
| const Index input_outer_dim_size = |
| input_block.block_sizes().TotalSize() / num_outputs_to_update; |
| for (Index input_outer_dim_index = 0; |
| input_outer_dim_index < input_outer_dim_size; |
| ++input_outer_dim_index) { |
| const Index input_outer_dim_base = input_outer_dim_index * |
| num_outputs_to_update; |
| for (Index i = 0; i < preserved_dim_vector_reducer_count; ++i) { |
| reducers[i].Reduce(input_outer_dim_base + i * PacketSize, |
| PacketSize, input_block.data()); |
| } |
| const Index scalar_reducer_base = input_outer_dim_base + |
| preserved_dim_vector_coeff_count; |
| for (Index i = preserved_dim_vector_reducer_count; |
| i < preserved_dim_reducer_limit; ++i) { |
| reducers[i].Reduce(scalar_reducer_base + i - |
| preserved_dim_vector_reducer_count, |
| 1, |
| input_block.data()); |
| } |
| } |
| } else { |
| for (Index i = 0; i < num_outputs_to_update; ++i) { |
| reducers[i].Reduce(i * num_values_to_reduce, |
| num_values_to_reduce, |
| input_block.data()); |
| } |
| } |
| } |
| |
| // Finalize all reducers for this output shard. |
| const Index output_base_index = |
| output_outer_index * output_block_inner_dim_size + |
| output_shard_index * output_shard_size; |
| if (inner_most_dim_preserved) { |
| EIGEN_ALIGN_DEFAULT typename internal::remove_const<CoeffReturnType>::type values[PacketSize]; |
| for (Index i = 0; i < preserved_dim_vector_reducer_count; ++i) { |
| const Index reducer_base = output_base_index + i * PacketSize; |
| internal::pstore<CoeffReturnType, PacketReturnType>( |
| values, reducers[i].FinalizePacket()); |
| for (Index j = 0; j < PacketSize; ++j) { |
| output_block->data()[reducer_base + j] = values[j]; |
| } |
| } |
| const Index scalar_reducer_base = output_base_index + |
| preserved_dim_vector_coeff_count; |
| |
| for (Index i = preserved_dim_vector_reducer_count; |
| i < preserved_dim_reducer_limit; ++i) { |
| output_block->data()[ |
| scalar_reducer_base + i - preserved_dim_vector_reducer_count] = |
| reducers[i].Finalize(); |
| } |
| } else { |
| for (int i = 0; i < num_outputs_to_update; ++i) { |
| output_block->data()[output_base_index + i] = |
| reducers[i].Finalize(); |
| } |
| } |
| |
| // Update 'tensor_slice_offsets' by num outputs for this output shard. |
| tensor_slice_offsets[first_preserved_dim_input_index] += |
| num_outputs_to_update; |
| } |
| // Update slice offset for inner preserved dim. |
| tensor_slice_offsets[first_preserved_dim_input_index] -= |
| output_block_inner_dim_size; |
| // Update slice offsets for remaining output dims. |
| for (int i = 0; i < NumOutputDims - 1; ++i) { |
| BlockIteratorState& b = block_iter_state[i]; |
| if (++b.output_count < b.output_size) { |
| ++tensor_slice_offsets[b.input_dim]; |
| break; |
| } |
| b.output_count = 0; |
| tensor_slice_offsets[b.input_dim] -= b.output_size - 1; |
| } |
| } |
| |
| // Free memory. |
| m_device.deallocate(input_block_data); |
| m_device.deallocate(reducers); |
| } |
| |
| EIGEN_DEVICE_FUNC Scalar* data() const { return NULL; } |
| |
| private: |
| template <int, typename, typename> friend struct internal::GenericDimReducer; |
| template <typename, typename, bool> friend struct internal::InnerMostDimReducer; |
| template <int, typename, typename, bool> friend struct internal::InnerMostDimPreserver; |
| template <typename S, typename O, typename D, bool V> friend struct internal::FullReducer; |
| #ifdef EIGEN_USE_THREADS |
| template <typename S, typename O, bool V> friend struct internal::FullReducerShard; |
| #endif |
| #ifdef EIGEN_HAS_CUDA_FP16 |
| template <typename S, typename R, typename I> friend void internal::ReductionInitKernelHalfFloat(R, const S, I, half*); |
| template <typename S, typename R, typename I> friend void internal::ReductionInitFullReduxKernelHalfFloat(R, const S, I, half2*); |
| template <int B, int N, typename S, typename R, typename I> friend void internal::FullReductionKernelHalfFloat(R, const S, I, half*, half2*); |
| |
| #endif |
| |
| #if defined(EIGEN_USE_GPU) && defined(__CUDACC__) |
| template <int B, int N, typename S, typename R, typename I> KERNEL_FRIEND void internal::FullReductionKernel(R, const S, I, typename S::CoeffReturnType*); |
| template <int N, typename S, typename R, typename I> KERNEL_FRIEND void internal::InnerReductionKernel(R, const S, I, I, typename S::CoeffReturnType*); |
| template <int N, typename S, typename R, typename I> KERNEL_FRIEND void internal::OuterReductionKernel(R, const S, I, I, typename S::CoeffReturnType*); |
| #endif |
| |
| struct BlockIteratorState { |
| Index input_dim; |
| Index output_size; |
| Index output_count; |
| }; |
| |
| // Returns the Index in the input tensor of the first value that needs to be |
| // used to compute the reduction at output index "index". |
| EIGEN_DEVICE_FUNC EIGEN_STRONG_INLINE Index firstInput(Index index) const { |
| if (ReducingInnerMostDims) { |
| return index * m_numValuesToReduce; |
| } |
| Index startInput = 0; |
| if (static_cast<int>(Layout) == static_cast<int>(ColMajor)) { |
| for (int i = NumOutputDims - 1; i > 0; --i) { |
| // This is index_i in the output tensor. |
| const Index idx = index / m_fastOutputStrides[i]; |
| startInput += idx * m_preservedStrides[i]; |
| index -= idx * m_outputStrides[i]; |
| } |
| } else { |
| for (int i = 0; i < NumOutputDims - 1; ++i) { |
| // This is index_i in the output tensor. |
| const Index idx = index / m_fastOutputStrides[i]; |
| startInput += idx * m_preservedStrides[i]; |
| index -= idx * m_outputStrides[i]; |
| } |
| } |
| if (PreservingInnerMostDims) { |
| eigen_assert(m_numValuesToReduce == 1); |
| startInput += index; |
| } else { |
| startInput += index * m_numValuesToReduce; |
| } |
| return startInput; |
| } |
| |
| EIGEN_DEVICE_FUNC EIGEN_STRONG_INLINE void GetInputCoordsForOutputIndex( |
| Index index, |
| DSizes<Index, NumInputDims>* coords) const { |
| for (int i = 0; i < NumInputDims; ++i) { |
| (*coords)[i] = 0; |
| } |
| if (static_cast<int>(Layout) == static_cast<int>(ColMajor)) { |
| for (int i = NumOutputDims - 1; i > 0; --i) { |
| const Index idx = index / m_fastOutputStrides[i]; |
| (*coords)[m_output_to_input_dim_map[i]] = idx; |
| index -= idx * m_outputStrides[i]; |
| } |
| (*coords)[m_output_to_input_dim_map[0]] = index; |
| } else { |
| for (int i = 0; i < NumOutputDims - 1; ++i) { |
| const Index idx = index / m_fastOutputStrides[i]; |
| (*coords)[m_output_to_input_dim_map[i]] = idx; |
| index -= idx * m_outputStrides[i]; |
| } |
| (*coords)[m_output_to_input_dim_map[NumOutputDims-1]] = index; |
| } |
| } |
| |
| EIGEN_DEVICE_FUNC EIGEN_STRONG_INLINE void CalculateTargetInputBlockShape( |
| const Index max_coeff_count, |
| const DSizes<Index, NumInputDims>& input_slice_sizes, |
| DSizes<Index, NumInputDims>* target_input_block_sizes) const { |
| typedef typename internal::packet_traits<Scalar>::type Packet; |
| typedef internal::BlockReducer<Self, Op> BlockReducer; |
| // TODO(andydavis) Compute reducer overhead correctly for the case where |
| // we are preserving the inner most dimension, and a single reducer |
| // reduces a packet's worth of output coefficients. |
| const Index reducer_overhead = sizeof(BlockReducer) / sizeof(Scalar); |
| |
| Index coeff_to_allocate = max_coeff_count; |
| bool first_preserved_dim_allocated = false; |
| bool first_reduced_dim_allocated = false; |
| for (int i = 0; i < NumInputDims; ++i) { |
| const int dim = static_cast<int>(Layout) == static_cast<int>(ColMajor) |
| ? i : NumInputDims - i - 1; |
| (*target_input_block_sizes)[dim] = 1; |
| if (m_reduced_dim[dim]) { |
| // TODO(andydavis) Consider allocating to multiple reduced dimensions. |
| // Watch out for cases where reduced dimensions are not contiguous, |
| // which induces scattered reads. |
| if (!first_reduced_dim_allocated) { |
| (*target_input_block_sizes)[dim] = numext::mini(input_slice_sizes[dim], |
| coeff_to_allocate); |
| coeff_to_allocate /= (*target_input_block_sizes)[dim]; |
| first_reduced_dim_allocated = true; |
| } |
| } else if (!first_preserved_dim_allocated) { |
| // TODO(andydavis) Include output block size in this L1 working set |
| // calculation. |
| const Index allocated = max_coeff_count - coeff_to_allocate; |
| const Index alloc_size = numext::maxi(static_cast<Index>(1), |
| coeff_to_allocate / |
| reducer_overhead); |
| (*target_input_block_sizes)[dim] = numext::mini(input_slice_sizes[dim], |
| alloc_size); |
| coeff_to_allocate = numext::maxi( |
| static_cast<Index>(1), |
| coeff_to_allocate / ((*target_input_block_sizes)[dim] * |
| reducer_overhead)); |
| first_preserved_dim_allocated = true; |
| } |
| } |
| } |
| |
| // Bitmap indicating if an input dimension is reduced or not. |
| array<bool, NumInputDims> m_reduced_dim; |
| // Dimensions of the output of the operation. |
| Dimensions m_dimensions; |
| // Precomputed strides for the input tensor. |
| array<Index, NumInputDims> m_inputStrides; |
| // Precomputed strides for the output tensor. |
| array<Index, NumOutputDims> m_outputStrides; |
| array<internal::TensorIntDivisor<Index>, NumOutputDims> m_fastOutputStrides; |
| // Subset of strides of the input tensor for the non-reduced dimensions. |
| // Indexed by output dimensions. |
| array<Index, NumOutputDims> m_preservedStrides; |
| // Map from output to input dimension index. |
| array<Index, NumOutputDims> m_output_to_input_dim_map; |
| // How many values go into each reduction |
| Index m_numValuesToReduce; |
| |
| // Subset of strides of the input tensor for the reduced dimensions. |
| // Indexed by reduced dimensions. |
| array<Index, NumReducedDims> m_reducedStrides; |
| // Size of the input dimensions that are reduced. |
| // Indexed by reduced dimensions. |
| array<Index, NumReducedDims> m_reducedDims; |
| |
| // Evaluator for the input expression. |
| TensorEvaluator<ArgType, Device> m_impl; |
| |
| // Operation to apply for computing the reduction. |
| Op m_reducer; |
| |
| // For full reductions |
| #ifdef EIGEN_USE_GPU |
| static const bool RunningOnGPU = internal::is_same<Device, Eigen::GpuDevice>::value; |
| #else |
| static const bool RunningOnGPU = false; |
| #endif |
| CoeffReturnType* m_result; |
| std::size_t m_block_total_size_max; |
| |
| const Device& m_device; |
| }; |
| |
| } // end namespace Eigen |
| |
| #endif // EIGEN_CXX11_TENSOR_TENSOR_REDUCTION_H |
