diff --git a/llvm/include/llvm/Support/ThreadPool.h b/llvm/include/llvm/Support/ThreadPool.h --- a/llvm/include/llvm/Support/ThreadPool.h +++ b/llvm/include/llvm/Support/ThreadPool.h @@ -29,29 +29,94 @@ namespace llvm { -/// A ThreadPool for asynchronous parallel execution on a defined number of -/// threads. +/// A ThreadPoolWithResult for asynchronous parallel execution on a defined +/// number of threads. /// /// The pool keeps a vector of threads alive, waiting on a condition variable /// for some work to become available. -class ThreadPool { +template class ThreadPoolWithResult { public: - using TaskTy = std::function; - using PackagedTaskTy = std::packaged_task; + using TaskTy = std::function; + using PackagedTaskTy = std::packaged_task; /// Construct a pool using the hardware strategy \p S for mapping hardware /// execution resources (threads, cores, CPUs) /// Defaults to using the maximum execution resources in the system, but /// accounting for the affinity mask. - ThreadPool(ThreadPoolStrategy S = hardware_concurrency()); + ThreadPoolWithResult(ThreadPoolStrategy S = hardware_concurrency()) + : ThreadCount(S.compute_thread_count()) { +#ifdef LLVM_ENABLE_THREADS + // Create ThreadCount threads that will loop forever, wait on QueueCondition + // for tasks to be queued or the Pool to be destroyed. + Threads.reserve(ThreadCount); + for (unsigned ThreadID = 0; ThreadID < ThreadCount; ++ThreadID) { + Threads.emplace_back([S, ThreadID, this] { + S.apply_thread_strategy(ThreadID); + while (true) { + PackagedTaskTy Task; + { + std::unique_lock LockGuard(QueueLock); + // Wait for tasks to be pushed in the queue + QueueCondition.wait(LockGuard, + [&] { return !EnableFlag || !Tasks.empty(); }); + // Exit condition + if (!EnableFlag && Tasks.empty()) + return; + // Yeah, we have a task, grab it and release the lock on the queue + + // We first need to signal that we are active before popping the + // queue in order for wait() to properly detect that even if the + // queue is empty, there is still a task in flight. + ++ActiveThreads; + Task = std::move(Tasks.front()); + Tasks.pop(); + } + // Run the task we just grabbed + Task(); + + bool Notify; + { + // Adjust `ActiveThreads`, in case someone waits on + // ThreadPoolWithResult::wait() + std::lock_guard LockGuard(QueueLock); + --ActiveThreads; + Notify = workCompletedUnlocked(); + } + // Notify task completion if this is the last active thread, in case + // someone waits on ThreadPoolWithResult::wait(). + if (Notify) + CompletionCondition.notify_all(); + } + }); + } +#else // LLVM_ENABLE_THREADS Disabled + if (ThreadCount != 1) { + errs() << "Warning: request a ThreadPoolWithResult with " << ThreadCount + << " threads, but LLVM_ENABLE_THREADS has been turned off\n"; + } +#endif + } /// Blocking destructor: the pool will wait for all the threads to complete. - ~ThreadPool(); + ~ThreadPoolWithResult() { +#if LLVM_ENABLE_THREADS + { + std::unique_lock LockGuard(QueueLock); + EnableFlag = false; + } + QueueCondition.notify_all(); + for (auto &Worker : Threads) + Worker.join(); +#else // LLVM_ENABLE_THREADS Disabled + wait(); +#endif + } /// Asynchronous submission of a task to the pool. The returned future can be /// used to wait for the task to finish and is *non-blocking* on destruction. template - inline std::shared_future async(Function &&F, Args &&... ArgList) { + inline std::shared_future async(Function &&F, + Args &&...ArgList) { auto Task = std::bind(std::forward(F), std::forward(ArgList)...); return asyncImpl(std::move(Task)); @@ -60,25 +125,75 @@ /// Asynchronous submission of a task to the pool. The returned future can be /// used to wait for the task to finish and is *non-blocking* on destruction. template - inline std::shared_future async(Function &&F) { + inline std::shared_future async(Function &&F) { return asyncImpl(std::forward(F)); } /// Blocking wait for all the threads to complete and the queue to be empty. /// It is an error to try to add new tasks while blocking on this call. - void wait(); + void wait() { +#if LLVM_ENABLE_THREADS + // Wait for all threads to complete and the queue to be empty + std::unique_lock LockGuard(QueueLock); + CompletionCondition.wait(LockGuard, + [&] { return workCompletedUnlocked(); }); + +#else // LLVM_ENABLE_THREADS Disabled + + // Sequential implementation running the tasks + while (!Tasks.empty()) { + auto Task = std::move(Tasks.front()); + Tasks.pop(); + Task(); + } +#endif + } unsigned getThreadCount() const { return ThreadCount; } /// Returns true if the current thread is a worker thread of this thread pool. - bool isWorkerThread() const; + bool isWorkerThread() const { + llvm::thread::id CurrentThreadId = llvm::this_thread::get_id(); + for (const llvm::thread &Thread : Threads) + if (CurrentThreadId == Thread.get_id()) + return true; + return false; + } private: bool workCompletedUnlocked() { return !ActiveThreads && Tasks.empty(); } /// Asynchronous submission of a task to the pool. The returned future can be /// used to wait for the task to finish and is *non-blocking* on destruction. - std::shared_future asyncImpl(TaskTy F); + std::shared_future asyncImpl(TaskTy Task) { +#ifdef LLVM_ENABLE_THREADS + /// Wrap the Task in a packaged_task to return a future object. + PackagedTaskTy PackagedTask(std::move(Task)); + auto Future = PackagedTask.get_future(); + { + // Lock the queue and push the new task + std::unique_lock LockGuard(QueueLock); + + // Don't allow enqueueing after disabling the pool + assert(EnableFlag && + "Queuing a thread during ThreadPoolWithResult destruction"); + + Tasks.push(std::move(PackagedTask)); + } + QueueCondition.notify_one(); + return Future.share(); + +#else // LLVM_ENABLE_THREADS Disabled + + // Get a Future with launch::deferred execution using std::async + auto Future = std::async(std::launch::deferred, std::move(Task)).share(); + // Wrap the future so that both ThreadPoolWithResult::wait() can operate and + // the returned future can be sync'ed on. + PackagedTaskTy PackagedTask([Future]() { Future.get(); }); + Tasks.push(std::move(PackagedTask)); + return Future; +#endif + } /// Threads in flight std::vector Threads; @@ -103,6 +218,8 @@ unsigned ThreadCount; }; + +using ThreadPool = ThreadPoolWithResult; } #endif // LLVM_SUPPORT_THREADPOOL_H diff --git a/llvm/lib/Support/CMakeLists.txt b/llvm/lib/Support/CMakeLists.txt --- a/llvm/lib/Support/CMakeLists.txt +++ b/llvm/lib/Support/CMakeLists.txt @@ -205,7 +205,6 @@ SystemUtils.cpp TarWriter.cpp TargetParser.cpp - ThreadPool.cpp TimeProfiler.cpp Timer.cpp ToolOutputFile.cpp diff --git a/llvm/lib/Support/ThreadPool.cpp b/llvm/lib/Support/ThreadPool.cpp deleted file mode 100644 --- a/llvm/lib/Support/ThreadPool.cpp +++ /dev/null @@ -1,143 +0,0 @@ -//==-- llvm/Support/ThreadPool.cpp - A ThreadPool implementation -*- C++ -*-==// -// -// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions. -// See https://llvm.org/LICENSE.txt for license information. -// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception -// -//===----------------------------------------------------------------------===// -// -// This file implements a crude C++11 based thread pool. -// -//===----------------------------------------------------------------------===// - -#include "llvm/Support/ThreadPool.h" - -#include "llvm/Config/llvm-config.h" -#include "llvm/Support/Threading.h" -#include "llvm/Support/raw_ostream.h" - -using namespace llvm; - -#if LLVM_ENABLE_THREADS - -ThreadPool::ThreadPool(ThreadPoolStrategy S) - : ThreadCount(S.compute_thread_count()) { - // Create ThreadCount threads that will loop forever, wait on QueueCondition - // for tasks to be queued or the Pool to be destroyed. - Threads.reserve(ThreadCount); - for (unsigned ThreadID = 0; ThreadID < ThreadCount; ++ThreadID) { - Threads.emplace_back([S, ThreadID, this] { - S.apply_thread_strategy(ThreadID); - while (true) { - PackagedTaskTy Task; - { - std::unique_lock LockGuard(QueueLock); - // Wait for tasks to be pushed in the queue - QueueCondition.wait(LockGuard, - [&] { return !EnableFlag || !Tasks.empty(); }); - // Exit condition - if (!EnableFlag && Tasks.empty()) - return; - // Yeah, we have a task, grab it and release the lock on the queue - - // We first need to signal that we are active before popping the queue - // in order for wait() to properly detect that even if the queue is - // empty, there is still a task in flight. - ++ActiveThreads; - Task = std::move(Tasks.front()); - Tasks.pop(); - } - // Run the task we just grabbed - Task(); - - bool Notify; - { - // Adjust `ActiveThreads`, in case someone waits on ThreadPool::wait() - std::lock_guard LockGuard(QueueLock); - --ActiveThreads; - Notify = workCompletedUnlocked(); - } - // Notify task completion if this is the last active thread, in case - // someone waits on ThreadPool::wait(). - if (Notify) - CompletionCondition.notify_all(); - } - }); - } -} - -void ThreadPool::wait() { - // Wait for all threads to complete and the queue to be empty - std::unique_lock LockGuard(QueueLock); - CompletionCondition.wait(LockGuard, [&] { return workCompletedUnlocked(); }); -} - -bool ThreadPool::isWorkerThread() const { - llvm::thread::id CurrentThreadId = llvm::this_thread::get_id(); - for (const llvm::thread &Thread : Threads) - if (CurrentThreadId == Thread.get_id()) - return true; - return false; -} - -std::shared_future ThreadPool::asyncImpl(TaskTy Task) { - /// Wrap the Task in a packaged_task to return a future object. - PackagedTaskTy PackagedTask(std::move(Task)); - auto Future = PackagedTask.get_future(); - { - // Lock the queue and push the new task - std::unique_lock LockGuard(QueueLock); - - // Don't allow enqueueing after disabling the pool - assert(EnableFlag && "Queuing a thread during ThreadPool destruction"); - - Tasks.push(std::move(PackagedTask)); - } - QueueCondition.notify_one(); - return Future.share(); -} - -// The destructor joins all threads, waiting for completion. -ThreadPool::~ThreadPool() { - { - std::unique_lock LockGuard(QueueLock); - EnableFlag = false; - } - QueueCondition.notify_all(); - for (auto &Worker : Threads) - Worker.join(); -} - -#else // LLVM_ENABLE_THREADS Disabled - -// No threads are launched, issue a warning if ThreadCount is not 0 -ThreadPool::ThreadPool(ThreadPoolStrategy S) - : ThreadCount(S.compute_thread_count()) { - if (ThreadCount != 1) { - errs() << "Warning: request a ThreadPool with " << ThreadCount - << " threads, but LLVM_ENABLE_THREADS has been turned off\n"; - } -} - -void ThreadPool::wait() { - // Sequential implementation running the tasks - while (!Tasks.empty()) { - auto Task = std::move(Tasks.front()); - Tasks.pop(); - Task(); - } -} - -std::shared_future ThreadPool::asyncImpl(TaskTy Task) { - // Get a Future with launch::deferred execution using std::async - auto Future = std::async(std::launch::deferred, std::move(Task)).share(); - // Wrap the future so that both ThreadPool::wait() can operate and the - // returned future can be sync'ed on. - PackagedTaskTy PackagedTask([Future]() { Future.get(); }); - Tasks.push(std::move(PackagedTask)); - return Future; -} - -ThreadPool::~ThreadPool() { wait(); } - -#endif diff --git a/mlir/include/mlir/IR/MLIRContext.h b/mlir/include/mlir/IR/MLIRContext.h --- a/mlir/include/mlir/IR/MLIRContext.h +++ b/mlir/include/mlir/IR/MLIRContext.h @@ -16,7 +16,9 @@ #include namespace llvm { -class ThreadPool; +template +class ThreadPoolWithResult; +using ThreadPool = ThreadPoolWithResult; } // end namespace llvm namespace mlir {