Allow specifying inner & outer stride for CWiseUnaryView - fixes #2398

This commit is contained in:
Andrew Johnson
2022-01-05 19:24:46 +00:00
committed by Rasmus Munk Larsen
parent 27a78e4f96
commit a491c7f898
6 changed files with 79 additions and 27 deletions

View File

@@ -15,8 +15,8 @@
namespace Eigen {
namespace internal {
template<typename ViewOp, typename MatrixType>
struct traits<CwiseUnaryView<ViewOp, MatrixType> >
template<typename ViewOp, typename MatrixType, typename StrideType>
struct traits<CwiseUnaryView<ViewOp, MatrixType, StrideType> >
: traits<MatrixType>
{
typedef typename result_of<
@@ -30,17 +30,22 @@ struct traits<CwiseUnaryView<ViewOp, MatrixType> >
MatrixTypeInnerStride = inner_stride_at_compile_time<MatrixType>::ret,
// need to cast the sizeof's from size_t to int explicitly, otherwise:
// "error: no integral type can represent all of the enumerator values
InnerStrideAtCompileTime = MatrixTypeInnerStride == Dynamic
? int(Dynamic)
: int(MatrixTypeInnerStride) * int(sizeof(typename traits<MatrixType>::Scalar) / sizeof(Scalar)),
OuterStrideAtCompileTime = outer_stride_at_compile_time<MatrixType>::ret == Dynamic
? int(Dynamic)
: outer_stride_at_compile_time<MatrixType>::ret * int(sizeof(typename traits<MatrixType>::Scalar) / sizeof(Scalar))
InnerStrideAtCompileTime = StrideType::InnerStrideAtCompileTime == 0
? (MatrixTypeInnerStride == Dynamic
? int(Dynamic)
: int(MatrixTypeInnerStride) * int(sizeof(typename traits<MatrixType>::Scalar) / sizeof(Scalar)))
: int(StrideType::InnerStrideAtCompileTime),
OuterStrideAtCompileTime = StrideType::OuterStrideAtCompileTime == 0
? (outer_stride_at_compile_time<MatrixType>::ret == Dynamic
? int(Dynamic)
: outer_stride_at_compile_time<MatrixType>::ret * int(sizeof(typename traits<MatrixType>::Scalar) / sizeof(Scalar)))
: int(StrideType::OuterStrideAtCompileTime)
};
};
}
template<typename ViewOp, typename MatrixType, typename StorageKind>
template<typename ViewOp, typename MatrixType, typename StrideType, typename StorageKind>
class CwiseUnaryViewImpl;
/** \class CwiseUnaryView
@@ -56,12 +61,12 @@ class CwiseUnaryViewImpl;
*
* \sa MatrixBase::unaryViewExpr(const CustomUnaryOp &) const, class CwiseUnaryOp
*/
template<typename ViewOp, typename MatrixType>
class CwiseUnaryView : public CwiseUnaryViewImpl<ViewOp, MatrixType, typename internal::traits<MatrixType>::StorageKind>
template<typename ViewOp, typename MatrixType, typename StrideType>
class CwiseUnaryView : public CwiseUnaryViewImpl<ViewOp, MatrixType, StrideType, typename internal::traits<MatrixType>::StorageKind>
{
public:
typedef typename CwiseUnaryViewImpl<ViewOp, MatrixType,typename internal::traits<MatrixType>::StorageKind>::Base Base;
typedef typename CwiseUnaryViewImpl<ViewOp, MatrixType, StrideType, typename internal::traits<MatrixType>::StorageKind>::Base Base;
EIGEN_GENERIC_PUBLIC_INTERFACE(CwiseUnaryView)
typedef typename internal::ref_selector<MatrixType>::non_const_type MatrixTypeNested;
typedef typename internal::remove_all<MatrixType>::type NestedExpression;
@@ -93,22 +98,22 @@ class CwiseUnaryView : public CwiseUnaryViewImpl<ViewOp, MatrixType, typename in
};
// Generic API dispatcher
template<typename ViewOp, typename XprType, typename StorageKind>
template<typename ViewOp, typename XprType, typename StrideType, typename StorageKind>
class CwiseUnaryViewImpl
: public internal::generic_xpr_base<CwiseUnaryView<ViewOp, XprType> >::type
: public internal::generic_xpr_base<CwiseUnaryView<ViewOp, XprType, StrideType> >::type
{
public:
typedef typename internal::generic_xpr_base<CwiseUnaryView<ViewOp, XprType> >::type Base;
typedef typename internal::generic_xpr_base<CwiseUnaryView<ViewOp, XprType, StrideType> >::type Base;
};
template<typename ViewOp, typename MatrixType>
class CwiseUnaryViewImpl<ViewOp,MatrixType,Dense>
: public internal::dense_xpr_base< CwiseUnaryView<ViewOp, MatrixType> >::type
template<typename ViewOp, typename MatrixType, typename StrideType>
class CwiseUnaryViewImpl<ViewOp,MatrixType,StrideType,Dense>
: public internal::dense_xpr_base< CwiseUnaryView<ViewOp, MatrixType, StrideType> >::type
{
public:
typedef CwiseUnaryView<ViewOp, MatrixType> Derived;
typedef typename internal::dense_xpr_base< CwiseUnaryView<ViewOp, MatrixType> >::type Base;
typedef CwiseUnaryView<ViewOp, MatrixType,StrideType> Derived;
typedef typename internal::dense_xpr_base< CwiseUnaryView<ViewOp, MatrixType,StrideType> >::type Base;
EIGEN_DENSE_PUBLIC_INTERFACE(Derived)
EIGEN_INHERIT_ASSIGNMENT_OPERATORS(CwiseUnaryViewImpl)
@@ -118,12 +123,16 @@ class CwiseUnaryViewImpl<ViewOp,MatrixType,Dense>
EIGEN_DEVICE_FUNC EIGEN_CONSTEXPR inline Index innerStride() const
{
return derived().nestedExpression().innerStride() * sizeof(typename internal::traits<MatrixType>::Scalar) / sizeof(Scalar);
return StrideType::InnerStrideAtCompileTime != 0
? int(StrideType::InnerStrideAtCompileTime)
: derived().nestedExpression().innerStride() * sizeof(typename internal::traits<MatrixType>::Scalar) / sizeof(Scalar);
}
EIGEN_DEVICE_FUNC EIGEN_CONSTEXPR inline Index outerStride() const
{
return derived().nestedExpression().outerStride() * sizeof(typename internal::traits<MatrixType>::Scalar) / sizeof(Scalar);
return StrideType::OuterStrideAtCompileTime != 0
? int(StrideType::OuterStrideAtCompileTime)
: derived().nestedExpression().outerStride() * sizeof(typename internal::traits<MatrixType>::Scalar) / sizeof(Scalar);
}
protected:
EIGEN_DEFAULT_EMPTY_CONSTRUCTOR_AND_DESTRUCTOR(CwiseUnaryViewImpl)