mirror of
https://gitlab.com/libeigen/eigen.git
synced 2026-04-10 11:34:33 +08:00
Reworked the expression evaluation mechanism in order to make it possible to efficiently compute convolutions and contractions in the future:
* The scheduling of computation is moved out the the assignment code and into a new TensorExecutor class * The assignment itself is now a regular node on the expression tree * The expression evaluators start by recursively evaluating all their subexpressions if needed
This commit is contained in:
Benoit Steiner
@@ -10,10 +10,6 @@
|
||||
#ifndef EIGEN_CXX11_TENSOR_TENSOR_ASSIGN_H
|
||||
#define EIGEN_CXX11_TENSOR_TENSOR_ASSIGN_H
|
||||
|
||||
#ifdef EIGEN_USE_THREADS
|
||||
#include <future>
|
||||
#endif
|
||||
|
||||
namespace Eigen {
|
||||
|
||||
/** \class TensorAssign
|
||||
@@ -21,172 +17,134 @@ namespace Eigen {
|
||||
*
|
||||
* \brief The tensor assignment class.
|
||||
*
|
||||
* This class is responsible for triggering the evaluation of the expressions
|
||||
* used on the lhs and rhs of an assignment operator and copy the result of
|
||||
* the evaluation of the rhs expression at the address computed during the
|
||||
* evaluation lhs expression.
|
||||
*
|
||||
* TODO: vectorization. For now the code only uses scalars
|
||||
* TODO: parallelisation using multithreading on cpu, or kernels on gpu.
|
||||
* This class is represents the assignment of the values resulting from the evaluation of
|
||||
* the rhs expression to the memory locations denoted by the lhs expression.
|
||||
*/
|
||||
namespace internal {
|
||||
|
||||
// Default strategy: the expressions are evaluated with a single cpu thread.
|
||||
template<typename Derived1, typename Derived2, typename Device = DefaultDevice, bool Vectorizable = TensorEvaluator<Derived1, Device>::PacketAccess & TensorEvaluator<Derived2, Device>::PacketAccess>
|
||||
struct TensorAssign
|
||||
template<typename LhsXprType, typename RhsXprType>
|
||||
struct traits<TensorAssignOp<LhsXprType, RhsXprType> >
|
||||
{
|
||||
typedef typename Derived1::Index Index;
|
||||
EIGEN_DEVICE_FUNC
|
||||
static inline void run(Derived1& dst, const Derived2& src, const Device& device = Device())
|
||||
{
|
||||
TensorEvaluator<Derived1, Device> evalDst(dst, device);
|
||||
TensorEvaluator<Derived2, Device> evalSrc(src, device);
|
||||
const Index size = dst.size();
|
||||
for (Index i = 0; i < size; ++i) {
|
||||
evalDst.coeffRef(i) = evalSrc.coeff(i);
|
||||
}
|
||||
}
|
||||
typedef typename LhsXprType::Scalar Scalar;
|
||||
typedef typename internal::packet_traits<Scalar>::type Packet;
|
||||
typedef typename traits<LhsXprType>::StorageKind StorageKind;
|
||||
typedef typename promote_index_type<typename traits<LhsXprType>::Index,
|
||||
typename traits<RhsXprType>::Index>::type Index;
|
||||
typedef typename LhsXprType::Nested LhsNested;
|
||||
typedef typename RhsXprType::Nested RhsNested;
|
||||
typedef typename remove_reference<LhsNested>::type _LhsNested;
|
||||
typedef typename remove_reference<RhsNested>::type _RhsNested;
|
||||
|
||||
enum {
|
||||
Flags = 0,
|
||||
};
|
||||
};
|
||||
|
||||
|
||||
template<typename Derived1, typename Derived2, typename Device>
|
||||
struct TensorAssign<Derived1, Derived2, Device, true>
|
||||
template<typename LhsXprType, typename RhsXprType>
|
||||
struct eval<TensorAssignOp<LhsXprType, RhsXprType>, Eigen::Dense>
|
||||
{
|
||||
typedef typename Derived1::Index Index;
|
||||
static inline void run(Derived1& dst, const Derived2& src, const Device& device = Device())
|
||||
{
|
||||
TensorEvaluator<Derived1, Device> evalDst(dst, device);
|
||||
TensorEvaluator<Derived2, Device> evalSrc(src, device);
|
||||
const Index size = dst.size();
|
||||
|
||||
static const int LhsStoreMode = TensorEvaluator<Derived1, Device>::IsAligned ? Aligned : Unaligned;
|
||||
static const int RhsLoadMode = TensorEvaluator<Derived2, Device>::IsAligned ? Aligned : Unaligned;
|
||||
static const int PacketSize = unpacket_traits<typename TensorEvaluator<Derived1, Device>::PacketReturnType>::size;
|
||||
const int VectorizedSize = (size / PacketSize) * PacketSize;
|
||||
|
||||
for (Index i = 0; i < VectorizedSize; i += PacketSize) {
|
||||
evalDst.template writePacket<LhsStoreMode>(i, evalSrc.template packet<RhsLoadMode>(i));
|
||||
}
|
||||
for (Index i = VectorizedSize; i < size; ++i) {
|
||||
evalDst.coeffRef(i) = evalSrc.coeff(i);
|
||||
}
|
||||
}
|
||||
typedef const TensorAssignOp<LhsXprType, RhsXprType>& type;
|
||||
};
|
||||
|
||||
|
||||
|
||||
// Multicore strategy: the index space is partitioned and each core is assigned to a partition
|
||||
#ifdef EIGEN_USE_THREADS
|
||||
template <typename LhsEval, typename RhsEval, typename Index, bool Vectorizable = LhsEval::PacketAccess & RhsEval::PacketAccess>
|
||||
struct EvalRange {
|
||||
static void run(LhsEval& dst, const RhsEval& src, const Index first, const Index last) {
|
||||
eigen_assert(last > first);
|
||||
for (Index i = first; i < last; ++i) {
|
||||
dst.coeffRef(i) = src.coeff(i);
|
||||
}
|
||||
}
|
||||
};
|
||||
|
||||
template <typename LhsEval, typename RhsEval, typename Index>
|
||||
struct EvalRange<LhsEval, RhsEval, Index, true> {
|
||||
static void run(LhsEval& dst, const RhsEval& src, const Index first, const Index last) {
|
||||
eigen_assert(last > first);
|
||||
|
||||
Index i = first;
|
||||
static const int PacketSize = unpacket_traits<typename LhsEval::PacketReturnType>::size;
|
||||
if (last - first > PacketSize) {
|
||||
static const int LhsStoreMode = LhsEval::IsAligned ? Aligned : Unaligned;
|
||||
static const int RhsLoadMode = RhsEval::IsAligned ? Aligned : Unaligned;
|
||||
eigen_assert(first % PacketSize == 0);
|
||||
Index lastPacket = last - (last % PacketSize);
|
||||
for (; i < lastPacket; i += PacketSize) {
|
||||
dst.template writePacket<LhsStoreMode>(i, src.template packet<RhsLoadMode>(i));
|
||||
}
|
||||
}
|
||||
|
||||
for (; i < last; ++i) {
|
||||
dst.coeffRef(i) = src.coeff(i);
|
||||
}
|
||||
}
|
||||
};
|
||||
|
||||
template<typename Derived1, typename Derived2>
|
||||
struct TensorAssignMultiThreaded
|
||||
template<typename LhsXprType, typename RhsXprType>
|
||||
struct nested<TensorAssignOp<LhsXprType, RhsXprType>, 1, typename eval<TensorAssignOp<LhsXprType, RhsXprType> >::type>
|
||||
{
|
||||
typedef typename Derived1::Index Index;
|
||||
static inline void run(Derived1& dst, const Derived2& src, const ThreadPoolDevice& device)
|
||||
{
|
||||
TensorEvaluator<Derived1, DefaultDevice> evalDst(dst, DefaultDevice());
|
||||
TensorEvaluator<Derived2, DefaultDevice> evalSrc(src, Defaultevice());
|
||||
const Index size = dst.size();
|
||||
|
||||
static const bool Vectorizable = TensorEvaluator<Derived1, DefaultDevice>::PacketAccess & TensorEvaluator<Derived2, DefaultDevice>::PacketAccess;
|
||||
static const int PacketSize = Vectorizable ? unpacket_traits<typename TensorEvaluator<Derived1, DefaultDevice>::PacketReturnType>::size : 1;
|
||||
|
||||
int blocksz = static_cast<int>(ceil(static_cast<float>(size)/device.numThreads()) + PacketSize - 1);
|
||||
const Index blocksize = std::max<Index>(PacketSize, (blocksz - (blocksz % PacketSize)));
|
||||
const Index numblocks = size / blocksize;
|
||||
|
||||
Index i = 0;
|
||||
vector<std::future<void> > results;
|
||||
results.reserve(numblocks);
|
||||
for (int i = 0; i < numblocks; ++i) {
|
||||
results.push_back(std::async(std::launch::async, &EvalRange<TensorEvaluator<Derived1, DefaultDevice>, TensorEvaluator<Derived2, DefaultDevice>, Index>::run, evalDst, evalSrc, i*blocksize, (i+1)*blocksize));
|
||||
}
|
||||
|
||||
for (int i = 0; i < numblocks; ++i) {
|
||||
results[i].get();
|
||||
}
|
||||
|
||||
if (numblocks * blocksize < size) {
|
||||
EvalRange<TensorEvaluator<Derived1>, TensorEvaluator<Derived2>, Index>::run(evalDst, evalSrc, numblocks * blocksize, size);
|
||||
}
|
||||
}
|
||||
typedef TensorAssignOp<LhsXprType, RhsXprType> type;
|
||||
};
|
||||
#endif
|
||||
|
||||
} // end namespace internal
|
||||
|
||||
|
||||
// GPU: the evaluation of the expressions is offloaded to a GPU.
|
||||
#if defined(EIGEN_USE_GPU) && defined(__CUDACC__)
|
||||
template <typename LhsEvaluator, typename RhsEvaluator>
|
||||
__global__ void EigenMetaKernelNoCheck(LhsEvaluator evalDst, const RhsEvaluator evalSrc) {
|
||||
const int index = blockIdx.x * blockDim.x + threadIdx.x;
|
||||
evalDst.coeffRef(index) = evalSrc.coeff(index);
|
||||
}
|
||||
template <typename LhsEvaluator, typename RhsEvaluator>
|
||||
__global__ void EigenMetaKernelPeel(LhsEvaluator evalDst, const RhsEvaluator evalSrc, int peel_start_offset, int size) {
|
||||
const int index = peel_start_offset + blockIdx.x * blockDim.x + threadIdx.x;
|
||||
if (index < size) {
|
||||
evalDst.coeffRef(index) = evalSrc.coeff(index);
|
||||
|
||||
template<typename LhsXprType, typename RhsXprType>
|
||||
class TensorAssignOp : public TensorBase<TensorAssignOp<LhsXprType, RhsXprType> >
|
||||
{
|
||||
public:
|
||||
typedef typename Eigen::internal::traits<TensorAssignOp>::Scalar Scalar;
|
||||
typedef typename Eigen::internal::traits<TensorAssignOp>::Packet Packet;
|
||||
typedef typename Eigen::NumTraits<Scalar>::Real RealScalar;
|
||||
typedef typename LhsXprType::CoeffReturnType CoeffReturnType;
|
||||
typedef typename LhsXprType::PacketReturnType PacketReturnType;
|
||||
typedef typename Eigen::internal::nested<TensorAssignOp>::type Nested;
|
||||
typedef typename Eigen::internal::traits<TensorAssignOp>::StorageKind StorageKind;
|
||||
typedef typename Eigen::internal::traits<TensorAssignOp>::Index Index;
|
||||
|
||||
EIGEN_DEVICE_FUNC EIGEN_STRONG_INLINE TensorAssignOp(LhsXprType& lhs, const RhsXprType& rhs)
|
||||
: m_lhs_xpr(lhs), m_rhs_xpr(rhs) {}
|
||||
|
||||
/** \returns the nested expressions */
|
||||
EIGEN_DEVICE_FUNC
|
||||
typename internal::remove_all<typename LhsXprType::Nested>::type&
|
||||
lhsExpression() const { return *((typename internal::remove_all<typename LhsXprType::Nested>::type*)&m_lhs_xpr); }
|
||||
|
||||
EIGEN_DEVICE_FUNC
|
||||
const typename internal::remove_all<typename RhsXprType::Nested>::type&
|
||||
rhsExpression() const { return m_rhs_xpr; }
|
||||
|
||||
protected:
|
||||
typename internal::remove_all<typename LhsXprType::Nested>::type& m_lhs_xpr;
|
||||
const typename internal::remove_all<typename RhsXprType::Nested>::type& m_rhs_xpr;
|
||||
};
|
||||
|
||||
|
||||
template<typename LeftArgType, typename RightArgType, typename Device>
|
||||
struct TensorEvaluator<const TensorAssignOp<LeftArgType, RightArgType>, Device>
|
||||
{
|
||||
typedef TensorAssignOp<LeftArgType, RightArgType> XprType;
|
||||
|
||||
enum {
|
||||
IsAligned = TensorEvaluator<LeftArgType, Device>::IsAligned & TensorEvaluator<RightArgType, Device>::IsAligned,
|
||||
PacketAccess = TensorEvaluator<LeftArgType, Device>::PacketAccess & TensorEvaluator<RightArgType, Device>::PacketAccess,
|
||||
};
|
||||
|
||||
EIGEN_DEVICE_FUNC TensorEvaluator(const XprType& op, const Device& device) :
|
||||
m_leftImpl(op.lhsExpression(), device),
|
||||
m_rightImpl(op.rhsExpression(), device)
|
||||
{ }
|
||||
|
||||
typedef typename XprType::Index Index;
|
||||
typedef typename XprType::CoeffReturnType CoeffReturnType;
|
||||
typedef typename XprType::PacketReturnType PacketReturnType;
|
||||
typedef typename TensorEvaluator<RightArgType, Device>::Dimensions Dimensions;
|
||||
|
||||
EIGEN_DEVICE_FUNC const Dimensions& dimensions() const
|
||||
{
|
||||
// TODO: use left impl instead if right impl dimensions are known at compile time.
|
||||
return m_rightImpl.dimensions();
|
||||
}
|
||||
|
||||
EIGEN_DEVICE_FUNC EIGEN_STRONG_INLINE void evalSubExprsIfNeeded() {
|
||||
m_leftImpl.evalSubExprsIfNeeded();
|
||||
m_rightImpl.evalSubExprsIfNeeded();
|
||||
}
|
||||
EIGEN_DEVICE_FUNC EIGEN_STRONG_INLINE void cleanup() {
|
||||
m_leftImpl.cleanup();
|
||||
m_rightImpl.cleanup();
|
||||
}
|
||||
|
||||
EIGEN_DEVICE_FUNC EIGEN_STRONG_INLINE void evalScalar(Index i) {
|
||||
m_leftImpl.coeffRef(i) = m_rightImpl.coeff(i);
|
||||
}
|
||||
EIGEN_DEVICE_FUNC EIGEN_STRONG_INLINE void evalPacket(Index i) {
|
||||
static const int LhsStoreMode = TensorEvaluator<LeftArgType, Device>::IsAligned ? Aligned : Unaligned;
|
||||
static const int RhsLoadMode = TensorEvaluator<RightArgType, Device>::IsAligned ? Aligned : Unaligned;
|
||||
m_leftImpl.template writePacket<LhsStoreMode>(i, m_rightImpl.template packet<RhsLoadMode>(i));
|
||||
}
|
||||
EIGEN_DEVICE_FUNC CoeffReturnType coeff(Index index) const
|
||||
{
|
||||
return m_leftImpl.coeff(index);
|
||||
}
|
||||
template<int LoadMode>
|
||||
EIGEN_DEVICE_FUNC PacketReturnType packet(Index index) const
|
||||
{
|
||||
return m_leftImpl.template packet<LoadMode>(index);
|
||||
}
|
||||
|
||||
private:
|
||||
TensorEvaluator<LeftArgType, Device> m_leftImpl;
|
||||
TensorEvaluator<RightArgType, Device> m_rightImpl;
|
||||
};
|
||||
|
||||
}
|
||||
|
||||
template<typename Derived1, typename Derived2>
|
||||
struct TensorAssignGpu
|
||||
{
|
||||
typedef typename Derived1::Index Index;
|
||||
static inline void run(Derived1& dst, const Derived2& src, const GpuDevice& device)
|
||||
{
|
||||
TensorEvaluator<Derived1, GpuDevice> evalDst(dst, device);
|
||||
TensorEvaluator<Derived2, GpuDevice> evalSrc(src, device);
|
||||
const Index size = dst.size();
|
||||
const int block_size = std::min<int>(size, 32*32);
|
||||
const int num_blocks = size / block_size;
|
||||
EigenMetaKernelNoCheck<TensorEvaluator<Derived1, GpuDevice>, TensorEvaluator<Derived2, GpuDevice> > <<<num_blocks, block_size, 0, device.stream()>>>(evalDst, evalSrc);
|
||||
|
||||
const int remaining_items = size % block_size;
|
||||
if (remaining_items > 0) {
|
||||
const int peel_start_offset = num_blocks * block_size;
|
||||
const int peel_block_size = std::min<int>(size, 32);
|
||||
const int peel_num_blocks = (remaining_items + peel_block_size - 1) / peel_block_size;
|
||||
EigenMetaKernelPeel<TensorEvaluator<Derived1, GpuDevice>, TensorEvaluator<Derived2, GpuDevice> > <<<peel_num_blocks, peel_block_size, 0, device.stream()>>>(evalDst, evalSrc, peel_start_offset, size);
|
||||
}
|
||||
}
|
||||
};
|
||||
#endif
|
||||
|
||||
} // end namespace internal
|
||||
|
||||
} // end namespace Eigen
|
||||
|
||||
#endif // EIGEN_CXX11_TENSOR_TENSOR_ASSIGN_H
|
||||
|
||||
Reference in New Issue
Block a user