File tree Expand file tree Collapse file tree 2 files changed +13
-3
lines changed
kernels/portable/cpu/util Expand file tree Collapse file tree 2 files changed +13
-3
lines changed Original file line number Diff line number Diff line change 1212
1313#include < executorch/runtime/core/exec_aten/exec_aten.h>
1414#include < executorch/runtime/core/exec_aten/util/tensor_util.h>
15+ #include < executorch/runtime/kernel/thread_parallel_interface.h>
1516
1617namespace torch {
1718namespace executor {
@@ -53,9 +54,15 @@ inline void apply_unary_map_fn(
5354 CTYPE_OUT* const data_out,
5455 const int64_t size,
5556 const int64_t stride = 1 ) {
56- for (const auto i : c10::irange (size)) {
57- data_out[i * stride] = map_fun (data_in[i * stride]);
58- }
57+ executorch::extension::parallel_for (
58+ 0 ,
59+ size,
60+ ::executorch::extension::internal::GRAIN_SIZE,
61+ [&](const auto begin, const auto end) {
62+ for (const auto i : c10::irange (begin, end)) {
63+ data_out[i * stride] = map_fun (data_in[i * stride]);
64+ }
65+ });
5966}
6067
6168//
Original file line number Diff line number Diff line change @@ -244,6 +244,9 @@ def define_common_targets():
244244 name = "functional_util" ,
245245 srcs = [],
246246 exported_headers = ["functional_util.h" ],
247+ exported_deps = [
248+ "//executorch/runtime/kernel:thread_parallel_interface" ,
249+ ],
247250 deps = [
248251 "//executorch/runtime/kernel:kernel_includes" ,
249252 "//executorch/runtime/core/exec_aten/util:tensor_util" ,
You can’t perform that action at this time.
0 commit comments