From 372c62e9d5ba3c3d07250cc5b2cee63cdb1d719c Mon Sep 17 00:00:00 2001
From: Wei Su <wsu@meta.com>
Date: Mon, 10 Jul 2023 19:54:59 -0700
Subject: [PATCH] Block the combination of bf16 input and float output in
 EmbeddingSpMDM when running on Mac and Windows (#1865)

Summary:
Pull Request resolved: https://github.com/pytorch/FBGEMM/pull/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
---
 src/EmbeddingSpMDM.cc      | 4 ++--
 test/EmbeddingSpMDMTest.cc | 2 +-
 2 files changed, 3 insertions(+), 3 deletions(-)

diff --git a/src/EmbeddingSpMDM.cc b/src/EmbeddingSpMDM.cc
index 26c5fc45f..5e4e11b75 100644
--- 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
index d4dbac6b3..68962cff2 100644
--- 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;
   }