Speed up exp(x).

This commit is contained in:
Rasmus Munk Larsen
2024-11-19 17:50:34 +00:00
parent 93ec5450cb
commit 8ee6f8475a
5 changed files with 66 additions and 18 deletions

View File

@@ -274,22 +274,20 @@ EIGEN_STRONG_INLINE EIGEN_DEVICE_FUNC Packet pldexp_generic(const Packet& a, con
//
// Assumes IEEE floating point format
template <typename Packet>
struct pldexp_fast_impl {
EIGEN_STRONG_INLINE EIGEN_DEVICE_FUNC Packet pldexp_fast(const Packet& a, const Packet& exponent) {
typedef typename unpacket_traits<Packet>::integer_packet PacketI;
typedef typename unpacket_traits<Packet>::type Scalar;
typedef typename unpacket_traits<PacketI>::type ScalarI;
static constexpr int TotalBits = sizeof(Scalar) * CHAR_BIT, MantissaBits = numext::numeric_limits<Scalar>::digits - 1,
ExponentBits = TotalBits - MantissaBits - 1;
static EIGEN_STRONG_INLINE EIGEN_DEVICE_FUNC Packet run(const Packet& a, const Packet& exponent) {
const Packet bias = pset1<Packet>(Scalar((ScalarI(1) << (ExponentBits - 1)) - ScalarI(1))); // 127
const Packet limit = pset1<Packet>(Scalar((ScalarI(1) << ExponentBits) - ScalarI(1))); // 255
// restrict biased exponent between 0 and 255 for float.
const PacketI e = pcast<Packet, PacketI>(pmin(pmax(padd(exponent, bias), pzero(limit)), limit)); // exponent + 127
// return a * (2^e)
return pmul(a, preinterpret<Packet>(plogical_shift_left<MantissaBits>(e)));
}
};
const Packet bias = pset1<Packet>(Scalar((ScalarI(1) << (ExponentBits - 1)) - ScalarI(1))); // 127
const Packet limit = pset1<Packet>(Scalar((ScalarI(1) << ExponentBits) - ScalarI(1))); // 255
// restrict biased exponent between 0 and 255 for float.
const PacketI e = pcast<Packet, PacketI>(pmin(pmax(padd(exponent, bias), pzero(limit)), limit)); // exponent + 127
// return a * (2^e)
return pmul(a, preinterpret<Packet>(plogical_shift_left<MantissaBits>(e)));
}
// Natural or base 2 logarithm.
// Computes log(x) as log(2^e * m) = C*e + log(m), where the constant C =log(2)
@@ -549,7 +547,11 @@ EIGEN_DEFINE_FUNCTION_ALLOWING_MULTIPLE_DEFINITIONS Packet pexp_float(const Pack
y = pmadd(r2, y, p_low);
// Return 2^m * exp(r).
// TODO: replace pldexp with faster implementation since y in [-1, 1).
const Packet fast_pldexp_unsafe = pandnot(pcmp_lt(x, pset1<Packet>(-87.0)), zero_mask);
if (!predux_any(fast_pldexp_unsafe)) {
// For x >= -87, we can safely use the fast version of pldexp.
return pselect(zero_mask, cst_zero, pmax(pldexp_fast(y, m), _x));
}
return pselect(zero_mask, cst_zero, pmax(pldexp(y, m), _x));
}
@@ -562,7 +564,7 @@ EIGEN_DEFINE_FUNCTION_ALLOWING_MULTIPLE_DEFINITIONS Packet pexp_double(const Pac
const Packet cst_half = pset1<Packet>(0.5);
const Packet cst_exp_hi = pset1<Packet>(709.784);
const Packet cst_exp_lo = pset1<Packet>(-709.784);
const Packet cst_exp_lo = pset1<Packet>(-745.519);
const Packet cst_cephes_LOG2EF = pset1<Packet>(1.4426950408889634073599);
const Packet cst_cephes_exp_p0 = pset1<Packet>(1.26177193074810590878e-4);
@@ -616,7 +618,11 @@ EIGEN_DEFINE_FUNCTION_ALLOWING_MULTIPLE_DEFINITIONS Packet pexp_double(const Pac
// Construct the result 2^n * exp(g) = e * x. The max is used to catch
// non-finite values in the input.
// TODO: replace pldexp with faster implementation since x in [-1, 1).
const Packet fast_pldexp_unsafe = pandnot(pcmp_lt(_x, pset1<Packet>(-708.0)), zero_mask);
if (!predux_any(fast_pldexp_unsafe)) {
// For x >= -708, we can safely use the fast version of pldexp.
return pselect(zero_mask, cst_zero, pmax(pldexp_fast(x, fx), _x));
}
return pselect(zero_mask, cst_zero, pmax(pldexp(x, fx), _x));
}