Make BroadcastIndexesRange efficient if there is no broadcasting

swolchok · swolchok · commit c35fdea82bc3 · 2025-03-18T09:22:16.000-07:00
It seems to be OK to just check for broadcasting when creating the iterator and then leave a highly-predictable branch inside the loop. Avoids code near-duplication to handle both broadcast and non-broadcast cases. ghstack-source-id: 8be7c18 ghstack-comment-id: 2725945271 Pull Request resolved: #9298
diff --git a/kernels/optimized/cpu/op_where.cpp b/kernels/optimized/cpu/op_where.cpp
@@ -48,42 +48,24 @@ Tensor& opt_where_out(
       cond.scalar_type() == ScalarType::Bool) {
     auto out_numel = out.numel();
     ET_SWITCH_REALB_TYPES(compute_type, ctx, op_name, CTYPE_COMPUTE, [&]() {
-      const bool a_is_broadcasted = !out.sizes().equals(a.sizes());
-      const bool b_is_broadcasted = !out.sizes().equals(b.sizes());
-      const bool cond_is_broadcasted = !out.sizes().equals(cond.sizes());
-      const bool any_is_broadcasted =
-          (a_is_broadcasted || b_is_broadcasted || cond_is_broadcasted);
       const CTYPE_COMPUTE* const data_a = a.const_data_ptr<CTYPE_COMPUTE>();
       const CTYPE_COMPUTE* const data_b = b.const_data_ptr<CTYPE_COMPUTE>();
       const bool* const data_cond = cond.const_data_ptr<bool>();
       CTYPE_COMPUTE* const data_out = out.data_ptr<CTYPE_COMPUTE>();
-      if (any_is_broadcasted) {
-        executorch::extension::parallel_for(
-            0,
-            out_numel,
-            ::executorch::extension::internal::GRAIN_SIZE,
-            [&](const auto begin, const auto end) {
-              auto range = BroadcastIndexesRange<3>(out, a, b, cond);
-              auto begin_it = range.begin();
-              begin_it += begin;
-              for (; (*begin_it)[0] < end; ++begin_it) {
-                const auto [out_index, a_index, b_index, cond_index] =
-                    *begin_it;
-                data_out[out_index] =
-                    data_cond[cond_index] ? data_a[a_index] : data_b[b_index];
-              }
-            });
-      } else {
-        executorch::extension::parallel_for(
-            0,
-            out_numel,
-            ::executorch::extension::internal::GRAIN_SIZE,
-            [&](const auto begin, const auto end) {
-              for (const auto i : c10::irange(begin, end)) {
-                data_out[i] = data_cond[i] ? data_a[i] : data_b[i];
-              }
-            });
-      }
+      executorch::extension::parallel_for(
+          0,
+          out_numel,
+          ::executorch::extension::internal::GRAIN_SIZE,
+          [&](const auto begin, const auto end) {
+            auto range = BroadcastIndexesRange<3>(out, a, b, cond);
+            auto begin_it = range.begin();
+            begin_it += begin;
+            for (; (*begin_it)[0] < end; ++begin_it) {
+              const auto [out_index, a_index, b_index, cond_index] = *begin_it;
+              data_out[out_index] =
+                  data_cond[cond_index] ? data_a[a_index] : data_b[b_index];
+            }
+          });
     });
   } else {
     // Fall back for mixed dtype to keep code size and compile time
diff --git a/kernels/portable/cpu/util/broadcast_indexes_range.h b/kernels/portable/cpu/util/broadcast_indexes_range.h
@@ -34,14 +34,17 @@ class BroadcastIndexesIterator {
 
   template <typename... Args>
   explicit BroadcastIndexesIterator(const Tensor& output, const Args&... args)
-      : output_dim_(output.dim()),
-        output_shape_(output.sizes()),
-        effective_input_broadcast_strides_{
-            effective_input_broadcast_stride(output, args)...} {
+      : output_dim_or_zero_if_no_broadcasting_(
+            ((args.sizes() == output.sizes()) && ...) ? 0 : output.dim()),
+        output_shape_(output.sizes()) {
     static_assert(
         sizeof...(args) == kNumInputs && (std::is_same_v<Args, Tensor> && ...),
         "BroadcastIndexesIterator constructor requires kNumInputs input tensor"
         "arguments!");
+    if (output_dim_or_zero_if_no_broadcasting_ != 0) {
+      effective_input_broadcast_strides_ = {
+          effective_input_broadcast_stride(output, args)...};
+    }
   }
 
   struct make_end_t {
@@ -73,9 +76,14 @@ class BroadcastIndexesIterator {
 
   BroadcastIndexesIterator& operator++() {
     output_index()++;
+    if (output_dim_or_zero_if_no_broadcasting_ == 0) {
+      std::fill(
+          current_indexes_.begin() + 1, current_indexes_.end(), output_index());
+      return *this;
+    }
     // TODO: add optimization for particular input tensors not being
     // broadcasted?
-    for (auto ii = output_dim_ - 1; ii >= 0; --ii) {
+    for (auto ii = output_dim_or_zero_if_no_broadcasting_ - 1; ii >= 0; --ii) {
       // You might wonder what happens if output_shape_[ii] == 0. In
       // that case, output.numel() would be 0, and thus we would have
       // begin() == end() and no iteration.
@@ -121,7 +129,8 @@ class BroadcastIndexesIterator {
         delinearized_output_index_.size());
     for (const auto ii : c10::irange(1, kNumInputs + 1)) {
       current_indexes_[ii] = 0;
-      for (const auto jj : c10::irange(output_dim_)) {
+      for (const auto jj :
+           c10::irange(output_dim_or_zero_if_no_broadcasting_)) {
         current_indexes_[ii] += delinearized_output_index_[jj] *
             effective_input_broadcast_strides_[ii - 1][jj];
       }
@@ -180,7 +189,7 @@ class BroadcastIndexesIterator {
   // followed by kNumInputs input indexes.
   std::array<ssize_t, kNumInputs + 1> current_indexes_ = {0};
   ShapeType delinearized_output_index_ = {0};
-  ssize_t output_dim_;
+  ssize_t output_dim_or_zero_if_no_broadcasting_;
   ArrayRef<exec_aten::SizesType> output_shape_;
   // The linear index for a broadcast tensor is
   // sum(delinearized_output_index_[i] * input_stride_[i] if
@@ -189,8 +198,7 @@ class BroadcastIndexesIterator {
   // output_dim. This is straightforwardly implementable with an
   // adjusted stride array that contains 0s where the padded input
   // shape would contain 1s.
-  std::array<ShapeType, kNumInputs> effective_input_broadcast_strides_ = {
-      {{0}}};
+  std::array<ShapeType, kNumInputs> effective_input_broadcast_strides_;
 };
 } // namespace internal
 
diff --git a/kernels/portable/cpu/util/broadcast_util.h b/kernels/portable/cpu/util/broadcast_util.h
@@ -254,26 +254,13 @@ inline void apply_binary_elementwise_fn(
     const Tensor& a,
     const Tensor& b,
     const Tensor& out) {
-  const bool a_is_broadcasted = !out.sizes().equals(a.sizes());
-  const bool b_is_broadcasted = !out.sizes().equals(b.sizes());
-  const bool any_is_broadcasted = (a_is_broadcasted || b_is_broadcasted);
-
   const CTYPE_A* const data_a = a.const_data_ptr<CTYPE_A>();
   const CTYPE_B* const data_b = b.const_data_ptr<CTYPE_B>();
   CTYPE_OUT* const data_out = out.mutable_data_ptr<CTYPE_OUT>();
 
-  if (any_is_broadcasted) {
-    for (const auto [out_index, a_index, b_index] :
-         BroadcastIndexesRange<2>(out, a, b)) {
-      data_out[out_index] = compute_fun(data_a[a_index], data_b[b_index]);
-    }
-  } else {
-    for (const auto i : c10::irange(out.numel())) {
-      size_t a_linear_index = i;
-      size_t b_linear_index = i;
-
-      data_out[i] = compute_fun(data_a[a_linear_index], data_b[b_linear_index]);
-    }
+  for (const auto [out_index, a_index, b_index] :
+       BroadcastIndexesRange<2>(out, a, b)) {
+    data_out[out_index] = compute_fun(data_a[a_index], data_b[b_index]);
   }
 }
 
@@ -294,27 +281,15 @@ inline void apply_ternary_elementwise_fn(
     const Tensor& b,
     const Tensor& c,
     const Tensor& out) {
-  const bool a_is_broadcasted = !out.sizes().equals(a.sizes());
-  const bool b_is_broadcasted = !out.sizes().equals(b.sizes());
-  const bool c_is_broadcasted = !out.sizes().equals(c.sizes());
-  const bool any_is_broadcasted =
-      (a_is_broadcasted || b_is_broadcasted || c_is_broadcasted);
-
   const CTYPE_A* const data_a = a.const_data_ptr<CTYPE_A>();
   const CTYPE_B* const data_b = b.const_data_ptr<CTYPE_B>();
   const CTYPE_C* const data_c = c.const_data_ptr<CTYPE_C>();
   CTYPE_OUT* const data_out = out.mutable_data_ptr<CTYPE_OUT>();
 
-  if (any_is_broadcasted) {
-    for (const auto [out_index, a_index, b_index, c_index] :
-         BroadcastIndexesRange<3>(out, a, b, c)) {
-      data_out[out_index] =
-          compute_fun(data_a[a_index], data_b[b_index], data_c[c_index]);
-    }
-  } else {
-    for (const auto i : c10::irange(out.numel())) {
-      data_out[i] = compute_fun(data_a[i], data_b[i], data_c[i]);
-    }
+  for (const auto [out_index, a_index, b_index, c_index] :
+       BroadcastIndexesRange<3>(out, a, b, c)) {
+    data_out[out_index] =
+        compute_fun(data_a[a_index], data_b[b_index], data_c[c_index]);
   }
 }
 
diff --git a/kernels/portable/cpu/util/elementwise_util.h b/kernels/portable/cpu/util/elementwise_util.h
@@ -76,11 +76,6 @@ inline void apply_elementwise_fn(
           internal::check_tensor_dtype(out, out_dtypes, compute_type),
       InvalidArgument, );
 
-  bool any_is_broadcasted = false;
-  if constexpr (kNumInputs > 1) {
-    any_is_broadcasted = (!out.sizes().equals(inputs.first->sizes()) || ...);
-  }
-
   struct InputInfo {
     load_to_common_fn<CTYPE_COMMON> load_to_common;
     const char* data_ptr;
@@ -99,29 +94,16 @@ inline void apply_elementwise_fn(
   char* const data_out = reinterpret_cast<char*>(out.mutable_data_ptr());
   const auto out_element_size = out.element_size();
 
-  if (any_is_broadcasted) {
-    for (const auto& indexes :
-         BroadcastIndexesRange<kNumInputs>(out, (*inputs.first)...)) {
-      std::array<CTYPE_COMMON, kNumInputs> loaded_inputs;
-      for (const auto idx : c10::irange(kNumInputs)) {
-        const auto& input_info = inputs_info[idx];
-        loaded_inputs[idx] = input_info.load_to_common(
-            &input_info.data_ptr[indexes[idx + 1] * input_info.element_size]);
-      }
-      auto result = std::apply(compute_fun, loaded_inputs);
-      store_common_to_out(result, &data_out[indexes[0] * out_element_size]);
-    }
-  } else {
-    for (const auto i : c10::irange(out.numel())) {
-      std::array<CTYPE_COMMON, kNumInputs> loaded_inputs;
-      for (const auto idx : c10::irange(kNumInputs)) {
-        const auto& input_info = inputs_info[idx];
-        loaded_inputs[idx] = input_info.load_to_common(
-            &input_info.data_ptr[i * input_info.element_size]);
-      }
-      auto result = std::apply(compute_fun, loaded_inputs);
-      store_common_to_out(result, &data_out[i * out_element_size]);
+  for (const auto& indexes :
+       BroadcastIndexesRange<kNumInputs>(out, (*inputs.first)...)) {
+    std::array<CTYPE_COMMON, kNumInputs> loaded_inputs;
+    for (const auto idx : c10::irange(kNumInputs)) {
+      const auto& input_info = inputs_info[idx];
+      loaded_inputs[idx] = input_info.load_to_common(
+          &input_info.data_ptr[indexes[idx + 1] * input_info.element_size]);
     }
+    auto result = std::apply(compute_fun, loaded_inputs);
+    store_common_to_out(result, &data_out[indexes[0] * out_element_size]);
   }
 }
 } // namespace internal
