Rewrite logical op pattern. Binary ops: logical_and, logical_or, logical_xor (#6015)

Summary:
Pull Request resolved: #6015

Aggregate of binary logical ops:
- 680 K -> 5 K
ghstack-source-id: 246985134
exported-using-ghexport

Reviewed By: swolchok

Differential Revision: D63989999

fbshipit-source-id: e419a116b4fe7126ab2e499db64cb5e661cfd90e

Author: manuelcandales

kernels/portable/cpu/op_logical_and.cpp

@@ -6,7 +6,7 @@
  * LICENSE file in the root directory of this source tree.
  */
 
-#include <executorch/kernels/portable/cpu/pattern/pattern.h>
+#include <executorch/kernels/portable/cpu/pattern/logical_op.h>
 #include <executorch/runtime/kernel/kernel_includes.h>
 #include <cmath>
 
@@ -26,8 +26,9 @@ Tensor& logical_and_out(
     const Tensor& a,
     const Tensor& b,
     Tensor& out) {
-  return internal::binary_ufunc_realb_realb_to_realb_logical(
-      logical_and, ctx, a, b, out);
+  // @lint-ignore CLANGTIDY facebook-hte-CArray
+  static constexpr const char op_name[] = "logical_and.out";
+  return internal::logical_tensor_out<op_name>(logical_and, ctx, a, b, out);
 }
 
 } // namespace native

kernels/portable/cpu/op_logical_or.cpp

@@ -6,7 +6,7 @@
  * LICENSE file in the root directory of this source tree.
  */
 
-#include <executorch/kernels/portable/cpu/pattern/pattern.h>
+#include <executorch/kernels/portable/cpu/pattern/logical_op.h>
 #include <executorch/runtime/kernel/kernel_includes.h>
 #include <cmath>
 
@@ -26,8 +26,9 @@ Tensor& logical_or_out(
     const Tensor& a,
     const Tensor& b,
     Tensor& out) {
-  return internal::binary_ufunc_realb_realb_to_realb_logical(
-      logical_or, ctx, a, b, out);
+  // @lint-ignore CLANGTIDY facebook-hte-CArray
+  static constexpr const char op_name[] = "logical_or.out";
+  return internal::logical_tensor_out<op_name>(logical_or, ctx, a, b, out);
 }
 
 } // namespace native

kernels/portable/cpu/op_logical_xor.cpp

@@ -6,7 +6,7 @@
  * LICENSE file in the root directory of this source tree.
  */
 
-#include <executorch/kernels/portable/cpu/pattern/pattern.h>
+#include <executorch/kernels/portable/cpu/pattern/logical_op.h>
 #include <executorch/runtime/kernel/kernel_includes.h>
 #include <cmath>
 
@@ -26,8 +26,9 @@ Tensor& logical_xor_out(
     const Tensor& a,
     const Tensor& b,
     Tensor& out) {
-  return internal::binary_ufunc_realb_realb_to_realb_logical(
-      logical_xor, ctx, a, b, out);
+  // @lint-ignore CLANGTIDY facebook-hte-CArray
+  static constexpr const char op_name[] = "logical_xor.out";
+  return internal::logical_tensor_out<op_name>(logical_xor, ctx, a, b, out);
 }
 
 } // namespace native

kernels/portable/cpu/pattern/binary_ufunc_realb_realb_to_realb_logical.cpp

@@ -0,0 +1,53 @@
+/*
+ * 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/portable/cpu/util/elementwise_util.h>
+#include <executorch/runtime/kernel/kernel_includes.h>
+
+namespace torch {
+namespace executor {
+namespace native {
+namespace internal {
+
+/**
+ * Implements an op pattern for ops that take two broadcastable input tensors
+ * and performs an element-wise binary logical operation `fn`.
+ */
+template <const char* op_name>
+Tensor& logical_tensor_out(
+    bool (*fn)(bool, bool),
+    KernelRuntimeContext& ctx,
+    const Tensor& a,
+    const Tensor& b,
+    Tensor& out) {
+  ET_KERNEL_CHECK(
+      ctx, tensors_have_same_dim_order(a, b, out), InvalidArgument, out);
+
+  ET_KERNEL_CHECK(
+      ctx,
+      resize_to_broadcast_target_size(a, b, out) == Error::Ok,
+      InvalidArgument,
+      out);
+
+  utils::apply_bitensor_elementwise_fn<bool, op_name>(
+      fn,
+      ctx,
+      a,
+      utils::SupportedTensorDtypes::REALHBBF16,
+      b,
+      utils::SupportedTensorDtypes::REALHBBF16,
+      out,
+      utils::SupportedTensorDtypes::REALHBBF16);
+
+  return out;
+}
+
+} // namespace internal
+} // namespace native
+} // namespace executor
+} // namespace torch

kernels/portable/cpu/pattern/logical_op.h

@@ -89,19 +89,6 @@ Tensor& unary_ufunc_realhbbf16_to_floathbf16(
     const Tensor& in,
     Tensor& out);
 
-/**
- * Implements an op pattern for ops that take two broadcastable input tensors
- * of any realb dtype, no additional arguments, performs an element-wise binary
- * logical operation, and outputs a realb tensor. The function fn specifies the
- * binary logical operation which is applied to the input tensors element-wise.
- */
-Tensor& binary_ufunc_realb_realb_to_realb_logical(
-    bool (*fn)(bool, bool),
-    KernelRuntimeContext& ctx,
-    const Tensor& a,
-    const Tensor& b,
-    Tensor& out);
-
 } // namespace internal
 } // namespace native
 } // namespace executor

kernels/portable/cpu/pattern/pattern.h

@@ -15,6 +15,7 @@ def define_common_targets():
             "//executorch/kernels/portable/cpu/pattern:pattern",
             "//executorch/kernels/portable/cpu/pattern:bitwise_op",
             "//executorch/kernels/portable/cpu/pattern:comparison_op",
+            "//executorch/kernels/portable/cpu/pattern:logical_op"
         ],
         visibility = ["//executorch/...", "@EXECUTORCH_CLIENTS"],
     )
@@ -37,13 +38,21 @@ def define_common_targets():
         visibility = ["//executorch/kernels/portable/cpu/...", "//executorch/kernels/optimized/cpu/..."],
     )
 
+    runtime.cxx_library(
+        name = "logical_op",
+        exported_headers = [
+            "logical_op.h",
+        ],
+        compiler_flags = [],
+        visibility = ["//executorch/kernels/portable/cpu/...", "//executorch/kernels/optimized/cpu/..."],
+    )
+
     runtime.cxx_library(
         name = "pattern",
         srcs = [
             "unary_ufunc_realhb_to_bool.cpp",
             "unary_ufunc_realhbbf16_to_floathbf16.cpp",
             "unary_ufunc_realh.cpp",
-            "binary_ufunc_realb_realb_to_realb_logical.cpp",
         ],
         exported_headers = [
             "pattern.h",

kernels/portable/cpu/pattern/targets.bzl

@@ -699,7 +699,10 @@ ATEN_OPS = (
     op_target(
         name = "op_logical_and",
         deps = [
-            "//executorch/kernels/portable/cpu/pattern:pattern",
+            ":scalar_utils",
+            "//executorch/kernels/portable/cpu/pattern:logical_op",
+            "//executorch/kernels/portable/cpu/util:broadcast_util",
+            "//executorch/kernels/portable/cpu/util:elementwise_util",
         ],
     ),
     op_target(
@@ -712,13 +715,19 @@ ATEN_OPS = (
     op_target(
         name = "op_logical_or",
         deps = [
-            "//executorch/kernels/portable/cpu/pattern:pattern",
+            ":scalar_utils",
+            "//executorch/kernels/portable/cpu/pattern:logical_op",
+            "//executorch/kernels/portable/cpu/util:broadcast_util",
+            "//executorch/kernels/portable/cpu/util:elementwise_util",
         ],
     ),
     op_target(
         name = "op_logical_xor",
         deps = [
-            "//executorch/kernels/portable/cpu/pattern:pattern",
+            ":scalar_utils",
+            "//executorch/kernels/portable/cpu/pattern:logical_op",
+            "//executorch/kernels/portable/cpu/util:broadcast_util",
+            "//executorch/kernels/portable/cpu/util:elementwise_util",
         ],
     ),
     op_target(