155 lines
4.2 KiB
C++
155 lines
4.2 KiB
C++
// This file is part of Eigen, a lightweight C++ template library
|
|
// for linear algebra.
|
|
//
|
|
// Copyright (C) 2014 Benoit Steiner <benoit.steiner.goog@gmail.com>
|
|
//
|
|
// This Source Code Form is subject to the terms of the Mozilla
|
|
// Public License v. 2.0. If a copy of the MPL was not distributed
|
|
// with this file, You can obtain one at the mozilla.org home page
|
|
|
|
#ifndef EIGEN_CXX11_THREADPOOL_SIMPLE_THREAD_POOL_H
|
|
#define EIGEN_CXX11_THREADPOOL_SIMPLE_THREAD_POOL_H
|
|
|
|
namespace Eigen {
|
|
|
|
// The implementation of the ThreadPool type ensures that the Schedule method
|
|
// runs the functions it is provided in FIFO order when the scheduling is done
|
|
// by a single thread.
|
|
// Environment provides a way to create threads and also allows to intercept
|
|
// task submission and execution.
|
|
template <typename Environment>
|
|
class SimpleThreadPoolTempl : public ThreadPoolInterface {
|
|
public:
|
|
// Construct a pool that contains "num_threads" threads.
|
|
explicit SimpleThreadPoolTempl(int num_threads, Environment env = Environment())
|
|
: env_(env), threads_(num_threads), waiters_(num_threads) {
|
|
for (int i = 0; i < num_threads; i++) {
|
|
threads_.push_back(env.CreateThread([this, i]() { WorkerLoop(i); }));
|
|
}
|
|
}
|
|
|
|
// Wait until all scheduled work has finished and then destroy the
|
|
// set of threads.
|
|
~SimpleThreadPoolTempl() {
|
|
{
|
|
// Wait for all work to get done.
|
|
std::unique_lock<std::mutex> l(mu_);
|
|
while (!pending_.empty()) {
|
|
empty_.wait(l);
|
|
}
|
|
exiting_ = true;
|
|
|
|
// Wakeup all waiters.
|
|
for (auto w : waiters_) {
|
|
w->ready = true;
|
|
w->task.f = nullptr;
|
|
w->cv.notify_one();
|
|
}
|
|
}
|
|
|
|
// Wait for threads to finish.
|
|
for (auto t : threads_) {
|
|
delete t;
|
|
}
|
|
}
|
|
|
|
// Schedule fn() for execution in the pool of threads. The functions are
|
|
// executed in the order in which they are scheduled.
|
|
void Schedule(std::function<void()> fn) final {
|
|
Task t = env_.CreateTask(std::move(fn));
|
|
std::unique_lock<std::mutex> l(mu_);
|
|
if (waiters_.empty()) {
|
|
pending_.push_back(std::move(t));
|
|
} else {
|
|
Waiter* w = waiters_.back();
|
|
waiters_.pop_back();
|
|
w->ready = true;
|
|
w->task = std::move(t);
|
|
w->cv.notify_one();
|
|
}
|
|
}
|
|
|
|
int NumThreads() const final {
|
|
return static_cast<int>(threads_.size());
|
|
}
|
|
|
|
int CurrentThreadId() const final {
|
|
const PerThread* pt = this->GetPerThread();
|
|
if (pt->pool == this) {
|
|
return pt->thread_id;
|
|
} else {
|
|
return -1;
|
|
}
|
|
}
|
|
|
|
protected:
|
|
void WorkerLoop(int thread_id) {
|
|
std::unique_lock<std::mutex> l(mu_);
|
|
PerThread* pt = GetPerThread();
|
|
pt->pool = this;
|
|
pt->thread_id = thread_id;
|
|
Waiter w;
|
|
Task t;
|
|
while (!exiting_) {
|
|
if (pending_.empty()) {
|
|
// Wait for work to be assigned to me
|
|
w.ready = false;
|
|
waiters_.push_back(&w);
|
|
while (!w.ready) {
|
|
w.cv.wait(l);
|
|
}
|
|
t = w.task;
|
|
w.task.f = nullptr;
|
|
} else {
|
|
// Pick up pending work
|
|
t = std::move(pending_.front());
|
|
pending_.pop_front();
|
|
if (pending_.empty()) {
|
|
empty_.notify_all();
|
|
}
|
|
}
|
|
if (t.f) {
|
|
mu_.unlock();
|
|
env_.ExecuteTask(t);
|
|
t.f = nullptr;
|
|
mu_.lock();
|
|
}
|
|
}
|
|
}
|
|
|
|
private:
|
|
typedef typename Environment::Task Task;
|
|
typedef typename Environment::EnvThread Thread;
|
|
|
|
struct Waiter {
|
|
std::condition_variable cv;
|
|
Task task;
|
|
bool ready;
|
|
};
|
|
|
|
struct PerThread {
|
|
constexpr PerThread() : pool(NULL), thread_id(-1) { }
|
|
SimpleThreadPoolTempl* pool; // Parent pool, or null for normal threads.
|
|
int thread_id; // Worker thread index in pool.
|
|
};
|
|
|
|
Environment env_;
|
|
std::mutex mu_;
|
|
MaxSizeVector<Thread*> threads_; // All threads
|
|
MaxSizeVector<Waiter*> waiters_; // Stack of waiting threads.
|
|
std::deque<Task> pending_; // Queue of pending work
|
|
std::condition_variable empty_; // Signaled on pending_.empty()
|
|
bool exiting_ = false;
|
|
|
|
PerThread* GetPerThread() const {
|
|
EIGEN_THREAD_LOCAL PerThread per_thread;
|
|
return &per_thread;
|
|
}
|
|
};
|
|
|
|
typedef SimpleThreadPoolTempl<StlThreadEnvironment> SimpleThreadPool;
|
|
|
|
} // namespace Eigen
|
|
|
|
#endif // EIGEN_CXX11_THREADPOOL_SIMPLE_THREAD_POOL_H
|