WIP: Fix complex div broadcasting in CPU kernels

kernels/optimized/cpu/op_div.cpp

@@ -21,9 +21,6 @@ namespace native {
 namespace {
 
 ScalarType get_common_type(ScalarType a_type, ScalarType b_type) {
-  if (isComplexType(a_type) || isComplexType(b_type)) {
-    return promoteTypes(a_type, b_type);
-  }
   ET_CHECK(!isQIntType(a_type) && !isBitsType(a_type));
   ET_CHECK(!isQIntType(b_type) && !isBitsType(b_type));
 
@@ -64,14 +61,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/op_masked_scatter.cpp

@@ -49,18 +49,39 @@ Tensor& masked_scatter_out(
 
   ET_SWITCH_REALHBBF16_TYPES(in_type, ctx, name, CTYPE, [&]() {
     const CTYPE* const src_data = src.const_data_ptr<CTYPE>();
-    apply_binary_elementwise_fn<CTYPE, bool, CTYPE>(
-        [src_data, &idx, &src_numel, &src_numel_check](
+    const bool src_is_contiguous =
+        is_contiguous_dim_order(src.dim_order().data(), src.dim());
+    const auto compute =
+        [src_data, &src, src_is_contiguous, &idx, &src_numel, &src_numel_check](
             const CTYPE val_in, const bool val_mask) {
-          if (val_mask && idx >= src_numel) {
+          if (!val_mask) {
+            return val_in;
+          }
+          if (idx >= src_numel) {
             src_numel_check = false;
             return val_in;
           }
-          return val_mask ? src_data[idx++] : val_in;
-        },
-        in,
-        mask,
-        out);
+          const size_t src_index = idx++;
+          if (src_is_contiguous) {
+            return src_data[src_index];
+          }
+          size_t src_coord[kTensorDimensionLimit];
+          delinearize_index(src_index, src, src_coord, kTensorDimensionLimit);
+          return src_data[linearize_access_indexes(src_coord, src.dim(), src)];
+        };
+    if (is_contiguous_dim_order(out.dim_order().data(), out.dim())) {
+      apply_binary_elementwise_fn<CTYPE, bool, CTYPE>(compute, in, mask, out);
+    } else {
+      const CTYPE* const in_data = in.const_data_ptr<CTYPE>();
+      const bool* const mask_data = mask.const_data_ptr<bool>();
+      CTYPE* const out_data = out.mutable_data_ptr<CTYPE>();
+      // Source consumption must follow logical output order, even when all
+      // tensors have matching channels-last layouts.
+      for (const auto [unused_linear_index, in_index, mask_index, out_index] :
+           BroadcastIndexesRange<3, true>(out, in, mask, out)) {
+        out_data[out_index] = compute(in_data[in_index], mask_data[mask_index]);
+      }
+    }
   });
 
   ET_KERNEL_CHECK_MSG(

kernels/portable/cpu/util/broadcast_util.h

@@ -247,6 +247,15 @@ size_t linearize_access_indexes(
  * Useful for binary elementwise operators. For each element of the inputs,
  * perform a computation and write to the corresponding element of the output.
  * Tensor broadcasting is applied wherever it is required.
+ *
+ * The tensors must be dense in their respective dimension orders. `out` must
+ * already be resized to the broadcast target shape, and `a` and `b` must each
+ * have at most `out.dim()` dimensions and be broadcastable to `out`. The
+ * strides of `a`, `b`, and `out` are all respected. If all three tensors have
+ * identical sizes and strides, `compute_fun` is called in ascending physical
+ * offset order; otherwise, it is called in the logical row-major coordinate
+ * order of `out`. The call order is not guaranteed to stay the same across
+ * layouts or shapes.
  */
 template <typename CTYPE_A, typename CTYPE_B, typename CTYPE_OUT, typename Op>
 inline void apply_binary_elementwise_fn(
@@ -258,8 +267,19 @@ 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)) {
+  if (a.sizes() == out.sizes() && b.sizes() == out.sizes() &&
+      a.strides() == out.strides() && b.strides() == out.strides()) {
+    for (const auto i : c10::irange(out.numel())) {
+      data_out[i] = compute_fun(data_a[i], data_b[i]);
+    }
+    return;
+  }
+
+  // The first index is a logical linear index that ignores out's strides. As
+  // documented in BroadcastIndexesRange, pass out as an extra input to get its
+  // physical offset instead.
+  for (const auto [unused_linear_index, a_index, b_index, out_index] :
+       BroadcastIndexesRange<3, true>(out, a, b, out)) {
     data_out[out_index] = compute_fun(data_a[a_index], data_b[b_index]);
   }
 }

kernels/portable/cpu/util/test/broadcast_test.cpp

@@ -23,13 +23,139 @@ using executorch::aten::ScalarType;
 using executorch::aten::Tensor;
 using executorch::runtime::ArrayRef;
 using executorch::runtime::testing::TensorFactory;
+using torch::executor::apply_binary_elementwise_fn;
 using torch::executor::broadcast_tensor;
 using torch::executor::delinearize_index;
 using torch::executor::get_broadcast_target_size;
 using torch::executor::linearize_access_indexes;
 using torch::executor::tensor_is_broadcastable_to;
 using torch::executor::tensors_are_broadcastable_between;
 
+TEST(BroadcastUtilTest, BinaryElementwiseDefaultPhysicalOrder) {
+  TensorFactory<ScalarType::Int> tf;
+  Tensor a = tf.make({2, 2}, {1, 2, 3, 4});
+  Tensor b = tf.make({2, 2}, {11, 13, 17, 19});
+  Tensor out = tf.zeros({2, 2});
+  std::vector<int32_t> visited;
+  apply_binary_elementwise_fn<int32_t, int32_t, int32_t>(
+      [&visited](int32_t x, int32_t y) {
+        visited.push_back(x);
+        return x - y;
+      },
+      a,
+      b,
+      out);
+  EXPECT_EQ(visited, (std::vector<int32_t>{1, 2, 3, 4}));
+  EXPECT_TENSOR_EQ(out, tf.make({2, 2}, {-10, -11, -14, -15}));
+}
+
+TEST(BroadcastUtilTest, BinaryElementwiseChannelsLastPhysicalOrder) {
+  TensorFactory<ScalarType::Int> tf;
+  Tensor a = tf.channels_last_like(tf.make({1, 2, 1, 2}, {1, 2, 3, 4}));
+  Tensor b = tf.channels_last_like(tf.make({1, 2, 1, 2}, {11, 13, 17, 19}));
+  Tensor out = tf.channels_last_like(tf.zeros({1, 2, 1, 2}));
+  std::vector<int32_t> visited;
+  apply_binary_elementwise_fn<int32_t, int32_t, int32_t>(
+      [&visited](int32_t x, int32_t y) {
+        visited.push_back(x);
+        return x - y;
+      },
+      a,
+      b,
+      out);
+  EXPECT_EQ(visited, (std::vector<int32_t>{1, 3, 2, 4}));
+  EXPECT_TENSOR_EQ(
+      out, tf.channels_last_like(tf.make({1, 2, 1, 2}, {-10, -11, -14, -15})));
+}
+
+TEST(BroadcastUtilTest, BinaryElementwiseDifferentStrides) {
+  TensorFactory<ScalarType::Int> tf;
+  Tensor a = tf.make({1, 2, 1, 2}, {1, 2, 3, 4});
+  Tensor b = tf.make({1, 2, 1, 2}, {11, 13, 17, 19});
+  Tensor a_channels_last = tf.channels_last_like(a);
+  Tensor b_channels_last = tf.channels_last_like(b);
+  Tensor expected = tf.make({1, 2, 1, 2}, {-10, -11, -14, -15});
+  Tensor expected_channels_last = tf.channels_last_like(expected);
+  const auto subtract = [](int32_t x, int32_t y) { return x - y; };
+
+  Tensor out_channels_last_1 = tf.channels_last_like(tf.zeros({1, 2, 1, 2}));
+  apply_binary_elementwise_fn<int32_t, int32_t, int32_t>(
+      subtract, a, b_channels_last, out_channels_last_1);
+  EXPECT_TENSOR_EQ(out_channels_last_1, expected_channels_last);
+
+  Tensor out_channels_last_2 = tf.channels_last_like(tf.zeros({1, 2, 1, 2}));
+  apply_binary_elementwise_fn<int32_t, int32_t, int32_t>(
+      subtract, a_channels_last, b, out_channels_last_2);
+  EXPECT_TENSOR_EQ(out_channels_last_2, expected_channels_last);
+
+  Tensor out_contiguous = tf.zeros({1, 2, 1, 2});
+  apply_binary_elementwise_fn<int32_t, int32_t, int32_t>(
+      subtract, a_channels_last, b_channels_last, out_contiguous);
+  EXPECT_TENSOR_EQ(out_contiguous, expected);
+}
+
+TEST(BroadcastUtilTest, BinaryElementwiseChannelsLastBroadcast) {
+  TensorFactory<ScalarType::Int> tf;
+  Tensor a = tf.channels_last_like(tf.make({1, 2, 1, 2}, {1, 2, 3, 4}));
+  Tensor b = tf.channels_last_like(tf.make({1, 2, 1, 1}, {10, 20}));
+  Tensor out = tf.channels_last_like(tf.zeros({1, 2, 1, 2}));
+  const auto subtract = [](int32_t x, int32_t y) { return x - y; };
+
+  apply_binary_elementwise_fn<int32_t, int32_t, int32_t>(subtract, a, b, out);
+  EXPECT_TENSOR_EQ(
+      out, tf.channels_last_like(tf.make({1, 2, 1, 2}, {-9, -8, -17, -16})));
+  apply_binary_elementwise_fn<int32_t, int32_t, int32_t>(subtract, b, a, out);
+  EXPECT_TENSOR_EQ(
+      out, tf.channels_last_like(tf.make({1, 2, 1, 2}, {9, 8, 17, 16})));
+}
+
+TEST(BroadcastUtilTest, BinaryElementwiseMixedRank) {
+  TensorFactory<ScalarType::Int> tf;
+  for (bool channels_last : {false, true}) {
+    Tensor a = tf.make({2, 3, 1, 1}, {1, 2, 3, 4, 5, 6});
+    Tensor b = tf.make({1, 2, 3, 1, 1}, {10, 20, 30, 40, 50, 60});
+    Tensor out = tf.zeros({1, 2, 3, 1, 1});
+    Tensor expected = tf.make({1, 2, 3, 1, 1}, {-9, -18, -27, -36, -45, -54});
+    if (channels_last) {
+      a = tf.channels_last_like(a);
+      b = tf.channels_last_like(b);
+      out = tf.channels_last_like(out);
+      expected = tf.channels_last_like(expected);
+    }
+    apply_binary_elementwise_fn<int32_t, int32_t, int32_t>(
+        [](int32_t x, int32_t y) { return x - y; }, a, b, out);
+    EXPECT_TENSOR_EQ(out, expected);
+  }
+}
+
+TEST(BroadcastUtilTest, BinaryElementwiseBroadcastBothInputs) {
+  TensorFactory<ScalarType::Int> tf;
+  Tensor a = tf.make({2, 1}, {1, 4});
+  Tensor b = tf.make({1, 2}, {10, 20});
+  Tensor out = tf.zeros({2, 2});
+  apply_binary_elementwise_fn<int32_t, int32_t, int32_t>(
+      [](int32_t x, int32_t y) { return x - y; }, a, b, out);
+  EXPECT_TENSOR_EQ(out, tf.make({2, 2}, {-9, -19, -6, -16}));
+}
+
+TEST(BroadcastUtilTest, BinaryElementwiseEmpty) {
+  TensorFactory<ScalarType::Int> tf;
+  Tensor a = tf.zeros({0, 2});
+  Tensor b = tf.ones({1, 2});
+  Tensor out = tf.zeros({0, 2});
+  bool called = false;
+  apply_binary_elementwise_fn<int32_t, int32_t, int32_t>(
+      [&called](int32_t x, int32_t y) {
+        called = true;
+        return x - y;
+      },
+      a,
+      b,
+      out);
+  EXPECT_FALSE(called);
+  EXPECT_TENSOR_EQ(out, a);
+}
+
 TEST(BroadcastUtilTest, BroadcastTensor) {
   TensorFactory<ScalarType::Int> tf;
 

kernels/test/CMakeLists.txt

@@ -246,6 +246,7 @@ set(all_test_sources
     "op_logit_test.cpp"
     "op_lt_test.cpp"
     "op_masked_fill_test.cpp"
+    "op_masked_scatter_test.cpp"
     "op_masked_select_test.cpp"
     "op_max_test.cpp"
     "op_max_pool2d_with_indices_test.cpp"