From fddb6dcd3d03e5cc7bce35dd7900555ef953194c Mon Sep 17 00:00:00 2001 From: Rasmus Munk Larsen <4643818-rmlarsen1@users.noreply.gitlab.com> Date: Sun, 10 May 2026 16:45:48 -0700 Subject: [PATCH] Perf: tighter pldexp 4-way split (tree multiply, vpmovsxdq expansion, drop pnegate) libeigen/eigen!2527 Co-authored-by: Rasmus Munk Larsen --- Eigen/src/Core/arch/AVX/PacketMath.h | 49 ++++++++++--------- Eigen/src/Core/arch/AVX512/PacketMath.h | 29 +++++------ .../Default/GenericPacketMathFrexpLdexp.h | 34 ++++++++----- Eigen/src/Core/arch/SSE/PacketMath.h | 16 +++--- test/packetmath.cpp | 11 +++++ 5 files changed, 77 insertions(+), 62 deletions(-) diff --git a/Eigen/src/Core/arch/AVX/PacketMath.h b/Eigen/src/Core/arch/AVX/PacketMath.h index 12be8248d..6b4831643 100644 --- a/Eigen/src/Core/arch/AVX/PacketMath.h +++ b/Eigen/src/Core/arch/AVX/PacketMath.h @@ -1873,31 +1873,39 @@ EIGEN_STRONG_INLINE Packet8f pldexp(const Packet8f& a, const Packet8f& return pldexp_generic(a, exponent); } +// Build 2^k as Packet4d from a Packet4i holding the biased int32 exponent in +// each lane. AVX2 has a single-instruction widen+shift path; AVX-only must +// split the 128-bit input into two halves, widen+shift each separately with +// SSE intrinsics, and reassemble with vinsertf128. +EIGEN_STRONG_INLINE Packet4d pldexp_avx_pow2_from_biased(const Packet4i& biased) { +#ifdef EIGEN_VECTORIZE_AVX2 + return _mm256_castsi256_pd(_mm256_slli_epi64(_mm256_cvtepi32_epi64(biased), 52)); +#else + __m128i lo = _mm_cvtepi32_epi64(biased); // SSE4.1: lower 2 int32 -> 2 int64 + __m128i hi = _mm_cvtepi32_epi64(_mm_unpackhi_epi64(biased, biased)); // upper 2 int32 -> 2 int64 + lo = _mm_slli_epi64(lo, 52); + hi = _mm_slli_epi64(hi, 52); + return _mm256_castsi256_pd(_mm256_insertf128_si256(_mm256_castsi128_si256(lo), hi, 1)); +#endif +} + template <> EIGEN_STRONG_INLINE Packet4d pldexp(const Packet4d& a, const Packet4d& exponent) { // Clamp exponent to [-2099, 2099] const Packet4d max_exponent = pset1(2099.0); const Packet4i e = _mm256_cvtpd_epi32(pmin(pmax(exponent, pnegate(max_exponent)), max_exponent)); - // Split 2^e into four factors and multiply. + // 4-way split + depth-3 multiply tree; see pldexp_generic for derivation + // (including why the first multiply must be a*c1, not a*c2). const Packet4i bias = pset1(1023); - Packet4i b = parithmetic_shift_right<2>(e); // floor(e/4) + const Packet4i b = parithmetic_shift_right<2>(e); // floor(e/4) + const Packet4i b_remainder = psub(psub(e, b), padd(b, b)); // e - 3b (depth 2) + const Packet4d c1 = pldexp_avx_pow2_from_biased(padd(b, bias)); // 2^b + const Packet4d c2 = pldexp_avx_pow2_from_biased(padd(b_remainder, bias)); // 2^(e-3b) - // 2^b - Packet4i hi = vec4i_swizzle1(padd(b, bias), 0, 2, 1, 3); - Packet4i lo = _mm_slli_epi64(hi, 52); - hi = _mm_slli_epi64(_mm_srli_epi64(hi, 32), 52); - Packet4d c = _mm256_castsi256_pd(_mm256_insertf128_si256(_mm256_castsi128_si256(lo), hi, 1)); - Packet4d out = pmul(pmul(pmul(a, c), c), c); // a * 2^(3b) - - // 2^(e - 3b) - b = psub(psub(psub(e, b), b), b); // e - 3b - hi = vec4i_swizzle1(padd(b, bias), 0, 2, 1, 3); - lo = _mm_slli_epi64(hi, 52); - hi = _mm_slli_epi64(_mm_srli_epi64(hi, 32), 52); - c = _mm256_castsi256_pd(_mm256_insertf128_si256(_mm256_castsi128_si256(lo), hi, 1)); - out = pmul(out, c); // a * 2^e - return out; + const Packet4d c1_squared = pmul(c1, c1); + const Packet4d a_c1 = pmul(a, c1); + return pmul(pmul(a_c1, c1_squared), c2); // a * 2^e } template <> @@ -1908,12 +1916,7 @@ EIGEN_STRONG_INLINE Packet4d pldexp_fast(const Packet4d& a, const Pack const Packet4i e = _mm256_cvtpd_epi32(pmin(pmax(exponent, min_exponent), max_exponent)); const Packet4i bias = pset1(1023); - // 2^e - Packet4i hi = vec4i_swizzle1(padd(e, bias), 0, 2, 1, 3); - const Packet4i lo = _mm_slli_epi64(hi, 52); - hi = _mm_slli_epi64(_mm_srli_epi64(hi, 32), 52); - const Packet4d c = _mm256_castsi256_pd(_mm256_insertf128_si256(_mm256_castsi128_si256(lo), hi, 1)); - return pmul(a, c); // a * 2^e + return pmul(a, pldexp_avx_pow2_from_biased(padd(e, bias))); // a * 2^e } template <> diff --git a/Eigen/src/Core/arch/AVX512/PacketMath.h b/Eigen/src/Core/arch/AVX512/PacketMath.h index 0bc98195d..b2f7fef1a 100644 --- a/Eigen/src/Core/arch/AVX512/PacketMath.h +++ b/Eigen/src/Core/arch/AVX512/PacketMath.h @@ -1477,26 +1477,19 @@ EIGEN_STRONG_INLINE Packet8d pldexp(const Packet8d& a, const Packet8d& const Packet8d max_exponent = pset1(2099.0); const Packet8i e = _mm512_cvtpd_epi32(pmin(pmax(exponent, pnegate(max_exponent)), max_exponent)); - // Split 2^e into four factors and multiply. + // 4-way split + depth-3 multiply tree; see pldexp_generic for derivation. + // 2^b and 2^(e-3b) are built by widening the biased int32 exponent to int64 + // with vpmovsxdq and shifting into the double exponent field with vpsllq. const Packet8i bias = pset1(1023); - Packet8i b = parithmetic_shift_right<2>(e); // floor(e/4) + const Packet8i b = parithmetic_shift_right<2>(e); // floor(e/4) + const Packet8i b_remainder = psub(psub(e, b), padd(b, b)); // e - 3b (depth 2) + const Packet8d c1 = _mm512_castsi512_pd(_mm512_slli_epi64(_mm512_cvtepi32_epi64(padd(b, bias)), 52)); // 2^b + const Packet8d c2 = + _mm512_castsi512_pd(_mm512_slli_epi64(_mm512_cvtepi32_epi64(padd(b_remainder, bias)), 52)); // 2^(e-3b) - // 2^b - const Packet8i permute_idx = _mm256_setr_epi32(0, 4, 1, 5, 2, 6, 3, 7); - Packet8i hi = _mm256_permutevar8x32_epi32(padd(b, bias), permute_idx); - Packet8i lo = _mm256_slli_epi64(hi, 52); - hi = _mm256_slli_epi64(_mm256_srli_epi64(hi, 32), 52); - Packet8d c = _mm512_castsi512_pd(_mm512_inserti64x4(_mm512_castsi256_si512(lo), hi, 1)); - Packet8d out = pmul(pmul(pmul(a, c), c), c); // a * 2^(3b) - - // 2^(e - 3b) - b = psub(psub(psub(e, b), b), b); // e - 3b - hi = _mm256_permutevar8x32_epi32(padd(b, bias), permute_idx); - lo = _mm256_slli_epi64(hi, 52); - hi = _mm256_slli_epi64(_mm256_srli_epi64(hi, 32), 52); - c = _mm512_castsi512_pd(_mm512_inserti64x4(_mm512_castsi256_si512(lo), hi, 1)); - out = pmul(out, c); // a * 2^e - return out; + const Packet8d c1_squared = pmul(c1, c1); + const Packet8d a_c1 = pmul(a, c1); + return pmul(pmul(a_c1, c1_squared), c2); // a * 2^e } #ifdef EIGEN_VECTORIZE_AVX512DQ diff --git a/Eigen/src/Core/arch/Default/GenericPacketMathFrexpLdexp.h b/Eigen/src/Core/arch/Default/GenericPacketMathFrexpLdexp.h index cdf329103..632eca427 100644 --- a/Eigen/src/Core/arch/Default/GenericPacketMathFrexpLdexp.h +++ b/Eigen/src/Core/arch/Default/GenericPacketMathFrexpLdexp.h @@ -110,26 +110,34 @@ EIGEN_STRONG_INLINE EIGEN_DEVICE_FUNC Packet pldexp_generic(const Packet& a, con // // Set e = min(max(exponent, -278), 278); // b = floor(e/4); - // out = ((((a * 2^(b)) * 2^(b)) * 2^(b)) * 2^(e-3*b)) + // c1 = 2^b + // c2 = 2^(e - 3b) + // out = a * c1^3 * c2 (= a * 2^e) // - // This will avoid any intermediate overflows and correctly handle 0, inf, - // NaN cases. + // Re-associate as (a * c1) * (c1 * c1) * c2 so c1*c1 runs in parallel with + // a*c1; the dependent-multiply chain is 3 deep instead of 4. Multiplying + // by c1 (the downscale factor) first is required for safety: for negative + // exponents the remainder factor c2 = 2^(e-3b) can be > 1 (e.g. e=-1 -> + // b=-1 -> c2=4), so a*c2 would overflow at |a| near max even when the + // final ldexp result is finite. typedef typename unpacket_traits::integer_packet PacketI; typedef typename unpacket_traits::type Scalar; typedef typename unpacket_traits::type ScalarI; static constexpr int TotalBits = sizeof(Scalar) * CHAR_BIT, MantissaBits = numext::numeric_limits::digits - 1, ExponentBits = TotalBits - MantissaBits - 1; - const Packet max_exponent = pset1(Scalar((ScalarI(1) << ExponentBits) + ScalarI(MantissaBits - 1))); // 278 - const PacketI bias = pset1((ScalarI(1) << (ExponentBits - 1)) - ScalarI(1)); // 127 - const PacketI e = pcast(pmin(pmax(exponent, pnegate(max_exponent)), max_exponent)); - PacketI b = parithmetic_shift_right<2>(e); // floor(e/4); - Packet c = preinterpret(plogical_shift_left(padd(b, bias))); // 2^b - Packet out = pmul(pmul(pmul(a, c), c), c); // a * 2^(3b) - b = pnmadd(pset1(3), b, e); // e - 3b - c = preinterpret(plogical_shift_left(padd(b, bias))); // 2^(e-3*b) - out = pmul(out, c); - return out; + constexpr Scalar max_exp_value = Scalar((ScalarI(1) << ExponentBits) + ScalarI(MantissaBits - 1)); // 278 + const Packet max_exponent = pset1(max_exp_value); + const Packet neg_max_exponent = pset1(-max_exp_value); + const PacketI bias = pset1((ScalarI(1) << (ExponentBits - 1)) - ScalarI(1)); // 127 + const PacketI e = pcast(pmin(pmax(exponent, neg_max_exponent), max_exponent)); + const PacketI b = parithmetic_shift_right<2>(e); // floor(e/4); + const PacketI b_remainder = pnmadd(pset1(3), b, e); // e - 3b + const Packet c1 = preinterpret(plogical_shift_left(padd(b, bias))); // 2^b + const Packet c2 = preinterpret(plogical_shift_left(padd(b_remainder, bias))); // 2^(e-3*b) + const Packet c1_squared = pmul(c1, c1); + const Packet a_c1 = pmul(a, c1); + return pmul(pmul(a_c1, c1_squared), c2); } // Explicitly multiplies diff --git a/Eigen/src/Core/arch/SSE/PacketMath.h b/Eigen/src/Core/arch/SSE/PacketMath.h index 5d0d46bdb..7596e0bb3 100644 --- a/Eigen/src/Core/arch/SSE/PacketMath.h +++ b/Eigen/src/Core/arch/SSE/PacketMath.h @@ -1799,15 +1799,15 @@ EIGEN_STRONG_INLINE Packet2d pldexp(const Packet2d& a, const Packet2d& // Convert e to integer and swizzle to low-order bits. const Packet4i ei = vec4i_swizzle1(_mm_cvtpd_epi32(e), 0, 3, 1, 3); - // Split 2^e into four factors and multiply: + // 4-way split + depth-3 multiply tree; see pldexp_generic for derivation. const Packet4i bias = _mm_set_epi32(0, 1023, 0, 1023); - Packet4i b = parithmetic_shift_right<2>(ei); // floor(e/4) - Packet2d c = _mm_castsi128_pd(_mm_slli_epi64(padd(b, bias), 52)); // 2^b - Packet2d out = pmul(pmul(pmul(a, c), c), c); // a * 2^(3b) - b = psub(psub(psub(ei, b), b), b); // e - 3b - c = _mm_castsi128_pd(_mm_slli_epi64(padd(b, bias), 52)); // 2^(e - 3b) - out = pmul(out, c); // a * 2^e - return out; + const Packet4i b = parithmetic_shift_right<2>(ei); // floor(e/4) + const Packet4i b_remainder = psub(psub(ei, b), padd(b, b)); // e - 3b (depth 2) + const Packet2d c1 = _mm_castsi128_pd(_mm_slli_epi64(padd(b, bias), 52)); // 2^b + const Packet2d c2 = _mm_castsi128_pd(_mm_slli_epi64(padd(b_remainder, bias), 52)); // 2^(e - 3b) + const Packet2d c1_squared = pmul(c1, c1); + const Packet2d a_c1 = pmul(a, c1); + return pmul(pmul(a_c1, c1_squared), c2); // a * 2^e } // We specialize pldexp here, since the generic implementation uses Packet2l, which is not well diff --git a/test/packetmath.cpp b/test/packetmath.cpp index 3e4dcbbe9..7848370bb 100644 --- a/test/packetmath.cpp +++ b/test/packetmath.cpp @@ -1031,6 +1031,17 @@ void packetmath_real() { data1[0] = Scalar(std::ldexp(Scalar(1.0), NumTraits::max_exponent() - 1)); data1[PacketSize] = Scalar(+NumTraits::min_exponent() - NumTraits::max_exponent()); CHECK_CWISE2_IF(PacketTraits::HasExp, REF_LDEXP, internal::pldexp); + // Near-max magnitude with small negative exponents. Regression guard for + // the 4-way scale-factor split: the remainder factor c2 = 2^(e-3*floor(e/4)) + // is > 1 for e in {-1, -2, -5, -6, ...}, so the multiply tree must apply + // the downscale c1 before c2 -- otherwise (numext::abs(a)) * c2 spuriously + // overflows to inf for finite results like ldexp((numext::numeric_limits) + // ::max(), -1). + for (int i = 0; i < PacketSize; ++i) { + data1[i] = (numext::numeric_limits::max)(); + data1[i + PacketSize] = Scalar(-1 - (i % 8)); // -1, -2, ..., -8 + } + CHECK_CWISE2_IF(PacketTraits::HasExp, REF_LDEXP, internal::pldexp); } for (int i = 0; i < size; ++i) {