Update

[ghstack-poisoned]
pytorch · leslie-fang-intel · Apr 29, 2025 · Apr 29, 2025 · Apr 29, 2025 · Apr 30, 2025
commit 36960ae861195cb2492a6a427afedfeb7c63c67b
diff --git a/aten/src/ATen/cpu/vec/vec512/vec512_convert.h b/aten/src/ATen/cpu/vec/vec512/vec512_convert.h
@@ -291,6 +291,28 @@ struct VecConvert<
   }
 };
 
+template <>
+struct VecConvert<Float8_e4m3fn, 1, float, 1> {
+  static inline VectorizedN<Float8_e4m3fn, 1> apply(const VectorizedN<float, 1>& src_n) {
+    std::cout<<"---- hit specilize cvt fp32 to float8 ----"<<std::endl;
+    at::vec::Vectorized<float> src = src_n[0];
+    __m128i res128 = cvtfp32_fp8e4m3(src);
+    return at::vec::Vectorized<Float8_e4m3fn>(_mm512_castsi128_si512(res128));
+  }
+};
+
+template <>
+struct VecConvert<float, 1, Float8_e4m3fn, 1> {
+  static inline VectorizedN<float, 1> apply(const VectorizedN<Float8_e4m3fn, 1>& src_n) {
+    // cvt first 16x8 bits from Float8_e4m3fn to float
+    std::cout<<"---- hit specilize cvt float8 to fp32 ----"<<std::endl;
+    at::vec::Vectorized<Float8_e4m3fn> src = src_n[0];
+    __m512 result;
+    cvtfp8e4m3_fp32(_mm512_castsi512_si128(src), result);
+    return at::vec::Vectorized<float>(result);
+  }
+};
+
 #endif
 
 } // namespace CPU_CAPABILITY

diff --git a/aten/src/ATen/cpu/vec/vec_convert.h b/aten/src/ATen/cpu/vec/vec_convert.h
@@ -15,6 +15,8 @@ template <
 struct VecConvert {
   static inline VectorizedN<dst_t, dst_n> apply(
       const VectorizedN<src_t, src_n>& src) {
+
+    std::cout<<"--- hit the legacy cvt ----"<<std::endl;
     constexpr int count = std::min(
         VectorizedN<src_t, src_n>::size(), VectorizedN<dst_t, dst_n>::size());
     __at_align__ src_t src_buf[VectorizedN<src_t, src_n>::size()];
@@ -36,6 +38,7 @@ inline std::enable_if_t<std::is_same_v<dst_t, src_t>, Vectorized<src_t>> convert
 template <typename dst_t, typename src_t>
 inline std::enable_if_t<!std::is_same_v<dst_t, src_t>, Vectorized<dst_t>>
 convert(const Vectorized<src_t>& src) {
+  std::cout<<"--- should hit this convert ----"<<std::endl;
   return VecConvert<dst_t, 1, src_t, 1>::apply(src);
 }
 

diff --git a/torch/_inductor/codegen/cpp.py b/torch/_inductor/codegen/cpp.py
@@ -154,6 +154,7 @@ def get_export_declaration():
     torch.int8,
     torch.int32,
     torch.int64,
+    torch.float8_e4m3fn,
 ]
 
 MASKED_VECTORIZABLE_DTYPES: list[torch.dtype] = [
@@ -1599,6 +1600,7 @@ def to_dtype(x, dtype, src_dtype=None, use_compute_dtypes=True):
             torch.int8,
             torch.int32,
             torch.int64,
+            torch.float8_e4m3fn,
         ], f"{__name__} does not support {dtype}"
         assert isinstance(x, CppCSEVariable)
         src_dtype = x.dtype