Support Half/BFloat16 in relu (#7858)

Partial fix for #7748.
pytorch · Jan 28, 2025 · e2d23a5 · e2d23a5
1 parent a561e31
commit e2d23a5
Show file tree

Hide file tree

Showing 2 changed files with 14 additions and 2 deletions.
diff --git a/kernels/portable/cpu/op_relu.cpp b/kernels/portable/cpu/op_relu.cpp
@@ -33,12 +33,16 @@ Tensor& relu_out(KernelRuntimeContext& ctx, const Tensor& in, Tensor& out) {
   ET_KERNEL_CHECK(
       ctx, tensors_have_same_shape_and_dtype(in, out), InvalidArgument, out);
 
-  ET_KERNEL_CHECK(ctx, tensor_is_real_type(out), InvalidArgument, out);
+  ET_KERNEL_CHECK(
+      ctx,
+      executorch::runtime::tensor_is_realhbf16_type(out),
+      InvalidArgument,
+      out);
 
   ET_KERNEL_CHECK(
       ctx, tensors_have_same_dim_order(in, out), InvalidArgument, out);
 
-  ET_SWITCH_REAL_TYPES(in.scalar_type(), ctx, "relu.out", CTYPE, [&]() {
+  ET_SWITCH_REALHBF16_TYPES(in.scalar_type(), ctx, "relu.out", CTYPE, [&]() {
     apply_unary_map_fn(
         [](const CTYPE val_in) {
           return (std::isnan(val_in) || val_in >= CTYPE(0)) ? val_in : CTYPE(0);

diff --git a/kernels/test/op_relu_test.cpp b/kernels/test/op_relu_test.cpp
@@ -82,6 +82,14 @@ TEST_F(OpReluTest, DoubleTensors) {
   test_relu_execution_floats<ScalarType::Double>();
 }
 
+TEST_F(OpReluTest, HalfTensors) {
+  test_relu_execution_floats<ScalarType::Half>();
+}
+
+TEST_F(OpReluTest, BFloat16Tensors) {
+  test_relu_execution_floats<ScalarType::BFloat16>();
+}
+
 TEST_F(OpReluTest, ByteTensors) {
   TensorFactory<ScalarType::Byte> tf;