Fix pexp complex test edge-cases.

This commit is contained in:
Antonio Sánchez
2024-03-04 17:44:38 +00:00
committed by Rasmus Munk Larsen
parent 251ec42087
commit 38fcedaf8e
3 changed files with 140 additions and 47 deletions

View File

@@ -164,7 +164,7 @@ struct imag_ref_default_impl {
typedef typename NumTraits<Scalar>::Real RealScalar;
EIGEN_DEVICE_FUNC static inline RealScalar& run(Scalar& x) { return reinterpret_cast<RealScalar*>(&x)[1]; }
EIGEN_DEVICE_FUNC static inline const RealScalar& run(const Scalar& x) {
return reinterpret_cast<RealScalar*>(&x)[1];
return reinterpret_cast<const RealScalar*>(&x)[1];
}
};
@@ -1541,6 +1541,25 @@ EIGEN_DEVICE_FUNC EIGEN_ALWAYS_INLINE T exp(const T& x) {
return exp(x);
}
// MSVC screws up some edge-cases for std::exp(complex).
#ifdef EIGEN_COMP_MSVC
template <typename RealScalar>
EIGEN_DEVICE_FUNC EIGEN_ALWAYS_INLINE std::complex<RealScalar> exp(const std::complex<RealScalar>& x) {
EIGEN_USING_STD(exp);
// If z is (x,±∞) (for any finite x), the result is (NaN,NaN) and FE_INVALID is raised.
// If z is (x,NaN) (for any finite x), the result is (NaN,NaN) and FE_INVALID may be raised.
if ((isfinite)(real_ref(x)) && !(isfinite)(imag_ref(x))) {
return std::complex<RealScalar>(NumTraits<RealScalar>::quiet_NaN(), NumTraits<RealScalar>::quiet_NaN());
}
// If z is (+∞,±∞), the result is (±∞,NaN) and FE_INVALID is raised (the sign of the real part is unspecified)
// If z is (+∞,NaN), the result is (±∞,NaN) (the sign of the real part is unspecified)
if ((real_ref(x) == NumTraits<RealScalar>::infinity() && !(isfinite)(imag_ref(x)))) {
return std::complex<RealScalar>(NumTraits<RealScalar>::infinity(), NumTraits<RealScalar>::quiet_NaN());
}
return exp(x);
}
#endif
#if defined(SYCL_DEVICE_ONLY)
SYCL_SPECIALIZE_FLOATING_TYPES_UNARY(exp, exp)
#endif

View File

@@ -1068,40 +1068,39 @@ EIGEN_DEFINE_FUNCTION_ALLOWING_MULTIPLE_DEFINITIONS Packet pexp_complex(const Pa
typedef typename unpacket_traits<Packet>::type Scalar;
typedef typename Scalar::value_type RealScalar;
const RealPacket even_mask = peven_mask(a.v);
const Packet even_maskp = Packet(even_mask);
const RealPacket odd_mask = pcplxflip(Packet(even_mask)).v;
Packet p0y = Packet(pand(odd_mask, a.v));
Packet py0 = pcplxflip(p0y);
Packet pyy = padd(p0y, py0);
// Let a = x + iy.
// exp(a) = exp(x) * cis(y), plus some special edge-case handling.
RealPacket sincos = psincos_float<false, RealPacket, true>(pyy.v);
RealPacket cossin = pcplxflip(Packet(sincos)).v;
// exp(x):
RealPacket x = pand(a.v, even_mask);
x = por(x, pcplxflip(Packet(x)).v);
RealPacket expx = pexp(x); // exp(x);
// cis(y):
RealPacket y = pand(odd_mask, a.v);
y = por(y, pcplxflip(Packet(y)).v);
RealPacket cisy = psincos_float<false, RealPacket, true>(y);
cisy = pcplxflip(Packet(cisy)).v; // cos(y) + i * sin(y)
const RealPacket cst_pos_inf = pset1<RealPacket>(NumTraits<RealScalar>::infinity());
const RealPacket cst_neg_inf = pset1<RealPacket>(-NumTraits<RealScalar>::infinity());
Packet x_is_inf = Packet(pcmp_eq(a.v, cst_pos_inf));
Packet x_is_minf = Packet(pcmp_eq(a.v, cst_neg_inf));
Packet x_is_zero = Packet(pcmp_eq(pzero(a).v, a.v));
Packet x_real_is_inf = pand(even_maskp, x_is_inf);
Packet x_real_is_minf = pand(even_maskp, x_is_minf);
Packet inf0 = pset1<Packet>(Scalar(NumTraits<RealScalar>::infinity(), RealScalar(0)));
Packet x_is_inf0 = pand(x_real_is_inf, pcplxflip(x_is_zero));
x_is_inf0 = por(x_is_inf0, pcplxflip(x_is_inf0));
Packet x_imag_goes_zero = pand(por(x_is_minf, x_is_inf), pcplxflip(x_real_is_minf));
Packet x_is_nan = Packet(pisnan(a.v));
Packet x_real_goes_zero = pand(x_is_nan, pcplxflip(x_real_is_minf));
RealPacket pexp_real = pexp(a.v);
Packet pexp_half = Packet(pand(even_mask, pexp_real));
RealPacket xexp_flip_rp = pcplxflip(pexp_half).v;
RealPacket xexp = padd(pexp_half.v, xexp_flip_rp);
Packet result(pmul(cossin, xexp));
// If x is -inf, we know that cossin(y) is bounded,
// so the result is (0, +/-0), where the sign of the imaginary part comes
// from the sign of cossin(y).
RealPacket cisy_sign = por(pandnot(cisy, pabs(cisy)), pset1<RealPacket>(RealScalar(1)));
cisy = pselect(pcmp_eq(x, cst_neg_inf), cisy_sign, cisy);
result = pselect(x_is_inf0, inf0, result);
result = pselect(x_real_is_minf, pzero(a), result);
result = pselect(x_imag_goes_zero, pzero(a), result);
result = pselect(x_real_goes_zero, pzero(a), result);
// If x is inf, and cos(y) has unknown sign (y is inf or NaN), the result
// is (+/-inf, NaN), where the signs are undetermined (take the sign of y).
RealPacket y_sign = por(pandnot(y, pabs(y)), pset1<RealPacket>(RealScalar(1)));
cisy = pselect(pand(pcmp_eq(x, cst_pos_inf), pisnan(cisy)), pand(y_sign, even_mask), cisy);
Packet result = Packet(pmul(expx, cisy));
// If y is +/- 0, the input is real, so take the real result for consistency.
result = pselect(Packet(pcmp_eq(y, pzero(y))), Packet(por(pand(expx, even_mask), pand(y, odd_mask))), result);
return result;
}