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

Add reciprocal packet op and fast specializations for float with SSE, AVX, and AVX512.

Browse Source
This commit is contained in:
Rasmus Munk Larsen
2022-01-21 23:49:18 +00:00
parent 4b0926f99b
commit ea2c02060c
14 changed files with 95 additions and 63 deletions
Download Patch File Download Diff File Expand all files Collapse all files

21
Eigen/src/Core/GenericPacketMath.h
View File

@@ -65,6 +65,7 @@ struct default_packet_traits
HasCmp = 0,
HasDiv = 0,
HasReciprocal = 0,
HasSqrt = 0,
HasRsqrt = 0,
HasExp = 0,
@@ -816,13 +817,6 @@ Packet plog2(const Packet& a) {
template<typename Packet> EIGEN_DECLARE_FUNCTION_ALLOWING_MULTIPLE_DEFINITIONS
Packet psqrt(const Packet& a) { return numext::sqrt(a); }
/** \internal \returns the reciprocal square-root of \a a (coeff-wise) */
template<typename Packet> EIGEN_DECLARE_FUNCTION_ALLOWING_MULTIPLE_DEFINITIONS
Packet prsqrt(const Packet& a) {
typedef typename internal::unpacket_traits<Packet>::type Scalar;
return pdiv(pset1<Packet>(Scalar(1)), psqrt(a));
}
/** \internal \returns the rounded value of \a a (coeff-wise) */
template<typename Packet> EIGEN_DECLARE_FUNCTION_ALLOWING_MULTIPLE_DEFINITIONS
Packet pround(const Packet& a) { using numext::round; return round(a); }
@@ -1035,6 +1029,19 @@ pblend(const Selector<unpacket_traits<Packet>::size>& ifPacket, const Packet& th
return ifPacket.select[0] ? thenPacket : elsePacket;
}
/** \internal \returns 1 / a (coeff-wise) */
template <typename Packet>
EIGEN_DEVICE_FUNC inline Packet preciprocal(const Packet& a) {
using Scalar = typename unpacket_traits<Packet>::type;
return pdiv(pset1<Packet>(Scalar(1)), a);
}
/** \internal \returns the reciprocal square-root of \a a (coeff-wise) */
template<typename Packet> EIGEN_DECLARE_FUNCTION_ALLOWING_MULTIPLE_DEFINITIONS
Packet prsqrt(const Packet& a) {
return preciprocal<Packet>(psqrt(a));
}
} // end namespace internal
} // end namespace Eigen

29
Eigen/src/Core/MathFunctionsImpl.h
View File

@@ -17,6 +17,35 @@ namespace Eigen {
namespace internal {
/** \internal Fast reciprocal using Newton-Raphson's method.
We assume that the starting guess provided in approx_a_recip has at least
half the leading mantissa bits in the correct result, such that a single
Newton-Raphson step is sufficient to get within 1-2 ulps of the currect result.
*/
template <typename Packet, int Steps>
struct generic_reciprocal_newton_step {
static_assert(Steps > 0, "Steps must be at least 1.");
EIGEN_DEVICE_FUNC static EIGEN_STRONG_INLINE Packet
run(const Packet& a, const Packet& approx_a_recip) {
using Scalar = typename unpacket_traits<Packet>::type;
const Packet two = pset1<Packet>(Scalar(2));
const Packet neg_a = pnegate(a);
// Refine the approximation using one Newton-Raphson step:
// x_{i} = x_{i-1} * (2 - a * x_{i-1})
const Packet x =
generic_reciprocal_newton_step<Packet,Steps - 1>::run(a, approx_a_recip);
return pmul(x, pmadd(neg_a, x, two));
}
};
template<typename Packet>
struct generic_reciprocal_newton_step<Packet, 0> {
EIGEN_DEVICE_FUNC static EIGEN_STRONG_INLINE Packet
run(const Packet& /*unused*/, const Packet& approx_a_recip) {
return approx_a_recip;
}
};
/** \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],

15
Eigen/src/Core/arch/AVX/MathFunctions.h
View File

@@ -166,19 +166,12 @@ Packet8f prsqrt<Packet8f>(const Packet8f& _x) {
return pselect<Packet8f>(not_normal_finite_mask, y_approx, y_newton);
}
#else
template <> EIGEN_DEFINE_FUNCTION_ALLOWING_MULTIPLE_DEFINITIONS EIGEN_UNUSED
Packet8f prsqrt<Packet8f>(const Packet8f& _x) {
EIGEN_DECLARE_CONST_Packet8f(one, 1.0f);
return _mm256_div_ps(p8f_one, _mm256_sqrt_ps(_x));
template<> EIGEN_STRONG_INLINE Packet8f preciprocal<Packet8f>(const Packet8f& a) {
return generic_reciprocal_newton_step<Packet8f, /*Steps=*/1>::run(a, _mm256_rcp_ps(a));
}
#endif
template <> EIGEN_DEFINE_FUNCTION_ALLOWING_MULTIPLE_DEFINITIONS EIGEN_UNUSED
Packet4d prsqrt<Packet4d>(const Packet4d& _x) {
EIGEN_DECLARE_CONST_Packet4d(one, 1.0);
return _mm256_div_pd(p4d_one, _mm256_sqrt_pd(_x));
}
F16_PACKET_FUNCTION(Packet8f, Packet8h, psin)
F16_PACKET_FUNCTION(Packet8f, Packet8h, pcos)
@@ -190,6 +183,7 @@ F16_PACKET_FUNCTION(Packet8f, Packet8h, pexp)
F16_PACKET_FUNCTION(Packet8f, Packet8h, ptanh)
F16_PACKET_FUNCTION(Packet8f, Packet8h, psqrt)
F16_PACKET_FUNCTION(Packet8f, Packet8h, prsqrt)
F16_PACKET_FUNCTION(Packet8f, Packet8h, preciprocal)
template <>
EIGEN_STRONG_INLINE Packet8h pfrexp(const Packet8h& a, Packet8h& exponent) {
@@ -214,6 +208,7 @@ BF16_PACKET_FUNCTION(Packet8f, Packet8bf, pexp)
BF16_PACKET_FUNCTION(Packet8f, Packet8bf, ptanh)
BF16_PACKET_FUNCTION(Packet8f, Packet8bf, psqrt)
BF16_PACKET_FUNCTION(Packet8f, Packet8bf, prsqrt)
BF16_PACKET_FUNCTION(Packet8f, Packet8bf, preciprocal)
template <>
EIGEN_STRONG_INLINE Packet8bf pfrexp(const Packet8bf& a, Packet8bf& exponent) {

1
Eigen/src/Core/arch/AVX/PacketMath.h
View File

@@ -78,6 +78,7 @@ template<> struct packet_traits<float> : default_packet_traits
HasCmp = 1,
HasDiv = 1,
HasReciprocal = EIGEN_FAST_MATH,
HasSin = EIGEN_FAST_MATH,
HasCos = EIGEN_FAST_MATH,
HasLog = 1,

20
Eigen/src/Core/arch/AVX512/MathFunctions.h
View File

@@ -253,13 +253,6 @@ prsqrt<Packet16f>(const Packet16f& _x) {
// return rsqrt(+inf) = 0, rsqrt(x) = NaN if x < 0, and rsqrt(0) = +inf.
return _mm512_mask_blend_ps(not_finite_pos_mask, y_newton, y_approx);
}
#else
template <>
EIGEN_STRONG_INLINE Packet16f prsqrt<Packet16f>(const Packet16f& x) {
EIGEN_DECLARE_CONST_Packet16f(one, 1.0f);
return _mm512_div_ps(p16f_one, _mm512_sqrt_ps(x));
}
#endif
F16_PACKET_FUNCTION(Packet16f, Packet16h, prsqrt)
@@ -304,12 +297,17 @@ prsqrt<Packet8d>(const Packet8d& _x) {
// return rsqrt(+inf) = 0, rsqrt(x) = NaN if x < 0, and rsqrt(0) = +inf.
return _mm512_mask_blend_pd(not_finite_pos_mask, y_newton, y_approx);
}
template<> EIGEN_STRONG_INLINE Packet16f preciprocal<Packet16f>(const Packet16f& a) {
#ifdef EIGEN_VECTORIZE_AVX512ER
return _mm512_rcp28_ps(a));
#else
template <>
EIGEN_STRONG_INLINE Packet8d prsqrt<Packet8d>(const Packet8d& x) {
EIGEN_DECLARE_CONST_Packet8d(one, 1.0f);
return _mm512_div_pd(p8d_one, _mm512_sqrt_pd(x));
return generic_reciprocal_newton_step<Packet16f, /*Steps=*/1>::run(a, _mm512_rcp14_ps(a));
#endif
}
F16_PACKET_FUNCTION(Packet16f, Packet16h, preciprocal)
BF16_PACKET_FUNCTION(Packet16f, Packet16bf, preciprocal)
#endif
template<> EIGEN_DEFINE_FUNCTION_ALLOWING_MULTIPLE_DEFINITIONS EIGEN_UNUSED

1
Eigen/src/Core/arch/AVX512/PacketMath.h
View File

@@ -120,6 +120,7 @@ template<> struct packet_traits<float> : default_packet_traits
HasExp = 1,
HasSqrt = EIGEN_FAST_MATH,
HasRsqrt = EIGEN_FAST_MATH,
HasReciprocal = EIGEN_FAST_MATH,
HasTanh = EIGEN_FAST_MATH,
HasErf = EIGEN_FAST_MATH,
#endif

14
Eigen/src/Core/arch/SSE/MathFunctions.h
View File

@@ -148,21 +148,13 @@ Packet4f prsqrt<Packet4f>(const Packet4f& _x) {
return pselect<Packet4f>(not_normal_finite_mask, y_approx, y_newton);
}
#else
template<> EIGEN_DEFINE_FUNCTION_ALLOWING_MULTIPLE_DEFINITIONS EIGEN_UNUSED
Packet4f prsqrt<Packet4f>(const Packet4f& x) {
// Unfortunately we can't use the much faster mm_rsqrt_ps since it only provides an approximation.
return _mm_div_ps(pset1<Packet4f>(1.0f), _mm_sqrt_ps(x));
}
#endif
template<> EIGEN_DEFINE_FUNCTION_ALLOWING_MULTIPLE_DEFINITIONS EIGEN_UNUSED
Packet2d prsqrt<Packet2d>(const Packet2d& x) {
return _mm_div_pd(pset1<Packet2d>(1.0), _mm_sqrt_pd(x));
template<> EIGEN_STRONG_INLINE Packet4f preciprocal<Packet4f>(const Packet4f& a) {
return generic_reciprocal_newton_step<Packet4f, /*Steps=*/1>::run(a, _mm_rcp_ps(a));
}
// Hyperbolic Tangent function.
template <>
EIGEN_DEFINE_FUNCTION_ALLOWING_MULTIPLE_DEFINITIONS EIGEN_UNUSED Packet4f

1
Eigen/src/Core/arch/SSE/PacketMath.h
View File

@@ -136,6 +136,7 @@ struct packet_traits<float> : default_packet_traits {
HasCmp = 1,
HasDiv = 1,
HasReciprocal = EIGEN_FAST_MATH,
HasSin = EIGEN_FAST_MATH,
HasCos = EIGEN_FAST_MATH,
HasLog = 1,

11
Eigen/src/Core/functors/UnaryFunctors.h
View File

@@ -723,13 +723,18 @@ struct scalar_inverse_op {
EIGEN_DEVICE_FUNC inline Scalar operator() (const Scalar& a) const { return Scalar(1)/a; }
template<typename Packet>
EIGEN_DEVICE_FUNC inline const Packet packetOp(const Packet& a) const
{ return internal::pdiv(pset1<Packet>(Scalar(1)),a); }
{ return internal::preciprocal(a); }
};
template <typename Scalar>
struct functor_traits<scalar_inverse_op<Scalar> > {
enum {
PacketAccess = packet_traits<Scalar>::HasDiv,
Cost = scalar_div_cost<Scalar, PacketAccess>::value
// If packet_traits<Scalar>::HasReciprocal then the Estimated cost is that
// of computing an approximation plus a single Newton-Raphson step, which
// consists of 1 pmul + 1 pmadd.
Cost = (packet_traits<Scalar>::HasReciprocal
? 4 * NumTraits<Scalar>::MulCost
: scalar_div_cost<Scalar, PacketAccess>::value)
};
};
@@ -1033,7 +1038,7 @@ struct scalar_logistic_op {
}
};
// TODO(rmlarsen): Enable the following on host when integer_packet is defined
// TODO(rmlarsen): Enable the following on host when integer_packet is defined
// for the relevant packet types.
#ifdef EIGEN_GPU_CC

26
Eigen/src/LU/arch/InverseSize4.h
View File

@@ -58,10 +58,10 @@ struct compute_inverse_size4<Architecture::Target, float, MatrixType, ResultType
const float* data = matrix.data();
const Index stride = matrix.innerStride();
Packet4f L1_ = ploadt<Packet4f,MatrixAlignment>(data);
Packet4f L2_ = ploadt<Packet4f,MatrixAlignment>(data + stride*4);
Packet4f L3_ = ploadt<Packet4f,MatrixAlignment>(data + stride*8);
Packet4f L4_ = ploadt<Packet4f,MatrixAlignment>(data + stride*12);
Packet4f L1 = ploadt<Packet4f,MatrixAlignment>(data);
Packet4f L2 = ploadt<Packet4f,MatrixAlignment>(data + stride*4);
Packet4f L3 = ploadt<Packet4f,MatrixAlignment>(data + stride*8);
Packet4f L4 = ploadt<Packet4f,MatrixAlignment>(data + stride*12);
// Four 2x2 sub-matrices of the input matrix
// input = [[A, B],
@@ -70,17 +70,17 @@ struct compute_inverse_size4<Architecture::Target, float, MatrixType, ResultType
if (!StorageOrdersMatch)
{
A = vec4f_unpacklo(L1_, L2_);
B = vec4f_unpacklo(L3_, L4_);
C = vec4f_unpackhi(L1_, L2_);
D = vec4f_unpackhi(L3_, L4_);
A = vec4f_unpacklo(L1, L2);
B = vec4f_unpacklo(L3, L4);
C = vec4f_unpackhi(L1, L2);
D = vec4f_unpackhi(L3, L4);
}
else
{
A = vec4f_movelh(L1_, L2_);
B = vec4f_movehl(L2_, L1_);
C = vec4f_movelh(L3_, L4_);
D = vec4f_movehl(L4_, L3_);
A = vec4f_movelh(L1, L2);
B = vec4f_movehl(L2, L1);
C = vec4f_movelh(L3, L4);
D = vec4f_movehl(L4, L3);
}
Packet4f AB, DC;
@@ -120,7 +120,7 @@ struct compute_inverse_size4<Architecture::Target, float, MatrixType, ResultType
Packet4f det = vec4f_duplane(psub(padd(d1, d2), d), 0);
// reciprocal of the determinant of the input matrix, rd = 1/det
Packet4f rd = pdiv(pset1<Packet4f>(1.0f), det);
Packet4f rd = preciprocal(det);
// Four sub-matrices of the inverse
Packet4f iA, iB, iC, iD;