Block the combination of bf16 input and float output in EmbeddingSpMD…

…M when running on Mac and Windows (pytorch#1865) Summary: Pull Request resolved: pytorch#1865 Block the combination of bfloat16 input and float32 output in EmbeddingSpMDM API when running on Mac and Windows. This is to prevent a random test failure. Reviewed By: q10 Differential Revision: D47306447 fbshipit-source-id: 703c81820f90f8c50f496e9ed4d91d46592e4a04
q10 · Jul 11, 2023 · 372c62e · 372c62e
1 parent eea0116
commit 372c62e
Show file tree

Hide file tree

Showing 2 changed files with 3 additions and 3 deletions.
diff --git a/src/EmbeddingSpMDM.cc b/src/EmbeddingSpMDM.cc
@@ -1032,11 +1032,11 @@ typename EmbeddingSpMDMKernelSignature<inType, indxType, offsetType, outType>::
   if (!cpuinfo_initialize()) {
     throw std::runtime_error("Failed to initialize cpuinfo!");
   }
-#if defined(__APPLE__)
+#if defined(__APPLE__) || defined(_WIN32)
   if (std::is_same<inType, uint16_t>::value && is_bf16_in &&
       std::is_same<outType, float>::value) {
     throw std::runtime_error(
-        "Bfloat16 input with float32 output is not yet supported on Apple");
+        "Bfloat16 input with float32 output is not yet supported on Apple or Windows");
   }
 #endif
   if (output_stride == -1) {

diff --git a/test/EmbeddingSpMDMTest.cc b/test/EmbeddingSpMDMTest.cc
@@ -145,7 +145,7 @@ TEST_P(EmbeddingSpMDMTest, basicTest) {
     return;
   }
 
-#if defined(__APPLE__)
+#if defined(__APPLE__) || defined(_WIN32)
   if (in_type == BFLOAT16 && out_type == FLOAT) {
     return;
   }