From deae9b5f50af38c76d49109e2ed03b167e4102a0 Mon Sep 17 00:00:00 2001 From: River Riddle Date: Wed, 23 Jun 2021 01:16:55 +0000 Subject: [PATCH] [mlir] Add a ThreadPool to MLIRContext and refactor MLIR threading usage This revision refactors the usage of multithreaded utilities in MLIR to use a common thread pool within the MLIR context, in addition to a new utility that makes writing multi-threaded code in MLIR less error prone. Using a unified thread pool brings about several advantages: * Better thread usage and more control We currently use the static llvm threading utilities, which do not allow multiple levels of asynchronous scheduling (even if there are open threads). This is due to how the current TaskGroup structure works, which only allows one truly multithreaded instance at a time. By having our own ThreadPool we gain more control and flexibility over our job/thread scheduling, and in a followup can enable threading more parts of the compiler. * The static nature of TaskGroup causes issues in certain configurations Due to the static nature of TaskGroup, there have been quite a few problems related to destruction that have caused several downstream projects to disable threading. See D104207 for discussion on some related fallout. By having a ThreadPool scoped to the context, we don't have to worry about destruction and can ensure that any additional MLIR thread usage ends when the context is destroyed. Differential Revision: https://reviews.llvm.org/D104516 --- include/llvm/Support/ThreadPool.h | 3 +++ lib/Support/ThreadPool.cpp | 8 ++++++++ 2 files changed, 11 insertions(+) diff --git a/include/llvm/Support/ThreadPool.h b/include/llvm/Support/ThreadPool.h index 9d319eb71bb..4c41b88d604 100644 --- a/include/llvm/Support/ThreadPool.h +++ b/include/llvm/Support/ThreadPool.h @@ -70,6 +70,9 @@ public: unsigned getThreadCount() const { return ThreadCount; } + /// Returns true if the current thread is a worker thread of this thread pool. + bool isWorkerThread() const; + private: bool workCompletedUnlocked() { return !ActiveThreads && Tasks.empty(); } diff --git a/lib/Support/ThreadPool.cpp b/lib/Support/ThreadPool.cpp index 46a1990cd71..f442b3b0bc9 100644 --- a/lib/Support/ThreadPool.cpp +++ b/lib/Support/ThreadPool.cpp @@ -72,6 +72,14 @@ void ThreadPool::wait() { CompletionCondition.wait(LockGuard, [&] { return workCompletedUnlocked(); }); } +bool ThreadPool::isWorkerThread() const { + std::thread::id CurrentThreadId = std::this_thread::get_id(); + for (const std::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));