Add generic fast psqrt and prsqrt impls and make them correct for 0, +Inf, NaN, and negative arguments.

This commit is contained in:
Rasmus Munk Larsen
2022-02-05 00:20:13 +00:00
parent 4bffbe84f9
commit 979fdd58a4
5 changed files with 135 additions and 169 deletions

View File

@@ -75,30 +75,12 @@ Packet4f pcos<Packet4f>(const Packet4f& _x)
return pcos_float(_x);
}
// Functions for sqrt.
// The EIGEN_FAST_MATH version uses the _mm_rsqrt_ps approximation and one step
// of Newton's method, at a cost of 1-2 bits of precision as opposed to the
// exact solution. It does not handle +inf, or denormalized numbers correctly.
// The main advantage of this approach is not just speed, but also the fact that
// it can be inlined and pipelined with other computations, further reducing its
// effective latency. This is similar to Quake3's fast inverse square root.
// For detail see here: http://www.beyond3d.com/content/articles/8/
#if EIGEN_FAST_MATH
template<>
EIGEN_DEFINE_FUNCTION_ALLOWING_MULTIPLE_DEFINITIONS EIGEN_UNUSED
Packet4f psqrt<Packet4f>(const Packet4f& _x)
{
const Packet4f minus_half_x = pmul(_x, pset1<Packet4f>(-0.5f));
const Packet4f denormal_mask = pandnot(
pcmp_lt(_x, pset1<Packet4f>((std::numeric_limits<float>::min)())),
pcmp_lt(_x, pzero(_x)));
// Compute approximate reciprocal sqrt.
Packet4f x = _mm_rsqrt_ps(_x);
// Do a single step of Newton's iteration.
x = pmul(x, pmadd(minus_half_x, pmul(x,x), pset1<Packet4f>(1.5f)));
// Flush results for denormals to zero.
return pandnot(pmul(_x,x), denormal_mask);
return generic_sqrt_newton_step<Packet4f>::run(_x, _mm_rsqrt_ps(_x));
}
#else
@@ -117,43 +99,16 @@ Packet16b psqrt<Packet16b>(const Packet16b& x) { return x; }
#if EIGEN_FAST_MATH
template<> EIGEN_DEFINE_FUNCTION_ALLOWING_MULTIPLE_DEFINITIONS EIGEN_UNUSED
Packet4f prsqrt<Packet4f>(const Packet4f& _x) {
EIGEN_DECLARE_CONST_Packet4f(one_point_five, 1.5f);
EIGEN_DECLARE_CONST_Packet4f(minus_half, -0.5f);
EIGEN_DECLARE_CONST_Packet4f_FROM_INT(inf, 0x7f800000u);
EIGEN_DECLARE_CONST_Packet4f_FROM_INT(flt_min, 0x00800000u);
Packet4f neg_half = pmul(_x, p4f_minus_half);
// Identity infinite, zero, negative and denormal arguments.
Packet4f lt_min_mask = _mm_cmplt_ps(_x, p4f_flt_min);
Packet4f inf_mask = _mm_cmpeq_ps(_x, p4f_inf);
Packet4f not_normal_finite_mask = _mm_or_ps(lt_min_mask, inf_mask);
// Compute an approximate result using the rsqrt intrinsic.
Packet4f y_approx = _mm_rsqrt_ps(_x);
// Do a single step of Newton-Raphson iteration to improve the approximation.
// This uses the formula y_{n+1} = y_n * (1.5 - y_n * (0.5 * x) * y_n).
// It is essential to evaluate the inner term like this because forming
// y_n^2 may over- or underflow.
Packet4f y_newton = pmul(
y_approx, pmadd(y_approx, pmul(neg_half, y_approx), p4f_one_point_five));
// Select the result of the Newton-Raphson step for positive normal arguments.
// For other arguments, choose the output of the intrinsic. This will
// return rsqrt(+inf) = 0, rsqrt(x) = NaN if x < 0, and rsqrt(x) = +inf if
// x is zero or a positive denormalized float (equivalent to flushing positive
// denormalized inputs to zero).
return pselect<Packet4f>(not_normal_finite_mask, y_approx, y_newton);
Packet4f prsqrt<Packet4f>(const Packet4f& x) {
return generic_rsqrt_newton_step<Packet4f, /*Steps=*/1>::run(x, _mm_rsqrt_ps(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
// 30% faster.
template<> EIGEN_STRONG_INLINE Packet4f preciprocal<Packet4f>(const Packet4f& a) {
return generic_reciprocal_newton_step<Packet4f, /*Steps=*/1>::run(a, _mm_rcp_ps(a));
template<> EIGEN_STRONG_INLINE Packet4f preciprocal<Packet4f>(const Packet4f& x) {
return generic_reciprocal_newton_step<Packet4f, /*Steps=*/1>::run(x, _mm_rcp_ps(x));
}
#endif