refactor AssignmentFunctors.h, unify with existing scalar_op

This commit is contained in:
Charles Schlosser 2025-03-06 01:28:39 +00:00 committed by Rasmus Munk Larsen
parent 9a86214039
commit d28041ed5a

View File

@ -27,7 +27,7 @@ struct assign_op {
template <int Alignment, typename Packet> template <int Alignment, typename Packet>
EIGEN_STRONG_INLINE void assignPacket(DstScalar* a, const Packet& b) const { EIGEN_STRONG_INLINE void assignPacket(DstScalar* a, const Packet& b) const {
internal::pstoret<DstScalar, Packet, Alignment>(a, b); pstoret<DstScalar, Packet, Alignment>(a, b);
} }
}; };
@ -36,7 +36,7 @@ template <typename DstScalar>
struct assign_op<DstScalar, void> {}; struct assign_op<DstScalar, void> {};
template <typename DstScalar, typename SrcScalar> template <typename DstScalar, typename SrcScalar>
struct functor_traits<assign_op<DstScalar, SrcScalar> > { struct functor_traits<assign_op<DstScalar, SrcScalar>> {
enum { enum {
Cost = NumTraits<DstScalar>::ReadCost, Cost = NumTraits<DstScalar>::ReadCost,
PacketAccess = is_same<DstScalar, SrcScalar>::value && packet_traits<DstScalar>::Vectorizable && PacketAccess = is_same<DstScalar, SrcScalar>::value && packet_traits<DstScalar>::Vectorizable &&
@ -45,88 +45,70 @@ struct functor_traits<assign_op<DstScalar, SrcScalar> > {
}; };
/** \internal /** \internal
* \brief Template functor for scalar/packet assignment with addition * \brief Template functor for scalar/packet compound assignment
* *
*/ */
template <typename DstScalar, typename SrcScalar> template <typename DstScalar, typename SrcScalar, typename Func>
struct add_assign_op { struct compound_assign_op {
EIGEN_DEVICE_FUNC EIGEN_STRONG_INLINE void assignCoeff(DstScalar& a, const SrcScalar& b) const { a += b; } using traits = functor_traits<compound_assign_op<DstScalar, SrcScalar, Func>>;
EIGEN_DEVICE_FUNC EIGEN_STRONG_INLINE constexpr void assignCoeff(DstScalar& a, const SrcScalar& b) const {
assign_op<DstScalar, DstScalar>().assignCoeff(a, Func().operator()(a, b));
}
template <int Alignment, typename Packet> template <int Alignment, typename Packet>
EIGEN_STRONG_INLINE void assignPacket(DstScalar* a, const Packet& b) const { EIGEN_STRONG_INLINE void assignPacket(DstScalar* a, const Packet& b) const {
internal::pstoret<DstScalar, Packet, Alignment>(a, internal::padd(internal::ploadt<Packet, Alignment>(a), b)); assign_op<DstScalar, DstScalar>().template assignPacket<Alignment, Packet>(
a, Func().packetOp(ploadt<Packet, Alignment>(a), b));
} }
}; };
template <typename DstScalar, typename SrcScalar>
struct functor_traits<add_assign_op<DstScalar, SrcScalar> > { template <typename DstScalar, typename SrcScalar, typename Func>
struct functor_traits<compound_assign_op<DstScalar, SrcScalar, Func>> {
enum { enum {
Cost = NumTraits<DstScalar>::ReadCost + NumTraits<DstScalar>::AddCost, Cost = functor_traits<assign_op<DstScalar, DstScalar>>::Cost + functor_traits<Func>::Cost,
PacketAccess = is_same<DstScalar, SrcScalar>::value && packet_traits<DstScalar>::HasAdd PacketAccess = functor_traits<assign_op<DstScalar, DstScalar>>::PacketAccess && functor_traits<Func>::PacketAccess
}; };
}; };
/** \internal
* \brief Template functor for scalar/packet assignment with addition
*
*/
template <typename DstScalar, typename SrcScalar = DstScalar>
struct add_assign_op : compound_assign_op<DstScalar, SrcScalar, scalar_sum_op<DstScalar, SrcScalar>> {};
template <typename DstScalar, typename SrcScalar>
struct functor_traits<add_assign_op<DstScalar, SrcScalar>> : add_assign_op<DstScalar, SrcScalar>::traits {};
/** \internal /** \internal
* \brief Template functor for scalar/packet assignment with subtraction * \brief Template functor for scalar/packet assignment with subtraction
* *
*/ */
template <typename DstScalar, typename SrcScalar> template <typename DstScalar, typename SrcScalar = DstScalar>
struct sub_assign_op { struct sub_assign_op : compound_assign_op<DstScalar, SrcScalar, scalar_difference_op<DstScalar, SrcScalar>> {};
EIGEN_DEVICE_FUNC EIGEN_STRONG_INLINE void assignCoeff(DstScalar& a, const SrcScalar& b) const { a -= b; }
template <int Alignment, typename Packet>
EIGEN_STRONG_INLINE void assignPacket(DstScalar* a, const Packet& b) const {
internal::pstoret<DstScalar, Packet, Alignment>(a, internal::psub(internal::ploadt<Packet, Alignment>(a), b));
}
};
template <typename DstScalar, typename SrcScalar> template <typename DstScalar, typename SrcScalar>
struct functor_traits<sub_assign_op<DstScalar, SrcScalar> > { struct functor_traits<sub_assign_op<DstScalar, SrcScalar>> : sub_assign_op<DstScalar, SrcScalar>::traits {};
enum {
Cost = NumTraits<DstScalar>::ReadCost + NumTraits<DstScalar>::AddCost,
PacketAccess = is_same<DstScalar, SrcScalar>::value && packet_traits<DstScalar>::HasSub
};
};
/** \internal /** \internal
* \brief Template functor for scalar/packet assignment with multiplication * \brief Template functor for scalar/packet assignment with multiplication
* *
*/ */
template <typename DstScalar, typename SrcScalar = DstScalar> template <typename DstScalar, typename SrcScalar = DstScalar>
struct mul_assign_op { struct mul_assign_op : compound_assign_op<DstScalar, SrcScalar, scalar_product_op<DstScalar, SrcScalar>> {};
EIGEN_DEVICE_FUNC EIGEN_STRONG_INLINE void assignCoeff(DstScalar& a, const SrcScalar& b) const { a *= b; }
template <int Alignment, typename Packet>
EIGEN_STRONG_INLINE void assignPacket(DstScalar* a, const Packet& b) const {
internal::pstoret<DstScalar, Packet, Alignment>(a, internal::pmul(internal::ploadt<Packet, Alignment>(a), b));
}
};
template <typename DstScalar, typename SrcScalar> template <typename DstScalar, typename SrcScalar>
struct functor_traits<mul_assign_op<DstScalar, SrcScalar> > { struct functor_traits<mul_assign_op<DstScalar, SrcScalar>> : mul_assign_op<DstScalar, SrcScalar>::traits {};
enum {
Cost = NumTraits<DstScalar>::ReadCost + NumTraits<DstScalar>::MulCost,
PacketAccess = is_same<DstScalar, SrcScalar>::value && packet_traits<DstScalar>::HasMul
};
};
/** \internal /** \internal
* \brief Template functor for scalar/packet assignment with diviving * \brief Template functor for scalar/packet assignment with dividing
* *
*/ */
template <typename DstScalar, typename SrcScalar = DstScalar> template <typename DstScalar, typename SrcScalar = DstScalar>
struct div_assign_op { struct div_assign_op : compound_assign_op<DstScalar, SrcScalar, scalar_quotient_op<DstScalar, SrcScalar>> {};
EIGEN_DEVICE_FUNC EIGEN_STRONG_INLINE void assignCoeff(DstScalar& a, const SrcScalar& b) const { a /= b; }
template <int Alignment, typename Packet>
EIGEN_STRONG_INLINE void assignPacket(DstScalar* a, const Packet& b) const {
internal::pstoret<DstScalar, Packet, Alignment>(a, internal::pdiv(internal::ploadt<Packet, Alignment>(a), b));
}
};
template <typename DstScalar, typename SrcScalar> template <typename DstScalar, typename SrcScalar>
struct functor_traits<div_assign_op<DstScalar, SrcScalar> > { struct functor_traits<div_assign_op<DstScalar, SrcScalar>> : div_assign_op<DstScalar, SrcScalar>::traits {};
enum {
Cost = NumTraits<DstScalar>::ReadCost + NumTraits<DstScalar>::MulCost,
PacketAccess = is_same<DstScalar, SrcScalar>::value && packet_traits<DstScalar>::HasDiv
};
};
/** \internal /** \internal
* \brief Template functor for scalar/packet assignment with swapping * \brief Template functor for scalar/packet assignment with swapping
@ -158,7 +140,7 @@ struct swap_assign_op {
} }
}; };
template <typename Scalar> template <typename Scalar>
struct functor_traits<swap_assign_op<Scalar> > { struct functor_traits<swap_assign_op<Scalar>> {
enum { enum {
Cost = 3 * NumTraits<Scalar>::ReadCost, Cost = 3 * NumTraits<Scalar>::ReadCost,
PacketAccess = PacketAccess =