Remove assumption of std::complex for complex scalar types.

This commit is contained in:
Antonio Sanchez
2025-02-12 11:21:44 -08:00
parent 6b4881ad48
commit 22cd7307dd
21 changed files with 273 additions and 115 deletions

View File

@@ -182,6 +182,35 @@ struct imag_ref_retval {
typedef typename NumTraits<Scalar>::Real& type;
};
} // namespace internal
namespace numext {
template <typename Scalar>
EIGEN_DEVICE_FUNC inline EIGEN_MATHFUNC_RETVAL(real, Scalar) real(const Scalar& x) {
return EIGEN_MATHFUNC_IMPL(real, Scalar)::run(x);
}
template <typename Scalar>
EIGEN_DEVICE_FUNC inline internal::add_const_on_value_type_t<EIGEN_MATHFUNC_RETVAL(real_ref, Scalar)> real_ref(
const Scalar& x) {
return internal::real_ref_impl<Scalar>::run(x);
}
template <typename Scalar>
EIGEN_DEVICE_FUNC inline EIGEN_MATHFUNC_RETVAL(real_ref, Scalar) real_ref(Scalar& x) {
return EIGEN_MATHFUNC_IMPL(real_ref, Scalar)::run(x);
}
template <typename Scalar>
EIGEN_DEVICE_FUNC inline EIGEN_MATHFUNC_RETVAL(imag, Scalar) imag(const Scalar& x) {
return EIGEN_MATHFUNC_IMPL(imag, Scalar)::run(x);
}
} // namespace numext
namespace internal {
/****************************************************************************
* Implementation of conj *
****************************************************************************/
@@ -221,7 +250,9 @@ template <typename Scalar>
struct abs2_impl_default<Scalar, true> // IsComplex
{
typedef typename NumTraits<Scalar>::Real RealScalar;
EIGEN_DEVICE_FUNC static inline RealScalar run(const Scalar& x) { return x.real() * x.real() + x.imag() * x.imag(); }
EIGEN_DEVICE_FUNC static inline RealScalar run(const Scalar& x) {
return numext::real(x) * numext::real(x) + numext::imag(x) * numext::imag(x);
}
};
template <typename Scalar>
@@ -250,16 +281,14 @@ struct sqrt_impl {
};
// Complex sqrt defined in MathFunctionsImpl.h.
template <typename T>
EIGEN_DEVICE_FUNC std::complex<T> complex_sqrt(const std::complex<T>& a_x);
template <typename ComplexT>
EIGEN_DEVICE_FUNC ComplexT complex_sqrt(const ComplexT& a_x);
// Custom implementation is faster than `std::sqrt`, works on
// GPU, and correctly handles special cases (unlike MSVC).
template <typename T>
struct sqrt_impl<std::complex<T>> {
EIGEN_DEVICE_FUNC static EIGEN_ALWAYS_INLINE std::complex<T> run(const std::complex<T>& x) {
return complex_sqrt<T>(x);
}
EIGEN_DEVICE_FUNC static EIGEN_ALWAYS_INLINE std::complex<T> run(const std::complex<T>& x) { return complex_sqrt(x); }
};
template <typename Scalar>
@@ -272,13 +301,13 @@ template <typename T>
struct rsqrt_impl;
// Complex rsqrt defined in MathFunctionsImpl.h.
template <typename T>
EIGEN_DEVICE_FUNC std::complex<T> complex_rsqrt(const std::complex<T>& a_x);
template <typename ComplexT>
EIGEN_DEVICE_FUNC ComplexT complex_rsqrt(const ComplexT& a_x);
template <typename T>
struct rsqrt_impl<std::complex<T>> {
EIGEN_DEVICE_FUNC static EIGEN_ALWAYS_INLINE std::complex<T> run(const std::complex<T>& x) {
return complex_rsqrt<T>(x);
return complex_rsqrt(x);
}
};
@@ -299,7 +328,7 @@ struct norm1_default_impl<Scalar, true> {
typedef typename NumTraits<Scalar>::Real RealScalar;
EIGEN_DEVICE_FUNC static inline RealScalar run(const Scalar& x) {
EIGEN_USING_STD(abs);
return abs(x.real()) + abs(x.imag());
return abs(numext::real(x)) + abs(numext::imag(x));
}
};
@@ -469,8 +498,8 @@ struct expm1_retval {
****************************************************************************/
// Complex log defined in MathFunctionsImpl.h.
template <typename T>
EIGEN_DEVICE_FUNC std::complex<T> complex_log(const std::complex<T>& z);
template <typename ComplexT>
EIGEN_DEVICE_FUNC ComplexT complex_log(const ComplexT& z);
template <typename Scalar>
struct log_impl {
@@ -846,7 +875,7 @@ struct sign_impl<Scalar, true, IsInteger> {
real_type aa = abs(a);
if (aa == real_type(0)) return Scalar(0);
aa = real_type(1) / aa;
return Scalar(a.real() * aa, a.imag() * aa);
return Scalar(numext::real(a) * aa, numext::imag(a) * aa);
}
};
@@ -1042,27 +1071,6 @@ SYCL_SPECIALIZE_FLOATING_TYPES_BINARY(maxi, fmax)
#endif
template <typename Scalar>
EIGEN_DEVICE_FUNC inline EIGEN_MATHFUNC_RETVAL(real, Scalar) real(const Scalar& x) {
return EIGEN_MATHFUNC_IMPL(real, Scalar)::run(x);
}
template <typename Scalar>
EIGEN_DEVICE_FUNC inline internal::add_const_on_value_type_t<EIGEN_MATHFUNC_RETVAL(real_ref, Scalar)> real_ref(
const Scalar& x) {
return internal::real_ref_impl<Scalar>::run(x);
}
template <typename Scalar>
EIGEN_DEVICE_FUNC inline EIGEN_MATHFUNC_RETVAL(real_ref, Scalar) real_ref(Scalar& x) {
return EIGEN_MATHFUNC_IMPL(real_ref, Scalar)::run(x);
}
template <typename Scalar>
EIGEN_DEVICE_FUNC inline EIGEN_MATHFUNC_RETVAL(imag, Scalar) imag(const Scalar& x) {
return EIGEN_MATHFUNC_IMPL(imag, Scalar)::run(x);
}
template <typename Scalar>
EIGEN_DEVICE_FUNC inline EIGEN_MATHFUNC_RETVAL(arg, Scalar) arg(const Scalar& x) {
return EIGEN_MATHFUNC_IMPL(arg, Scalar)::run(x);

View File

@@ -171,8 +171,8 @@ struct hypot_impl {
// Generic complex sqrt implementation that correctly handles corner cases
// according to https://en.cppreference.com/w/cpp/numeric/complex/sqrt
template <typename T>
EIGEN_DEVICE_FUNC std::complex<T> complex_sqrt(const std::complex<T>& z) {
template <typename ComplexT>
EIGEN_DEVICE_FUNC ComplexT complex_sqrt(const ComplexT& z) {
// Computes the principal sqrt of the input.
//
// For a complex square root of the number x + i*y. We want to find real
@@ -194,21 +194,21 @@ EIGEN_DEVICE_FUNC std::complex<T> complex_sqrt(const std::complex<T>& z) {
// if x == 0: u = w, v = sign(y) * w
// if x > 0: u = w, v = y / (2 * w)
// if x < 0: u = |y| / (2 * w), v = sign(y) * w
using T = typename NumTraits<ComplexT>::Real;
const T x = numext::real(z);
const T y = numext::imag(z);
const T zero = T(0);
const T w = numext::sqrt(T(0.5) * (numext::abs(x) + numext::hypot(x, y)));
return (numext::isinf)(y) ? std::complex<T>(NumTraits<T>::infinity(), y)
: numext::is_exactly_zero(x) ? std::complex<T>(w, y < zero ? -w : w)
: x > zero ? std::complex<T>(w, y / (2 * w))
: std::complex<T>(numext::abs(y) / (2 * w), y < zero ? -w : w);
return (numext::isinf)(y) ? ComplexT(NumTraits<T>::infinity(), y)
: numext::is_exactly_zero(x) ? ComplexT(w, y < zero ? -w : w)
: x > zero ? ComplexT(w, y / (2 * w))
: ComplexT(numext::abs(y) / (2 * w), y < zero ? -w : w);
}
// Generic complex rsqrt implementation.
template <typename T>
EIGEN_DEVICE_FUNC std::complex<T> complex_rsqrt(const std::complex<T>& z) {
template <typename ComplexT>
EIGEN_DEVICE_FUNC ComplexT complex_rsqrt(const ComplexT& z) {
// Computes the principal reciprocal sqrt of the input.
//
// For a complex reciprocal square root of the number z = x + i*y. We want to
@@ -230,7 +230,7 @@ EIGEN_DEVICE_FUNC std::complex<T> complex_rsqrt(const std::complex<T>& z) {
// if x == 0: u = w / |z|, v = -sign(y) * w / |z|
// if x > 0: u = w / |z|, v = -y / (2 * w * |z|)
// if x < 0: u = |y| / (2 * w * |z|), v = -sign(y) * w / |z|
using T = typename NumTraits<ComplexT>::Real;
const T x = numext::real(z);
const T y = numext::imag(z);
const T zero = T(0);
@@ -239,20 +239,21 @@ EIGEN_DEVICE_FUNC std::complex<T> complex_rsqrt(const std::complex<T>& z) {
const T w = numext::sqrt(T(0.5) * (numext::abs(x) + abs_z));
const T woz = w / abs_z;
// Corner cases consistent with 1/sqrt(z) on gcc/clang.
return numext::is_exactly_zero(abs_z) ? std::complex<T>(NumTraits<T>::infinity(), NumTraits<T>::quiet_NaN())
: ((numext::isinf)(x) || (numext::isinf)(y)) ? std::complex<T>(zero, zero)
: numext::is_exactly_zero(x) ? std::complex<T>(woz, y < zero ? woz : -woz)
: x > zero ? std::complex<T>(woz, -y / (2 * w * abs_z))
: std::complex<T>(numext::abs(y) / (2 * w * abs_z), y < zero ? woz : -woz);
return numext::is_exactly_zero(abs_z) ? ComplexT(NumTraits<T>::infinity(), NumTraits<T>::quiet_NaN())
: ((numext::isinf)(x) || (numext::isinf)(y)) ? ComplexT(zero, zero)
: numext::is_exactly_zero(x) ? ComplexT(woz, y < zero ? woz : -woz)
: x > zero ? ComplexT(woz, -y / (2 * w * abs_z))
: ComplexT(numext::abs(y) / (2 * w * abs_z), y < zero ? woz : -woz);
}
template <typename T>
EIGEN_DEVICE_FUNC std::complex<T> complex_log(const std::complex<T>& z) {
template <typename ComplexT>
EIGEN_DEVICE_FUNC ComplexT complex_log(const ComplexT& z) {
// Computes complex log.
using T = typename NumTraits<ComplexT>::Real;
T a = numext::abs(z);
EIGEN_USING_STD(atan2);
T b = atan2(z.imag(), z.real());
return std::complex<T>(numext::log(a), b);
return ComplexT(numext::log(a), b);
}
} // end namespace internal

View File

@@ -112,7 +112,7 @@ class MatrixBase : public DenseBase<Derived> {
ConstTransposeReturnType>
AdjointReturnType;
/** \internal Return type of eigenvalues() */
typedef Matrix<std::complex<RealScalar>, internal::traits<Derived>::ColsAtCompileTime, 1, ColMajor>
typedef Matrix<internal::make_complex_t<Scalar>, internal::traits<Derived>::ColsAtCompileTime, 1, ColMajor>
EigenvaluesReturnType;
/** \internal the return type of identity */
typedef CwiseNullaryOp<internal::scalar_identity_op<Scalar>, PlainObject> IdentityReturnType;
@@ -468,7 +468,7 @@ class MatrixBase : public DenseBase<Derived> {
EIGEN_MATRIX_FUNCTION(MatrixSquareRootReturnValue, sqrt, square root)
EIGEN_MATRIX_FUNCTION(MatrixLogarithmReturnValue, log, logarithm)
EIGEN_MATRIX_FUNCTION_1(MatrixPowerReturnValue, pow, power to \c p, const RealScalar& p)
EIGEN_MATRIX_FUNCTION_1(MatrixComplexPowerReturnValue, pow, power to \c p, const std::complex<RealScalar>& p)
EIGEN_MATRIX_FUNCTION_1(MatrixComplexPowerReturnValue, pow, power to \c p, const internal::make_complex_t<Scalar>& p)
protected:
EIGEN_DEFAULT_COPY_CONSTRUCTOR(MatrixBase)

View File

@@ -497,7 +497,7 @@ class MatrixComplexPowerReturnValue;
namespace internal {
template <typename Scalar>
struct stem_function {
typedef std::complex<typename NumTraits<Scalar>::Real> ComplexScalar;
typedef internal::make_complex_t<Scalar> ComplexScalar;
typedef ComplexScalar type(ComplexScalar, int);
};
} // namespace internal

View File

@@ -745,6 +745,9 @@ using std::is_constant_evaluated;
constexpr bool is_constant_evaluated() { return false; }
#endif
template <typename Scalar>
using make_complex_t = std::conditional_t<NumTraits<Scalar>::IsComplex, Scalar, std::complex<Scalar>>;
} // end namespace internal
} // end namespace Eigen

View File

@@ -885,8 +885,12 @@ struct scalar_div_cost {
};
template <typename T, bool Vectorized>
struct scalar_div_cost<std::complex<T>, Vectorized> {
enum { value = 2 * scalar_div_cost<T>::value + 6 * NumTraits<T>::MulCost + 3 * NumTraits<T>::AddCost };
struct scalar_div_cost<T, Vectorized, std::enable_if_t<NumTraits<T>::IsComplex>> {
using RealScalar = typename NumTraits<T>::Real;
enum {
value =
2 * scalar_div_cost<RealScalar>::value + 6 * NumTraits<RealScalar>::MulCost + 3 * NumTraits<RealScalar>::AddCost
};
};
template <bool Vectorized>

View File

@@ -70,7 +70,7 @@ class ComplexEigenSolver {
* \c float or \c double) and just \c Scalar if #Scalar is
* complex.
*/
typedef std::complex<RealScalar> ComplexScalar;
typedef internal::make_complex_t<Scalar> ComplexScalar;
/** \brief Type for vector of eigenvalues as returned by eigenvalues().
*

View File

@@ -75,7 +75,7 @@ class ComplexSchur {
* \c float or \c double) and just \c Scalar if #Scalar is
* complex.
*/
typedef std::complex<RealScalar> ComplexScalar;
typedef internal::make_complex_t<Scalar> ComplexScalar;
/** \brief Type for the matrices in the Schur decomposition.
*

View File

@@ -89,7 +89,7 @@ class EigenSolver {
* \c float or \c double) and just \c Scalar if #Scalar is
* complex.
*/
typedef std::complex<RealScalar> ComplexScalar;
typedef internal::make_complex_t<Scalar> ComplexScalar;
/** \brief Type for vector of eigenvalues as returned by eigenvalues().
*

View File

@@ -83,7 +83,7 @@ class GeneralizedEigenSolver {
* \c float or \c double) and just \c Scalar if #Scalar is
* complex.
*/
typedef std::complex<RealScalar> ComplexScalar;
typedef internal::make_complex_t<Scalar> ComplexScalar;
/** \brief Type for vector of real scalar values eigenvalues as returned by betas().
*

View File

@@ -69,7 +69,7 @@ class RealQZ {
MaxColsAtCompileTime = MatrixType::MaxColsAtCompileTime
};
typedef typename MatrixType::Scalar Scalar;
typedef std::complex<typename NumTraits<Scalar>::Real> ComplexScalar;
typedef internal::make_complex_t<Scalar> ComplexScalar;
typedef Eigen::Index Index; ///< \deprecated since Eigen 3.3
typedef Matrix<ComplexScalar, ColsAtCompileTime, 1, Options & ~RowMajor, MaxColsAtCompileTime, 1> EigenvalueType;

View File

@@ -66,7 +66,7 @@ class RealSchur {
MaxColsAtCompileTime = MatrixType::MaxColsAtCompileTime
};
typedef typename MatrixType::Scalar Scalar;
typedef std::complex<typename NumTraits<Scalar>::Real> ComplexScalar;
typedef internal::make_complex_t<Scalar> ComplexScalar;
typedef Eigen::Index Index; ///< \deprecated since Eigen 3.3
typedef Matrix<ComplexScalar, ColsAtCompileTime, 1, Options & ~RowMajor, MaxColsAtCompileTime, 1> EigenvalueType;