File tree 2 files changed +13
-3
lines changed
kernels/portable/cpu/util
2 files changed +13
-3
lines changed Original file line number Diff line number Diff line change 12
12
13
13
#include < executorch/runtime/core/exec_aten/exec_aten.h>
14
14
#include < executorch/runtime/core/exec_aten/util/tensor_util.h>
15
+ #include < executorch/runtime/kernel/thread_parallel_interface.h>
15
16
16
17
namespace torch {
17
18
namespace executor {
@@ -53,9 +54,15 @@ inline void apply_unary_map_fn(
53
54
CTYPE_OUT* const data_out,
54
55
const int64_t size,
55
56
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
+ });
59
66
}
60
67
61
68
//
Original file line number Diff line number Diff line change @@ -244,6 +244,9 @@ def define_common_targets():
244
244
name = "functional_util" ,
245
245
srcs = [],
246
246
exported_headers = ["functional_util.h" ],
247
+ exported_deps = [
248
+ "//executorch/runtime/kernel:thread_parallel_interface" ,
249
+ ],
247
250
deps = [
248
251
"//executorch/runtime/kernel:kernel_includes" ,
249
252
"//executorch/runtime/core/exec_aten/util:tensor_util" ,
You can’t perform that action at this time.
0 commit comments