devtools/eigen
1
0
Fork 0
mirror of https://gitlab.com/libeigen/eigen.git synced 2026-04-10 11:34:33 +08:00
Code Issues Packages Projects Releases Wiki Activity

Vectorize cbrt for float and double.

Browse Source
This commit is contained in:
Rasmus Munk Larsen
2025-04-17 23:31:20 +00:00
parent 5330960900
commit 33f5f59614
12 changed files with 183 additions and 11 deletions
Download Patch File Download Diff File Expand all files Collapse all files

139
Eigen/src/Core/arch/Default/GenericPacketMathFunctions.h
View File

@@ -289,6 +289,143 @@ EIGEN_STRONG_INLINE EIGEN_DEVICE_FUNC Packet pldexp_fast(const Packet& a, const
return pmul(a, preinterpret<Packet>(plogical_shift_left<MantissaBits>(e)));
}
// This function implements a single step of Halley's iteration for
// computing x = y^(1/3):
// x_{k+1} = x_k - (x_k^3 - y) x_k / (2x_k^3 + y)
template <typename Packet>
EIGEN_DEFINE_FUNCTION_ALLOWING_MULTIPLE_DEFINITIONS Packet cbrt_halley_iteration_step(const Packet& x_k,
const Packet& y) {
typedef typename unpacket_traits<Packet>::type Scalar;
Packet x_k_cb = pmul(x_k, pmul(x_k, x_k));
Packet denom = pmadd(pset1<Packet>(Scalar(2)), x_k_cb, y);
Packet num = psub(x_k_cb, y);
Packet r = pdiv(num, denom);
return pnmadd(x_k, r, x_k);
}
// Decompose the input such that x^(1/3) = y^(1/3) * 2^e_div3, and y is in the
// interval [0.125,1].
template <typename Packet>
EIGEN_DEFINE_FUNCTION_ALLOWING_MULTIPLE_DEFINITIONS Packet cbrt_decompose(const Packet& x, Packet& e_div3) {
typedef typename unpacket_traits<Packet>::type Scalar;
// Extract the significant s in the range [0.5,1) and exponent e, such that
// x = 2^e * s.
Packet e, s;
s = pfrexp(x, e);
// Split the exponent into a part divisible by 3 and the remainder.
// e = 3*e_div3 + e_mod3.
constexpr Scalar kOneThird = Scalar(1) / 3;
e_div3 = pceil(pmul(e, pset1<Packet>(kOneThird)));
Packet e_mod3 = pnmadd(pset1<Packet>(Scalar(3)), e_div3, e);
// Replace s by y = (s * 2^e_mod3).
return pldexp_fast(s, e_mod3);
}
template <typename Packet>
EIGEN_DEFINE_FUNCTION_ALLOWING_MULTIPLE_DEFINITIONS Packet cbrt_special_cases_and_sign(const Packet& x,
const Packet& abs_root) {
typedef typename unpacket_traits<Packet>::type Scalar;
// Set sign.
const Packet sign_mask = pset1<Packet>(Scalar(-0.0));
const Packet x_sign = pand(sign_mask, x);
Packet root = por(x_sign, abs_root);
// Handle non-finite and zero values of x.
// constexpr Scalar kInf = NumTraits<Scalar>::infinity();
const Packet is_not_finite = psub(x,x);;
const Packet is_zero = pcmp_eq(pzero(x), x);
const Packet use_root = por(is_not_finite, is_zero);
return pselect(use_root, x, root);
}
// Generic implementation of cbrt(x) for float.
//
// The algorithm computes the cubic root of the input by first
// decomposing it into a exponent and significant
// x = s * 2^e.
//
// We can then write the cube root as
//
// x^(1/3) = 2^(e/3) * s^(1/3)
// = 2^((3*e_div3 + e_mod3)/3) * s^(1/3)
// = 2^(e_div3) * 2^(e_mod3/3) * s^(1/3)
// = 2^(e_div3) * (s * 2^e_mod3)^(1/3)
//
// where e_div3 = ceil(e/3) and e_mod3 = e - 3*e_div3.
//
// The cube root of the second term y = (s * 2^e_mod3)^(1/3) is coarsely
// approximated using a cubic polynomial and subsequently refined using a
// single step of Halley's iteration, and finally the two terms are combined
// using pldexp_fast.
//
// Note: Many alternatives exist for implementing cbrt. See, for example,
// the excellent discussion in Kahan's note:
// https://csclub.uwaterloo.ca/~pbarfuss/qbrt.pdf
// This particular implementation was found to be very fast and accurate
// among several alternatives tried, but is probably not "optimal" on all
// platforms.
//
// This is accurate to 2 ULP.
template <typename Packet>
EIGEN_DEFINE_FUNCTION_ALLOWING_MULTIPLE_DEFINITIONS Packet pcbrt_float(const Packet& x) {
typedef typename unpacket_traits<Packet>::type Scalar;
static_assert(std::is_same<Scalar, float>::value, "Scalar type must be float");
// Decompose the input such that x^(1/3) = y^(1/3) * 2^e_div3, and y is in the
// interval [0.125,1].
Packet e_div3;
const Packet y = cbrt_decompose(pabs(x), e_div3);
// Compute initial approximation accurate to 5.22e-3.
// The polynomial was computed using Rminimax.
constexpr float alpha[] = {5.9220016002655029296875e-01f, -1.3859539031982421875e+00f, 1.4581282138824462890625e+00f,
3.408401906490325927734375e-01f};
Packet r = ppolevl<Packet, 3>::run(y, alpha);
// Take one step of Halley's iteration.
r = cbrt_halley_iteration_step(r, y);
// Finally multiply by 2^(e_div3)
r = pldexp_fast(r, e_div3);
return cbrt_special_cases_and_sign(x, r);
}
// Generic implementation of cbrt(x) for double.
//
// The algorithm is identical to the one for float except that a different initial
// approximation is used for y^(1/3) and two Halley iteration steps are peformed.
//
// This is accurate to 1 ULP.
template <typename Packet>
EIGEN_DEFINE_FUNCTION_ALLOWING_MULTIPLE_DEFINITIONS Packet pcbrt_double(const Packet& x) {
typedef typename unpacket_traits<Packet>::type Scalar;
static_assert(std::is_same<Scalar, double>::value, "Scalar type must be double");
// Decompose the input such that x^(1/3) = y^(1/3) * 2^e_div3, and y is in the
// interval [0.125,1].
Packet e_div3;
const Packet y = cbrt_decompose(pabs(x), e_div3);
// Compute initial approximation accurate to 0.016.
// The polynomial was computed using Rminimax.
constexpr double alpha[] = {-4.69470621553356115551736138513660989701747894287109375e-01,
1.072314636518546304699839311069808900356292724609375e+00,
3.81249427609571867048288140722434036433696746826171875e-01};
Packet r = ppolevl<Packet, 2>::run(y, alpha);
// Take two steps of Halley's iteration.
r = cbrt_halley_iteration_step(r, y);
r = cbrt_halley_iteration_step(r, y);
// Finally multiply by 2^(e_div3).
r = pldexp_fast(r, e_div3);
return cbrt_special_cases_and_sign(x, r);
}
// Natural or base 2 logarithm.
// Computes log(x) as log(2^e * m) = C*e + log(m), where the constant C =log(2)
// and m is in the range [sqrt(1/2),sqrt(2)). In this range, the logarithm can
@@ -1123,7 +1260,7 @@ EIGEN_DEFINE_FUNCTION_ALLOWING_MULTIPLE_DEFINITIONS Packet generic_atan(const Pa
constexpr Scalar kPiOverTwo = static_cast<Scalar>(EIGEN_PI / 2);
const Packet cst_signmask = pset1<Packet>(-Scalar(0));
const Packet cst_signmask = pset1<Packet>(Scalar(-0.0));
const Packet cst_one = pset1<Packet>(Scalar(1));
const Packet cst_pi_over_two = pset1<Packet>(kPiOverTwo);