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
15namespace taichi {
16
17using RangeForTaskFunc = void(void *, int thread_id, int i);
18using ParallelFor = void(int n, int num_threads, void *, RangeForTaskFunc func);
19
20class 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