Fix lower precision check for MKLDNN on Windows (#122645)

Fixes #120788 Pull Request resolved: #121618 Approved by: https://github.com/xuhancn, https://github.com/jgong5, https://github.com/mingfeima, https://github.com/seemethere (cherry picked from commit 0371743) Co-authored-by: CaoE <[email protected]>
pytorch · Mar 25, 2024 · 39901f2 · 39901f2
1 parent 9e6f42d
commit 39901f2
Showing 1 changed file with 2 additions and 2 deletions.
diff --git a/aten/src/ATen/native/mkldnn/Utils.h b/aten/src/ATen/native/mkldnn/Utils.h
@@ -97,7 +97,7 @@ constexpr bool mkldnn_bf16_device_check_arm() {
 
 #if AT_MKLDNN_ENABLED()
 inline bool mkldnn_bf16_device_check() {
-#if defined(__x86_64__)
+#if defined(__x86_64__) || (defined(_M_X64) && !defined(_M_ARM64EC))
   // Use ideep to check bf16 on X64 as cpuinfo has no avx_ne_convert check.
   return ideep::has_bf16_type_support();
 #else
@@ -106,7 +106,7 @@ inline bool mkldnn_bf16_device_check() {
 }
 
 inline bool mkldnn_fp16_device_check() {
-#if defined(__x86_64__)
+#if defined(__x86_64__) || (defined(_M_X64) && !defined(_M_ARM64EC))
   return ideep::has_fp16_type_support();
 #else
   return false;