Fix optimized binary op broadcast rank mismatch (pytorch#21052)

### Summary

The optimized last-dimension and planned broadcast paths resize `out` to
one input's shape after the operator has already selected the broadcast
target. When leading dimensions of size one make the target rank larger,
that second resize tries to change the immutable rank and fails at
runtime.

Resize `out` to the broadcast target in both paths and add regression
coverage for add and mul in both argument orders.

Fixes pytorch#10959

### Test plan

`./cmake-out/kernels/test/optimized_kernels_test --gtest_brief=1`

`lintrunner kernels/optimized/cpu/binary_ops.cpp
kernels/optimized/cpu/binary_ops.h kernels/test/op_add_test.cpp
kernels/test/op_mul_test.cpp`

cc @larryliu0820 @manuelcandales

Author: Kilbex

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);
+}