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 tensor.isnan() by using typed predicates.

Browse Source
This commit is contained in:
Rasmus Munk Larsen
2023-03-16 04:04:22 +00:00
parent f02856c640
commit 0488b708b4
8 changed files with 110 additions and 25 deletions
Download Patch File Download Diff File Expand all files Collapse all files

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

@@ -634,6 +634,7 @@ template<> EIGEN_STRONG_INLINE Packet8f pcmp_le(const Packet8f& a, const Packet8
template<> EIGEN_STRONG_INLINE Packet8f pcmp_lt(const Packet8f& a, const Packet8f& b) { return _mm256_cmp_ps(a,b,_CMP_LT_OQ); }
template<> EIGEN_STRONG_INLINE Packet8f pcmp_lt_or_nan(const Packet8f& a, const Packet8f& b) { return _mm256_cmp_ps(a, b, _CMP_NGE_UQ); }
template<> EIGEN_STRONG_INLINE Packet8f pcmp_eq(const Packet8f& a, const Packet8f& b) { return _mm256_cmp_ps(a,b,_CMP_EQ_OQ); }
template<> EIGEN_STRONG_INLINE Packet8f pisnan(const Packet8f& a) { return _mm256_cmp_ps(a,a,_CMP_UNORD_Q); }
template<> EIGEN_STRONG_INLINE Packet4d pcmp_le(const Packet4d& a, const Packet4d& b) { return _mm256_cmp_pd(a,b,_CMP_LE_OQ); }
template<> EIGEN_STRONG_INLINE Packet4d pcmp_lt(const Packet4d& a, const Packet4d& b) { return _mm256_cmp_pd(a,b,_CMP_LT_OQ); }

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

@@ -353,7 +353,7 @@ EIGEN_STRONG_INLINE Packet8d pnegate(const Packet8d& a) {
}
template <>
EIGEN_STRONG_INLINE Packet16i pnegate(const Packet16i& a) {
return _mm512_sub_epi32(_mm512_set1_epi32(0), a);
return _mm512_sub_epi32(_mm512_setzero_si512(), a);
}
template <>
@@ -580,66 +580,72 @@ EIGEN_STRONG_INLINE __m256i Pack32To16(Packet16f rf) {
return _mm256_insertf128_si256(_mm256_castsi128_si256(result_lo), result_hi, 1);
}
template <>
EIGEN_STRONG_INLINE Packet16f pisnan(const Packet16f& a) {
__mmask16 mask = _mm512_cmp_ps_mask(a, a, _CMP_UNORD_Q);
return _mm512_castsi512_ps(_mm512_maskz_set1_epi32(mask, 0xffffffffu));
}
template <>
EIGEN_STRONG_INLINE Packet16f pcmp_eq(const Packet16f& a, const Packet16f& b) {
__mmask16 mask = _mm512_cmp_ps_mask(a, b, _CMP_EQ_OQ);
return _mm512_castsi512_ps(
_mm512_mask_set1_epi32(_mm512_set1_epi32(0), mask, 0xffffffffu));
_mm512_mask_set1_epi32(_mm512_setzero_epi32(), mask, 0xffffffffu));
}
template<> EIGEN_STRONG_INLINE Packet16f pcmp_le(const Packet16f& a, const Packet16f& b) {
__mmask16 mask = _mm512_cmp_ps_mask(a, b, _CMP_LE_OQ);
return _mm512_castsi512_ps(
_mm512_mask_set1_epi32(_mm512_set1_epi32(0), mask, 0xffffffffu));
_mm512_mask_set1_epi32(_mm512_setzero_epi32(), mask, 0xffffffffu));
}
template<> EIGEN_STRONG_INLINE Packet16f pcmp_lt(const Packet16f& a, const Packet16f& b) {
__mmask16 mask = _mm512_cmp_ps_mask(a, b, _CMP_LT_OQ);
return _mm512_castsi512_ps(
_mm512_mask_set1_epi32(_mm512_set1_epi32(0), mask, 0xffffffffu));
_mm512_mask_set1_epi32(_mm512_setzero_epi32(), mask, 0xffffffffu));
}
template<> EIGEN_STRONG_INLINE Packet16f pcmp_lt_or_nan(const Packet16f& a, const Packet16f& b) {
__mmask16 mask = _mm512_cmp_ps_mask(a, b, _CMP_NGE_UQ);
return _mm512_castsi512_ps(
_mm512_mask_set1_epi32(_mm512_set1_epi32(0), mask, 0xffffffffu));
_mm512_mask_set1_epi32(_mm512_setzero_epi32(), mask, 0xffffffffu));
}
template<> EIGEN_STRONG_INLINE Packet16i pcmp_eq(const Packet16i& a, const Packet16i& b) {
__mmask16 mask = _mm512_cmp_epi32_mask(a, b, _MM_CMPINT_EQ);
return _mm512_mask_set1_epi32(_mm512_set1_epi32(0), mask, 0xffffffffu);
return _mm512_mask_set1_epi32(_mm512_setzero_epi32(), mask, 0xffffffffu);
}
template<> EIGEN_STRONG_INLINE Packet16i pcmp_le(const Packet16i& a, const Packet16i& b) {
__mmask16 mask = _mm512_cmp_epi32_mask(a, b, _MM_CMPINT_LE);
return _mm512_mask_set1_epi32(_mm512_set1_epi32(0), mask, 0xffffffffu);
return _mm512_mask_set1_epi32(_mm512_setzero_epi32(), mask, 0xffffffffu);
}
template<> EIGEN_STRONG_INLINE Packet16i pcmp_lt(const Packet16i& a, const Packet16i& b) {
__mmask16 mask = _mm512_cmp_epi32_mask(a, b, _MM_CMPINT_LT);
return _mm512_mask_set1_epi32(_mm512_set1_epi32(0), mask, 0xffffffffu);
return _mm512_mask_set1_epi32(_mm512_setzero_epi32(), mask, 0xffffffffu);
}
template <>
EIGEN_STRONG_INLINE Packet8d pcmp_eq(const Packet8d& a, const Packet8d& b) {
__mmask8 mask = _mm512_cmp_pd_mask(a, b, _CMP_EQ_OQ);
return _mm512_castsi512_pd(
_mm512_mask_set1_epi64(_mm512_set1_epi64(0), mask, 0xffffffffffffffffu));
_mm512_mask_set1_epi64(_mm512_setzero_epi32(), mask, 0xffffffffffffffffu));
}
template <>
EIGEN_STRONG_INLINE Packet8d pcmp_le(const Packet8d& a, const Packet8d& b) {
__mmask8 mask = _mm512_cmp_pd_mask(a, b, _CMP_LE_OQ);
return _mm512_castsi512_pd(
_mm512_mask_set1_epi64(_mm512_set1_epi64(0), mask, 0xffffffffffffffffu));
_mm512_mask_set1_epi64(_mm512_setzero_epi32(), mask, 0xffffffffffffffffu));
}
template <>
EIGEN_STRONG_INLINE Packet8d pcmp_lt(const Packet8d& a, const Packet8d& b) {
__mmask8 mask = _mm512_cmp_pd_mask(a, b, _CMP_LT_OQ);
return _mm512_castsi512_pd(
_mm512_mask_set1_epi64(_mm512_set1_epi64(0), mask, 0xffffffffffffffffu));
_mm512_mask_set1_epi64(_mm512_setzero_epi32(), mask, 0xffffffffffffffffu));
}
template <>
EIGEN_STRONG_INLINE Packet8d pcmp_lt_or_nan(const Packet8d& a, const Packet8d& b) {
__mmask8 mask = _mm512_cmp_pd_mask(a, b, _CMP_NGE_UQ);
return _mm512_castsi512_pd(
_mm512_mask_set1_epi64(_mm512_set1_epi64(0), mask, 0xffffffffffffffffu));
_mm512_mask_set1_epi64(_mm512_setzero_epi32(), mask, 0xffffffffffffffffu));
}
template<> EIGEN_STRONG_INLINE Packet16f print<Packet16f>(const Packet16f& a) { return _mm512_roundscale_ps(a, _MM_FROUND_CUR_DIRECTION); }

27
Eigen/src/Core/arch/AVX512/TypeCasting.h
View File

@@ -16,6 +16,33 @@ namespace Eigen {
namespace internal {
template <>
struct type_casting_traits<float, bool> {
enum {
VectorizedCast = 1,
SrcCoeffRatio = 1,
TgtCoeffRatio = 1
};
};
template <>
struct type_casting_traits<bool,float> {
enum {
VectorizedCast = 1,
SrcCoeffRatio = 1,
TgtCoeffRatio = 1
};
};
template<> EIGEN_STRONG_INLINE Packet16b pcast<Packet16f, Packet16b>(const Packet16f& a) {
__mmask16 mask = _mm512_cmpneq_ps_mask(a, pzero(a));
return _mm512_maskz_cvtepi32_epi8(mask, _mm512_set1_epi32(1));
}
template<> EIGEN_STRONG_INLINE Packet16f pcast<Packet16b, Packet16f>(const Packet16b& a) {
return _mm512_cvtepi32_ps(_mm512_cvtepi8_epi32(a));
}
template<> EIGEN_STRONG_INLINE Packet16i pcast<Packet16f, Packet16i>(const Packet16f& a) {
return _mm512_cvttps_epi32(a);
}

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

@@ -1124,7 +1124,7 @@ Packet psqrt_complex(const Packet& a) {
Packet imag_inf_result;
imag_inf_result.v = por(pand(cst_pos_inf, real_mask), pandnot(a.v, real_mask));
// unless otherwise specified, if either the real or imaginary component is nan, the entire result is nan
Packet result_is_nan = pandnot(ptrue(result), pcmp_eq(result, result));
Packet result_is_nan = pisnan(result);
result = por(result_is_nan, result);
return pselect(is_imag_inf, imag_inf_result, pselect(is_real_inf, real_inf_result, result));
@@ -1796,7 +1796,7 @@ EIGEN_DEFINE_FUNCTION_ALLOWING_MULTIPLE_DEFINITIONS Packet generic_pow(const Pac
const Packet abs_x_is_lt_one = pcmp_lt(abs_x, cst_one);
const Packet x_is_one = pandnot(abs_x_is_one, x_is_neg);
const Packet x_is_neg_one = pand(abs_x_is_one, x_is_neg);
const Packet x_is_nan = pandnot(ptrue(x), pcmp_eq(x, x));
const Packet x_is_nan = pisnan(x);
// Predicates for sign and magnitude of y.
const Packet abs_y = pabs(y);
@@ -1804,7 +1804,7 @@ EIGEN_DEFINE_FUNCTION_ALLOWING_MULTIPLE_DEFINITIONS Packet generic_pow(const Pac
const Packet abs_y_is_zero = pcmp_eq(abs_y, cst_zero);
const Packet y_is_neg = pcmp_lt(y, cst_zero);
const Packet y_is_pos = pandnot(ptrue(y), por(abs_y_is_zero, y_is_neg));
const Packet y_is_nan = pandnot(ptrue(y), pcmp_eq(y, y));
const Packet y_is_nan = pisnan(y);
const Packet abs_y_is_inf = pcmp_eq(abs_y, cst_pos_inf);
EIGEN_CONSTEXPR Scalar huge_exponent =
(NumTraits<Scalar>::max_exponent() * Scalar(EIGEN_LN2)) / NumTraits<Scalar>::epsilon();