Modify pselect and various masks to use Scalar(1) for true.

This commit is contained in:
Antonio Sánchez
2025-06-20 22:40:46 +00:00
parent 6de0515fa6
commit db8bd5b825
6 changed files with 103 additions and 113 deletions

View File

@@ -375,7 +375,7 @@ EIGEN_DEVICE_FUNC inline bool pdiv(const bool& a, const bool& b) {
return a && b;
}
// In the generic case, memset to all one bits.
// In the generic packet case, memset to all one bits.
template <typename Packet, typename EnableIf = void>
struct ptrue_impl {
static EIGEN_DEVICE_FUNC inline Packet run(const Packet& /*a*/) {
@@ -385,19 +385,16 @@ struct ptrue_impl {
}
};
// Use a value of one for scalars.
template <typename Scalar>
struct ptrue_impl<Scalar, std::enable_if_t<is_scalar<Scalar>::value>> {
static EIGEN_DEVICE_FUNC inline Scalar run(const Scalar&) { return Scalar(1); }
};
// For booleans, we can only directly set a valid `bool` value to avoid UB.
template <>
struct ptrue_impl<bool, void> {
static EIGEN_DEVICE_FUNC inline bool run(const bool& /*a*/) { return true; }
};
// For non-trivial scalars, set to Scalar(1) (i.e. a non-zero value).
// Although this is technically not a valid bitmask, the scalar path for pselect
// uses a comparison to zero, so this should still work in most cases. We don't
// have another option, since the scalar type requires initialization.
template <typename T>
struct ptrue_impl<T, std::enable_if_t<is_scalar<T>::value && NumTraits<T>::RequireInitialization>> {
static EIGEN_DEVICE_FUNC inline T run(const T& /*a*/) { return T(1); }
static EIGEN_DEVICE_FUNC inline bool run(const bool&) { return true; }
};
/** \internal \returns one bits. */
@@ -406,7 +403,7 @@ EIGEN_DEVICE_FUNC inline Packet ptrue(const Packet& a) {
return ptrue_impl<Packet>::run(a);
}
// In the general case, memset to zero.
// In the general packet case, memset to zero.
template <typename Packet, typename EnableIf = void>
struct pzero_impl {
static EIGEN_DEVICE_FUNC inline Packet run(const Packet& /*a*/) {
@@ -875,17 +872,29 @@ EIGEN_DEVICE_FUNC EIGEN_STRONG_INLINE Packet plset(const typename unpacket_trait
return a;
}
template <typename Packet, typename EnableIf = void>
struct peven_mask_impl {
static EIGEN_DEVICE_FUNC EIGEN_STRONG_INLINE Packet run(const Packet&) {
typedef typename unpacket_traits<Packet>::type Scalar;
const size_t n = unpacket_traits<Packet>::size;
EIGEN_ALIGN_TO_BOUNDARY(sizeof(Packet)) Scalar elements[n];
for (size_t i = 0; i < n; ++i) {
memset(elements + i, ((i & 1) == 0 ? 0xff : 0), sizeof(Scalar));
}
return ploadu<Packet>(elements);
}
};
template <typename Scalar>
struct peven_mask_impl<Scalar, std::enable_if_t<is_scalar<Scalar>::value>> {
static EIGEN_DEVICE_FUNC EIGEN_STRONG_INLINE Scalar run(const Scalar&) { return Scalar(1); }
};
/** \internal \returns a packet with constant coefficients \a a, e.g.: (x, 0, x, 0),
where x is the value of all 1-bits. */
template <typename Packet>
EIGEN_DEVICE_FUNC inline Packet peven_mask(const Packet& /*a*/) {
typedef typename unpacket_traits<Packet>::type Scalar;
const size_t n = unpacket_traits<Packet>::size;
EIGEN_ALIGN_TO_BOUNDARY(sizeof(Packet)) Scalar elements[n];
for (size_t i = 0; i < n; ++i) {
memset(elements + i, ((i & 1) == 0 ? 0xff : 0), sizeof(Scalar));
}
return ploadu<Packet>(elements);
EIGEN_DEVICE_FUNC inline Packet peven_mask(const Packet& a) {
return peven_mask_impl<Packet>::run(a);
}
/** \internal copy the packet \a from to \a *to, \a to must be properly aligned */