[ExecuTorch] Handle rank 0 tensors correctly in optimized add/sub/div/mul

Differential Revision: D62310838

Pull Request resolved: #5140

Author: swolchok

kernels/optimized/cpu/op_add.cpp

@@ -85,13 +85,12 @@ Tensor& opt_add_out(
   if (b.numel() == 1) {
     if (a_type == b_type && a_type == out_type && a_type != ScalarType::Half &&
         a_type != ScalarType::BFloat16) {
-      auto error = resize_tensor(out, a.sizes());
-      ET_KERNEL_CHECK_MSG(
+      ET_KERNEL_CHECK(
           ctx,
-          error == Error::Ok,
+          resize_to_broadcast_target_size(a, b, out) == Error::Ok,
           InvalidArgument,
-          out,
-          "Failed to resize output tensor.");
+          out);
+
       ET_SWITCH_REALB_TYPES(a_type, ctx, "add.out", CTYPE, [&]() {
         ET_SWITCH_REALB_TYPES(b_type, ctx, "add.out", CTYPE_B, [&]() {
           CTYPE alpha_val;

kernels/optimized/cpu/op_div.cpp

@@ -66,13 +66,11 @@ Tensor& opt_div_out(
         scalar = &b;
         scalar_type = b_type;
       }
-      auto error = resize_tensor(out, tensor->sizes());
-      ET_KERNEL_CHECK_MSG(
+      ET_KERNEL_CHECK(
           ctx,
-          error == Error::Ok,
+          resize_to_broadcast_target_size(a, b, out) == Error::Ok,
           InvalidArgument,
-          out,
-          "Failed to resize output tensor.");
+          out);
       ET_SWITCH_REALB_TYPES(tensor_type, ctx, "div.out", CTYPE, [&]() {
         ET_SWITCH_REALB_TYPES(scalar_type, ctx, "div.out", CTYPE_SCALAR, [&]() {
           CTYPE_SCALAR scalar_val = *scalar->const_data_ptr<CTYPE_SCALAR>();

kernels/optimized/cpu/op_mul.cpp

@@ -82,13 +82,12 @@ Tensor& opt_mul_out(
   if (b.numel() == 1) {
     if (a_type == b_type && a_type == out_type && a_type != ScalarType::Half &&
         a_type != ScalarType::BFloat16) {
-      auto error = resize_tensor(out, a.sizes());
-      ET_KERNEL_CHECK_MSG(
+      ET_KERNEL_CHECK(
           ctx,
-          error == Error::Ok,
+          resize_to_broadcast_target_size(a, b, out) == Error::Ok,
           InvalidArgument,
-          out,
-          "Failed to resize output tensor.");
+          out);
+
       ET_SWITCH_REALB_TYPES(a_type, ctx, "mul.out", CTYPE, [&]() {
         ET_SWITCH_REALB_TYPES(b_type, ctx, "mul.out", CTYPE_B, [&]() {
           CTYPE_B b_val = *b.const_data_ptr<CTYPE_B>();

kernels/optimized/cpu/op_sub.cpp

@@ -101,13 +101,11 @@ Tensor& opt_sub_out(
         scalar = &b;
         scalar_type = b_type;
       }
-      auto error = resize_tensor(out, tensor->sizes());
-      ET_KERNEL_CHECK_MSG(
+      ET_KERNEL_CHECK(
           ctx,
-          error == Error::Ok,
+          resize_to_broadcast_target_size(a, b, out) == Error::Ok,
           InvalidArgument,
-          out,
-          "Failed to resize output tensor.");
+          out);
       ET_SWITCH_REAL_TYPES(tensor_type, ctx, "sub.out", CTYPE, [&]() {
         ET_SWITCH_REAL_TYPES(scalar_type, ctx, "sub.out", CTYPE_SCALAR, [&]() {
           CTYPE alpha_val;

kernels/test/op_add_test.cpp

@@ -352,6 +352,23 @@ TEST_F(OpAddOutKernelTest, BroadcastOneElementTensorTypePromotion) {
   EXPECT_TENSOR_EQ(out, expected);
 }
 
+TEST_F(OpAddOutKernelTest, BroadcastOneElementRank0Tensor) {
+  TensorFactory<ScalarType::Float> tf;
+
+  Tensor a = tf.make({1}, {5});
+  Tensor b = tf.make({}, {2});
+
+  Tensor out = tf.zeros({1});
+
+  op_add_out(a, b, 1, out);
+
+  Tensor ret = tf.make({1}, {7});
+  EXPECT_TENSOR_EQ(out, ret);
+
+  op_add_out(b, a, 1, out);
+  EXPECT_TENSOR_EQ(out, ret);
+}
+
 //
 // Death Tests
 //

kernels/test/op_div_test.cpp

@@ -237,6 +237,25 @@ TEST_F(OpDivOutTest, BroadcastScalarSupported2) {
   EXPECT_TENSOR_EQ(out, ret);
 }
 
+TEST_F(OpDivOutTest, BroadcastScalarRank0Supported) {
+  TensorFactory<ScalarType::Float> tf;
+
+  Tensor a = tf.make({1}, {8});
+  Tensor b = tf.make({}, {2});
+
+  Tensor out = tf.zeros({1});
+
+  op_div_out(a, b, out);
+
+  Tensor ret = tf.make({1}, {4});
+  EXPECT_TENSOR_EQ(out, ret);
+
+  op_div_out(b, a, out);
+
+  ret = tf.make({1}, {0.25});
+  EXPECT_TENSOR_EQ(out, ret);
+}
+
 TEST_F(OpDivOutTest, BroadcastDimSizeIsOneAB) {
   TensorFactory<ScalarType::Float> tf;
 

kernels/test/op_mul_test.cpp

@@ -182,6 +182,23 @@ class OpMulOutTest : public OperatorTest {
     EXPECT_TENSOR_CLOSE(op_mul_out(a, b, out), expected);
     EXPECT_TENSOR_CLOSE(op_mul_out(b, a, out), expected);
   }
+
+  template <ScalarType DTYPE>
+  void test_both_scalar_input_broadcast() {
+    TensorFactory<DTYPE> tf_a;
+
+    // a is a rank-1 scalar and b is a rank-0 scalar
+    Tensor a = tf_a.make({1}, /*data=*/{2});
+    Tensor b = tf_a.make({}, /*data=*/{2});
+
+    // Destination for output of mul.
+    Tensor out = tf_a.make({1}, /*data=*/{2});
+    Tensor expected = tf_a.make({1}, /*data=*/{4});
+
+    // Check that it matches the expected output.
+    EXPECT_TENSOR_CLOSE(op_mul_out(a, b, out), expected);
+    EXPECT_TENSOR_CLOSE(op_mul_out(b, a, out), expected);
+  }
 };
 
 class OpMulScalarOutTest : public OperatorTest {
@@ -309,6 +326,12 @@ TEST_F(OpMulOutTest, ScalarInputBroadcastTest) {
   test_scalar_input_broadcast<ScalarType::BFloat16>();
 }
 
+TEST_F(OpMulOutTest, BothScalarInputBroadcastTest) {
+  test_both_scalar_input_broadcast<ScalarType::Int>();
+  test_both_scalar_input_broadcast<ScalarType::Half>();
+  test_both_scalar_input_broadcast<ScalarType::BFloat16>();
+}
+
 TEST_F(OpMulOutTest, MismatchedOutputShapesDies) {
   if (SupportedFeatures::get()->is_aten) {
     GTEST_SKIP() << "ATen currently supports mismatched shapes";

kernels/test/op_sub_test.cpp

@@ -206,6 +206,25 @@ TEST_F(OpSubOutTest, BroadcastScalarSupported2) {
   EXPECT_TENSOR_EQ(out, ret);
 }
 
+TEST_F(OpSubOutTest, BroadcastScalarRank0Supported) {
+  TensorFactory<ScalarType::Float> tf;
+
+  Tensor a = tf.make({1}, {5});
+  Tensor b = tf.make({}, {2});
+
+  Tensor out = tf.zeros({1});
+
+  op_sub_out(a, b, 1, out);
+
+  Tensor ret = tf.make({1}, {3});
+  EXPECT_TENSOR_EQ(out, ret);
+
+  op_sub_out(b, a, 1, out);
+
+  ret = tf.make({1}, {-3});
+  EXPECT_TENSOR_EQ(out, ret);
+}
+
 //
 // Death Tests
 //