Add parallel_for_each_reduce_over_dim_list_output_index (#9143)

Author: swolchok

kernels/portable/cpu/util/reduce_util.h

@@ -823,10 +823,30 @@ template <typename Func>
     executorch::aten::optional<int64_t> dim,
     const Tensor& out,
     const Func& func) {
-  const int64_t reduction_size = get_reduced_dim_product(in, dim);
+  const ssize_t reduction_size = get_reduced_dim_product(in, dim);
   const auto grain_size = std::max(
-      static_cast<int64_t>(1),
-      executorch::extension::internal::GRAIN_SIZE / reduction_size);
+      static_cast<ssize_t>(1),
+      static_cast<ssize_t>(executorch::extension::internal::GRAIN_SIZE) /
+          reduction_size);
+  return executorch::extension::parallel_for(0, out.numel(), grain_size, func);
+}
+
+/**
+ * parallel_for wrapper for reductions that call reduce_over_dim_list or
+ * map_reduce_over_dim_list for each output element. Automatically
+ * calculates appropriate grain size.
+ */
+template <typename Func>
+[[nodiscard]] bool parallel_for_each_reduce_over_dim_list_output_index(
+    const Tensor& in,
+    optional<ArrayRef<int64_t>> dim_list,
+    const Tensor& out,
+    const Func& func) {
+  const ssize_t reduction_size = get_reduced_dim_product(in, dim_list);
+  const auto grain_size = std::max(
+      static_cast<ssize_t>(1),
+      static_cast<ssize_t>(executorch::extension::internal::GRAIN_SIZE) /
+          reduction_size);
   return executorch::extension::parallel_for(0, out.numel(), grain_size, func);
 }