[ExecuTorch] Support bf16 for binary logical ops

kernels/portable/cpu/pattern/binary_ufunc_realb_realb_to_realb_logical.cpp

@@ -34,9 +34,9 @@ Tensor& binary_ufunc_realb_realb_to_realb_logical(
   ScalarType b_type = b.scalar_type();
   ScalarType out_type = out.scalar_type();
 
-  ET_SWITCH_REAL_TYPES_AND(Bool, a_type, ctx, __func__, CTYPE_A, [&]() {
-    ET_SWITCH_REAL_TYPES_AND(Bool, b_type, ctx, __func__, CTYPE_B, [&]() {
-      ET_SWITCH_REAL_TYPES_AND(Bool, out_type, ctx, __func__, CTYPE_OUT, [&]() {
+  ET_SWITCH_REALHBBF16_TYPES(a_type, ctx, __func__, CTYPE_A, [&]() {
+    ET_SWITCH_REALHBBF16_TYPES(b_type, ctx, __func__, CTYPE_B, [&]() {
+      ET_SWITCH_REALHBBF16_TYPES(out_type, ctx, __func__, CTYPE_OUT, [&]() {
         apply_binary_elementwise_fn<CTYPE_A, CTYPE_B, CTYPE_OUT>(
             [fn](const CTYPE_A val_a, const CTYPE_B val_b) {
               bool a_casted = static_cast<bool>(val_a);

kernels/test/BinaryLogicalOpTest.cpp

@@ -0,0 +1,27 @@
+/*
+ * Copyright (c) Meta Platforms, Inc. and affiliates.
+ * All rights reserved.
+ *
+ * This source code is licensed under the BSD-style license found in the
+ * LICENSE file in the root directory of this source tree.
+ */
+
+#include <executorch/kernels/test/BinaryLogicalOpTest.h>
+
+namespace torch::executor::testing {
+
+void BinaryLogicalOpTest::test_all_dtypes() {
+#define TEST_ENTRY(ctype, dtype) \
+  test_op_out<ScalarType::dtype, ScalarType::Double, ScalarType::Double>();
+  ET_FORALL_REALHBF16_TYPES(TEST_ENTRY);
+#undef TEST_ENTRY
+#define TEST_ENTRY(ctype, dtype) \
+  test_op_out<ScalarType::Double, ScalarType::dtype, ScalarType::Double>();
+  ET_FORALL_REALHBF16_TYPES(TEST_ENTRY);
+#undef TEST_ENTRY
+#define TEST_ENTRY(ctype, dtype) \
+  test_op_out<ScalarType::Double, ScalarType::Double, ScalarType::dtype>();
+  ET_FORALL_REALHBF16_TYPES(TEST_ENTRY);
+#undef TEST_ENTRY
+}
+} // namespace torch::executor::testing

kernels/test/BinaryLogicalOpTest.h

@@ -0,0 +1,72 @@
+/*
+ * Copyright (c) Meta Platforms, Inc. and affiliates.
+ * All rights reserved.
+ *
+ * This source code is licensed under the BSD-style license found in the
+ * LICENSE file in the root directory of this source tree.
+ */
+
+#pragma once
+
+#include <executorch/kernels/test/TestUtil.h>
+#include <executorch/runtime/core/exec_aten/exec_aten.h>
+#include <executorch/runtime/core/exec_aten/testing_util/tensor_factory.h>
+#include <executorch/runtime/core/exec_aten/testing_util/tensor_util.h>
+
+namespace torch::executor::testing {
+class BinaryLogicalOpTest : public OperatorTest {
+ protected:
+  // Implement this to call the torch::executor::aten::op_outf function for the
+  // op.
+  virtual exec_aten::Tensor& op_out(
+      const exec_aten::Tensor& lhs,
+      const exec_aten::Tensor& rhs,
+      exec_aten::Tensor& out) = 0;
+
+  // Scalar reference implementation of the function in question for testing.
+  virtual double op_reference(double x, double y) const = 0;
+
+  template <
+      exec_aten::ScalarType IN_DTYPE,
+      exec_aten::ScalarType IN_DTYPE2,
+      exec_aten::ScalarType OUT_DTYPE>
+  void test_op_out() {
+    TensorFactory<IN_DTYPE> tf_in;
+    TensorFactory<IN_DTYPE2> tf_in2;
+    TensorFactory<OUT_DTYPE> tf_out;
+
+    exec_aten::Tensor out = tf_out.zeros({1, 4});
+
+    using CTYPE1 = typename decltype(tf_in)::ctype;
+    std::vector<CTYPE1> test_vector1 = {0, CTYPE1(-1), CTYPE1(0), CTYPE1(31)};
+
+    using CTYPE2 = typename decltype(tf_in2)::ctype;
+    std::vector<CTYPE2> test_vector2 = {
+        CTYPE2(0),
+        CTYPE2(0),
+        CTYPE2(15),
+        CTYPE2(12),
+    };
+
+    std::vector<typename decltype(tf_out)::ctype> expected_vector;
+    for (int ii = 0; ii < test_vector1.size(); ++ii) {
+      expected_vector.push_back(
+          op_reference(test_vector1[ii], test_vector2[ii]));
+    }
+
+    op_out(
+        tf_in.make({1, 4}, test_vector1),
+        tf_in2.make({1, 4}, test_vector2),
+        out);
+
+    EXPECT_TENSOR_CLOSE(out, tf_out.make({1, 4}, expected_vector));
+  }
+
+  void test_all_dtypes();
+};
+
+#define IMPLEMENT_BINARY_LOGICAL_OP_TEST(TestName) \
+  TEST_F(TestName, SimpleTestAllTypes) {           \
+    test_all_dtypes();                             \
+  }
+} // namespace torch::executor::testing

kernels/test/CMakeLists.txt

@@ -68,6 +68,7 @@ add_custom_target(
 )
 
 set(all_test_sources
+    "BinaryLogicalOpTest.cpp"
     "op__to_dim_order_copy_test.cpp"
     "op_abs_test.cpp"
     "op_acos_test.cpp"

kernels/test/op_logical_and_test.cpp

@@ -6,23 +6,26 @@
  * LICENSE file in the root directory of this source tree.
  */
 
+#include <executorch/kernels/test/BinaryLogicalOpTest.h>
 #include <executorch/kernels/test/FunctionHeaderWrapper.h> // Declares the operator
-#include <executorch/kernels/test/TestUtil.h>
-#include <executorch/runtime/core/exec_aten/exec_aten.h>
-#include <executorch/runtime/core/exec_aten/testing_util/tensor_factory.h>
-#include <executorch/runtime/core/exec_aten/testing_util/tensor_util.h>
 
 #include <gtest/gtest.h>
 
-using namespace ::testing;
-using exec_aten::ScalarType;
 using exec_aten::Tensor;
-using torch::executor::testing::TensorFactory;
 
-class OpLogicalAndTest : public OperatorTest {
+class OpLogicalAndTest : public torch::executor::testing::BinaryLogicalOpTest {
  protected:
-  Tensor&
-  op_logical_and_out(const Tensor& self, const Tensor& other, Tensor& out) {
+  Tensor& op_out(const Tensor& self, const Tensor& other, Tensor& out)
+      override {
     return torch::executor::aten::logical_and_outf(context_, self, other, out);
   }
+
+  double op_reference(double x, double y) const override {
+    uint64_t lhs, rhs;
+    std::memcpy(&lhs, &x, sizeof(lhs));
+    std::memcpy(&rhs, &y, sizeof(rhs));
+    return lhs && rhs;
+  }
 };
+
+IMPLEMENT_BINARY_LOGICAL_OP_TEST(OpLogicalAndTest)

kernels/test/op_logical_or_test.cpp

@@ -6,23 +6,26 @@
  * LICENSE file in the root directory of this source tree.
  */
 
+#include <executorch/kernels/test/BinaryLogicalOpTest.h>
 #include <executorch/kernels/test/FunctionHeaderWrapper.h> // Declares the operator
-#include <executorch/kernels/test/TestUtil.h>
-#include <executorch/runtime/core/exec_aten/exec_aten.h>
-#include <executorch/runtime/core/exec_aten/testing_util/tensor_factory.h>
-#include <executorch/runtime/core/exec_aten/testing_util/tensor_util.h>
 
 #include <gtest/gtest.h>
 
-using namespace ::testing;
-using exec_aten::ScalarType;
 using exec_aten::Tensor;
-using torch::executor::testing::TensorFactory;
 
-class OpLogicalOrTest : public OperatorTest {
+class OpLogicalOrTest : public torch::executor::testing::BinaryLogicalOpTest {
  protected:
-  Tensor&
-  op_logical_or_out(const Tensor& self, const Tensor& other, Tensor& out) {
+  Tensor& op_out(const Tensor& self, const Tensor& other, Tensor& out)
+      override {
     return torch::executor::aten::logical_or_outf(context_, self, other, out);
   }
+
+  double op_reference(double x, double y) const override {
+    uint64_t lhs, rhs;
+    std::memcpy(&lhs, &x, sizeof(lhs));
+    std::memcpy(&rhs, &y, sizeof(rhs));
+    return lhs || rhs;
+  }
 };
+
+IMPLEMENT_BINARY_LOGICAL_OP_TEST(OpLogicalOrTest)

kernels/test/op_logical_xor_test.cpp

@@ -6,23 +6,26 @@
  * LICENSE file in the root directory of this source tree.
  */
 
+#include <executorch/kernels/test/BinaryLogicalOpTest.h>
 #include <executorch/kernels/test/FunctionHeaderWrapper.h> // Declares the operator
-#include <executorch/kernels/test/TestUtil.h>
-#include <executorch/runtime/core/exec_aten/exec_aten.h>
-#include <executorch/runtime/core/exec_aten/testing_util/tensor_factory.h>
-#include <executorch/runtime/core/exec_aten/testing_util/tensor_util.h>
 
 #include <gtest/gtest.h>
 
-using namespace ::testing;
-using exec_aten::ScalarType;
 using exec_aten::Tensor;
-using torch::executor::testing::TensorFactory;
 
-class OpLogicalXorTest : public OperatorTest {
+class OpLogicalXorTest : public torch::executor::testing::BinaryLogicalOpTest {
  protected:
-  Tensor&
-  op_logical_xor_out(const Tensor& self, const Tensor& other, Tensor& out) {
+  Tensor& op_out(const Tensor& self, const Tensor& other, Tensor& out)
+      override {
     return torch::executor::aten::logical_xor_outf(context_, self, other, out);
   }
+
+  double op_reference(double x, double y) const override {
+    uint64_t lhs, rhs;
+    std::memcpy(&lhs, &x, sizeof(lhs));
+    std::memcpy(&rhs, &y, sizeof(rhs));
+    return bool(lhs) != bool(rhs);
+  }
 };
+
+IMPLEMENT_BINARY_LOGICAL_OP_TEST(OpLogicalXorTest)

kernels/test/targets.bzl

@@ -44,9 +44,11 @@ def define_common_targets():
         runtime.cxx_library(
             name = "test_util" + aten_suffix,
             srcs = [
+                "BinaryLogicalOpTest.cpp",
                 "UnaryUfuncRealHBBF16ToFloatHBF16Test.cpp",
             ],
             exported_headers = [
+                "BinaryLogicalOpTest.h",
                 "TestUtil.h",
                 "UnaryUfuncRealHBBF16ToFloatHBF16Test.h",
             ],