xref: /aosp_15_r20/external/executorch/extension/parallel/thread_parallel.cpp (revision 523fa7a60841cd1ecfb9cc4201f1ca8b03ed023a)
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