Fix optimized binary op broadcast rank mismatch

kernels/optimized/cpu/binary_ops.cpp

@@ -32,10 +32,9 @@ std::optional<BroadcastElementwisePlan> plan_broadcast_elementwise(
     plan.lhs = &a;
     plan.rhs = &b;
   }
-  auto error = resize_tensor(out, plan.lhs->sizes());
   ET_KERNEL_CHECK_MSG(
       ctx,
-      error == Error::Ok,
+      resize_to_broadcast_target_size(a, b, out) == Error::Ok,
       InvalidArgument,
       std::nullopt,
       "Failed to resize output tensor.");

kernels/optimized/cpu/binary_ops.h

@@ -11,6 +11,7 @@
 #include <executorch/kernels/optimized/vec/functional.h>
 #include <executorch/kernels/portable/cpu/scalar_utils.h>
 #include <executorch/kernels/portable/cpu/util/broadcast_indexes_range.h>
+#include <executorch/kernels/portable/cpu/util/broadcast_util.h>
 #include <executorch/runtime/kernel/kernel_includes.h>
 
 #include <optional>
@@ -189,10 +190,9 @@ Tensor& handle_last_dim_broadcast_elementwise(
     lhs = &a;
     rhs = &b;
   }
-  auto error = resize_tensor(out, lhs->sizes());
   ET_KERNEL_CHECK_MSG(
       ctx,
-      error == Error::Ok,
+      resize_to_broadcast_target_size(a, b, out) == Error::Ok,
       InvalidArgument,
       out,
       "Failed to resize output tensor.");

kernels/test/op_add_test.cpp

@@ -604,6 +604,29 @@ TEST_F(OpAddOutKernelTest, BroadcastBToA) {
   EXPECT_TENSOR_CLOSE(op_add_out(a, b, 1.0, out), expected);
 }
 
+TEST_F(OpAddOutKernelTest, BroadcastLastDimRankMismatch) {
+  TensorFactory<ScalarType::Float> tf;
+  Tensor a = tf.make({3, 4}, /*data=*/{1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12});
+  Tensor b = tf.make({1, 3, 1}, /*data=*/{10, 20, 30});
+  Tensor out = tf.zeros({1, 3, 4});
+
+  Tensor expected = tf.make(
+      {1, 3, 4}, /*data=*/{11, 12, 13, 14, 25, 26, 27, 28, 39, 40, 41, 42});
+  EXPECT_TENSOR_CLOSE(op_add_out(a, b, 1.0, out), expected);
+  EXPECT_TENSOR_CLOSE(op_add_out(b, a, 1.0, out), expected);
+}
+
+TEST_F(OpAddOutKernelTest, Broadcast2dBy1dRankMismatch) {
+  TensorFactory<ScalarType::Float> tf;
+  Tensor a = tf.make({2, 3}, /*data=*/{1, 2, 3, 4, 5, 6});
+  Tensor b = tf.make({1, 1, 3}, /*data=*/{10, 20, 30});
+  Tensor out = tf.zeros({1, 2, 3});
+
+  Tensor expected = tf.make({1, 2, 3}, /*data=*/{11, 22, 33, 14, 25, 36});
+  EXPECT_TENSOR_CLOSE(op_add_out(a, b, 1.0, out), expected);
+  EXPECT_TENSOR_CLOSE(op_add_out(b, a, 1.0, out), expected);
+}
+
 //
 // Death Tests
 //

kernels/test/op_mul_test.cpp

@@ -923,3 +923,27 @@ TEST_F(OpMulOutTest, BroadcastDimensionMismatchWithDifferentTypes) {
     EXPECT_TENSOR_EQ(result, expected);
   }
 }
+
+TEST_F(OpMulOutTest, BroadcastLastDimRankMismatch) {
+  TensorFactory<ScalarType::Float> tf;
+  Tensor a = tf.make({3, 4}, /*data=*/{1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12});
+  Tensor b = tf.make({1, 3, 1}, /*data=*/{10, 20, 30});
+  Tensor out = tf.zeros({1, 3, 4});
+
+  Tensor expected = tf.make(
+      {1, 3, 4},
+      /*data=*/{10, 20, 30, 40, 100, 120, 140, 160, 270, 300, 330, 360});
+  EXPECT_TENSOR_CLOSE(op_mul_out(a, b, out), expected);
+  EXPECT_TENSOR_CLOSE(op_mul_out(b, a, out), expected);
+}
+
+TEST_F(OpMulOutTest, Broadcast2dBy1dRankMismatch) {
+  TensorFactory<ScalarType::Float> tf;
+  Tensor a = tf.make({2, 3}, /*data=*/{1, 2, 3, 4, 5, 6});
+  Tensor b = tf.make({1, 1, 3}, /*data=*/{10, 20, 30});
+  Tensor out = tf.zeros({1, 2, 3});
+
+  Tensor expected = tf.make({1, 2, 3}, /*data=*/{10, 40, 90, 40, 100, 180});
+  EXPECT_TENSOR_CLOSE(op_mul_out(a, b, out), expected);
+  EXPECT_TENSOR_CLOSE(op_mul_out(b, a, out), expected);
+}