Fix MSVC complex sqrt and packetmath test.

MSVC incorrectly handles `inf` cases for `std::sqrt<std::complex<T>>`.
Here we replace it with a custom version (currently used on GPU).

Also fixed the `packetmath` test, which previously skipped several
corner cases since `CHECK_CWISE1` only tests the first `PacketSize`
elements.
This commit is contained in:
Antonio Sanchez
2021-01-07 09:39:05 -08:00
committed by Rasmus Munk Larsen
parent 8d9cfba799
commit f149e0ebc3
5 changed files with 84 additions and 47 deletions

View File

@@ -338,6 +338,22 @@ struct sqrt_impl
}
};
// Complex sqrt defined in MathFunctionsImpl.h.
template<typename T> std::complex<T> complex_sqrt(const std::complex<T>& a_x);
// MSVC incorrectly handles inf cases.
#if EIGEN_COMP_MSVC > 0
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);
}
};
#endif
template<typename Scalar>
struct sqrt_retval
{

View File

@@ -99,6 +99,50 @@ 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) {
// Computes the principal sqrt of the input.
//
// For a complex square root of the number x + i*y. We want to find real
// numbers u and v such that
// (u + i*v)^2 = x + i*y <=>
// u^2 - v^2 + i*2*u*v = x + i*v.
// By equating the real and imaginary parts we get:
// u^2 - v^2 = x
// 2*u*v = y.
//
// For x >= 0, this has the numerically stable solution
// u = sqrt(0.5 * (x + sqrt(x^2 + y^2)))
// v = y / (2 * u)
// and for x < 0,
// v = sign(y) * sqrt(0.5 * (-x + sqrt(x^2 + y^2)))
// u = y / (2 * v)
//
// Letting w = sqrt(0.5 * (|x| + |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
const T x = numext::real(z);
const T y = numext::imag(z);
const T zero = T(0);
const T cst_half = T(0.5);
// Special case of isinf(y)
if ((numext::isinf)(y)) {
const T inf = std::numeric_limits<T>::infinity();
return std::complex<T>(inf, y);
}
T w = numext::sqrt(cst_half * (numext::abs(x) + numext::abs(z)));
return
x == zero ? 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 );
}
} // end namespace internal
} // end namespace Eigen

View File

@@ -95,46 +95,12 @@ template<typename T> struct scalar_quotient_op<const std::complex<T>, const std:
template<typename T> struct scalar_quotient_op<std::complex<T>, std::complex<T> > : scalar_quotient_op<const std::complex<T>, const std::complex<T> > {};
template<typename T>
struct sqrt_impl<std::complex<T>> {
static EIGEN_DEVICE_FUNC std::complex<T> run(const std::complex<T>& z) {
// Computes the principal sqrt of the input.
//
// For a complex square root of the number x + i*y. We want to find real
// numbers u and v such that
// (u + i*v)^2 = x + i*y <=>
// u^2 - v^2 + i*2*u*v = x + i*v.
// By equating the real and imaginary parts we get:
// u^2 - v^2 = x
// 2*u*v = y.
//
// For x >= 0, this has the numerically stable solution
// u = sqrt(0.5 * (x + sqrt(x^2 + y^2)))
// v = y / (2 * u)
// and for x < 0,
// v = sign(y) * sqrt(0.5 * (-x + sqrt(x^2 + y^2)))
// u = y / (2 * v)
//
// Letting w = sqrt(0.5 * (|x| + |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
const T x = numext::real(z);
const T y = numext::imag(z);
const T zero = T(0);
const T cst_half = T(0.5);
// Special case of isinf(y)
if ((numext::isinf)(y)) {
const T inf = std::numeric_limits<T>::infinity();
return std::complex<T>(inf, y);
}
T w = numext::sqrt(cst_half * (numext::abs(x) + numext::abs(z)));
return
x == zero ? 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 );
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);
}
};