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

@ -16,54 +16,56 @@ namespace Eigen {
namespace internal {
template<typename Derived, int UnrollCount, int Rows>
template<typename Derived, int UnrollCount, int InnerSize>
struct all_unroller
{
enum {
col = (UnrollCount-1) / Rows,
row = (UnrollCount-1) % Rows
IsRowMajor = (Derived::Flags & RowMajor),
i = (UnrollCount-1) / InnerSize,
j = (UnrollCount-1) % InnerSize
};
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>
struct all_unroller<Derived, 0, Rows>
template<typename Derived, int InnerSize>
struct all_unroller<Derived, 0, InnerSize>
{
EIGEN_DEVICE_FUNC static inline bool run(const Derived &/*mat*/) { return true; }
};
template<typename Derived, int Rows>
struct all_unroller<Derived, Dynamic, Rows>
template<typename Derived, int InnerSize>
struct all_unroller<Derived, Dynamic, InnerSize>
{
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
{
enum {
col = (UnrollCount-1) / Rows,
row = (UnrollCount-1) % Rows
IsRowMajor = (Derived::Flags & RowMajor),
i = (UnrollCount-1) / InnerSize,
j = (UnrollCount-1) % InnerSize
};
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>
struct any_unroller<Derived, 0, Rows>
template<typename Derived, int InnerSize>
struct any_unroller<Derived, 0, InnerSize>
{
EIGEN_DEVICE_FUNC static inline bool run(const Derived & /*mat*/) { return false; }
};
template<typename Derived, int Rows>
struct any_unroller<Derived, Dynamic, Rows>
template<typename Derived, int InnerSize>
struct any_unroller<Derived, Dynamic, InnerSize>
{
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;
enum {
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());
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
{
for(Index j = 0; j < cols(); ++j)
for(Index i = 0; i < rows(); ++i)
if (!evaluator.coeff(i, j)) return false;
for(Index i = 0; i < derived().outerSize(); ++i)
for(Index j = 0; j < derived().innerSize(); ++j)
if (!evaluator.coeff(IsRowMajor ? i : j, IsRowMajor ? j : i)) return false;
return true;
}
}
@ -107,16 +111,18 @@ EIGEN_DEVICE_FUNC inline bool DenseBase<Derived>::any() const
typedef internal::evaluator<Derived> Evaluator;
enum {
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());
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
{
for(Index j = 0; j < cols(); ++j)
for(Index i = 0; i < rows(); ++i)
if (evaluator.coeff(i, j)) return true;
for(Index i = 0; i < derived().outerSize(); ++i)
for(Index j = 0; j < derived().innerSize(); ++j)
if (evaluator.coeff(IsRowMajor ? i : j, IsRowMajor ? j : i)) return true;
return false;
}
}