Add missing x86 pcasts

2026-04-10 11:34:33 +08:00 · 2023-07-28 23:41:38 +00:00
parent 24d15e086f
commit 5527e78a64
5 changed files with 177 additions and 238 deletions
--- a/Eigen/src/Core/GenericPacketMath.h
+++ b/Eigen/src/Core/GenericPacketMath.h
@@ -206,6 +206,17 @@ struct type_casting_traits {
  };
 };

+// provides a succint template to define vectorized casting traits with respect to the largest accessible packet types
+template <typename Src, typename Tgt>
+struct vectorized_type_casting_traits {
+  enum : int {
+    DefaultSrcPacketSize = packet_traits<Src>::size,
+    DefaultTgtPacketSize = packet_traits<Tgt>::size,
+    VectorizedCast = 1,
+    SrcCoeffRatio = plain_enum_max(DefaultTgtPacketSize / DefaultSrcPacketSize, 1),
+    TgtCoeffRatio = plain_enum_max(DefaultSrcPacketSize / DefaultTgtPacketSize, 1)
+  };
+};

 /** \internal Wrapper to ensure that multiple packet types can map to the same
    same underlying vector type. */
--- a/Eigen/src/Core/arch/AVX/TypeCasting.h
+++ b/Eigen/src/Core/arch/AVX/TypeCasting.h
@@ -17,76 +17,24 @@ namespace Eigen {
 namespace internal {

 #ifndef EIGEN_VECTORIZE_AVX512
-template <>
-struct type_casting_traits<Eigen::half, float> {
-  enum {
-    VectorizedCast = 1,
-    SrcCoeffRatio = 1,
-    TgtCoeffRatio = 1
-  };
-};
+template<> struct type_casting_traits<float, bool> : vectorized_type_casting_traits<float, bool> {};
+template<> struct type_casting_traits<bool, float> : vectorized_type_casting_traits<bool, float> {};

+template<> struct type_casting_traits<float, int> : vectorized_type_casting_traits<float, int> {};
+template<> struct type_casting_traits<int, float> : vectorized_type_casting_traits<int, float> {};

-template <>
-struct type_casting_traits<float, Eigen::half> {
-  enum {
-    VectorizedCast = 1,
-    SrcCoeffRatio = 1,
-    TgtCoeffRatio = 1
-  };
-};
+template<> struct type_casting_traits<float, double> : vectorized_type_casting_traits<float, double> {};
+template<> struct type_casting_traits<double, float> : vectorized_type_casting_traits<double, float> {};

-template <>
-struct type_casting_traits<bfloat16, float> {
-  enum {
-    VectorizedCast = 1,
-    SrcCoeffRatio = 1,
-    TgtCoeffRatio = 1
-  };
-};
+template<> struct type_casting_traits<double, int> : vectorized_type_casting_traits<double, int> {};
+template<> struct type_casting_traits<int, double> : vectorized_type_casting_traits<int, double> {};

-template <>
-struct type_casting_traits<float, bfloat16> {
-  enum {
-    VectorizedCast = 1,
-    SrcCoeffRatio = 1,
-    TgtCoeffRatio = 1
-  };
-};
+template<> struct type_casting_traits<half, float> : vectorized_type_casting_traits<half, float> {};
+template<> struct type_casting_traits<float, half> : vectorized_type_casting_traits<float, half> {};

-template <>
-struct type_casting_traits<float, bool> {
-  enum {
-    VectorizedCast = 1,
-    SrcCoeffRatio = 2,
-    TgtCoeffRatio = 1
-  };
-};
-#endif  // EIGEN_VECTORIZE_AVX512
-
-template<> EIGEN_STRONG_INLINE Packet8i pcast<Packet8f, Packet8i>(const Packet8f& a) {
-  return _mm256_cvttps_epi32(a);
-}
-
-template<> EIGEN_STRONG_INLINE Packet8f pcast<Packet8i, Packet8f>(const Packet8i& a) {
-  return _mm256_cvtepi32_ps(a);
-}
-
-template<> EIGEN_STRONG_INLINE Packet8f pcast<Packet4d, Packet8f>(const Packet4d& a, const Packet4d& b) {
-  return _mm256_set_m128(_mm256_cvtpd_ps(b), _mm256_cvtpd_ps(a));
-}
-
-template<> EIGEN_STRONG_INLINE Packet8i pcast<Packet4d, Packet8i>(const Packet4d& a, const Packet4d& b) {
-  return _mm256_set_m128i(_mm256_cvttpd_epi32(b), _mm256_cvttpd_epi32(a));
-}
-
-template <> EIGEN_STRONG_INLINE Packet4f pcast<Packet4d, Packet4f>(const Packet4d& a) {
-  return _mm256_cvtpd_ps(a);
-}
-
-template <> EIGEN_STRONG_INLINE Packet4i pcast<Packet4d, Packet4i>(const Packet4d& a) {
-  return _mm256_cvttpd_epi32(a);
-}
+template<> struct type_casting_traits<bfloat16, float> : vectorized_type_casting_traits<bfloat16, float> {};
+template<> struct type_casting_traits<float, bfloat16> : vectorized_type_casting_traits<float, bfloat16> {};
+#endif

 template <>
 EIGEN_STRONG_INLINE Packet16b pcast<Packet8f, Packet16b>(const Packet8f& a,
@@ -118,6 +66,63 @@ EIGEN_STRONG_INLINE Packet16b pcast<Packet8f, Packet16b>(const Packet8f& a,
 #endif
 }

+template <>
+EIGEN_STRONG_INLINE Packet8f pcast<Packet16b, Packet8f>(const Packet16b& a) {
+  const __m256 cst_one = _mm256_set1_ps(1.0f);
+  #ifdef EIGEN_VECTORIZE_AVX2
+  __m256i a_extended = _mm256_cvtepi8_epi32(a);
+  __m256i abcd_efgh = _mm256_cmpeq_epi32(a_extended, _mm256_setzero_si256());
+  #else
+  __m128i abcd_efhg_ijkl_mnop = _mm_cmpeq_epi8(a, _mm_setzero_si128());
+  __m128i aabb_ccdd_eeff_gghh = _mm_unpacklo_epi8(abcd_efhg_ijkl_mnop, abcd_efhg_ijkl_mnop);
+  __m128i aaaa_bbbb_cccc_dddd = _mm_unpacklo_epi8(aabb_ccdd_eeff_gghh, aabb_ccdd_eeff_gghh);
+  __m128i eeee_ffff_gggg_hhhh = _mm_unpackhi_epi8(aabb_ccdd_eeff_gghh, aabb_ccdd_eeff_gghh);
+  __m256i abcd_efgh = _mm256_setr_m128i(aaaa_bbbb_cccc_dddd, eeee_ffff_gggg_hhhh);
+  #endif
+  __m256 result = _mm256_andnot_ps(_mm256_castsi256_ps(abcd_efgh), cst_one);
+  return result;
+}
+
+template<> EIGEN_STRONG_INLINE Packet8i pcast<Packet8f, Packet8i>(const Packet8f& a) {
+  return _mm256_cvttps_epi32(a);
+}
+
+template<> EIGEN_STRONG_INLINE Packet8i pcast<Packet4d, Packet8i>(const Packet4d& a, const Packet4d& b) {
+  return _mm256_set_m128i(_mm256_cvttpd_epi32(b), _mm256_cvttpd_epi32(a));
+}
+
+template <> EIGEN_STRONG_INLINE Packet4i pcast<Packet4d, Packet4i>(const Packet4d& a) {
+  return _mm256_cvttpd_epi32(a);
+}
+
+template<> EIGEN_STRONG_INLINE Packet8f pcast<Packet8i, Packet8f>(const Packet8i& a) {
+  return _mm256_cvtepi32_ps(a);
+}
+
+template<> EIGEN_STRONG_INLINE Packet8f pcast<Packet4d, Packet8f>(const Packet4d& a, const Packet4d& b) {
+  return _mm256_set_m128(_mm256_cvtpd_ps(b), _mm256_cvtpd_ps(a));
+}
+
+template <> EIGEN_STRONG_INLINE Packet4f pcast<Packet4d, Packet4f>(const Packet4d& a) {
+  return _mm256_cvtpd_ps(a);
+}
+
+template <> EIGEN_STRONG_INLINE Packet4d pcast<Packet8i, Packet4d>(const Packet8i& a) {
+  return _mm256_cvtepi32_pd(_mm256_castsi256_si128(a));
+}
+
+template <> EIGEN_STRONG_INLINE Packet4d pcast<Packet4i, Packet4d>(const Packet4i& a) {
+  return _mm256_cvtepi32_pd(a);
+}
+
+template <> EIGEN_STRONG_INLINE Packet4d pcast<Packet8f, Packet4d>(const Packet8f& a) {
+  return _mm256_cvtps_pd(_mm256_castps256_ps128(a));
+}
+
+template <> EIGEN_STRONG_INLINE Packet4d pcast<Packet4f, Packet4d>(const Packet4f& a) {
+  return _mm256_cvtps_pd(a);
+}
+
 template<> EIGEN_STRONG_INLINE Packet8i preinterpret<Packet8i,Packet8f>(const Packet8f& a) {
  return _mm256_castps_si256(a);
 }
--- a/Eigen/src/Core/arch/AVX512/TypeCasting.h
+++ b/Eigen/src/Core/arch/AVX512/TypeCasting.h
@@ -16,23 +16,23 @@ namespace Eigen {

 namespace internal {

-template <>
-struct type_casting_traits<float, bool> {
-  enum {
-    VectorizedCast = 1,
-    SrcCoeffRatio = 1,
-    TgtCoeffRatio = 1
-  };
-};
+template<> struct type_casting_traits<float, bool> : vectorized_type_casting_traits<float, bool> {};
+template<> struct type_casting_traits<bool, float> : vectorized_type_casting_traits<bool, float> {};

-template <>
-struct type_casting_traits<bool, float> {
-  enum {
-    VectorizedCast = 1,
-    SrcCoeffRatio = 1,
-    TgtCoeffRatio = 1
-  };
-};
+template<> struct type_casting_traits<float, int> : vectorized_type_casting_traits<float, int> {};
+template<> struct type_casting_traits<int, float> : vectorized_type_casting_traits<int, float> {};
+
+template<> struct type_casting_traits<float, double> : vectorized_type_casting_traits<float, double> {};
+template<> struct type_casting_traits<double, float> : vectorized_type_casting_traits<double, float> {};
+
+template<> struct type_casting_traits<double, int> : vectorized_type_casting_traits<double, int> {};
+template<> struct type_casting_traits<int, double> : vectorized_type_casting_traits<int, double> {};
+
+template<> struct type_casting_traits<half, float> : vectorized_type_casting_traits<half, float> {};
+template<> struct type_casting_traits<float, half> : vectorized_type_casting_traits<float, half> {};
+
+template<> struct type_casting_traits<bfloat16, float> : vectorized_type_casting_traits<bfloat16, float> {};
+template<> struct type_casting_traits<float, bfloat16> : vectorized_type_casting_traits<float, bfloat16> {};

 template<> EIGEN_STRONG_INLINE Packet16b pcast<Packet16f, Packet16b>(const Packet16f& a) {
  __mmask16 mask = _mm512_cmpneq_ps_mask(a, pzero(a));
@@ -47,10 +47,26 @@ template<> EIGEN_STRONG_INLINE Packet16i pcast<Packet16f, Packet16i>(const Packe
  return _mm512_cvttps_epi32(a);
 }

+template<> EIGEN_STRONG_INLINE Packet8d pcast<Packet16f, Packet8d>(const Packet16f& a) {
+  return _mm512_cvtps_pd(_mm512_castps512_ps256(a));
+}
+
+template<> EIGEN_STRONG_INLINE Packet8d pcast<Packet8f, Packet8d>(const Packet8f& a) {
+  return _mm512_cvtps_pd(a);
+}
+
 template<> EIGEN_STRONG_INLINE Packet16f pcast<Packet16i, Packet16f>(const Packet16i& a) {
  return _mm512_cvtepi32_ps(a);
 }

+template<> EIGEN_STRONG_INLINE Packet8d pcast<Packet16i, Packet8d>(const Packet16i& a) {
+  return _mm512_cvtepi32_pd(_mm512_castsi512_si256(a));
+}
+
+template<> EIGEN_STRONG_INLINE Packet8d pcast<Packet8i, Packet8d>(const Packet8i& a) {
+  return _mm512_cvtepi32_pd(a);
+}
+
 template<> EIGEN_STRONG_INLINE Packet16f pcast<Packet8d, Packet16f>(const Packet8d& a, const Packet8d& b) {
  return  cat256(_mm512_cvtpd_ps(a), _mm512_cvtpd_ps(b));
 }
@@ -131,80 +147,26 @@ template<> EIGEN_STRONG_INLINE Packet8bf preinterpret<Packet8bf, Packet16bf>(con

 #ifndef EIGEN_VECTORIZE_AVX512FP16

-template <>
-struct type_casting_traits<half, float> {
-  enum {
-    VectorizedCast = 1,
-    SrcCoeffRatio = 1,
-    TgtCoeffRatio = 1
-  };
-};
-
 template<> EIGEN_STRONG_INLINE Packet16f pcast<Packet16h, Packet16f>(const Packet16h& a) {
  return half2float(a);
 }

-template <>
-struct type_casting_traits<float, half> {
-  enum {
-    VectorizedCast = 1,
-    SrcCoeffRatio = 1,
-    TgtCoeffRatio = 1
-  };
-};
-
 template<> EIGEN_STRONG_INLINE Packet16h pcast<Packet16f, Packet16h>(const Packet16f& a) {
  return float2half(a);
 }

 #endif

-template <>
-struct type_casting_traits<bfloat16, float> {
-  enum {
-    VectorizedCast = 1,
-    SrcCoeffRatio = 1,
-    TgtCoeffRatio = 1
-  };
-};
-
 template<> EIGEN_STRONG_INLINE Packet16f pcast<Packet16bf, Packet16f>(const Packet16bf& a) {
  return Bf16ToF32(a);
 }

-template <>
-struct type_casting_traits<float, bfloat16> {
-  enum {
-    VectorizedCast = 1,
-    SrcCoeffRatio = 1,
-    TgtCoeffRatio = 1
-  };
-};
-
 template<> EIGEN_STRONG_INLINE Packet16bf pcast<Packet16f, Packet16bf>(const Packet16f& a) {
  return F32ToBf16(a);
 }

 #ifdef EIGEN_VECTORIZE_AVX512FP16

-template <>
-struct type_casting_traits<half, float> {
-  enum {
-    VectorizedCast = 1,
-    SrcCoeffRatio = 1,
-    TgtCoeffRatio = 2
-  };
-};
-
-template <>
-struct type_casting_traits<float, half> {
-  enum {
-    VectorizedCast = 1,
-    SrcCoeffRatio = 2,
-    TgtCoeffRatio = 1
-  };
-};
-
 template<> EIGEN_STRONG_INLINE Packet16h preinterpret<Packet16h, Packet32h>(const Packet32h& a) {
  return _mm256_castpd_si256(_mm512_extractf64x4_pd(_mm512_castph_pd(a), 0));
 }
@@ -257,7 +219,7 @@ EIGEN_STRONG_INLINE Packet8h pcast<Packet4f, Packet8h>(const Packet4f& a, const
  __m256 result = _mm256_undefined_ps();
  result = _mm256_insertf128_ps(result, a, 0);
  result = _mm256_insertf128_ps(result, b, 1);
-  return _mm256_cvtps_ph(result, _MM_FROUND_TO_NEAREST_INT|_MM_FROUND_NO_EXC);
+  return _mm256_cvtps_ph(result, _MM_FROUND_TO_NEAREST_INT);
 }


--- a/Eigen/src/Core/arch/SSE/TypeCasting.h
+++ b/Eigen/src/Core/arch/SSE/TypeCasting.h
@@ -17,61 +17,19 @@ namespace Eigen {
 namespace internal {

 #ifndef EIGEN_VECTORIZE_AVX
-template <>
-struct type_casting_traits<float, bool> {
-  enum {
-    VectorizedCast = 1,
-    SrcCoeffRatio = 4,
-    TgtCoeffRatio = 1
-  };
-};
+template<> struct type_casting_traits<float, bool> : vectorized_type_casting_traits<float, bool> {};
+template<> struct type_casting_traits<bool, float> : vectorized_type_casting_traits<bool, float> {};

-template <>
-struct type_casting_traits<float, double> {
-  enum {
-    VectorizedCast = 1,
-    SrcCoeffRatio = 1,
-    TgtCoeffRatio = 2
-  };
-};
+template<> struct type_casting_traits<float, int> : vectorized_type_casting_traits<float, int> {};
+template<> struct type_casting_traits<int, float> : vectorized_type_casting_traits<int, float> {};
+
+template<> struct type_casting_traits<float, double> : vectorized_type_casting_traits<float, double> {};
+template<> struct type_casting_traits<double, float> : vectorized_type_casting_traits<double, float> {};
+
+template<> struct type_casting_traits<double, int> : vectorized_type_casting_traits<double, int> {};
+template<> struct type_casting_traits<int, double> : vectorized_type_casting_traits<int, double> {};
 #endif

-template <>
-struct type_casting_traits<int, float> {
-  enum {
-    VectorizedCast = 1,
-    SrcCoeffRatio = 1,
-    TgtCoeffRatio = 1
-  };
-};
-
-template <>
-struct type_casting_traits<float, int> {
-  enum {
-    VectorizedCast = 1,
-    SrcCoeffRatio = 1,
-    TgtCoeffRatio = 1
-  };
-};
-
-template <>
-struct type_casting_traits<double, int> {
-  enum {
-    VectorizedCast = 1,
-    SrcCoeffRatio = 2,
-    TgtCoeffRatio = 1
-  };
-};
-
-template <>
-struct type_casting_traits<double, float> {
-  enum {
-    VectorizedCast = 1,
-    SrcCoeffRatio = 2,
-    TgtCoeffRatio = 1
-  };
-};
-
 template <>
 EIGEN_STRONG_INLINE Packet16b pcast<Packet4f, Packet16b>(const Packet4f& a,
                                                         const Packet4f& b,
@@ -88,10 +46,31 @@ EIGEN_STRONG_INLINE Packet16b pcast<Packet4f, Packet16b>(const Packet4f& a,
  return _mm_and_si128(merged, _mm_set1_epi8(1));
 }

+template <>
+EIGEN_STRONG_INLINE Packet4f pcast<Packet16b, Packet4f>(const Packet16b& a) {
+  const __m128 cst_one = _mm_set_ps1(1.0f);
+  #ifdef EIGEN_VECTORIZE_SSE4_1
+  __m128i a_extended = _mm_cvtepi8_epi32(a);
+  __m128i abcd = _mm_cmpeq_epi32(a_extended, _mm_setzero_si128());
+  #else
+  __m128i abcd_efhg_ijkl_mnop = _mm_cmpeq_epi8(a, _mm_setzero_si128());
+  __m128i aabb_ccdd_eeff_gghh = _mm_unpacklo_epi8(abcd_efhg_ijkl_mnop, abcd_efhg_ijkl_mnop);
+  __m128i abcd = _mm_unpacklo_epi8(aabb_ccdd_eeff_gghh, aabb_ccdd_eeff_gghh);
+  #endif
+  __m128 result = _mm_andnot_ps(_mm_castsi128_ps(abcd), cst_one);
+  return result;
+}
+
 template<> EIGEN_STRONG_INLINE Packet4i pcast<Packet4f, Packet4i>(const Packet4f& a) {
  return _mm_cvttps_epi32(a);
 }

+template<> EIGEN_STRONG_INLINE Packet4i pcast<Packet2d, Packet4i>(const Packet2d& a, const Packet2d& b) {
+  return _mm_castps_si128(_mm_shuffle_ps(_mm_castsi128_ps(_mm_cvttpd_epi32(a)),
+                                         _mm_castsi128_ps(_mm_cvttpd_epi32(b)),
+                                         (1 << 2) | (1 << 6)));
+}
+
 template<> EIGEN_STRONG_INLINE Packet4f pcast<Packet4i, Packet4f>(const Packet4i& a) {
  return _mm_cvtepi32_ps(a);
 }
@@ -100,10 +79,9 @@ template<> EIGEN_STRONG_INLINE Packet4f pcast<Packet2d, Packet4f>(const Packet2d
  return _mm_shuffle_ps(_mm_cvtpd_ps(a), _mm_cvtpd_ps(b), (1 << 2) | (1 << 6));
 }

-template<> EIGEN_STRONG_INLINE Packet4i pcast<Packet2d, Packet4i>(const Packet2d& a, const Packet2d& b) {
-  return _mm_castps_si128(_mm_shuffle_ps(_mm_castsi128_ps(_mm_cvttpd_epi32(a)),
-                                         _mm_castsi128_ps(_mm_cvttpd_epi32(b)),
-                                         (1 << 2) | (1 << 6)));
+template<> EIGEN_STRONG_INLINE Packet2d pcast<Packet4i, Packet2d>(const Packet4i& a) {
+  // Simply discard the second half of the input
+  return _mm_cvtepi32_pd(a);
 }

 template<> EIGEN_STRONG_INLINE Packet2d pcast<Packet4f, Packet2d>(const Packet4f& a) {