File tree 1 file changed +8
-0
lines changed
kernels/portable/cpu/util
1 file changed +8
-0
lines changed Original file line number Diff line number Diff line change @@ -823,11 +823,15 @@ template <typename Func>
823
823
executorch::aten::optional<int64_t > dim,
824
824
const Tensor& out,
825
825
const Func& func) {
826
+ #ifdef ET_USE_THREADPOOL
826
827
const ssize_t reduction_size = get_reduced_dim_product (in, dim);
827
828
const auto grain_size = std::max (
828
829
static_cast <ssize_t >(1 ),
829
830
static_cast <ssize_t >(executorch::extension::internal::GRAIN_SIZE) /
830
831
reduction_size);
832
+ #else // ET_USE_THREADPOOL
833
+ const auto grain_size = 1 ;
834
+ #endif // ET_USE_THREADPOOL
831
835
return executorch::extension::parallel_for (0 , out.numel (), grain_size, func);
832
836
}
833
837
@@ -842,11 +846,15 @@ template <typename Func>
842
846
optional<ArrayRef<int64_t >> dim_list,
843
847
const Tensor& out,
844
848
const Func& func) {
849
+ #ifdef ET_UE_THREADPOOL
845
850
const ssize_t reduction_size = get_reduced_dim_product (in, dim_list);
846
851
const auto grain_size = std::max (
847
852
static_cast <ssize_t >(1 ),
848
853
static_cast <ssize_t >(executorch::extension::internal::GRAIN_SIZE) /
849
854
reduction_size);
855
+ #else // ET_USE_THREADPOOL
856
+ const auto grain_size = 1 ;
857
+ #endif // ET_USE_THREADPOOL
850
858
return executorch::extension::parallel_for (0 , out.numel (), grain_size, func);
851
859
}
852
860
You can’t perform that action at this time.
0 commit comments