Enable Half/BF16: abs, full, gelu, neg (#5856)

Summary: Pull Request resolved: #5856

Differential Revision: D63863399

Author: manuelcandales

kernels/portable/cpu/op_abs.cpp

@@ -31,7 +31,7 @@ Tensor& abs_out(KernelRuntimeContext& ctx, const Tensor& in, Tensor& out) {
   ET_KERNEL_CHECK(
       ctx, tensors_have_same_dim_order(in, out), InvalidArgument, out);
 
-  ET_SWITCH_REAL_TYPES(in.scalar_type(), ctx, "abs.out", CTYPE, [&] {
+  ET_SWITCH_REALHBF16_TYPES(in.scalar_type(), ctx, "abs.out", CTYPE, [&] {
     apply_unary_map_fn(
         [](const CTYPE val_in) {
           if (val_in < 0) {

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/portable/cpu/op_gelu.cpp

@@ -37,7 +37,7 @@ Tensor& gelu_out(
   ET_KERNEL_CHECK(
       ctx, tensors_have_same_dim_order(in, out), InvalidArgument, out);
 
-  ET_SWITCH_FLOAT_TYPES(in.scalar_type(), ctx, "gelu.out", CTYPE, [&]() {
+  ET_SWITCH_FLOATH_TYPES(in.scalar_type(), ctx, "gelu.out", CTYPE, [&]() {
     if (approximate == "tanh") {
       apply_unary_map_fn(
           [](const CTYPE x) {

kernels/portable/cpu/op_neg.cpp

@@ -33,7 +33,7 @@ Tensor& neg_out(KernelRuntimeContext& ctx, const Tensor& in, Tensor& out) {
   ET_KERNEL_CHECK(
       ctx, tensors_have_same_dim_order(in, out), InvalidArgument, out);
 
-  ET_SWITCH_REAL_TYPES(in.scalar_type(), ctx, "neg.out", CTYPE, [&] {
+  ET_SWITCH_REALHBF16_TYPES(in.scalar_type(), ctx, "neg.out", CTYPE, [&] {
     apply_unary_map_fn(
         [](const CTYPE val_in) { return static_cast<CTYPE>(-val_in); },
         in.const_data_ptr<CTYPE>(),

kernels/test/op_abs_test.cpp

@@ -24,8 +24,44 @@ class OpAbsTest : public OperatorTest {
   Tensor& op_abs_out(const Tensor& self, Tensor& out) {
     return torch::executor::aten::abs_outf(context_, self, out);
   }
+
+  template <ScalarType DTYPE>
+  void test_dtype() {
+    TensorFactory<DTYPE> tf;
+
+    Tensor in = tf.make({2, 3}, {-3, -2, -1, 0, 1, 2});
+    Tensor out = tf.zeros({2, 3});
+    Tensor expected = tf.make({2, 3}, {3, 2, 1, 0, 1, 2});
+
+    Tensor ret = op_abs_out(in, out);
+
+    EXPECT_TENSOR_EQ(out, ret);
+    EXPECT_TENSOR_EQ(out, expected);
+  }
+
+  template <>
+  void test_dtype<ScalarType::Byte>() {
+    TensorFactory<ScalarType::Byte> tf;
+
+    Tensor in = tf.make({2, 3}, {253, 254, 255, 0, 1, 2});
+    Tensor out = tf.zeros({2, 3});
+    Tensor expected = tf.make({2, 3}, {253, 254, 255, 0, 1, 2});
+
+    Tensor ret = op_abs_out(in, out);
+
+    EXPECT_TENSOR_EQ(out, ret);
+    EXPECT_TENSOR_EQ(out, expected);
+  }
 };
 
+TEST_F(OpAbsTest, AllRealHBF16Input) {
+#define TEST_KERNEL(INPUT_CTYPE, INPUT_DTYPE) \
+  test_dtype<ScalarType::INPUT_DTYPE>();
+
+  ET_FORALL_REALHBF16_TYPES(TEST_KERNEL);
+#undef TEST_KERNEL
+}
+
 TEST_F(OpAbsTest, SanityCheck) {
   TensorFactory<ScalarType::Float> tf;
 

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));
+}

kernels/test/op_gelu_test.cpp

@@ -66,6 +66,10 @@ class OpGeluTest : public OperatorTest {
   }
 };
 
+TEST_F(OpGeluTest, HalfTensors) {
+  test_gelu_execution<ScalarType::Half>();
+}
+
 TEST_F(OpGeluTest, FloatTensors) {
   test_gelu_execution<ScalarType::Float>();
 }

kernels/test/op_neg_test.cpp

@@ -24,8 +24,44 @@ class OpNegTest : public OperatorTest {
   Tensor& op_neg_out(const Tensor& self, Tensor& out) {
     return torch::executor::aten::neg_outf(context_, self, out);
   }
+
+  template <ScalarType DTYPE>
+  void test_dtype() {
+    TensorFactory<DTYPE> tf;
+
+    Tensor in = tf.make({2, 3}, {-3, -2, -1, 0, 1, 2});
+    Tensor out = tf.zeros({2, 3});
+    Tensor expected = tf.make({2, 3}, {3, 2, 1, 0, -1, -2});
+
+    Tensor ret = op_neg_out(in, out);
+
+    EXPECT_TENSOR_EQ(out, ret);
+    EXPECT_TENSOR_EQ(out, expected);
+  }
+
+  template <>
+  void test_dtype<ScalarType::Byte>() {
+    TensorFactory<ScalarType::Byte> tf;
+
+    Tensor in = tf.make({2, 3}, {253, 254, 255, 0, 1, 2});
+    Tensor out = tf.zeros({2, 3});
+    Tensor expected = tf.make({2, 3}, {3, 2, 1, 0, 255, 254});
+
+    Tensor ret = op_neg_out(in, out);
+
+    EXPECT_TENSOR_EQ(out, ret);
+    EXPECT_TENSOR_EQ(out, expected);
+  }
 };
 
+TEST_F(OpNegTest, AllRealHBF16Input) {
+#define TEST_KERNEL(INPUT_CTYPE, INPUT_DTYPE) \
+  test_dtype<ScalarType::INPUT_DTYPE>();
+
+  ET_FORALL_REALHBF16_TYPES(TEST_KERNEL);
+#undef TEST_KERNEL
+}
+
 TEST_F(OpNegTest, SanityCheck) {
   TensorFactory<ScalarType::Float> tf;