1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
|
#pragma once
#include <atomic>
#include <condition_variable>
#include <cstddef>
#include <functional>
#include <mutex>
#include <queue>
#include <thread>
#include <utility>
#include <vector>
#include <c10/macros/Export.h>
#include <c10/util/Registry.h>
#include <c10/util/numa.h>
#include <c10/util/thread_name.h>
namespace c10 {
class C10_API TaskThreadPoolBase {
public:
virtual void run(std::function<void()> func) = 0;
virtual size_t size() const = 0;
/**
* The number of available (i.e. idle) threads in this thread pool.
*/
virtual size_t numAvailable() const = 0;
/**
* Check if the current thread is from the thread pool.
*/
virtual bool inThreadPool() const = 0;
virtual ~TaskThreadPoolBase() noexcept = default;
static size_t defaultNumThreads();
};
class C10_API ThreadPool : public c10::TaskThreadPoolBase {
protected:
struct task_element_t {
bool run_with_id;
// NOLINTNEXTLINE(cppcoreguidelines-avoid-const-or-ref-data-members)
const std::function<void()> no_id;
// NOLINTNEXTLINE(cppcoreguidelines-avoid-const-or-ref-data-members)
const std::function<void(std::size_t)> with_id;
explicit task_element_t(std::function<void()> f)
: run_with_id(false), no_id(std::move(f)), with_id(nullptr) {}
explicit task_element_t(std::function<void(std::size_t)> f)
: run_with_id(true), no_id(nullptr), with_id(std::move(f)) {}
};
std::queue<task_element_t> tasks_;
std::vector<std::thread> threads_;
mutable std::mutex mutex_;
std::condition_variable condition_;
std::condition_variable completed_;
std::atomic_bool running_;
bool complete_;
std::size_t available_;
std::size_t total_;
int numa_node_id_;
public:
ThreadPool() = delete;
explicit ThreadPool(
int pool_size,
int numa_node_id = -1,
const std::function<void()>& init_thread = nullptr);
~ThreadPool() override;
size_t size() const override;
size_t numAvailable() const override;
bool inThreadPool() const override;
void run(std::function<void()> func) override;
template <typename Task>
void runTaskWithID(Task task) {
std::unique_lock<std::mutex> lock(mutex_);
// Set task and signal condition variable so that a worker thread will
// wake up and use the task.
tasks_.emplace(static_cast<std::function<void(std::size_t)>>(task));
complete_ = false;
condition_.notify_one();
}
/// @brief Wait for queue to be empty
void waitWorkComplete();
private:
// @brief Entry point for pool threads.
void main_loop(std::size_t index);
};
class C10_API TaskThreadPool : public c10::ThreadPool {
public:
explicit TaskThreadPool(int pool_size, int numa_node_id = -1)
: ThreadPool(pool_size, numa_node_id, [numa_node_id]() {
setThreadName("CaffeTaskThread");
NUMABind(numa_node_id);
}) {}
};
C10_DECLARE_SHARED_REGISTRY(
ThreadPoolRegistry,
TaskThreadPoolBase,
int,
int,
bool);
} // namespace c10
|
