1292915Sdim//===-- llvm/Support/ThreadPool.h - A ThreadPool implementation -*- C++ -*-===//
2292915Sdim//
3353358Sdim// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
4353358Sdim// See https://llvm.org/LICENSE.txt for license information.
5353358Sdim// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
6292915Sdim//
7292915Sdim//===----------------------------------------------------------------------===//
8292915Sdim//
9292915Sdim// This file defines a crude C++11 based thread pool.
10292915Sdim//
11292915Sdim//===----------------------------------------------------------------------===//
12292915Sdim
13292915Sdim#ifndef LLVM_SUPPORT_THREAD_POOL_H
14292915Sdim#define LLVM_SUPPORT_THREAD_POOL_H
15292915Sdim
16341825Sdim#include "llvm/Config/llvm-config.h"
17292915Sdim#include "llvm/Support/thread.h"
18292915Sdim
19292915Sdim#include <future>
20292915Sdim
21297244Sdim#include <atomic>
22292915Sdim#include <condition_variable>
23292915Sdim#include <functional>
24292915Sdim#include <memory>
25292915Sdim#include <mutex>
26292915Sdim#include <queue>
27292915Sdim#include <utility>
28292915Sdim
29292915Sdimnamespace llvm {
30292915Sdim
31292915Sdim/// A ThreadPool for asynchronous parallel execution on a defined number of
32292915Sdim/// threads.
33292915Sdim///
34292915Sdim/// The pool keeps a vector of threads alive, waiting on a condition variable
35292915Sdim/// for some work to become available.
36292915Sdimclass ThreadPool {
37292915Sdimpublic:
38292915Sdim  using TaskTy = std::function<void()>;
39292915Sdim  using PackagedTaskTy = std::packaged_task<void()>;
40292915Sdim
41327952Sdim  /// Construct a pool with the number of threads found by
42327952Sdim  /// hardware_concurrency().
43292915Sdim  ThreadPool();
44292915Sdim
45292915Sdim  /// Construct a pool of \p ThreadCount threads
46292915Sdim  ThreadPool(unsigned ThreadCount);
47292915Sdim
48292915Sdim  /// Blocking destructor: the pool will wait for all the threads to complete.
49292915Sdim  ~ThreadPool();
50292915Sdim
51292915Sdim  /// Asynchronous submission of a task to the pool. The returned future can be
52292915Sdim  /// used to wait for the task to finish and is *non-blocking* on destruction.
53292915Sdim  template <typename Function, typename... Args>
54321369Sdim  inline std::shared_future<void> async(Function &&F, Args &&... ArgList) {
55292915Sdim    auto Task =
56292915Sdim        std::bind(std::forward<Function>(F), std::forward<Args>(ArgList)...);
57292915Sdim    return asyncImpl(std::move(Task));
58292915Sdim  }
59292915Sdim
60292915Sdim  /// Asynchronous submission of a task to the pool. The returned future can be
61292915Sdim  /// used to wait for the task to finish and is *non-blocking* on destruction.
62292915Sdim  template <typename Function>
63321369Sdim  inline std::shared_future<void> async(Function &&F) {
64292915Sdim    return asyncImpl(std::forward<Function>(F));
65292915Sdim  }
66292915Sdim
67292915Sdim  /// Blocking wait for all the threads to complete and the queue to be empty.
68292915Sdim  /// It is an error to try to add new tasks while blocking on this call.
69292915Sdim  void wait();
70292915Sdim
71292915Sdimprivate:
72292915Sdim  /// Asynchronous submission of a task to the pool. The returned future can be
73292915Sdim  /// used to wait for the task to finish and is *non-blocking* on destruction.
74321369Sdim  std::shared_future<void> asyncImpl(TaskTy F);
75292915Sdim
76292915Sdim  /// Threads in flight
77292915Sdim  std::vector<llvm::thread> Threads;
78292915Sdim
79292915Sdim  /// Tasks waiting for execution in the pool.
80292915Sdim  std::queue<PackagedTaskTy> Tasks;
81292915Sdim
82292915Sdim  /// Locking and signaling for accessing the Tasks queue.
83292915Sdim  std::mutex QueueLock;
84292915Sdim  std::condition_variable QueueCondition;
85292915Sdim
86292915Sdim  /// Locking and signaling for job completion
87292915Sdim  std::mutex CompletionLock;
88292915Sdim  std::condition_variable CompletionCondition;
89292915Sdim
90292915Sdim  /// Keep track of the number of thread actually busy
91292915Sdim  std::atomic<unsigned> ActiveThreads;
92292915Sdim
93292915Sdim#if LLVM_ENABLE_THREADS // avoids warning for unused variable
94292915Sdim  /// Signal for the destruction of the pool, asking thread to exit.
95292915Sdim  bool EnableFlag;
96292915Sdim#endif
97292915Sdim};
98292915Sdim}
99292915Sdim
100292915Sdim#endif // LLVM_SUPPORT_THREAD_POOL_H
101