diff --git a/kernels/portable/cpu/util/reduce_util.h b/kernels/portable/cpu/util/reduce_util.h index 1c6a6de4101..ff7589af4f5 100644 --- a/kernels/portable/cpu/util/reduce_util.h +++ b/kernels/portable/cpu/util/reduce_util.h @@ -823,11 +823,15 @@ template executorch::aten::optional dim, const Tensor& out, const Func& func) { +#ifdef ET_USE_THREADPOOL const ssize_t reduction_size = get_reduced_dim_product(in, dim); const auto grain_size = std::max( static_cast(1), static_cast(executorch::extension::internal::GRAIN_SIZE) / reduction_size); +#else // ET_USE_THREADPOOL + const auto grain_size = 1; +#endif // ET_USE_THREADPOOL return executorch::extension::parallel_for(0, out.numel(), grain_size, func); } @@ -842,11 +846,15 @@ template optional> dim_list, const Tensor& out, const Func& func) { +#ifdef ET_UE_THREADPOOL const ssize_t reduction_size = get_reduced_dim_product(in, dim_list); const auto grain_size = std::max( static_cast(1), static_cast(executorch::extension::internal::GRAIN_SIZE) / reduction_size); +#else // ET_USE_THREADPOOL + const auto grain_size = 1; +#endif // ET_USE_THREADPOOL return executorch::extension::parallel_for(0, out.numel(), grain_size, func); }