288 lines
9.7 KiB
C++
288 lines
9.7 KiB
C++
// This file is part of Eigen, a lightweight C++ template library
|
|
// for linear algebra.
|
|
//
|
|
// Copyright (C) 2016 Igor Babuschkin <igor@babuschk.in>
|
|
//
|
|
// 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 the mozilla.org home page
|
|
|
|
#ifndef EIGEN_CXX11_TENSOR_TENSOR_SCAN_H
|
|
#define EIGEN_CXX11_TENSOR_TENSOR_SCAN_H
|
|
|
|
namespace Eigen {
|
|
|
|
namespace internal {
|
|
|
|
template <typename Op, typename XprType>
|
|
struct traits<TensorScanOp<Op, XprType> >
|
|
: public traits<XprType> {
|
|
typedef typename XprType::Scalar Scalar;
|
|
typedef traits<XprType> XprTraits;
|
|
typedef typename XprTraits::StorageKind StorageKind;
|
|
typedef typename XprType::Nested Nested;
|
|
typedef typename remove_reference<Nested>::type _Nested;
|
|
static const int NumDimensions = XprTraits::NumDimensions;
|
|
static const int Layout = XprTraits::Layout;
|
|
};
|
|
|
|
template<typename Op, typename XprType>
|
|
struct eval<TensorScanOp<Op, XprType>, Eigen::Dense>
|
|
{
|
|
typedef const TensorScanOp<Op, XprType>& type;
|
|
};
|
|
|
|
template<typename Op, typename XprType>
|
|
struct nested<TensorScanOp<Op, XprType>, 1,
|
|
typename eval<TensorScanOp<Op, XprType> >::type>
|
|
{
|
|
typedef TensorScanOp<Op, XprType> type;
|
|
};
|
|
} // end namespace internal
|
|
|
|
/** \class TensorScan
|
|
* \ingroup CXX11_Tensor_Module
|
|
*
|
|
* \brief Tensor scan class.
|
|
*/
|
|
template <typename Op, typename XprType>
|
|
class TensorScanOp
|
|
: public TensorBase<TensorScanOp<Op, XprType>, ReadOnlyAccessors> {
|
|
public:
|
|
typedef typename Eigen::internal::traits<TensorScanOp>::Scalar Scalar;
|
|
typedef typename Eigen::NumTraits<Scalar>::Real RealScalar;
|
|
typedef typename XprType::CoeffReturnType CoeffReturnType;
|
|
typedef typename Eigen::internal::nested<TensorScanOp>::type Nested;
|
|
typedef typename Eigen::internal::traits<TensorScanOp>::StorageKind StorageKind;
|
|
typedef typename Eigen::internal::traits<TensorScanOp>::Index Index;
|
|
|
|
EIGEN_DEVICE_FUNC EIGEN_STRONG_INLINE TensorScanOp(
|
|
const XprType& expr, const Index& axis, bool exclusive = false, const Op& op = Op())
|
|
: m_expr(expr), m_axis(axis), m_accumulator(op), m_exclusive(exclusive) {}
|
|
|
|
EIGEN_DEVICE_FUNC EIGEN_STRONG_INLINE
|
|
const Index axis() const { return m_axis; }
|
|
EIGEN_DEVICE_FUNC EIGEN_STRONG_INLINE
|
|
const XprType& expression() const { return m_expr; }
|
|
EIGEN_DEVICE_FUNC EIGEN_STRONG_INLINE
|
|
const Op accumulator() const { return m_accumulator; }
|
|
EIGEN_DEVICE_FUNC EIGEN_STRONG_INLINE
|
|
bool exclusive() const { return m_exclusive; }
|
|
|
|
protected:
|
|
typename XprType::Nested m_expr;
|
|
const Index m_axis;
|
|
const Op m_accumulator;
|
|
const bool m_exclusive;
|
|
};
|
|
|
|
template <typename Self, typename Reducer, typename Device>
|
|
struct ScanLauncher;
|
|
|
|
// Eval as rvalue
|
|
template <typename Op, typename ArgType, typename Device>
|
|
struct TensorEvaluator<const TensorScanOp<Op, ArgType>, Device> {
|
|
|
|
typedef TensorScanOp<Op, ArgType> XprType;
|
|
typedef typename XprType::Index Index;
|
|
static const int NumDims = internal::array_size<typename TensorEvaluator<ArgType, Device>::Dimensions>::value;
|
|
typedef DSizes<Index, NumDims> Dimensions;
|
|
typedef typename internal::remove_const<typename XprType::Scalar>::type Scalar;
|
|
typedef typename XprType::CoeffReturnType CoeffReturnType;
|
|
typedef typename PacketType<CoeffReturnType, Device>::type PacketReturnType;
|
|
typedef TensorEvaluator<const TensorScanOp<Op, ArgType>, Device> Self;
|
|
|
|
enum {
|
|
IsAligned = false,
|
|
PacketAccess = (internal::unpacket_traits<PacketReturnType>::size > 1),
|
|
BlockAccess = false,
|
|
Layout = TensorEvaluator<ArgType, Device>::Layout,
|
|
CoordAccess = false,
|
|
RawAccess = true
|
|
};
|
|
|
|
EIGEN_DEVICE_FUNC EIGEN_STRONG_INLINE TensorEvaluator(const XprType& op,
|
|
const Device& device)
|
|
: m_impl(op.expression(), device),
|
|
m_device(device),
|
|
m_exclusive(op.exclusive()),
|
|
m_accumulator(op.accumulator()),
|
|
m_size(m_impl.dimensions()[op.axis()]),
|
|
m_stride(1),
|
|
m_output(NULL) {
|
|
|
|
// Accumulating a scalar isn't supported.
|
|
EIGEN_STATIC_ASSERT((NumDims > 0), YOU_MADE_A_PROGRAMMING_MISTAKE);
|
|
eigen_assert(op.axis() >= 0 && op.axis() < NumDims);
|
|
|
|
// Compute stride of scan axis
|
|
const Dimensions& dims = m_impl.dimensions();
|
|
if (static_cast<int>(Layout) == static_cast<int>(ColMajor)) {
|
|
for (int i = 0; i < op.axis(); ++i) {
|
|
m_stride = m_stride * dims[i];
|
|
}
|
|
} else {
|
|
for (int i = NumDims - 1; i > op.axis(); --i) {
|
|
m_stride = m_stride * dims[i];
|
|
}
|
|
}
|
|
}
|
|
|
|
EIGEN_DEVICE_FUNC EIGEN_STRONG_INLINE const Dimensions& dimensions() const {
|
|
return m_impl.dimensions();
|
|
}
|
|
|
|
EIGEN_DEVICE_FUNC EIGEN_STRONG_INLINE const Index& stride() const {
|
|
return m_stride;
|
|
}
|
|
|
|
EIGEN_DEVICE_FUNC EIGEN_STRONG_INLINE const Index& size() const {
|
|
return m_size;
|
|
}
|
|
|
|
EIGEN_DEVICE_FUNC EIGEN_STRONG_INLINE const Op& accumulator() const {
|
|
return m_accumulator;
|
|
}
|
|
|
|
EIGEN_DEVICE_FUNC EIGEN_STRONG_INLINE bool exclusive() const {
|
|
return m_exclusive;
|
|
}
|
|
|
|
EIGEN_DEVICE_FUNC EIGEN_STRONG_INLINE const TensorEvaluator<ArgType, Device>& inner() const {
|
|
return m_impl;
|
|
}
|
|
|
|
EIGEN_DEVICE_FUNC EIGEN_STRONG_INLINE const Device& device() const {
|
|
return m_device;
|
|
}
|
|
|
|
EIGEN_STRONG_INLINE bool evalSubExprsIfNeeded(Scalar* data) {
|
|
m_impl.evalSubExprsIfNeeded(NULL);
|
|
ScanLauncher<Self, Op, Device> launcher;
|
|
if (data) {
|
|
launcher(*this, data);
|
|
return false;
|
|
}
|
|
|
|
const Index total_size = internal::array_prod(dimensions());
|
|
m_output = static_cast<CoeffReturnType*>(m_device.allocate(total_size * sizeof(Scalar)));
|
|
launcher(*this, m_output);
|
|
return true;
|
|
}
|
|
|
|
template<int LoadMode>
|
|
EIGEN_DEVICE_FUNC PacketReturnType packet(Index index) const {
|
|
return internal::ploadt<PacketReturnType, LoadMode>(m_output + index);
|
|
}
|
|
|
|
EIGEN_DEVICE_FUNC EIGEN_STRONG_INLINE CoeffReturnType* data() const
|
|
{
|
|
return m_output;
|
|
}
|
|
|
|
EIGEN_DEVICE_FUNC EIGEN_STRONG_INLINE CoeffReturnType coeff(Index index) const
|
|
{
|
|
return m_output[index];
|
|
}
|
|
|
|
EIGEN_DEVICE_FUNC EIGEN_STRONG_INLINE TensorOpCost costPerCoeff(bool) const {
|
|
return TensorOpCost(sizeof(CoeffReturnType), 0, 0);
|
|
}
|
|
|
|
EIGEN_DEVICE_FUNC EIGEN_STRONG_INLINE void cleanup() {
|
|
if (m_output != NULL) {
|
|
m_device.deallocate(m_output);
|
|
m_output = NULL;
|
|
}
|
|
m_impl.cleanup();
|
|
}
|
|
|
|
protected:
|
|
TensorEvaluator<ArgType, Device> m_impl;
|
|
const Device& m_device;
|
|
const bool m_exclusive;
|
|
Op m_accumulator;
|
|
const Index m_size;
|
|
Index m_stride;
|
|
CoeffReturnType* m_output;
|
|
};
|
|
|
|
// CPU implementation of scan
|
|
// TODO(ibab) This single-threaded implementation should be parallelized,
|
|
// at least by running multiple scans at the same time.
|
|
template <typename Self, typename Reducer, typename Device>
|
|
struct ScanLauncher {
|
|
void operator()(Self& self, typename Self::CoeffReturnType *data) {
|
|
Index total_size = internal::array_prod(self.dimensions());
|
|
|
|
// We fix the index along the scan axis to 0 and perform a
|
|
// scan per remaining entry. The iteration is split into two nested
|
|
// loops to avoid an integer division by keeping track of each idx1 and idx2.
|
|
for (Index idx1 = 0; idx1 < total_size; idx1 += self.stride() * self.size()) {
|
|
for (Index idx2 = 0; idx2 < self.stride(); idx2++) {
|
|
// Calculate the starting offset for the scan
|
|
Index offset = idx1 + idx2;
|
|
|
|
// Compute the scan along the axis, starting at the calculated offset
|
|
typename Self::CoeffReturnType accum = self.accumulator().initialize();
|
|
for (Index idx3 = 0; idx3 < self.size(); idx3++) {
|
|
Index curr = offset + idx3 * self.stride();
|
|
|
|
if (self.exclusive()) {
|
|
data[curr] = self.accumulator().finalize(accum);
|
|
self.accumulator().reduce(self.inner().coeff(curr), &accum);
|
|
} else {
|
|
self.accumulator().reduce(self.inner().coeff(curr), &accum);
|
|
data[curr] = self.accumulator().finalize(accum);
|
|
}
|
|
}
|
|
}
|
|
}
|
|
}
|
|
};
|
|
|
|
#if defined(EIGEN_USE_GPU) && defined(__CUDACC__)
|
|
|
|
// GPU implementation of scan
|
|
// TODO(ibab) This placeholder implementation performs multiple scans in
|
|
// parallel, but it would be better to use a parallel scan algorithm and
|
|
// optimize memory access.
|
|
template <typename Self, typename Reducer>
|
|
__global__ void ScanKernel(Self self, Index total_size, typename Self::CoeffReturnType* data) {
|
|
// Compute offset as in the CPU version
|
|
Index val = threadIdx.x + blockIdx.x * blockDim.x;
|
|
Index offset = (val / self.stride()) * self.stride() * self.size() + val % self.stride();
|
|
|
|
if (offset + (self.size() - 1) * self.stride() < total_size) {
|
|
// Compute the scan along the axis, starting at the calculated offset
|
|
typename Self::CoeffReturnType accum = self.accumulator().initialize();
|
|
for (Index idx = 0; idx < self.size(); idx++) {
|
|
Index curr = offset + idx * self.stride();
|
|
if (self.exclusive()) {
|
|
data[curr] = self.accumulator().finalize(accum);
|
|
self.accumulator().reduce(self.inner().coeff(curr), &accum);
|
|
} else {
|
|
self.accumulator().reduce(self.inner().coeff(curr), &accum);
|
|
data[curr] = self.accumulator().finalize(accum);
|
|
}
|
|
}
|
|
}
|
|
__syncthreads();
|
|
|
|
}
|
|
|
|
template <typename Self, typename Reducer>
|
|
struct ScanLauncher<Self, Reducer, GpuDevice> {
|
|
void operator()(const Self& self, typename Self::CoeffReturnType* data) {
|
|
Index total_size = internal::array_prod(self.dimensions());
|
|
Index num_blocks = (total_size / self.size() + 63) / 64;
|
|
Index block_size = 64;
|
|
LAUNCH_CUDA_KERNEL((ScanKernel<Self, Reducer>), num_blocks, block_size, 0, self.device(), self, total_size, data);
|
|
}
|
|
};
|
|
#endif // EIGEN_USE_GPU && __CUDACC__
|
|
|
|
} // end namespace Eigen
|
|
|
|
#endif // EIGEN_CXX11_TENSOR_TENSOR_SCAN_H
|