Fix bug in le broadcast for single element broadcast

kernels/optimized/cpu/op_le.cpp

@@ -30,57 +30,8 @@ Tensor& opt_le_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 = at::vec::Vectorized<CTYPE>;
-        if (a.numel() == 1) {
-          at::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 {
-          at::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;
-  }
-
   // Check for optimized broadcast paths
   auto selected_optimized_path = select_optimized_path(a, b, out);
   if (selected_optimized_path == ElementwiseOptimizedPath::kTreatAs1d) {

kernels/test/op_le_test.cpp

@@ -1112,3 +1112,70 @@ TEST_F(OpLeTensorOutTest, Broadcast22dBy1dReverseTest) {
 
   EXPECT_TENSOR_EQ(out, tf_bool.make({3, 4}, expected_data));
 }
+
+TEST_F(OpLeTensorOutTest, MonotonicIncreasingVsScalarBroadcastTest) {
+  TensorFactory<ScalarType::Int> tf;
+  TensorFactory<ScalarType::Bool> tf_bool;
+
+  // Test case: 1D tensor [0, 1, 2, ..., 63] vs 2D tensor [1, 1] with value 2
+  std::vector<int32_t> lhs_data;
+  for (int i = 0; i < 64; ++i) {
+    lhs_data.push_back(i);
+  }
+
+  Tensor lhs = tf.make({64}, lhs_data);
+  Tensor rhs = tf.make({1, 1}, {2});
+  Tensor out = tf_bool.zeros({1, 64});
+
+  op_le_tensor_out(lhs, rhs, out);
+
+  // Expected: [0, 1, 2] <= 2 should be [true, true, true], rest false
+  using ctype =
+      executorch::runtime::testing::internal::ScalarTypeToCppTypeWrapper<
+          ScalarType::Bool>::ctype;
+  std::vector<ctype> expected_data;
+  for (int i = 0; i < 64; ++i) {
+    expected_data.push_back(i <= 2);
+  }
+
+  EXPECT_TENSOR_EQ(out, tf_bool.make({1, 64}, expected_data));
+
+  // Test with rhs value 4
+  rhs = tf.make({1, 1}, {4});
+  out = tf_bool.zeros({1, 64});
+
+  op_le_tensor_out(lhs, rhs, out);
+
+  expected_data.clear();
+  for (int i = 0; i < 64; ++i) {
+    expected_data.push_back(i <= 4);
+  }
+
+  EXPECT_TENSOR_EQ(out, tf_bool.make({1, 64}, expected_data));
+
+  // Test with rhs value 10
+  rhs = tf.make({1, 1}, {10});
+  out = tf_bool.zeros({1, 64});
+
+  op_le_tensor_out(lhs, rhs, out);
+
+  expected_data.clear();
+  for (int i = 0; i < 64; ++i) {
+    expected_data.push_back(i <= 10);
+  }
+
+  EXPECT_TENSOR_EQ(out, tf_bool.make({1, 64}, expected_data));
+
+  // Test with rhs value 32
+  rhs = tf.make({1, 1}, {32});
+  out = tf_bool.zeros({1, 64});
+
+  op_le_tensor_out(lhs, rhs, out);
+
+  expected_data.clear();
+  for (int i = 0; i < 64; ++i) {
+    expected_data.push_back(i <= 32);
+  }
+
+  EXPECT_TENSOR_EQ(out, tf_bool.make({1, 64}, expected_data));
+}