Fix complex div broadcasting in CPU kernels

kernels/optimized/cpu/op_div.cpp

@@ -64,14 +64,14 @@ Tensor& opt_div_out(
 
   // Handle complex types
   if (isComplexType(a_type) || isComplexType(b_type)) {
-    ScalarType common_type = get_common_type(a_type, b_type);
-    ET_SWITCH_COMPLEX_TYPES(common_type, ctx, op_name, CTYPE, [&]() {
-      const CTYPE* a_data = a.const_data_ptr<CTYPE>();
-      const CTYPE* b_data = b.const_data_ptr<CTYPE>();
-      CTYPE* out_data = out.mutable_data_ptr<CTYPE>();
-      for (size_t i = 0; i < out.numel(); ++i) {
-        out_data[i] = a_data[i] / b_data[i];
-      }
+    ET_KERNEL_CHECK(
+        ctx, a_type == b_type && a_type == out_type, InvalidArgument, out);
+    ET_SWITCH_COMPLEX_TYPES(out_type, ctx, op_name, CTYPE, [&]() {
+      apply_binary_elementwise_fn<CTYPE, CTYPE, CTYPE>(
+          [](const CTYPE val_a, const CTYPE val_b) { return val_a / val_b; },
+          a,
+          b,
+          out);
     });
     return out;
   }

kernels/portable/cpu/op_div.cpp

@@ -58,13 +58,18 @@ Tensor& div_out(
   ET_DEFINE_OPERATOR_NAME(op_name, "div.out");
 
   if (executorch::runtime::isComplexType(common_type)) {
-    ET_SWITCH_COMPLEX_TYPES(common_type, ctx, op_name, CTYPE, [&]() {
-      const CTYPE* a_data = a.const_data_ptr<CTYPE>();
-      const CTYPE* b_data = b.const_data_ptr<CTYPE>();
-      CTYPE* out_data = out.mutable_data_ptr<CTYPE>();
-      for (ssize_t i = 0; i < out.numel(); ++i) {
-        out_data[i] = a_data[i] / b_data[i];
-      }
+    ET_KERNEL_CHECK(
+        ctx,
+        a.scalar_type() == b.scalar_type() &&
+            a.scalar_type() == out.scalar_type(),
+        InvalidArgument,
+        out);
+    ET_SWITCH_COMPLEX_TYPES(out.scalar_type(), ctx, op_name, CTYPE, [&]() {
+      apply_binary_elementwise_fn<CTYPE, CTYPE, CTYPE>(
+          [](const CTYPE val_a, const CTYPE val_b) { return val_a / val_b; },
+          a,
+          b,
+          out);
     });
   } else {
     // Compute Dtype for real types

kernels/portable/cpu/util/broadcast_util.h

@@ -258,8 +258,9 @@ inline void apply_binary_elementwise_fn(
   const CTYPE_B* const data_b = b.const_data_ptr<CTYPE_B>();
   CTYPE_OUT* const data_out = out.mutable_data_ptr<CTYPE_OUT>();
 
-  for (const auto [out_index, a_index, b_index] :
-       BroadcastIndexesRange<2>(out, a, b)) {
+  // Use an input index to respect output strides during broadcasting.
+  for (const auto [_, a_index, b_index, out_index] :
+       BroadcastIndexesRange<3>(out, a, b, out)) {
     data_out[out_index] = compute_fun(data_a[a_index], data_b[b_index]);
   }
 }

kernels/test/op_div_test.cpp

@@ -31,6 +31,80 @@ class OpDivOutTest : public OperatorTest {
     return torch::executor::aten::div_outf(context_, a, b, out);
   }
 
+  template <typename CTYPE, ScalarType DTYPE>
+  void test_complex_scalar_broadcast() {
+    TensorFactory<DTYPE> tf;
+    Tensor a =
+        tf.make({2, 2}, {CTYPE(1, 1), CTYPE(2, -2), CTYPE(4, 4), CTYPE(8, -8)});
+    Tensor right_expected = tf.make(
+        {2, 2}, {CTYPE(0.5, 0), CTYPE(0, -1), CTYPE(2, 0), CTYPE(0, -4)});
+    Tensor left_expected = tf.make(
+        {2, 2}, {CTYPE(2, 0), CTYPE(0, 1), CTYPE(0.5, 0), CTYPE(0, 0.25)});
+    const std::vector<std::vector<int32_t>> scalar_sizes = {{}, {1}, {1, 1}};
+    for (const auto& sizes : scalar_sizes) {
+      Tensor b = tf.make(sizes, {CTYPE(2, 2)});
+      Tensor out = tf.full({2, 2}, CTYPE(0, 0));
+      EXPECT_TENSOR_CLOSE(op_div_out(a, b, out), right_expected);
+      EXPECT_TENSOR_CLOSE(op_div_out(b, a, out), left_expected);
+    }
+
+    Tensor scalar_a = tf.make({}, {CTYPE(1, 1)});
+    Tensor scalar_b = tf.make({}, {CTYPE(2, 2)});
+    Tensor out = tf.full({}, CTYPE(0, 0));
+    EXPECT_TENSOR_CLOSE(
+        op_div_out(scalar_a, scalar_b, out), tf.make({}, {CTYPE(0.5, 0)}));
+  }
+
+  template <typename CTYPE, ScalarType DTYPE>
+  void test_complex_broadcast() {
+    TensorFactory<DTYPE> tf;
+    Tensor a = tf.make({2, 1}, {CTYPE(2, 2), CTYPE(4, -4)});
+    Tensor b = tf.make({1, 2}, {CTYPE(1, 1), CTYPE(2, -2)});
+    Tensor out = tf.full({2, 2}, CTYPE(0, 0));
+    EXPECT_TENSOR_CLOSE(
+        op_div_out(a, b, out),
+        tf.make({2, 2}, {CTYPE(2, 0), CTYPE(0, 1), CTYPE(0, -4), CTYPE(2, 0)}));
+
+    Tensor empty_a = tf.full({0, 2}, CTYPE(0, 0));
+    Tensor empty_out = tf.full({0, 2}, CTYPE(0, 0));
+    EXPECT_TENSOR_CLOSE(op_div_out(empty_a, b, empty_out), empty_a);
+  }
+
+  template <typename CTYPE, ScalarType DTYPE>
+  void test_complex_channels_last_broadcast() {
+    TensorFactory<DTYPE> tf;
+    Tensor a = tf.channels_last_like(tf.make(
+        {1, 2, 1, 2},
+        {CTYPE(2, 2), CTYPE(4, -4), CTYPE(8, 8), CTYPE(16, -16)}));
+    Tensor b = tf.channels_last_like(
+        tf.make({1, 2, 1, 1}, {CTYPE(1, 1), CTYPE(2, 2)}));
+    Tensor out = tf.channels_last_like(tf.full({1, 2, 1, 2}, CTYPE(0, 0)));
+    Tensor right_expected = tf.channels_last_like(tf.make(
+        {1, 2, 1, 2}, {CTYPE(2, 0), CTYPE(0, -4), CTYPE(4, 0), CTYPE(0, -8)}));
+    Tensor left_expected = tf.channels_last_like(tf.make(
+        {1, 2, 1, 2},
+        {CTYPE(0.5, 0), CTYPE(0, 0.25), CTYPE(0.25, 0), CTYPE(0, 0.125)}));
+
+    EXPECT_TENSOR_CLOSE(op_div_out(a, b, out), right_expected);
+    EXPECT_TENSOR_CLOSE(op_div_out(b, a, out), left_expected);
+  }
+
+  template <ScalarType DTYPE_A, ScalarType DTYPE_B, ScalarType DTYPE_OUT>
+  void test_mismatched_complex_dtypes_dies() {
+    ET_SKIP_IF(
+        torch::executor::testing::SupportedFeatures::get()->is_aten,
+        "ATen supports mixed complex dtypes");
+    TensorFactory<DTYPE_A> tf_a;
+    TensorFactory<DTYPE_B> tf_b;
+    TensorFactory<DTYPE_OUT> tf_out;
+
+    Tensor a = tf_a.ones({2});
+    Tensor b = tf_b.ones({2});
+    Tensor out = tf_out.zeros({2});
+
+    ET_EXPECT_KERNEL_FAILURE(context_, op_div_out(a, b, out));
+  }
+
   template <ScalarType DTYPE_A, ScalarType DTYPE_B, ScalarType DTYPE_OUT>
   void test_div() {
     TensorFactory<DTYPE_A> tf_a;
@@ -607,6 +681,55 @@ TEST_F(OpDivScalarOutTest, OptimizedSanityCheck) {
 // Complex Type Tests
 //
 
+TEST_F(OpDivOutTest, ComplexScalarBroadcast) {
+#define TEST_ENTRY(ctype, dtype) \
+  test_complex_scalar_broadcast<ctype, ScalarType::dtype>();
+  ET_FORALL_COMPLEX_TYPES(TEST_ENTRY);
+#undef TEST_ENTRY
+}
+
+TEST_F(OpDivOutTest, ComplexBroadcast) {
+#define TEST_ENTRY(ctype, dtype) \
+  test_complex_broadcast<ctype, ScalarType::dtype>();
+  ET_FORALL_COMPLEX_TYPES(TEST_ENTRY);
+#undef TEST_ENTRY
+}
+
+TEST_F(OpDivOutTest, ComplexChannelsLastBroadcast) {
+#define TEST_ENTRY(ctype, dtype) \
+  test_complex_channels_last_broadcast<ctype, ScalarType::dtype>();
+  ET_FORALL_COMPLEX_TYPES(TEST_ENTRY);
+#undef TEST_ENTRY
+}
+
+TEST_F(OpDivOutTest, ComplexMismatchedInputDtypesDies) {
+  test_mismatched_complex_dtypes_dies<
+      ScalarType::ComplexFloat,
+      ScalarType::ComplexDouble,
+      ScalarType::ComplexDouble>();
+}
+
+TEST_F(OpDivOutTest, ComplexMismatchedOutputDtypeDies) {
+  test_mismatched_complex_dtypes_dies<
+      ScalarType::ComplexFloat,
+      ScalarType::ComplexFloat,
+      ScalarType::ComplexDouble>();
+}
+
+TEST_F(OpDivOutTest, ComplexAndRealInputsDies) {
+  test_mismatched_complex_dtypes_dies<
+      ScalarType::ComplexFloat,
+      ScalarType::Float,
+      ScalarType::ComplexFloat>();
+}
+
+TEST_F(OpDivOutTest, RealAndComplexInputsDies) {
+  test_mismatched_complex_dtypes_dies<
+      ScalarType::Float,
+      ScalarType::ComplexFloat,
+      ScalarType::ComplexFloat>();
+}
+
 TEST_F(OpDivOutTest, ComplexFloatBasic) {
   TensorFactory<ScalarType::ComplexFloat> tf;