ThreadPool: grow the pool only as needed

[Benoit]

Summary

On my 96-core cloudtop 'machine', it seems unnecessary to always start
96 threads upfront... particularly as the ThreadPool is created even
with -mlir-disable-threading. Things like the resuling spew in GDB and
the obfuscated output of (gdb) info threads are my motivation here,
but it probably also doesn't hurt for at least some efficiency metrics to
avoid creating many threads upfront.

[Benoit]

off by one mistake in max number of worker threads

[Benoit]

llvm/lib/Support/ThreadPool.cpp

24

Here actually I'm really not sure: what should be the max number of threads in Threads: should it be Strategy.compute_thread_count() or should it be that minus one to account for the main thread? it looks like the existing code was doing the former but (as actually confirmed by GDB) that meant it had one more thread running than the detected hardware concurrency. Which could be fine if we expect that the main thread would be idle most of the time? What was the intent here?

[Benoit]

llvm/include/llvm/Support/ThreadPool.h

69–72

Please decide what we want here? The immediate problem that led me to add +1 here was that when Threads was empty, I was now returning 0, and the caller was clearly expecting at least 1:

At mlir/include/mlir/IR/Threading.h:75

size_t numActions = std::min(numElements, threadPool.getThreadCount());
SmallVector<std::shared_future<void>> threadFutures;
threadFutures.reserve(numActions - 1);
for (unsigned i = 1; i < numActions; ++i)
  threadFutures.emplace_back(threadPool.async(processFn));

[Benoit]

Got it, getThreadCount() was really meant to return the max number of threads, not the current number. Those two numbers used to be the same before this patch.

[Benoit]

llvm/include/llvm/Support/ThreadPool.h

69–72

Self-replying: this call site really shows that getThreadCount was expected to return the potential max number of threads, which incidentally was equal to the current number of threads before this diff, but that's what is changing here. To avoid breaking existing users, I reverted getThreadCount to this behavior (I believe that the earlier state of this diff would have effectively kept the number of threads to 1. Now I've checked in GDB that we do create a dozen threads for a simple lit test).

llvm/lib/Support/ThreadPool.cpp

24

Got it, nevermind, updated this diff (see reply to the other comment thread).

[mehdi_amini]

particularly as the ThreadPool is created even with -mlir-disable-threading.

I'm fairly sure I fixed this one or two months ago, at least the C++ API allows to setup a project without starting a ThreadPool at all. This was a bottleneck in TensorFlow for some small MLIR work we're doing there.

There might be some similar plumbing we could do in how we handle --mlir-disable-threading with mlir-opt as well, since it is a testing tool we haven't tried to "optimize" this kind of things (you're making a good case for it though!).

That said, I like this change to lazily create thread as we go: so that even in cases where I want parallelism, I won't create more threads than actually needed!

llvm/lib/Support/ThreadPool.cpp

27

Nit: remove trivial braces

[mehdi_amini]

llvm/include/llvm/Support/ThreadPool.h

71

I'm not fond of keeping the API name as-is with a new semantics.

What about removing it and using two APIs instead:

• getMaxThreadCount()
• getAvailableThreadCount() => return Threads.size()

[Benoit]

Address review comments.

[Benoit]

In D115019#3168840, @mehdi_amini wrote:

I'm fairly sure I fixed this one or two months ago, at least the C++ API allows to setup a project without starting a ThreadPool at all. This was a bottleneck in TensorFlow for some small MLIR work we're doing there.

There might be some similar plumbing we could do in how we handle --mlir-disable-threading with mlir-opt as well, since it is a testing tool we haven't tried to "optimize" this kind of things (you're making a good case for it though!).

Indeed, it is in mlir-opt that I was observing that behavior this week.

llvm/include/llvm/Support/ThreadPool.h

71

There are 3 call sites in one Clang file, and 2 call sites in 2 MLIR files. That's a little more than I feel save updating all in one shot. Are you OK to keep the current name for now, and defer to a follow-up? At least the behavior isn't changing, these existing users are still going to get effectively the same result value.

For the new methods: I agree that the current getThreadCount should get renamed getMaxThreadCount.

I don't feel comfortable with getAvailableThreadCount() => return Threads.size() because that would conflict with the meaning of "available" in the existing member AvailableThreads. How about getCurrentThreadCount ?

[mehdi_amini]

llvm/include/llvm/Support/ThreadPool.h

149

Nit: it won't have any impact right now, but can you guard this declaration with LLVM_ENABLE_THREADS?

That would lead to compile time failures instead of link-time failures if someone changes something incorrectly in the future.

[mehdi_amini]

LG, thanks!

llvm/include/llvm/Support/ThreadPool.h

71

getCurrentThreadCount() LGTM

[Benoit]

review comment

[Benoit]

In D115019#3168875, @mehdi_amini wrote:

LG, thanks!

Thanks for the quick review! I don't have permissions to push by myself, FWIW.
I'm also running this through google global presubmits tonight... let's wait until tomorrow to submit this.

[mehdi_amini]

Sure, ping me tomorrow!

[Benoit]

Wait a bit more before submitting... we're still debugging CI failures.

[Benoit]

fix preexisting race condition in isWorkerThread found by MLIR tests with TSan

[Benoit]

This should be good to submit now.

[kda]

Looks like this is breaking the buildbot:
https://lab.llvm.org/buildbot/#/builders/37/builds/8903

[kda]

salient log extract:

[127/558] Building CXX object lib/Support/CMakeFiles/LLVMSupport.dir/ThreadPool.cpp.o
FAILED: lib/Support/CMakeFiles/LLVMSupport.dir/ThreadPool.cpp.o 
/b/sanitizer-x86_64-linux/build/llvm_build64/bin/clang++  -DGTEST_HAS_RTTI=0 -D_GNU_SOURCE -D__STDC_CONSTANT_MACROS -D__STDC_FORMAT_MACROS -D__STDC_LIMIT_MACROS -Ilib/Support -I/b/sanitizer-x86_64-linux/build/llvm-project/llvm/lib/Support -Iinclude -I/b/sanitizer-x86_64-linux/build/llvm-project/llvm/include -fPIC -fno-semantic-interposition -fvisibility-inlines-hidden -Werror=date-time -Werror=unguarded-availability-new -Wall -Wextra -Wno-unused-parameter -Wwrite-strings -Wcast-qual -Wmissing-field-initializers -pedantic -Wno-long-long -Wc++98-compat-extra-semi -Wimplicit-fallthrough -Wcovered-switch-default -Wno-noexcept-type -Wnon-virtual-dtor -Wdelete-non-virtual-dtor -Wsuggest-override -Wstring-conversion -Wmisleading-indentation -fdiagnostics-color -ffunction-sections -fdata-sections -Werror=global-constructors -m32 -fPIC -flto -Os -g0 -DNDEBUG -fno-rtti -fno-exceptions -nostdinc++ -I/b/sanitizer-x86_64-linux/build/symbolizer_build32/symbolizer/zlib -isystem /b/sanitizer-x86_64-linux/build/symbolizer_build32/symbolizer/libcxx/include/x86_64-unknown-linux-gnu/c++/v1 -isystem /b/sanitizer-x86_64-linux/build/symbolizer_build32/symbolizer/libcxx/include/c++/v1 -Wno-error=global-constructors   -std=c++14  -fno-exceptions -fno-rtti -MD -MT lib/Support/CMakeFiles/LLVMSupport.dir/ThreadPool.cpp.o -MF lib/Support/CMakeFiles/LLVMSupport.dir/ThreadPool.cpp.o.d -o lib/Support/CMakeFiles/LLVMSupport.dir/ThreadPool.cpp.o -c /b/sanitizer-x86_64-linux/build/llvm-project/llvm/lib/Support/ThreadPool.cpp
/b/sanitizer-x86_64-linux/build/llvm-project/llvm/lib/Support/ThreadPool.cpp:97:13: error: redefinition of 'ThreadPool'
ThreadPool::ThreadPool(ThreadPoolStrategy S)
            ^
/b/sanitizer-x86_64-linux/build/llvm-project/llvm/include/llvm/Support/ThreadPool.h:43:3: note: previous definition is here
  ThreadPool(ThreadPoolStrategy S = hardware_concurrency())
  ^
/b/sanitizer-x86_64-linux/build/llvm-project/llvm/lib/Support/ThreadPool.cpp:98:7: error: initializer 'ThreadCount' does not name a non-static data member or base class; did you mean the member 'MaxThreadCount'?
    : ThreadCount(S.compute_thread_count()) {
      ^~~~~~~~~~~
      MaxThreadCount
/b/sanitizer-x86_64-linux/build/llvm-project/llvm/include/llvm/Support/ThreadPool.h:177:18: note: 'MaxThreadCount' declared here
  const unsigned MaxThreadCount;
                 ^
/b/sanitizer-x86_64-linux/build/llvm-project/llvm/lib/Support/ThreadPool.cpp:99:7: error: use of undeclared identifier 'ThreadCount'
  if (ThreadCount != 1) {
      ^
/b/sanitizer-x86_64-linux/build/llvm-project/llvm/lib/Support/ThreadPool.cpp:100:56: error: use of undeclared identifier 'ThreadCount'; did you mean 'MaxThreadCount'?
    errs() << "Warning: request a ThreadPool with " << ThreadCount
                                                       ^~~~~~~~~~~
                                                       MaxThreadCount
/b/sanitizer-x86_64-linux/build/llvm-project/llvm/include/llvm/Support/ThreadPool.h:177:18: note: 'MaxThreadCount' declared here
  const unsigned MaxThreadCount;
                 ^
4 errors generated.

[mehdi_amini]

Sorry! Pushed a fix in b28f317c8156 and will monitor the bot tonight.

[mehdi_amini]

https://lab.llvm.org/buildbot/ is broken for me right now though (502 Bad Gateway), does it work for you?

[Benoit]

In D115019#3171261, @mehdi_amini wrote:

Sorry! Pushed a fix in b28f317c8156 and will monitor the bot tonight.

Sorry, all, for the breakage, and thanks for handling it while I was away!!
It looks like I never paid attention to the not-LLVM_ENABLE_THREADS path.

llvm/include/llvm/Support/ThreadPool.h

 /// The pool keeps a vector of threads alive, waiting on a condition variable
 /// for some work to become available.
 class ThreadPool {
 public:
   /// 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());
+  ThreadPool(ThreadPoolStrategy S = hardware_concurrency())
+      : Strategy(S), MaxThreadCount(S.compute_thread_count()) {}
 
   /// Blocking destructor: the pool will wait for all the threads to complete.
   ~ThreadPool();
 
   /// 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 <typename Function, typename... Args>
   inline auto async(Function &&F, Args &&...ArgList) {
     auto Task =
         std::bind(std::forward<Function>(F), std::forward<Args>(ArgList)...);
     return async(std::move(Task));
   }
 
   /// 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 <typename Func>
   auto async(Func &&F) -> std::shared_future<decltype(F())> {
     return asyncImpl(std::function<decltype(F())()>(std::forward<Func>(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();
 
-  unsigned getThreadCount() const { return ThreadCount; }
+  // Returns maximum number of worker threads in the pool (not the current
+  // number of threads).
+  unsigned getThreadCount() const { return MaxThreadCount; }
 
   /// Returns true if the current thread is a worker thread of this thread pool.
   bool isWorkerThread() const;
 
 private:
   /// Helpers to create a promise and a callable wrapper of \p Task that sets
   /// the result of the promise. Returns the callable and a future to access the
   /// result.
   template <typename ResTy>
#if LLVM_ENABLE_THREADS
 
     {
       // Lock the queue and push the new task
       std::unique_lock<std::mutex> LockGuard(QueueLock);
 
       // Don't allow enqueueing after disabling the pool
       assert(EnableFlag && "Queuing a thread during ThreadPool destruction");
       Tasks.push(std::move(R.first));
+      grow();
     }
     QueueCondition.notify_one();
     return R.second.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 ThreadPool::wait() can operate and the
     // returned future can be sync'ed on.
     Tasks.push([Future]() { Future.get(); });
     return Future;
 #endif
   }
 
+  // Maybe create a new thread and add it to Threads.
+  //
+  // Requirements:
+  //   * this->QueueLock should be owned by the calling thread prior to
+  //     calling this function. It will neither lock it nor unlock it.
+  //     Calling this function without owning QueueLock would result in data
+  //     races as this function reads Tasks and ActiveThreads.
+  //   * this->Tasks should be populated with any pending tasks. This function
+  //     uses Tasks.size() to determine whether it needs to create a new thread.
+  //   * this->ActiveThreads should be up to date as it is also used to
+  //     determine whether to create a new thread.
+  void grow();
+
   /// Threads in flight
   std::vector<llvm::thread> Threads;
 
   /// Tasks waiting for execution in the pool.
   std::queue<std::function<void()>> Tasks;
 
   /// Locking and signaling for accessing the Tasks queue.
   std::mutex QueueLock;
   std::condition_variable QueueCondition;
 
   /// Signaling for job completion
   std::condition_variable CompletionCondition;
 
   /// Keep track of the number of thread actually busy
   unsigned ActiveThreads = 0;
 
 #if LLVM_ENABLE_THREADS // avoids warning for unused variable
   /// Signal for the destruction of the pool, asking thread to exit.
   bool EnableFlag = true;
 #endif
 
-  unsigned ThreadCount;
+  const ThreadPoolStrategy Strategy;
+
+  /// Maximum number of threads to potentially grow this pool to.
+  const unsigned MaxThreadCount;
 };
 }
 
 #endif // LLVM_SUPPORT_THREADPOOL_H

llvm/lib/Support/ThreadPool.cpp

 #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) {
-        std::function<void()> Task;
-        {
-          std::unique_lock<std::mutex> 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
+void ThreadPool::grow() {
+  if (Threads.size() >= MaxThreadCount) {
+    // Already hit the max thread pool size.
+    return;
+  }
+  if (ActiveThreads + Tasks.size() <= Threads.size()) {
+    // We have enough threads for now.
+    return;
+  }
+  int ThreadID = Threads.size();
+  Threads.emplace_back([this, ThreadID] {
+    Strategy.apply_thread_strategy(ThreadID);
 
-          // 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<std::mutex> 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<std::mutex> LockGuard(QueueLock);
   CompletionCondition.wait(LockGuard, [&] { return workCompletedUnlocked(); });
 }
 
 bool ThreadPool::isWorkerThread() const {