add EqualSpaced / setEqualSpaced

This commit is contained in:
Charles Schlosser
2022-12-13 00:54:57 +00:00
committed by Rasmus Munk Larsen
parent 273f803846
commit 2004831941
4 changed files with 78 additions and 6 deletions

View File

@@ -145,6 +145,39 @@ template <typename Scalar> struct linspaced_op
const linspaced_op_impl<Scalar,NumTraits<Scalar>::IsInteger> impl;
};
template <typename Scalar>
struct equalspaced_op {
typedef typename NumTraits<Scalar>::Real RealScalar;
EIGEN_DEVICE_FUNC equalspaced_op(const Scalar& start, const Scalar& step) : m_start(start), m_step(step) {}
template <typename IndexType>
EIGEN_DEVICE_FUNC EIGEN_STRONG_INLINE const Scalar operator()(IndexType i) const {
return m_start + m_step * static_cast<Scalar>(i);
}
template <typename Packet, typename IndexType>
EIGEN_DEVICE_FUNC EIGEN_STRONG_INLINE Packet packetOp(IndexType i) const {
const Packet cst_start = pset1<Packet>(m_start);
const Packet cst_step = pset1<Packet>(m_step);
const Packet cst_lin0 = plset<Packet>(Scalar(0));
const Packet cst_offset = pmadd(cst_lin0, cst_step, cst_start);
Packet istep = pset1<Packet>(static_cast<Scalar>(i) * m_step);
return padd(cst_offset, istep);
}
const Scalar m_start;
const Scalar m_step;
};
template <typename Scalar>
struct functor_traits<equalspaced_op<Scalar> > {
enum {
Cost = NumTraits<Scalar>::AddCost + NumTraits<Scalar>::MulCost,
PacketAccess =
packet_traits<Scalar>::HasSetLinear && packet_traits<Scalar>::HasMul && packet_traits<Scalar>::HasAdd,
IsRepeatable = true
};
};
// Linear access is automatically determined from the operator() prototypes available for the given functor.
// If it exposes an operator()(i,j), then we assume the i and j coefficients are required independently
// and linear access is not possible. In all other cases, linear access is enabled.