Enable optimized op_le broadcast against 1 element tensor

Differential Revision: D71652593

Pull Request resolved: #9507

Author: manuelcandales

kernels/optimized/cpu/op_le.cpp

@@ -9,6 +9,7 @@
 #include <executorch/kernels/optimized/vec/functional.h>
 #include <executorch/kernels/optimized/vec/vec.h>
 #include <executorch/kernels/portable/cpu/scalar_utils.h>
+#include <executorch/kernels/portable/cpu/util/broadcast_util.h>
 #include <executorch/runtime/kernel/kernel_includes.h>
 #include <executorch/runtime/platform/assert.h>
 
@@ -26,6 +27,58 @@ Tensor& opt_le_tensor_out(
     Tensor& out) {
   (void)ctx;
 
+  ScalarType a_type = a.scalar_type();
+  ScalarType b_type = b.scalar_type();
+  ScalarType out_type = out.scalar_type();
+
+  if (a.numel() == 1 || b.numel() == 1) {
+    const Tensor* tensor;
+    const Tensor* scalar;
+    ScalarType tensor_type;
+    ScalarType scalar_type;
+    if (a.numel() == 1) {
+      tensor = &b;
+      tensor_type = b_type;
+      scalar = &a;
+      scalar_type = a_type;
+    } else {
+      tensor = &a;
+      tensor_type = a_type;
+      scalar = &b;
+      scalar_type = b_type;
+    }
+    ET_KERNEL_CHECK(
+        ctx,
+        resize_to_broadcast_target_size(a, b, out) == Error::Ok,
+        InvalidArgument,
+        out);
+
+    constexpr auto name = "le.Tensor_out";
+
+    ET_SWITCH_REALB_TYPES(tensor_type, ctx, name, CTYPE, [&]() {
+      ET_SWITCH_REALB_TYPES(scalar_type, ctx, name, CTYPE_SCALAR, [&]() {
+        CTYPE_SCALAR scalar_val = *scalar->const_data_ptr<CTYPE_SCALAR>();
+        CTYPE scalar_casted = static_cast<CTYPE>(scalar_val);
+
+        using Vec = executorch::vec::Vectorized<CTYPE>;
+        if (a.numel() == 1) {
+          executorch::vec::map<CTYPE>(
+              [scalar_casted](Vec x) { return Vec(scalar_casted).le(x); },
+              out.mutable_data_ptr<CTYPE>(),
+              tensor->const_data_ptr<CTYPE>(),
+              out.numel());
+        } else {
+          executorch::vec::map<CTYPE>(
+              [scalar_casted](Vec x) { return x.le(Vec(scalar_casted)); },
+              out.mutable_data_ptr<CTYPE>(),
+              tensor->const_data_ptr<CTYPE>(),
+              out.numel());
+        }
+      });
+    });
+    return out;
+  }
+
   ET_KERNEL_CHECK(ctx, tensors_have_same_shape(a, b), InvalidArgument, out);
 
   // Resize for dynamic shape
@@ -37,10 +90,6 @@ Tensor& opt_le_tensor_out(
       out,
       "Failed to resize output tensor.");
 
-  ScalarType a_type = a.scalar_type();
-  ScalarType b_type = b.scalar_type();
-  ScalarType out_type = out.scalar_type();
-
   if (a_type == b_type && a_type == out_type) {
     ET_SWITCH_REAL_TYPES_AND(
         Bool, out_type, ctx, "le.Tensor_out", CTYPE, [&]() {

kernels/optimized/cpu/targets.bzl

@@ -46,6 +46,7 @@ _OPTIMIZED_ATEN_OPS = (
         name = "op_le",
         deps = [
             "//executorch/kernels/portable/cpu:scalar_utils",
+            "//executorch/kernels/portable/cpu/util:broadcast_util",
         ],
     ),
     op_target(

kernels/test/op_le_test.cpp

@@ -174,3 +174,15 @@ TEST_F(OpLeTensorOutTest, DynamicOutShapeTest) {
   op_le_tensor_out(a, b, out);
   EXPECT_TENSOR_EQ(out, tf.make({2, 2}, {false, true, true, false}));
 }
+
+TEST_F(OpLeTensorOutTest, BroadcastTest) {
+  TensorFactory<ScalarType::Int> tf;
+
+  Tensor a = tf.make(/*sizes=*/{4}, /*data=*/{2, 3, 2, 4});
+  Tensor b = tf.make({1, 1}, {3});
+
+  Tensor out = tf.zeros({1, 4});
+
+  op_le_tensor_out(a, b, out);
+  EXPECT_TENSOR_EQ(out, tf.make({1, 4}, {true, true, true, false}));
+}