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

AVX512 and AVX2 support for Packet16i and Packet8i added

Browse Source
This commit is contained in:
Jakub Lichman
2021-08-25 19:38:23 +00:00
committed by Rasmus Munk Larsen
parent ab28419298
commit dc5b1f7d75
4 changed files with 439 additions and 22 deletions
Download Patch File Download Diff File Expand all files Collapse all files

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

@@ -196,23 +196,21 @@ struct packet_traits<bfloat16> : default_packet_traits {
HasNdtri = 1
};
};
#endif
template<> struct scalar_div_cost<float,true> { enum { value = 14 }; };
template<> struct scalar_div_cost<double,true> { enum { value = 16 }; };
/* Proper support for integers is only provided by AVX2. In the meantime, we'll
use SSE instructions and packets to deal with integers.
template<> struct packet_traits<int> : default_packet_traits
template<> struct packet_traits<int> : default_packet_traits
{
typedef Packet8i type;
typedef Packet4i half;
enum {
Vectorizable = 1,
AlignedOnScalar = 1,
size=8
};
};
*/
#endif
template<> struct scalar_div_cost<float,true> { enum { value = 14 }; };
template<> struct scalar_div_cost<double,true> { enum { value = 16 }; };
template<> struct unpacket_traits<Packet8f> {
typedef float type;
@@ -226,8 +224,16 @@ template<> struct unpacket_traits<Packet4d> {
typedef Packet2d half;
enum {size=4, alignment=Aligned32, vectorizable=true, masked_load_available=false, masked_store_available=false};
};
template<> struct unpacket_traits<Packet8i> { typedef int type; typedef Packet4i half; enum {size=8, alignment=Aligned32, vectorizable=false, masked_load_available=false, masked_store_available=false}; };
template<> struct unpacket_traits<Packet8bf> { typedef bfloat16 type; typedef Packet8bf half; enum {size=8, alignment=Aligned16, vectorizable=true, masked_load_available=false, masked_store_available=false}; };
template<> struct unpacket_traits<Packet8i> {
typedef int type;
typedef Packet4i half;
enum {size=8, alignment=Aligned32, vectorizable=true, masked_load_available=false, masked_store_available=false};
};
template<> struct unpacket_traits<Packet8bf> {
typedef bfloat16 type;
typedef Packet8bf half;
enum {size=8, alignment=Aligned16, vectorizable=true, masked_load_available=false, masked_store_available=false};
};
// Helper function for bit packing snippet of low precision comparison.
// It packs the flags from 16x16 to 8x16.
@@ -258,6 +264,7 @@ template<> EIGEN_STRONG_INLINE Packet4d pload1<Packet4d>(const double* from) { r
template<> EIGEN_STRONG_INLINE Packet8f plset<Packet8f>(const float& a) { return _mm256_add_ps(_mm256_set1_ps(a), _mm256_set_ps(7.0,6.0,5.0,4.0,3.0,2.0,1.0,0.0)); }
template<> EIGEN_STRONG_INLINE Packet4d plset<Packet4d>(const double& a) { return _mm256_add_pd(_mm256_set1_pd(a), _mm256_set_pd(3.0,2.0,1.0,0.0)); }
template<> EIGEN_STRONG_INLINE Packet8i plset<Packet8i>(const int& a) { return _mm256_add_epi32(_mm256_set1_epi32(a), _mm256_set_epi32(7,6,5,4,3,2,1,0)); }
template<> EIGEN_STRONG_INLINE Packet8f padd<Packet8f>(const Packet8f& a, const Packet8f& b) { return _mm256_add_ps(a,b); }
template<> EIGEN_STRONG_INLINE Packet4d padd<Packet4d>(const Packet4d& a, const Packet4d& b) { return _mm256_add_pd(a,b); }
@@ -291,6 +298,10 @@ template<> EIGEN_STRONG_INLINE Packet4d pnegate(const Packet4d& a)
{
return _mm256_sub_pd(_mm256_set1_pd(0.0),a);
}
template<> EIGEN_STRONG_INLINE Packet8i pnegate(const Packet8i& a)
{
return _mm256_sub_epi32(_mm256_set1_epi32(0), a);
}
template<> EIGEN_STRONG_INLINE Packet8f pconj(const Packet8f& a) { return a; }
template<> EIGEN_STRONG_INLINE Packet4d pconj(const Packet4d& a) { return a; }
@@ -352,7 +363,26 @@ template<> EIGEN_STRONG_INLINE Packet4d pcmp_lt(const Packet4d& a, const Packet4
template<> EIGEN_STRONG_INLINE Packet4d pcmp_lt_or_nan(const Packet4d& a, const Packet4d& b) { return _mm256_cmp_pd(a, b, _CMP_NGE_UQ); }
template<> EIGEN_STRONG_INLINE Packet4d pcmp_eq(const Packet4d& a, const Packet4d& b) { return _mm256_cmp_pd(a,b,_CMP_EQ_OQ); }
template<> EIGEN_STRONG_INLINE Packet8i pcmp_le(const Packet8i& a, const Packet8i& b) {
#ifdef EIGEN_VECTORIZE_AVX2
return _mm256_xor_si256(_mm256_cmpgt_epi32(a,b), _mm256_set1_epi32(-1));
#else
__m128i lo = _mm_cmpgt_epi32(_mm256_extractf128_si256(a, 0), _mm256_extractf128_si256(b, 0));
lo = _mm_xor_si128(lo, _mm_set1_epi32(-1));
__m128i hi = _mm_cmpgt_epi32(_mm256_extractf128_si256(a, 1), _mm256_extractf128_si256(b, 1));
hi = _mm_xor_si128(hi, _mm_set1_epi32(-1));
return _mm256_insertf128_si256(_mm256_castsi128_si256(lo), (hi), 1);
#endif
}
template<> EIGEN_STRONG_INLINE Packet8i pcmp_lt(const Packet8i& a, const Packet8i& b) {
#ifdef EIGEN_VECTORIZE_AVX2
return _mm256_cmpgt_epi32(b,a);
#else
__m128i lo = _mm_cmpgt_epi32(_mm256_extractf128_si256(b, 0), _mm256_extractf128_si256(a, 0));
__m128i hi = _mm_cmpgt_epi32(_mm256_extractf128_si256(b, 1), _mm256_extractf128_si256(a, 1));
return _mm256_insertf128_si256(_mm256_castsi128_si256(lo), (hi), 1);
#endif
}
template<> EIGEN_STRONG_INLINE Packet8i pcmp_eq(const Packet8i& a, const Packet8i& b) {
#ifdef EIGEN_VECTORIZE_AVX2
return _mm256_cmpeq_epi32(a,b);
@@ -388,6 +418,9 @@ template<> EIGEN_STRONG_INLINE Packet4d pmin<Packet4d>(const Packet4d& a, const
return _mm256_min_pd(b,a);
#endif
}
template<> EIGEN_STRONG_INLINE Packet8i pmin<Packet8i>(const Packet8i& a, const Packet8i& b) {
return _mm256_min_epi32(a, b);
}
template<> EIGEN_STRONG_INLINE Packet8f pmax<Packet8f>(const Packet8f& a, const Packet8f& b) {
#if EIGEN_COMP_GNUC && EIGEN_COMP_GNUC < 63
@@ -411,6 +444,9 @@ template<> EIGEN_STRONG_INLINE Packet4d pmax<Packet4d>(const Packet4d& a, const
return _mm256_max_pd(b,a);
#endif
}
template<> EIGEN_STRONG_INLINE Packet8i pmax<Packet8i>(const Packet8i& a, const Packet8i& b) {
return _mm256_max_epi32(a, b);
}
// Add specializations for min/max with prescribed NaN progation.
template<>
@@ -611,6 +647,12 @@ template<> EIGEN_STRONG_INLINE Packet4d ploaddup<Packet4d>(const double* from)
Packet4d tmp = _mm256_broadcast_pd((const __m128d*)(const void*)from);
return _mm256_permute_pd(tmp, 3<<2);
}
// Loads 4 integers from memory a returns the packet {a0, a0 a1, a1, a2, a2, a3, a3}
template<> EIGEN_STRONG_INLINE Packet8i ploaddup<Packet8i>(const int* from)
{
Packet8i a = _mm256_castsi128_si256(pload<Packet4i>(from));
return _mm256_permutevar8x32_epi32(a, _mm256_setr_epi32(0, 0, 1, 1, 2, 2, 3, 3));
}
// Loads 2 floats from memory a returns the packet {a0, a0 a0, a0, a1, a1, a1, a1}
template<> EIGEN_STRONG_INLINE Packet8f ploadquad<Packet8f>(const float* from)
@@ -618,6 +660,10 @@ template<> EIGEN_STRONG_INLINE Packet8f ploadquad<Packet8f>(const float* from)
Packet8f tmp = _mm256_castps128_ps256(_mm_broadcast_ss(from));
return _mm256_insertf128_ps(tmp, _mm_broadcast_ss(from+1), 1);
}
template<> EIGEN_STRONG_INLINE Packet8i ploadquad<Packet8i>(const int* from)
{
return _mm256_inserti128_si256(_mm256_set1_epi32(*from), _mm_set1_epi32(*(from+1)), 1);
}
template<> EIGEN_STRONG_INLINE void pstore<float>(float* to, const Packet8f& from) { EIGEN_DEBUG_ALIGNED_STORE _mm256_store_ps(to, from); }
template<> EIGEN_STRONG_INLINE void pstore<double>(double* to, const Packet4d& from) { EIGEN_DEBUG_ALIGNED_STORE _mm256_store_pd(to, from); }
@@ -646,6 +692,11 @@ template<> EIGEN_DEVICE_FUNC inline Packet4d pgather<double, Packet4d>(const dou
{
return _mm256_set_pd(from[3*stride], from[2*stride], from[1*stride], from[0*stride]);
}
template<> EIGEN_DEVICE_FUNC inline Packet8i pgather<int, Packet8i>(const int* from, Index stride)
{
return _mm256_set_epi32(from[7*stride], from[6*stride], from[5*stride], from[4*stride],
from[3*stride], from[2*stride], from[1*stride], from[0*stride]);
}
template<> EIGEN_DEVICE_FUNC inline void pscatter<float, Packet8f>(float* to, const Packet8f& from, Index stride)
{
@@ -670,6 +721,20 @@ template<> EIGEN_DEVICE_FUNC inline void pscatter<double, Packet4d>(double* to,
to[stride*2] = _mm_cvtsd_f64(high);
to[stride*3] = _mm_cvtsd_f64(_mm_shuffle_pd(high, high, 1));
}
template<> EIGEN_DEVICE_FUNC inline void pscatter<int, Packet8i>(int* to, const Packet8i& from, Index stride)
{
__m128i low = _mm256_extracti128_si256(from, 0);
to[stride*0] = _mm_extract_epi32(low, 0);
to[stride*1] = _mm_extract_epi32(low, 1);
to[stride*2] = _mm_extract_epi32(low, 2);
to[stride*3] = _mm_extract_epi32(low, 3);
__m128i high = _mm256_extracti128_si256(from, 1);
to[stride*4] = _mm_extract_epi32(high, 0);
to[stride*5] = _mm_extract_epi32(high, 1);
to[stride*6] = _mm_extract_epi32(high, 2);
to[stride*7] = _mm_extract_epi32(high, 3);
}
template<> EIGEN_STRONG_INLINE void pstore1<Packet8f>(float* to, const float& a)
{
@@ -720,6 +785,10 @@ template<> EIGEN_STRONG_INLINE Packet4d preverse(const Packet4d& a)
return _mm256_permute_pd(swap_halves,5);
#endif
}
template<> EIGEN_STRONG_INLINE Packet8i preverse(const Packet8i& a)
{
return _mm256_castps_si256(preverse(_mm256_castsi256_ps(a)));
}
// pabs should be ok
template<> EIGEN_STRONG_INLINE Packet8f pabs(const Packet8f& a)
@@ -732,6 +801,10 @@ template<> EIGEN_STRONG_INLINE Packet4d pabs(const Packet4d& a)
const Packet4d mask = _mm256_castsi256_pd(_mm256_setr_epi32(0xFFFFFFFF,0x7FFFFFFF,0xFFFFFFFF,0x7FFFFFFF,0xFFFFFFFF,0x7FFFFFFF,0xFFFFFFFF,0x7FFFFFFF));
return _mm256_and_pd(a,mask);
}
template<> EIGEN_STRONG_INLINE Packet8i pabs(const Packet8i& a)
{
return _mm256_abs_epi32(a);
}
template<> EIGEN_STRONG_INLINE Packet8f pfrexp<Packet8f>(const Packet8f& a, Packet8f& exponent) {
return pfrexp_generic(a,exponent);
@@ -803,11 +876,19 @@ template<> EIGEN_STRONG_INLINE double predux<Packet4d>(const Packet4d& a)
{
return predux(Packet2d(_mm_add_pd(_mm256_castpd256_pd128(a),_mm256_extractf128_pd(a,1))));
}
template<> EIGEN_STRONG_INLINE int predux<Packet8i>(const Packet8i& a)
{
return predux(Packet4i(_mm_add_epi32(_mm256_castsi256_si128(a),_mm256_extractf128_si256(a,1))));
}
template<> EIGEN_STRONG_INLINE Packet4f predux_half_dowto4<Packet8f>(const Packet8f& a)
{
return _mm_add_ps(_mm256_castps256_ps128(a),_mm256_extractf128_ps(a,1));
}
template<> EIGEN_STRONG_INLINE Packet4i predux_half_dowto4<Packet8i>(const Packet8i& a)
{
return _mm_add_epi32(_mm256_castsi256_si128(a),_mm256_extractf128_si256(a,1));
}
template<> EIGEN_STRONG_INLINE float predux_mul<Packet8f>(const Packet8f& a)
{
@@ -905,6 +986,55 @@ ptranspose(PacketBlock<Packet8f,4>& kernel) {
kernel.packet[3] = _mm256_permute2f128_ps(S2, S3, 0x31);
}
#define MM256_SHUFFLE_EPI32(A, B, M) \
_mm256_castps_si256(_mm256_shuffle_ps(_mm256_castsi256_ps(A), _mm256_castsi256_ps(B), M))
EIGEN_DEVICE_FUNC inline void
ptranspose(PacketBlock<Packet8i,8>& kernel) {
__m256i T0 = _mm256_unpacklo_epi32(kernel.packet[0], kernel.packet[1]);
__m256i T1 = _mm256_unpackhi_epi32(kernel.packet[0], kernel.packet[1]);
__m256i T2 = _mm256_unpacklo_epi32(kernel.packet[2], kernel.packet[3]);
__m256i T3 = _mm256_unpackhi_epi32(kernel.packet[2], kernel.packet[3]);
__m256i T4 = _mm256_unpacklo_epi32(kernel.packet[4], kernel.packet[5]);
__m256i T5 = _mm256_unpackhi_epi32(kernel.packet[4], kernel.packet[5]);
__m256i T6 = _mm256_unpacklo_epi32(kernel.packet[6], kernel.packet[7]);
__m256i T7 = _mm256_unpackhi_epi32(kernel.packet[6], kernel.packet[7]);
__m256i S0 = MM256_SHUFFLE_EPI32(T0,T2,_MM_SHUFFLE(1,0,1,0));
__m256i S1 = MM256_SHUFFLE_EPI32(T0,T2,_MM_SHUFFLE(3,2,3,2));
__m256i S2 = MM256_SHUFFLE_EPI32(T1,T3,_MM_SHUFFLE(1,0,1,0));
__m256i S3 = MM256_SHUFFLE_EPI32(T1,T3,_MM_SHUFFLE(3,2,3,2));
__m256i S4 = MM256_SHUFFLE_EPI32(T4,T6,_MM_SHUFFLE(1,0,1,0));
__m256i S5 = MM256_SHUFFLE_EPI32(T4,T6,_MM_SHUFFLE(3,2,3,2));
__m256i S6 = MM256_SHUFFLE_EPI32(T5,T7,_MM_SHUFFLE(1,0,1,0));
__m256i S7 = MM256_SHUFFLE_EPI32(T5,T7,_MM_SHUFFLE(3,2,3,2));
kernel.packet[0] = _mm256_permute2f128_si256(S0, S4, 0x20);
kernel.packet[1] = _mm256_permute2f128_si256(S1, S5, 0x20);
kernel.packet[2] = _mm256_permute2f128_si256(S2, S6, 0x20);
kernel.packet[3] = _mm256_permute2f128_si256(S3, S7, 0x20);
kernel.packet[4] = _mm256_permute2f128_si256(S0, S4, 0x31);
kernel.packet[5] = _mm256_permute2f128_si256(S1, S5, 0x31);
kernel.packet[6] = _mm256_permute2f128_si256(S2, S6, 0x31);
kernel.packet[7] = _mm256_permute2f128_si256(S3, S7, 0x31);
}
EIGEN_DEVICE_FUNC inline void
ptranspose(PacketBlock<Packet8i,4>& kernel) {
__m256i T0 = _mm256_unpacklo_epi32(kernel.packet[0], kernel.packet[1]);
__m256i T1 = _mm256_unpackhi_epi32(kernel.packet[0], kernel.packet[1]);
__m256i T2 = _mm256_unpacklo_epi32(kernel.packet[2], kernel.packet[3]);
__m256i T3 = _mm256_unpackhi_epi32(kernel.packet[2], kernel.packet[3]);
__m256i S0 = MM256_SHUFFLE_EPI32(T0,T2,_MM_SHUFFLE(1,0,1,0));
__m256i S1 = MM256_SHUFFLE_EPI32(T0,T2,_MM_SHUFFLE(3,2,3,2));
__m256i S2 = MM256_SHUFFLE_EPI32(T1,T3,_MM_SHUFFLE(1,0,1,0));
__m256i S3 = MM256_SHUFFLE_EPI32(T1,T3,_MM_SHUFFLE(3,2,3,2));
kernel.packet[0] = _mm256_permute2f128_si256(S0, S1, 0x20);
kernel.packet[1] = _mm256_permute2f128_si256(S2, S3, 0x20);
kernel.packet[2] = _mm256_permute2f128_si256(S0, S1, 0x31);
kernel.packet[3] = _mm256_permute2f128_si256(S2, S3, 0x31);
}
EIGEN_DEVICE_FUNC inline void
ptranspose(PacketBlock<Packet4d,4>& kernel) {
__m256d T0 = _mm256_shuffle_pd(kernel.packet[0], kernel.packet[1], 15);