10 #if defined(EIGEN_USE_THREADS) && !defined(EIGEN_CXX11_TENSOR_TENSOR_DEVICE_THREAD_POOL_H) 11 #define EIGEN_CXX11_TENSOR_TENSOR_DEVICE_THREAD_POOL_H 17 class ThreadPoolInterface {
19 virtual void Schedule(std::function<
void()> fn) = 0;
21 virtual ~ThreadPoolInterface() {}
27 class ThreadPool :
public ThreadPoolInterface {
30 explicit ThreadPool(
int num_threads) {
31 for (
int i = 0; i < num_threads; i++) {
32 threads_.push_back(
new std::thread([
this]() { WorkerLoop(); }));
42 std::unique_lock<std::mutex> l(mu_);
43 empty_.wait(l, [
this]() {
return pending_.empty(); });
47 for (
auto w : waiters_) {
55 for (
auto t : threads_) {
63 void Schedule(std::function<
void()> fn) {
64 std::unique_lock<std::mutex> l(mu_);
65 if (waiters_.empty()) {
66 pending_.push_back(fn);
68 Waiter* w = waiters_.back();
78 std::unique_lock<std::mutex> l(mu_);
81 std::function<void()> fn;
82 if (pending_.empty()) {
85 waiters_.push_back(&w);
86 w.cv.wait(l, [&w]() {
return w.ready; });
91 fn = pending_.front();
93 if (pending_.empty()) {
107 std::condition_variable cv;
108 std::function<void()> work;
113 std::vector<std::thread*> threads_;
114 std::vector<Waiter*> waiters_;
115 std::deque<std::function<void()>> pending_;
116 std::condition_variable empty_;
117 bool exiting_ =
false;
128 Notification() : notified_(false) {}
132 std::unique_lock<std::mutex> l(mu_);
133 eigen_assert(!notified_);
138 void WaitForNotification() {
139 std::unique_lock<std::mutex> l(mu_);
140 cv_.wait(l, [
this]() {
return notified_; } );
145 std::condition_variable cv_;
151 template <
typename Function,
typename... Args>
struct FunctionWrapper
153 static void run(Notification* n, Function f, Args... args) {
159 static EIGEN_STRONG_INLINE
void wait_until_ready(Notification* n) {
161 n->WaitForNotification();
167 struct ThreadPoolDevice {
169 ThreadPoolDevice(ThreadPoolInterface* pool,
size_t num_cores) : pool_(pool), num_threads_(num_cores) { }
171 EIGEN_STRONG_INLINE
void* allocate(
size_t num_bytes)
const {
172 return internal::aligned_malloc(num_bytes);
175 EIGEN_STRONG_INLINE
void deallocate(
void* buffer)
const {
176 internal::aligned_free(buffer);
179 EIGEN_STRONG_INLINE
void memcpy(
void* dst,
const void* src,
size_t n)
const {
180 ::memcpy(dst, src, n);
182 EIGEN_STRONG_INLINE
void memcpyHostToDevice(
void* dst,
const void* src,
size_t n)
const {
185 EIGEN_STRONG_INLINE
void memcpyDeviceToHost(
void* dst,
const void* src,
size_t n)
const {
189 EIGEN_STRONG_INLINE
void memset(
void* buffer,
int c,
size_t n)
const {
190 ::memset(buffer, c, n);
193 EIGEN_STRONG_INLINE
size_t numThreads()
const {
197 EIGEN_DEVICE_FUNC EIGEN_STRONG_INLINE
int majorDeviceVersion()
const {
202 template <
class Function,
class... Args>
203 EIGEN_STRONG_INLINE Notification* enqueue(Function&& f, Args&&... args)
const {
204 Notification* n =
new Notification();
205 std::function<void()> func =
206 std::bind(&FunctionWrapper<Function, Args...>::run, n, f, args...);
207 pool_->Schedule(func);
210 template <
class Function,
class... Args>
211 EIGEN_STRONG_INLINE
void enqueueNoNotification(Function&& f, Args&&... args)
const {
212 std::function<void()> func = std::bind(f, args...);
213 pool_->Schedule(func);
217 ThreadPoolInterface* pool_;
224 #endif // EIGEN_CXX11_TENSOR_TENSOR_DEVICE_THREAD_POOL_H Namespace containing all symbols from the Eigen library.
Definition: CXX11Meta.h:13