Add linear redux evaluators

This commit is contained in:
Charles Schlosser 2023-05-24 17:07:25 +00:00 committed by Rasmus Munk Larsen
parent 67a1e881d9
commit da6a71faf0

View File

@ -44,9 +44,10 @@ public:
};
enum {
MayLinearize = (int(Evaluator::Flags) & LinearAccessBit),
MightVectorize = (int(Evaluator::Flags)&ActualPacketAccessBit)
&& (functor_traits<Func>::PacketAccess),
MayLinearVectorize = bool(MightVectorize) && (int(Evaluator::Flags)&LinearAccessBit),
MayLinearVectorize = bool(MightVectorize) && bool(MayLinearize),
MaySliceVectorize = bool(MightVectorize) && (int(SliceVectorizedWork)==Dynamic || int(SliceVectorizedWork)>=3)
};
@ -54,6 +55,7 @@ public:
enum {
Traversal = int(MayLinearVectorize) ? int(LinearVectorizedTraversal)
: int(MaySliceVectorize) ? int(SliceVectorizedTraversal)
: int(MayLinearize) ? int(LinearTraversal)
: int(DefaultTraversal)
};
@ -142,6 +144,46 @@ struct redux_novec_unroller<Func, Evaluator, Start, 0>
static EIGEN_STRONG_INLINE Scalar run(const Evaluator&, const Func&) { return Scalar(); }
};
template<typename Func, typename Evaluator, int Start, int Length>
struct redux_novec_linear_unroller
{
enum {
HalfLength = Length/2
};
typedef typename Evaluator::Scalar Scalar;
EIGEN_DEVICE_FUNC
static EIGEN_STRONG_INLINE Scalar run(const Evaluator &eval, const Func& func)
{
return func(redux_novec_linear_unroller<Func, Evaluator, Start, HalfLength>::run(eval,func),
redux_novec_linear_unroller<Func, Evaluator, Start+HalfLength, Length-HalfLength>::run(eval,func));
}
};
template<typename Func, typename Evaluator, int Start>
struct redux_novec_linear_unroller<Func, Evaluator, Start, 1>
{
typedef typename Evaluator::Scalar Scalar;
EIGEN_DEVICE_FUNC
static EIGEN_STRONG_INLINE Scalar run(const Evaluator &eval, const Func&)
{
return eval.coeff(Start);
}
};
// This is actually dead code and will never be called. It is required
// to prevent false warnings regarding failed inlining though
// for 0 length run() will never be called at all.
template<typename Func, typename Evaluator, int Start>
struct redux_novec_linear_unroller<Func, Evaluator, Start, 0>
{
typedef typename Evaluator::Scalar Scalar;
EIGEN_DEVICE_FUNC
static EIGEN_STRONG_INLINE Scalar run(const Evaluator&, const Func&) { return Scalar(); }
};
/*** vectorization ***/
template<typename Func, typename Evaluator, int Start, int Length>
@ -180,6 +222,40 @@ struct redux_vec_unroller<Func, Evaluator, Start, 1>
}
};
template<typename Func, typename Evaluator, int Start, int Length>
struct redux_vec_linear_unroller
{
template<typename PacketType>
EIGEN_DEVICE_FUNC
static EIGEN_STRONG_INLINE PacketType run(const Evaluator &eval, const Func& func)
{
enum {
PacketSize = unpacket_traits<PacketType>::size,
HalfLength = Length/2
};
return func.packetOp(
redux_vec_linear_unroller<Func, Evaluator, Start, HalfLength>::template run<PacketType>(eval,func),
redux_vec_linear_unroller<Func, Evaluator, Start+HalfLength, Length-HalfLength>::template run<PacketType>(eval,func) );
}
};
template<typename Func, typename Evaluator, int Start>
struct redux_vec_linear_unroller<Func, Evaluator, Start, 1>
{
template<typename PacketType>
EIGEN_DEVICE_FUNC
static EIGEN_STRONG_INLINE PacketType run(const Evaluator &eval, const Func&)
{
enum {
PacketSize = unpacket_traits<PacketType>::size,
index = Start * PacketSize,
alignment = Evaluator::Alignment
};
return eval.template packet<alignment,PacketType>(index);
}
};
/***************************************************************************
* Part 3 : implementation of all cases
***************************************************************************/
@ -210,6 +286,23 @@ struct redux_impl<Func, Evaluator, DefaultTraversal, NoUnrolling>
}
};
template<typename Func, typename Evaluator>
struct redux_impl<Func, Evaluator, LinearTraversal, NoUnrolling>
{
typedef typename Evaluator::Scalar Scalar;
template<typename XprType>
EIGEN_DEVICE_FUNC static EIGEN_STRONG_INLINE
Scalar run(const Evaluator &eval, const Func& func, const XprType& xpr)
{
eigen_assert(xpr.size()>0 && "you are using an empty matrix");
Scalar res = eval.coeff(0);
for(Index k = 1; k < xpr.size(); ++k)
res = func(res, eval.coeff(k));
return res;
}
};
template<typename Func, typename Evaluator>
struct redux_impl<Func,Evaluator, DefaultTraversal, CompleteUnrolling>
: redux_novec_unroller<Func,Evaluator, 0, Evaluator::SizeAtCompileTime>
@ -224,6 +317,20 @@ struct redux_impl<Func,Evaluator, DefaultTraversal, CompleteUnrolling>
}
};
template<typename Func, typename Evaluator>
struct redux_impl<Func,Evaluator, LinearTraversal, CompleteUnrolling>
: redux_novec_linear_unroller<Func,Evaluator, 0, Evaluator::SizeAtCompileTime>
{
typedef redux_novec_linear_unroller<Func,Evaluator, 0, Evaluator::SizeAtCompileTime> Base;
typedef typename Evaluator::Scalar Scalar;
template<typename XprType>
EIGEN_DEVICE_FUNC static EIGEN_STRONG_INLINE
Scalar run(const Evaluator &eval, const Func& func, const XprType& /*xpr*/)
{
return Base::run(eval,func);
}
};
template<typename Func, typename Evaluator>
struct redux_impl<Func, Evaluator, LinearVectorizedTraversal, NoUnrolling>
{
@ -342,13 +449,13 @@ struct redux_impl<Func, Evaluator, LinearVectorizedTraversal, CompleteUnrolling>
EIGEN_ONLY_USED_FOR_DEBUG(xpr)
eigen_assert(xpr.rows()>0 && xpr.cols()>0 && "you are using an empty matrix");
if (VectorizedSize > 0) {
Scalar res = func.predux(redux_vec_unroller<Func, Evaluator, 0, Size / PacketSize>::template run<PacketType>(eval,func));
Scalar res = func.predux(redux_vec_linear_unroller<Func, Evaluator, 0, Size / PacketSize>::template run<PacketType>(eval,func));
if (VectorizedSize != Size)
res = func(res,redux_novec_unroller<Func, Evaluator, VectorizedSize, Size-VectorizedSize>::run(eval,func));
res = func(res,redux_novec_linear_unroller<Func, Evaluator, VectorizedSize, Size-VectorizedSize>::run(eval,func));
return res;
}
else {
return redux_novec_unroller<Func, Evaluator, 0, Size>::run(eval,func);
return redux_novec_linear_unroller<Func, Evaluator, 0, Size>::run(eval,func);
}
}
};