Fix any/all reduction in the case of row-major layout

This commit is contained in:
Yury Gitman 2022-03-01 05:27:50 +00:00 committed by Antonio Sánchez
parent f03df0df53
commit bf6726a0c6

View File

@ -12,58 +12,60 @@
#include "./InternalHeaderCheck.h" #include "./InternalHeaderCheck.h"
namespace Eigen { namespace Eigen {
namespace internal { namespace internal {
template<typename Derived, int UnrollCount, int Rows> template<typename Derived, int UnrollCount, int InnerSize>
struct all_unroller struct all_unroller
{ {
enum { enum {
col = (UnrollCount-1) / Rows, IsRowMajor = (Derived::Flags & RowMajor),
row = (UnrollCount-1) % Rows i = (UnrollCount-1) / InnerSize,
j = (UnrollCount-1) % InnerSize
}; };
EIGEN_DEVICE_FUNC static inline bool run(const Derived &mat) EIGEN_DEVICE_FUNC static inline bool run(const Derived &mat)
{ {
return all_unroller<Derived, UnrollCount-1, Rows>::run(mat) && mat.coeff(row, col); return all_unroller<Derived, UnrollCount-1, InnerSize>::run(mat) && mat.coeff(IsRowMajor ? i : j, IsRowMajor ? j : i);
} }
}; };
template<typename Derived, int Rows> template<typename Derived, int InnerSize>
struct all_unroller<Derived, 0, Rows> struct all_unroller<Derived, 0, InnerSize>
{ {
EIGEN_DEVICE_FUNC static inline bool run(const Derived &/*mat*/) { return true; } EIGEN_DEVICE_FUNC static inline bool run(const Derived &/*mat*/) { return true; }
}; };
template<typename Derived, int Rows> template<typename Derived, int InnerSize>
struct all_unroller<Derived, Dynamic, Rows> struct all_unroller<Derived, Dynamic, InnerSize>
{ {
EIGEN_DEVICE_FUNC static inline bool run(const Derived &) { return false; } EIGEN_DEVICE_FUNC static inline bool run(const Derived &) { return false; }
}; };
template<typename Derived, int UnrollCount, int Rows> template<typename Derived, int UnrollCount, int InnerSize>
struct any_unroller struct any_unroller
{ {
enum { enum {
col = (UnrollCount-1) / Rows, IsRowMajor = (Derived::Flags & RowMajor),
row = (UnrollCount-1) % Rows i = (UnrollCount-1) / InnerSize,
j = (UnrollCount-1) % InnerSize
}; };
EIGEN_DEVICE_FUNC static inline bool run(const Derived &mat) EIGEN_DEVICE_FUNC static inline bool run(const Derived &mat)
{ {
return any_unroller<Derived, UnrollCount-1, Rows>::run(mat) || mat.coeff(row, col); return any_unroller<Derived, UnrollCount-1, InnerSize>::run(mat) || mat.coeff(IsRowMajor ? i : j, IsRowMajor ? j : i);
} }
}; };
template<typename Derived, int Rows> template<typename Derived, int InnerSize>
struct any_unroller<Derived, 0, Rows> struct any_unroller<Derived, 0, InnerSize>
{ {
EIGEN_DEVICE_FUNC static inline bool run(const Derived & /*mat*/) { return false; } EIGEN_DEVICE_FUNC static inline bool run(const Derived & /*mat*/) { return false; }
}; };
template<typename Derived, int Rows> template<typename Derived, int InnerSize>
struct any_unroller<Derived, Dynamic, Rows> struct any_unroller<Derived, Dynamic, InnerSize>
{ {
EIGEN_DEVICE_FUNC static inline bool run(const Derived &) { return false; } EIGEN_DEVICE_FUNC static inline bool run(const Derived &) { return false; }
}; };
@ -83,16 +85,18 @@ EIGEN_DEVICE_FUNC inline bool DenseBase<Derived>::all() const
typedef internal::evaluator<Derived> Evaluator; typedef internal::evaluator<Derived> Evaluator;
enum { enum {
unroll = SizeAtCompileTime != Dynamic unroll = SizeAtCompileTime != Dynamic
&& SizeAtCompileTime * (int(Evaluator::CoeffReadCost) + int(NumTraits<Scalar>::AddCost)) <= EIGEN_UNROLLING_LIMIT && SizeAtCompileTime * (int(Evaluator::CoeffReadCost) + int(NumTraits<Scalar>::AddCost)) <= EIGEN_UNROLLING_LIMIT,
IsRowMajor = (internal::traits<Derived>::Flags & RowMajor),
InnerSizeAtCompileTime = IsRowMajor ? internal::traits<Derived>::ColsAtCompileTime : internal::traits<Derived>::RowsAtCompileTime
}; };
Evaluator evaluator(derived()); Evaluator evaluator(derived());
if(unroll) if(unroll)
return internal::all_unroller<Evaluator, unroll ? int(SizeAtCompileTime) : Dynamic, internal::traits<Derived>::RowsAtCompileTime>::run(evaluator); return internal::all_unroller<Evaluator, unroll ? int(SizeAtCompileTime) : Dynamic, InnerSizeAtCompileTime>::run(evaluator);
else else
{ {
for(Index j = 0; j < cols(); ++j) for(Index i = 0; i < derived().outerSize(); ++i)
for(Index i = 0; i < rows(); ++i) for(Index j = 0; j < derived().innerSize(); ++j)
if (!evaluator.coeff(i, j)) return false; if (!evaluator.coeff(IsRowMajor ? i : j, IsRowMajor ? j : i)) return false;
return true; return true;
} }
} }
@ -107,16 +111,18 @@ EIGEN_DEVICE_FUNC inline bool DenseBase<Derived>::any() const
typedef internal::evaluator<Derived> Evaluator; typedef internal::evaluator<Derived> Evaluator;
enum { enum {
unroll = SizeAtCompileTime != Dynamic unroll = SizeAtCompileTime != Dynamic
&& SizeAtCompileTime * (int(Evaluator::CoeffReadCost) + int(NumTraits<Scalar>::AddCost)) <= EIGEN_UNROLLING_LIMIT && SizeAtCompileTime * (int(Evaluator::CoeffReadCost) + int(NumTraits<Scalar>::AddCost)) <= EIGEN_UNROLLING_LIMIT,
IsRowMajor = (internal::traits<Derived>::Flags & RowMajor),
InnerSizeAtCompileTime = IsRowMajor ? internal::traits<Derived>::ColsAtCompileTime : internal::traits<Derived>::RowsAtCompileTime
}; };
Evaluator evaluator(derived()); Evaluator evaluator(derived());
if(unroll) if(unroll)
return internal::any_unroller<Evaluator, unroll ? int(SizeAtCompileTime) : Dynamic, internal::traits<Derived>::RowsAtCompileTime>::run(evaluator); return internal::any_unroller<Evaluator, unroll ? int(SizeAtCompileTime) : Dynamic, InnerSizeAtCompileTime>::run(evaluator);
else else
{ {
for(Index j = 0; j < cols(); ++j) for(Index i = 0; i < derived().outerSize(); ++i)
for(Index i = 0; i < rows(); ++i) for(Index j = 0; j < derived().innerSize(); ++j)
if (evaluator.coeff(i, j)) return true; if (evaluator.coeff(IsRowMajor ? i : j, IsRowMajor ? j : i)) return true;
return false; return false;
} }
} }
@ -158,7 +164,7 @@ EIGEN_DEVICE_FUNC inline bool DenseBase<Derived>::allFinite() const
return !((derived()-derived()).hasNaN()); return !((derived()-derived()).hasNaN());
#endif #endif
} }
} // end namespace Eigen } // end namespace Eigen
#endif // EIGEN_ALLANDANY_H #endif // EIGEN_ALLANDANY_H