masked load/store framework

This commit is contained in:
Charles Schlosser
2025-04-12 00:31:10 +00:00
parent cebe09110c
commit 28c3b26d53
17 changed files with 1363 additions and 177 deletions

View File

@@ -29,6 +29,11 @@ struct assign_op {
EIGEN_STRONG_INLINE void assignPacket(DstScalar* a, const Packet& b) const {
pstoret<DstScalar, Packet, Alignment>(a, b);
}
template <int Alignment, typename Packet>
EIGEN_STRONG_INLINE void assignPacketSegment(DstScalar* a, const Packet& b, Index begin, Index count) const {
pstoretSegment<DstScalar, Packet, Alignment>(a, b, begin, count);
}
};
// Empty overload for void type (used by PermutationMatrix)
@@ -60,6 +65,12 @@ struct compound_assign_op {
assign_op<DstScalar, DstScalar>().template assignPacket<Alignment, Packet>(
a, Func().packetOp(ploadt<Packet, Alignment>(a), b));
}
template <int Alignment, typename Packet>
EIGEN_STRONG_INLINE void assignPacketSegment(DstScalar* a, const Packet& b, Index begin, Index count) const {
assign_op<DstScalar, DstScalar>().template assignPacketSegment<Alignment, Packet>(
a, Func().packetOp(ploadtSegment<Packet, Alignment>(a, begin, count), b), begin, count);
}
};
template <typename DstScalar, typename SrcScalar, typename Func>

View File

@@ -438,7 +438,6 @@ struct scalar_quotient_op : binary_op_base<LhsScalar, RhsScalar> {
}
template <typename Packet>
EIGEN_DEVICE_FUNC EIGEN_STRONG_INLINE const Packet packetOp(const Packet& a, const Packet& b) const {
maybe_raise_div_by_zero<Packet>::run(b);
return internal::pdiv(a, b);
}
};

View File

@@ -28,7 +28,7 @@ struct scalar_constant_op {
const Scalar m_other;
};
template <typename Scalar>
struct functor_traits<scalar_constant_op<Scalar> > {
struct functor_traits<scalar_constant_op<Scalar>> {
enum {
Cost = 0 /* as the constant value should be loaded in register only once for the whole expression */,
PacketAccess = packet_traits<Scalar>::Vectorizable,
@@ -56,7 +56,7 @@ struct scalar_identity_op {
}
};
template <typename Scalar>
struct functor_traits<scalar_identity_op<Scalar> > {
struct functor_traits<scalar_identity_op<Scalar>> {
enum { Cost = NumTraits<Scalar>::AddCost, PacketAccess = false, IsRepeatable = true };
};
@@ -86,18 +86,19 @@ struct linspaced_op_impl<Scalar, /*IsInteger*/ false> {
EIGEN_DEVICE_FUNC EIGEN_STRONG_INLINE const Packet packetOp(IndexType i) const {
// Principle:
// [low, ..., low] + ( [step, ..., step] * ( [i, ..., i] + [0, ..., size] ) )
Packet low = pset1<Packet>(m_low);
Packet high = pset1<Packet>(m_high);
Packet step = pset1<Packet>(m_step);
if (m_flip) {
Packet pi = plset<Packet>(Scalar(i - m_size1));
Packet res = padd(pset1<Packet>(m_high), pmul(pset1<Packet>(m_step), pi));
if (EIGEN_PREDICT_TRUE(i != 0)) return res;
Packet mask = pcmp_lt(pset1<Packet>(0), plset<Packet>(0));
return pselect<Packet>(mask, res, pset1<Packet>(m_low));
Packet res = pmadd(step, pi, high);
Packet mask = pcmp_lt(pzero(res), plset<Packet>(Scalar(i)));
return pselect<Packet>(mask, res, low);
} else {
Packet pi = plset<Packet>(Scalar(i));
Packet res = padd(pset1<Packet>(m_low), pmul(pset1<Packet>(m_step), pi));
if (EIGEN_PREDICT_TRUE(i != m_size1 - unpacket_traits<Packet>::size + 1)) return res;
Packet mask = pcmp_lt(plset<Packet>(0), pset1<Packet>(unpacket_traits<Packet>::size - 1));
return pselect<Packet>(mask, res, pset1<Packet>(m_high));
Packet res = pmadd(step, pi, low);
Packet mask = pcmp_lt(pi, pset1<Packet>(Scalar(m_size1)));
return pselect<Packet>(mask, res, high);
}
}
@@ -139,7 +140,7 @@ struct linspaced_op_impl<Scalar, /*IsInteger*/ true> {
template <typename Scalar>
struct linspaced_op;
template <typename Scalar>
struct functor_traits<linspaced_op<Scalar> > {
struct functor_traits<linspaced_op<Scalar>> {
enum {
Cost = 1,
PacketAccess = (!NumTraits<Scalar>::IsInteger) && packet_traits<Scalar>::HasSetLinear,
@@ -192,7 +193,7 @@ struct equalspaced_op {
};
template <typename Scalar>
struct functor_traits<equalspaced_op<Scalar> > {
struct functor_traits<equalspaced_op<Scalar>> {
enum {
Cost = NumTraits<Scalar>::AddCost + NumTraits<Scalar>::MulCost,
PacketAccess =