Enable Half & Bfloat16: op_eq & op_full

Differential Revision: D63863399

Author: manuelcandales

kernels/portable/cpu/op_eq.cpp

@@ -37,28 +37,27 @@ Tensor& eq_tensor_out(
   ET_KERNEL_CHECK(
       ctx, tensors_have_same_dim_order(a, b, out), InvalidArgument, out);
 
-  ET_SWITCH_REAL_TYPES_AND(Bool, a_type, ctx, "eq.Scalar_out", CTYPE_A, [&]() {
-    ET_SWITCH_REAL_TYPES_AND(
-        Bool, b_type, ctx, "eq.Scalar_out", CTYPE_B, [&]() {
-          using CTYPE_IN =
-              typename torch::executor::promote_types<CTYPE_A, CTYPE_B>::type;
-          ET_DCHECK(
-              CppTypeToScalarType<CTYPE_IN>::value ==
-              promoteTypes(a_type, b_type));
-          ET_SWITCH_REAL_TYPES_AND(
-              Bool, out_type, ctx, "eq.Scalar_out", CTYPE_OUT, [&]() {
-                apply_binary_elementwise_fn<CTYPE_A, CTYPE_B, CTYPE_OUT>(
-                    [](const CTYPE_A val_a, const CTYPE_B val_b) {
-                      CTYPE_IN a_casted = static_cast<CTYPE_IN>(val_a);
-                      CTYPE_IN b_casted = static_cast<CTYPE_IN>(val_b);
-                      bool value = a_casted == b_casted;
-                      return static_cast<CTYPE_OUT>(value);
-                    },
-                    a,
-                    b,
-                    out);
-              });
-        });
+  constexpr auto name = "eq.Tensor_out";
+
+  ET_SWITCH_REALHBBF16_TYPES(a_type, ctx, name, CTYPE_A, [&]() {
+    ET_SWITCH_REALHBBF16_TYPES(b_type, ctx, name, CTYPE_B, [&]() {
+      using CTYPE_IN =
+          typename torch::executor::promote_types<CTYPE_A, CTYPE_B>::type;
+      ET_DCHECK(
+          CppTypeToScalarType<CTYPE_IN>::value == promoteTypes(a_type, b_type));
+      ET_SWITCH_REALHBBF16_TYPES(out_type, ctx, name, CTYPE_OUT, [&]() {
+        apply_binary_elementwise_fn<CTYPE_A, CTYPE_B, CTYPE_OUT>(
+            [](const CTYPE_A val_a, const CTYPE_B val_b) {
+              CTYPE_IN a_casted = static_cast<CTYPE_IN>(val_a);
+              CTYPE_IN b_casted = static_cast<CTYPE_IN>(val_b);
+              bool value = a_casted == b_casted;
+              return static_cast<CTYPE_OUT>(value);
+            },
+            a,
+            b,
+            out);
+      });
+    });
   });
 
   return out;
@@ -86,27 +85,28 @@ Tensor& eq_scalar_out(
   ET_KERNEL_CHECK(
       ctx, tensors_have_same_dim_order(a, out), InvalidArgument, out);
 
-  ET_SWITCH_REAL_TYPES_AND(Bool, a_type, ctx, "eq.Scalar_out", CTYPE_A, [&]() {
-    ET_SWITCH_SCALAR_OBJ_TYPES(b_type, ctx, "eq.Scalar_out", CTYPE_B, [&]() {
+  constexpr auto name = "eq.Scalar_out";
+
+  ET_SWITCH_REALHBBF16_TYPES(a_type, ctx, name, CTYPE_A, [&]() {
+    ET_SWITCH_SCALAR_OBJ_TYPES(b_type, ctx, name, CTYPE_B, [&]() {
       using CTYPE_IN =
           typename torch::executor::promote_types<CTYPE_A, CTYPE_B>::type;
       ET_DCHECK(
           CppTypeToScalarType<CTYPE_IN>::value == promoteTypes(a_type, b_type));
-      ET_SWITCH_REAL_TYPES_AND(
-          Bool, out_type, ctx, "eq.Scalar_out", CTYPE_OUT, [&]() {
-            CTYPE_B val_b = 0;
-            utils::extract_scalar(b, &val_b);
-            apply_unary_map_fn(
-                [val_b](const CTYPE_A val_a) {
-                  CTYPE_IN a_casted = static_cast<CTYPE_IN>(val_a);
-                  CTYPE_IN b_casted = static_cast<CTYPE_IN>(val_b);
-                  bool value = a_casted == b_casted;
-                  return static_cast<CTYPE_OUT>(value);
-                },
-                a.const_data_ptr<CTYPE_A>(),
-                out.mutable_data_ptr<CTYPE_OUT>(),
-                out.numel());
-          });
+      ET_SWITCH_REALHBBF16_TYPES(out_type, ctx, name, CTYPE_OUT, [&]() {
+        CTYPE_B val_b = 0;
+        utils::extract_scalar(b, &val_b);
+        apply_unary_map_fn(
+            [val_b](const CTYPE_A val_a) {
+              CTYPE_IN a_casted = static_cast<CTYPE_IN>(val_a);
+              CTYPE_IN b_casted = static_cast<CTYPE_IN>(val_b);
+              bool value = a_casted == b_casted;
+              return static_cast<CTYPE_OUT>(value);
+            },
+            a.const_data_ptr<CTYPE_A>(),
+            out.mutable_data_ptr<CTYPE_OUT>(),
+            out.numel());
+      });
     });
   });
 

kernels/portable/cpu/op_full.cpp

@@ -40,7 +40,7 @@ Tensor& full_out(
     CTYPE_VAL val;
     utils::extract_scalar(fill_value, &val);
 
-    ET_SWITCH_REALHB_TYPES(out_type, ctx, name, CTYPE_OUT, [&] {
+    ET_SWITCH_REALHBBF16_TYPES(out_type, ctx, name, CTYPE_OUT, [&] {
       CTYPE_OUT val_casted = static_cast<CTYPE_OUT>(val);
       auto data_out = out.mutable_data_ptr<CTYPE_OUT>();
       for (size_t i = 0; i < out.numel(); ++i) {

kernels/test/op_eq_test.cpp

@@ -63,7 +63,7 @@ class OpEqScalarOutTest : public OperatorTest {
 
 TEST_F(OpEqScalarOutTest, AllRealInputBoolOutputSupport) {
 #define TEST_ENTRY(ctype, dtype) test_eq_scalar_out<ScalarType::dtype>();
-  ET_FORALL_REAL_TYPES(TEST_ENTRY);
+  ET_FORALL_REALHBF16_TYPES(TEST_ENTRY);
 #undef TEST_ENTRY
 }
 

kernels/test/op_full_test.cpp

@@ -122,3 +122,26 @@ TEST_F(OpFullOutTest, ZeroDim) {
   op_full_out(sizes, true, out);
   EXPECT_TENSOR_EQ(out, tf.ones(sizes_in32_t_vec));
 }
+
+TEST_F(OpFullOutTest, BFloat16Support) {
+  TensorFactory<ScalarType::BFloat16> tf;
+
+  std::vector<int64_t> sizes_int64_t_vec = {2, 3};
+  std::vector<int32_t> sizes_in32_t_vec = {2, 3};
+  auto sizes = IntArrayRef(sizes_int64_t_vec.data(), sizes_int64_t_vec.size());
+
+  // Boolean Scalar
+  Tensor out = tf.zeros(sizes_in32_t_vec);
+  op_full_out(sizes, true, out);
+  EXPECT_TENSOR_EQ(out, tf.ones(sizes_in32_t_vec));
+
+  // Integral Scalar
+  out = tf.zeros(sizes_in32_t_vec);
+  op_full_out(sizes, 1, out);
+  EXPECT_TENSOR_EQ(out, tf.ones(sizes_in32_t_vec));
+
+  // Floating Point Scalar
+  out = tf.zeros(sizes_in32_t_vec);
+  op_full_out(sizes, 3.1415926535, out);
+  EXPECT_TENSOR_EQ(out, tf.full(sizes_in32_t_vec, 3.1415926535));
+}