Changes to fast SQRT/RSQRT

This commit is contained in:
Rasmus Munk Larsen
2022-02-23 17:32:21 +00:00
committed by Antonio Sánchez
parent f9b7564faa
commit 8b875dbef1
4 changed files with 92 additions and 67 deletions

View File

@@ -84,19 +84,19 @@ struct generic_rsqrt_newton_step {
const Packet one_point_five = pset1<Packet>(Scalar(1.5));
const Packet minus_half = pset1<Packet>(Scalar(-0.5));
const Packet minus_half_a = pmul(minus_half, a);
const Packet neg_mask = pcmp_lt(a, pzero(a));
const Scalar norm_min = (std::numeric_limits<Scalar>::min)();
const Packet denorm_mask = pcmp_lt(a, pset1<Packet>(norm_min));
Packet x =
generic_rsqrt_newton_step<Packet,Steps - 1>::run(a, approx_rsqrt);
const Packet tmp = pmul(minus_half_a, x);
// If tmp is NaN, it means that a is either 0 or Inf.
// In this case return the approximation directly.
const Packet is_not_nan = pcmp_eq(tmp, tmp);
// If a is negative, return NaN.
x = por(x, neg_mask);
// In this case return the approximation directly. Do the same for
// positive subnormals. Otherwise return the Newton iterate.
const Packet return_x_newton = pandnot(pcmp_eq(tmp, tmp), denorm_mask);
// Refine the approximation using one Newton-Raphson step:
// x_{n+1} = x_n * (1.5 - x_n * ((0.5 * a) * x_n)).
const Packet x_newton = pmul(x, pmadd(tmp, x, one_point_five));
return pselect(is_not_nan, x_newton, x);
return pselect(return_x_newton, x_newton, x);
}
};
@@ -133,9 +133,11 @@ struct generic_sqrt_newton_step {
using Scalar = typename unpacket_traits<Packet>::type;
const Packet one_point_five = pset1<Packet>(Scalar(1.5));
const Packet negative_mask = pcmp_lt(a, pzero(a));
const Packet minus_half_a = pmul(a, pset1<Packet>(Scalar(-0.5)));
// Set negative arguments to NaN.
const Packet a_poisoned = por(a, negative_mask);
const Scalar norm_min = (std::numeric_limits<Scalar>::min)();
const Packet denorm_mask = pcmp_lt(a, pset1<Packet>(norm_min));
// Set negative arguments to NaN and positive subnormals to zero.
const Packet a_poisoned = por(pandnot(a, denorm_mask), negative_mask);
const Packet minus_half_a = pmul(a_poisoned, pset1<Packet>(Scalar(-0.5)));
// Do a single step of Newton's iteration for reciprocal square root:
// x_{n+1} = x_n * (1.5 - x_n * ((0.5 * a) * x_n)).
@@ -150,12 +152,10 @@ struct generic_sqrt_newton_step {
// Return sqrt(x) = x * rsqrt(x) for non-zero finite positive arguments.
// Return a itself for 0 or +inf, NaN for negative arguments.
return pselect(return_rsqrt, pmul(a_poisoned, rsqrt), por(a, negative_mask));
return pselect(return_rsqrt, pmul(a_poisoned, rsqrt), a_poisoned);
}
};
/** \internal \returns the hyperbolic tan of \a a (coeff-wise)
Doesn't do anything fancy, just a 13/6-degree rational interpolant which
is accurate up to a couple of ulps in the (approximate) range [-8, 8],

View File

@@ -89,28 +89,22 @@ pexp<Packet4d>(const Packet4d& _x) {
return pexp_double(_x);
}
#if EIGEN_FAST_MATH
template <>
EIGEN_DEFINE_FUNCTION_ALLOWING_MULTIPLE_DEFINITIONS EIGEN_UNUSED
Packet8f psqrt<Packet8f>(const Packet8f& _x) {
return generic_sqrt_newton_step<Packet8f>::run(_x, _mm256_rsqrt_ps(_x));
}
#else
// Notice that for newer processors, it is counterproductive to use Newton
// iteration for square root. In particular, Skylake and Zen2 processors
// have approximately doubled throughput of the _mm_sqrt_ps instruction
// compared to their predecessors.
template <> EIGEN_DEFINE_FUNCTION_ALLOWING_MULTIPLE_DEFINITIONS EIGEN_UNUSED
Packet8f psqrt<Packet8f>(const Packet8f& _x) {
return _mm256_sqrt_ps(_x);
}
#endif
template <> EIGEN_DEFINE_FUNCTION_ALLOWING_MULTIPLE_DEFINITIONS EIGEN_UNUSED
Packet4d psqrt<Packet4d>(const Packet4d& _x) {
return _mm256_sqrt_pd(_x);
}
// Even on Skylake, using Newton iteration is a win for reciprocal square root.
#if EIGEN_FAST_MATH
template<> EIGEN_DEFINE_FUNCTION_ALLOWING_MULTIPLE_DEFINITIONS EIGEN_UNUSED
Packet8f prsqrt<Packet8f>(const Packet8f& a) {

View File

@@ -75,29 +75,19 @@ Packet4f pcos<Packet4f>(const Packet4f& _x)
return pcos_float(_x);
}
#if EIGEN_FAST_MATH
template<>
EIGEN_DEFINE_FUNCTION_ALLOWING_MULTIPLE_DEFINITIONS EIGEN_UNUSED
Packet4f psqrt<Packet4f>(const Packet4f& _x)
{
return generic_sqrt_newton_step<Packet4f>::run(_x, _mm_rsqrt_ps(_x));
}
#else
// Notice that for newer processors, it is counterproductive to use Newton
// iteration for square root. In particular, Skylake and Zen2 processors
// have approximately doubled throughput of the _mm_sqrt_ps instruction
// compared to their predecessors.
template<>EIGEN_DEFINE_FUNCTION_ALLOWING_MULTIPLE_DEFINITIONS EIGEN_UNUSED
Packet4f psqrt<Packet4f>(const Packet4f& x) { return _mm_sqrt_ps(x); }
#endif
template<> EIGEN_DEFINE_FUNCTION_ALLOWING_MULTIPLE_DEFINITIONS EIGEN_UNUSED
Packet2d psqrt<Packet2d>(const Packet2d& x) { return _mm_sqrt_pd(x); }
template<> EIGEN_DEFINE_FUNCTION_ALLOWING_MULTIPLE_DEFINITIONS EIGEN_UNUSED
Packet16b psqrt<Packet16b>(const Packet16b& x) { return x; }
#if EIGEN_FAST_MATH
// Even on Skylake, using Newton iteration is a win for reciprocal square root.
template<> EIGEN_DEFINE_FUNCTION_ALLOWING_MULTIPLE_DEFINITIONS EIGEN_UNUSED
Packet4f prsqrt<Packet4f>(const Packet4f& x) {
return generic_rsqrt_newton_step<Packet4f, /*Steps=*/1>::run(x, _mm_rsqrt_ps(x));
@@ -105,7 +95,7 @@ Packet4f prsqrt<Packet4f>(const Packet4f& x) {
#ifdef EIGEN_VECTORIZE_FMA
// Trying to speed up reciprocal using Newton-Raphson is counterproductive
// unless FMA is available. Without FMA pdiv(pset1<Packet>(Scalar(1),a) is
// unless FMA is available. Without FMA pdiv(pset1<Packet>(Scalar(1),a)) is
// 30% faster.
template<> EIGEN_STRONG_INLINE Packet4f preciprocal<Packet4f>(const Packet4f& x) {
return generic_reciprocal_newton_step<Packet4f, /*Steps=*/1>::run(x, _mm_rcp_ps(x));