Support Half/BFloat16 in abs/neg (pytorch#7760)

Partial fix for pytorch#7748.

Author: swolchok

kernels/optimized/cpu/op_neg.cpp

@@ -26,7 +26,7 @@ Tensor& opt_neg_out(KernelRuntimeContext& ctx, const Tensor& in, Tensor& out) {
       out,
       "Failed to resize output tensor.");
 
-  ET_SWITCH_REAL_TYPES(in.scalar_type(), ctx, "neg.out", CTYPE, [&] {
+  ET_SWITCH_REALHBF16_TYPES(in.scalar_type(), ctx, "neg.out", CTYPE, [&] {
     using Vec = executorch::vec::Vectorized<CTYPE>;
     executorch::vec::map<CTYPE>(
         [](Vec x) { return x.neg(); },

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_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,19 +24,27 @@ class OpAbsTest : public OperatorTest {
   Tensor& op_abs_out(const Tensor& self, Tensor& out) {
     return torch::executor::aten::abs_outf(context_, self, out);
   }
-};
 
-TEST_F(OpAbsTest, SanityCheck) {
-  TensorFactory<ScalarType::Float> tf;
+  template <ScalarType DTYPE>
+  void run_smoke_test() {
+    TensorFactory<DTYPE> tf;
+
+    Tensor in = tf.make({1, 7}, {-3.0, -2.5, -1.01, 0.0, 1.01, 2.5, 3.0});
+    Tensor out = tf.zeros({1, 7});
+    Tensor expected = tf.make({1, 7}, {3.0, 2.5, 1.01, 0.0, 1.01, 2.5, 3.0});
 
-  Tensor in = tf.make({1, 7}, {-3.0, -2.5, -1.01, 0.0, 1.01, 2.5, 3.0});
-  Tensor out = tf.zeros({1, 7});
-  Tensor expected = tf.make({1, 7}, {3.0, 2.5, 1.01, 0.0, 1.01, 2.5, 3.0});
+    Tensor ret = op_abs_out(in, out);
 
-  Tensor ret = op_abs_out(in, out);
+    EXPECT_TENSOR_EQ(out, ret);
+    EXPECT_TENSOR_EQ(out, expected);
+  }
+};
 
-  EXPECT_TENSOR_EQ(out, ret);
-  EXPECT_TENSOR_EQ(out, expected);
+TEST_F(OpAbsTest, SmokeTest) {
+#define RUN_SMOKE_TEST(ctype, dtype) run_smoke_test<ScalarType::dtype>();
+  // TODO: cover all REALHBF16 types with generalized unary function test
+  // harness.
+  ET_FORALL_FLOATHBF16_TYPES(RUN_SMOKE_TEST);
 }
 
 TEST_F(OpAbsTest, MemoryFormatCheck) {

kernels/test/op_neg_test.cpp

@@ -24,17 +24,25 @@ class OpNegTest : public OperatorTest {
   Tensor& op_neg_out(const Tensor& self, Tensor& out) {
     return torch::executor::aten::neg_outf(context_, self, out);
   }
-};
 
-TEST_F(OpNegTest, SanityCheck) {
-  TensorFactory<ScalarType::Float> tf;
+  template <ScalarType DTYPE>
+  void run_smoke_test() {
+    TensorFactory<DTYPE> tf;
+
+    Tensor in = tf.make({1, 7}, {-3.0, -2.5, -1.01, 0.0, 1.01, 2.5, 3.0});
+    Tensor out = tf.zeros({1, 7});
+    Tensor expected = tf.make({1, 7}, {3.0, 2.5, 1.01, 0.0, -1.01, -2.5, -3.0});
 
-  Tensor in = tf.make({1, 7}, {-3.0, -2.5, -1.01, 0.0, 1.01, 2.5, 3.0});
-  Tensor out = tf.zeros({1, 7});
-  Tensor expected = tf.make({1, 7}, {3.0, 2.5, 1.01, 0.0, -1.01, -2.5, -3.0});
+    Tensor ret = op_neg_out(in, out);
 
-  Tensor ret = op_neg_out(in, out);
+    EXPECT_TENSOR_EQ(out, ret);
+    EXPECT_TENSOR_EQ(out, expected);
+  }
+};
 
-  EXPECT_TENSOR_EQ(out, ret);
-  EXPECT_TENSOR_EQ(out, expected);
+TEST_F(OpNegTest, SmokeTest) {
+#define RUN_SMOKE_TEST(ctype, dtype) run_smoke_test<ScalarType::dtype>();
+  // TODO: cover all REALHBF16 types with generalized unary function test
+  // harness.
+  ET_FORALL_FLOATHBF16_TYPES(RUN_SMOKE_TEST);
 }