Replace calls to numext::fma with numext:madd.

This commit is contained in:
Antonio Sánchez
2025-08-28 21:40:19 +00:00
committed by Rasmus Munk Larsen
parent 52f570a409
commit 2e8cc042a1
7 changed files with 88 additions and 44 deletions

View File

@@ -941,23 +941,44 @@ struct nearest_integer_impl<Scalar, true> {
static EIGEN_DEVICE_FUNC EIGEN_STRONG_INLINE Scalar run_trunc(const Scalar& x) { return x; }
};
// Extra namespace to prevent leaking std::fma into Eigen::internal.
namespace has_fma_detail {
template <typename T, typename EnableIf = void>
struct has_fma_impl : public std::false_type {};
using std::fma;
template <typename T>
struct has_fma_impl<
T, std::enable_if_t<std::is_same<T, decltype(fma(std::declval<T>(), std::declval<T>(), std::declval<T>()))>::value>>
: public std::true_type {};
} // namespace has_fma_detail
template <typename T>
struct has_fma : public has_fma_detail::has_fma_impl<T> {};
// Default implementation.
template <typename Scalar, typename Enable = void>
template <typename T, typename Enable = void>
struct fma_impl {
static EIGEN_DEVICE_FUNC EIGEN_STRONG_INLINE Scalar run(const Scalar& a, const Scalar& b, const Scalar& c) {
return a * b + c;
static_assert(has_fma<T>::value, "No function fma(...) for type. Please provide an implementation.");
};
// STD or ADL version if it exists.
template <typename T>
struct fma_impl<T, std::enable_if_t<has_fma<T>::value>> {
static T run(const T& a, const T& b, const T& c) {
using std::fma;
return fma(a, b, c);
}
};
// ADL version if it exists.
template <typename T>
struct fma_impl<
T,
std::enable_if_t<std::is_same<T, decltype(fma(std::declval<T>(), std::declval<T>(), std::declval<T>()))>::value>> {
static T run(const T& a, const T& b, const T& c) { return fma(a, b, c); }
};
#if defined(EIGEN_GPUCC)
template <>
struct has_fma<float> : public true_type {
}
template <>
struct fma_impl<float, void> {
static EIGEN_DEVICE_FUNC EIGEN_STRONG_INLINE float run(const float& a, const float& b, const float& c) {
@@ -965,6 +986,10 @@ struct fma_impl<float, void> {
}
};
template <>
struct has_fma<double> : public true_type {
}
template <>
struct fma_impl<double, void> {
static EIGEN_DEVICE_FUNC EIGEN_STRONG_INLINE double run(const double& a, const double& b, const double& c) {
@@ -973,6 +998,24 @@ struct fma_impl<double, void> {
};
#endif
// Basic multiply-add.
template <typename Scalar, typename EnableIf = void>
struct madd_impl {
static EIGEN_DEVICE_FUNC EIGEN_STRONG_INLINE Scalar run(const Scalar& x, const Scalar& y, const Scalar& z) {
return x * y + z;
}
};
// Use FMA if there is a single CPU instruction.
#ifdef EIGEN_VECTORIZE_FMA
template <typename Scalar>
struct madd_impl<Scalar, std::enable_if_t<has_fma<Scalar>::value>> {
static EIGEN_DEVICE_FUNC EIGEN_STRONG_INLINE Scalar run(const Scalar& x, const Scalar& y, const Scalar& z) {
return fma_impl<Scalar>::run(x, y, z);
}
};
#endif
} // end namespace internal
/****************************************************************************
@@ -1886,15 +1929,18 @@ EIGEN_DEVICE_FUNC EIGEN_STRONG_INLINE Scalar arithmetic_shift_right(const Scalar
return bit_cast<Scalar, SignedScalar>(bit_cast<SignedScalar, Scalar>(a) >> n);
}
// Use std::fma if available.
using std::fma;
// Otherwise, rely on template implementation.
template <typename Scalar>
EIGEN_DEVICE_FUNC EIGEN_STRONG_INLINE Scalar fma(const Scalar& x, const Scalar& y, const Scalar& z) {
return internal::fma_impl<Scalar>::run(x, y, z);
}
// Multiply-add.
template <typename Scalar>
EIGEN_DEVICE_FUNC EIGEN_STRONG_INLINE Scalar madd(const Scalar& x, const Scalar& y, const Scalar& z) {
return internal::madd_impl<Scalar>::run(x, y, z);
}
} // end namespace numext
namespace internal {