1*523fa7a6SAndroid Build Coastguard Worker /*
2*523fa7a6SAndroid Build Coastguard Worker * Copyright (c) Meta Platforms, Inc. and affiliates.
3*523fa7a6SAndroid Build Coastguard Worker * All rights reserved.
4*523fa7a6SAndroid Build Coastguard Worker *
5*523fa7a6SAndroid Build Coastguard Worker * This source code is licensed under the BSD-style license found in the
6*523fa7a6SAndroid Build Coastguard Worker * LICENSE file in the root directory of this source tree.
7*523fa7a6SAndroid Build Coastguard Worker */
8*523fa7a6SAndroid Build Coastguard Worker
9*523fa7a6SAndroid Build Coastguard Worker #include <tuple>
10*523fa7a6SAndroid Build Coastguard Worker
11*523fa7a6SAndroid Build Coastguard Worker #include <executorch/extension/parallel/thread_parallel.h>
12*523fa7a6SAndroid Build Coastguard Worker #include <executorch/extension/threadpool/threadpool.h>
13*523fa7a6SAndroid Build Coastguard Worker #include <executorch/runtime/core/exec_aten/util/tensor_util.h>
14*523fa7a6SAndroid Build Coastguard Worker #include <executorch/runtime/platform/assert.h>
15*523fa7a6SAndroid Build Coastguard Worker
16*523fa7a6SAndroid Build Coastguard Worker namespace executorch {
17*523fa7a6SAndroid Build Coastguard Worker namespace extension {
18*523fa7a6SAndroid Build Coastguard Worker
19*523fa7a6SAndroid Build Coastguard Worker namespace {
20*523fa7a6SAndroid Build Coastguard Worker thread_local int64_t thread_num_ = 0;
21*523fa7a6SAndroid Build Coastguard Worker }
22*523fa7a6SAndroid Build Coastguard Worker
23*523fa7a6SAndroid Build Coastguard Worker using namespace ::executorch::extension::threadpool;
24*523fa7a6SAndroid Build Coastguard Worker
divup(int64_t x,int64_t y)25*523fa7a6SAndroid Build Coastguard Worker inline int64_t divup(int64_t x, int64_t y) {
26*523fa7a6SAndroid Build Coastguard Worker return (x + y - 1) / y;
27*523fa7a6SAndroid Build Coastguard Worker }
28*523fa7a6SAndroid Build Coastguard Worker
get_thread_num()29*523fa7a6SAndroid Build Coastguard Worker int64_t get_thread_num() {
30*523fa7a6SAndroid Build Coastguard Worker return thread_num_;
31*523fa7a6SAndroid Build Coastguard Worker }
32*523fa7a6SAndroid Build Coastguard Worker
set_thread_num(int64_t thread_num)33*523fa7a6SAndroid Build Coastguard Worker void set_thread_num(int64_t thread_num) {
34*523fa7a6SAndroid Build Coastguard Worker thread_num_ = thread_num;
35*523fa7a6SAndroid Build Coastguard Worker }
36*523fa7a6SAndroid Build Coastguard Worker
37*523fa7a6SAndroid Build Coastguard Worker inline std::tuple<int64_t, int64_t>
calc_num_tasks_and_chunk_size(int64_t begin,int64_t end,int64_t grain_size)38*523fa7a6SAndroid Build Coastguard Worker calc_num_tasks_and_chunk_size(int64_t begin, int64_t end, int64_t grain_size) {
39*523fa7a6SAndroid Build Coastguard Worker if ((end - begin) < grain_size) {
40*523fa7a6SAndroid Build Coastguard Worker return std::make_tuple(1, std::max((int64_t)0, end - begin));
41*523fa7a6SAndroid Build Coastguard Worker }
42*523fa7a6SAndroid Build Coastguard Worker // Choose number of tasks based on grain size and number of threads.
43*523fa7a6SAndroid Build Coastguard Worker int64_t chunk_size =
44*523fa7a6SAndroid Build Coastguard Worker divup((end - begin), get_threadpool()->get_thread_count());
45*523fa7a6SAndroid Build Coastguard Worker // Make sure each task is at least grain_size size.
46*523fa7a6SAndroid Build Coastguard Worker chunk_size = std::max(grain_size, chunk_size);
47*523fa7a6SAndroid Build Coastguard Worker int64_t num_tasks = divup((end - begin), chunk_size);
48*523fa7a6SAndroid Build Coastguard Worker return std::make_tuple(num_tasks, chunk_size);
49*523fa7a6SAndroid Build Coastguard Worker }
50*523fa7a6SAndroid Build Coastguard Worker
parallel_for(const int64_t begin,const int64_t end,const int64_t grain_size,const std::function<void (int64_t,int64_t)> & f)51*523fa7a6SAndroid Build Coastguard Worker bool parallel_for(
52*523fa7a6SAndroid Build Coastguard Worker const int64_t begin,
53*523fa7a6SAndroid Build Coastguard Worker const int64_t end,
54*523fa7a6SAndroid Build Coastguard Worker const int64_t grain_size,
55*523fa7a6SAndroid Build Coastguard Worker const std::function<void(int64_t, int64_t)>& f) {
56*523fa7a6SAndroid Build Coastguard Worker ET_LOG_AND_RETURN_IF_FALSE(begin >= 0 && end >= 0);
57*523fa7a6SAndroid Build Coastguard Worker ET_LOG_AND_RETURN_IF_FALSE(end >= begin);
58*523fa7a6SAndroid Build Coastguard Worker ET_LOG_AND_RETURN_IF_FALSE(grain_size > 0);
59*523fa7a6SAndroid Build Coastguard Worker int64_t num_tasks = 0, chunk_size = 0;
60*523fa7a6SAndroid Build Coastguard Worker std::tie(num_tasks, chunk_size) =
61*523fa7a6SAndroid Build Coastguard Worker calc_num_tasks_and_chunk_size(begin, end, grain_size);
62*523fa7a6SAndroid Build Coastguard Worker
63*523fa7a6SAndroid Build Coastguard Worker auto task = [f, begin, end, chunk_size](size_t task_id) {
64*523fa7a6SAndroid Build Coastguard Worker set_thread_num(task_id);
65*523fa7a6SAndroid Build Coastguard Worker int64_t local_start = begin + static_cast<int64_t>(task_id) * chunk_size;
66*523fa7a6SAndroid Build Coastguard Worker if (local_start < end) {
67*523fa7a6SAndroid Build Coastguard Worker int64_t local_end = std::min(end, (int64_t)(chunk_size + local_start));
68*523fa7a6SAndroid Build Coastguard Worker f(local_start, local_end);
69*523fa7a6SAndroid Build Coastguard Worker }
70*523fa7a6SAndroid Build Coastguard Worker };
71*523fa7a6SAndroid Build Coastguard Worker
72*523fa7a6SAndroid Build Coastguard Worker // Per protocol from threadpool (pthreadpool), when this returns, all tasks
73*523fa7a6SAndroid Build Coastguard Worker // are executed, so this is synchronous.
74*523fa7a6SAndroid Build Coastguard Worker get_threadpool()->run(task, num_tasks);
75*523fa7a6SAndroid Build Coastguard Worker return true;
76*523fa7a6SAndroid Build Coastguard Worker }
77*523fa7a6SAndroid Build Coastguard Worker
78*523fa7a6SAndroid Build Coastguard Worker } // namespace extension
79*523fa7a6SAndroid Build Coastguard Worker } // namespace executorch
80