Use parallel_for in elementwise_util (#9243)

More straightforward rollout.

(Parallelizing over BroadcastIndexesRange is ugly, but so far we have
exactly two instances of it. If it keeps cropping up, we can add a
utility to make it nicer.)

Author: swolchok

kernels/portable/cpu/util/elementwise_util.h

@@ -13,6 +13,7 @@
 #include <executorch/kernels/portable/cpu/util/broadcast_util.h>
 #include <executorch/kernels/portable/cpu/util/dtype_util.h>
 #include <executorch/runtime/kernel/kernel_runtime_context.h>
+#include <executorch/runtime/kernel/thread_parallel_interface.h>
 
 #include <array>
 #include <utility>
@@ -94,17 +95,28 @@ inline void apply_elementwise_fn(
   char* const data_out = reinterpret_cast<char*>(out.mutable_data_ptr());
   const auto out_element_size = 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]);
-  }
+  ::executorch::extension::parallel_for(
+      0,
+      out.numel(),
+      ::executorch::extension::internal::GRAIN_SIZE,
+      [&](const auto begin, const auto end) {
+        const auto range =
+            BroadcastIndexesRange<kNumInputs>(out, (*inputs.first)...);
+        auto begin_it = range.begin();
+        begin_it += begin;
+        for (; (*begin_it)[0] < end; ++begin_it) {
+          const auto& indexes = *begin_it;
+          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
 

kernels/portable/cpu/util/targets.bzl

@@ -111,6 +111,7 @@ def define_common_targets():
             ":broadcast_util",
             ":dtype_util",
             "//executorch/runtime/kernel:kernel_runtime_context",
+            "//executorch/runtime/kernel:thread_parallel_interface",
         ],
         deps = [
             "//executorch/kernels/portable/cpu:scalar_utils",