1 | /******************************************************************************* |
2 | Copyright (c) The Taichi Authors (2016- ). All Rights Reserved. |
3 | The use of this software is governed by the LICENSE file. |
4 | *******************************************************************************/ |
5 | |
6 | #pragma once |
7 | |
8 | #include "taichi/common/core.h" |
9 | |
10 | #include <atomic> |
11 | #include <condition_variable> |
12 | #include <functional> |
13 | #include <thread> |
14 | |
15 | namespace taichi { |
16 | |
17 | using RangeForTaskFunc = void(void *, int thread_id, int i); |
18 | using ParallelFor = void(int n, int num_threads, void *, RangeForTaskFunc func); |
19 | |
20 | class ThreadPool { |
21 | public: |
22 | std::vector<std::thread> threads; |
23 | std::condition_variable slave_cv; |
24 | std::condition_variable master_cv; |
25 | std::mutex mutex; |
26 | std::atomic<int> task_head; |
27 | int task_tail; |
28 | int running_threads; |
29 | int max_num_threads; |
30 | int desired_num_threads; |
31 | uint64 timestamp; |
32 | uint64 last_finished; |
33 | bool started; |
34 | bool exiting; |
35 | RangeForTaskFunc *func; |
36 | void *range_for_task_context; // Note: this is a pointer to a |
37 | // range_task_helper_context defined in the |
38 | // LLVM runtime, which is different from |
39 | // taichi::lang::Context. |
40 | int thread_counter; |
41 | |
42 | explicit ThreadPool(int max_num_threads); |
43 | |
44 | void run(int splits, |
45 | int desired_num_threads, |
46 | void *range_for_task_context, |
47 | RangeForTaskFunc *func); |
48 | |
49 | static void static_run(ThreadPool *pool, |
50 | int splits, |
51 | int desired_num_threads, |
52 | void *range_for_task_context, |
53 | RangeForTaskFunc *func) { |
54 | return pool->run(splits, desired_num_threads, range_for_task_context, func); |
55 | } |
56 | |
57 | void target(); |
58 | |
59 | ~ThreadPool(); |
60 | }; |
61 | |
62 | } // namespace taichi |
63 | |