[ThreadPool] add ability to group tasks into separate groups

[llunak]

Summary

This is needed for parallelizing of loading modules symbols in LLDB (D122975). Currently LLDB can parallelize indexing symbols when loading a module, but modules are loaded sequentially. If LLDB index cache is enabled, this means that the cache loading is not parallelized, even though it could. However doing that creates a threadpool-within-threadpool situation, so the number of threads would not be properly limited.

This change adds ThreadPoolTaskGroup as a tag type that can be used with ThreadPool calls to put tasks into groups that can be independently waited for (even recursively from within a task) but still run in the same thread pool.

[aganea]

Not sure @chandlerc will be able to review your patch, adding some people that have been working in this space recently.

I think this is an interesting change, but I'm a bit worried that it adds complexity to the the thread task loop. I am wondering if this problem couldn't be solved by packaging the TaskGroup logic in a lambda. In essence the call stack would be: the thread loop in ThreadPool::grow() calls Task() which would call the TaskGroup logic lambda, which would call the user lambda. Regular non-TaskGroup would not go through that logic.

llvm/include/llvm/Support/ThreadPool.h

52

Can you move this above the ThreadPool? It'll be easier for future code readers I think.

61

Do you think we could have all these TaskGroup-specific functions inside the TaskGroup class instead? As an alternative, given your current usage, the tasks could be even queued in a std::vector prior, and passed to the TaskGroup constructor.

llvm/tools/llvm-profdata/llvm-profdata.cpp

41

Is this related to the current patch?

[clayborg]

I like the overall direction of this patch. I do see value in aganea's comments about possibly adding more methods to TaskGroup

llvm/include/llvm/Support/ThreadPool.h

61

I do like that idea, but if we do that it might be best to have TaskGroup take a "ThreadPool &" as a constructor argument so that it can ensure we always use the right ThreadPool if we ask the TaskGroup to wait.

Do we need to lock down the task group once "wait()" is called with a TaskGroup or can users keep adding things to the TaskGroup even if work is currently going on? Should we freeze the contents of a TaskGroup once we start waiting on this?

[mehdi_amini]

That's a pretty nice improvement! :)

Reading the patch, I'm not sure I have a good grasp about all the interactions of this with the existing invariants of this class: concurrency is always complex...

In essence the call stack would be: the thread loop in ThreadPool::grow() calls Task() which would call the TaskGroup logic lambda, which would call the user lambda. Regular non-TaskGroup would not go through that logic.

I don't quite get what you mean? Can you elaborate a bit?
I agree that the complexity increase is worrisome, so any idea to manage it is welcome :)

llvm/include/llvm/Support/ThreadPool.h

89

Do we have a way to assert on this?

101

I think you need more documentation for the groups, at minima:

• in particular ownership model / lifetime expectation for the TaskGroup objects.
• the existing APIs that don't take group should like be updated to be documented with respect to groups (is there a concept of a "default group" represented by the nullptr?).

103

Why the change to deque in this patch?

131–133

Can you document it please?

llvm/lib/Support/ThreadPool.cpp

98

Seems like if GroupOfTask is non-null you're calling workCompletedUnlocked twice, why?

104

" notify also threads waiting" ?

105

What does "this function" means here?

118

Nit: seems like The find_if(...) == end() could be replaced by !llvm::any_of(...)

llvm/tools/llvm-profdata/llvm-profdata.cpp

41

I suspect this file was depending on this header transitively and the include was removed from the ThreadPool.h header.

[aganea]

In D123225#3435142, @mehdi_amini wrote:

In essence the call stack would be: the thread loop in ThreadPool::grow() calls Task() which would call the TaskGroup logic lambda, which would call the user lambda. Regular non-TaskGroup would not go through that logic.

I don't quite get what you mean? Can you elaborate a bit?
I agree that the complexity increase is worrisome, so any idea to manage it is welcome :)

I'm just saying that we shouldn't be modifying ThreadPool::grow and instead implement the TaskGroup logic into a intermediate function, if that's possible. Instead of this currently:

template <typename Func>
auto async(TaskGroup &Group, Func &&F) -> std::shared_future<decltype(F())> {
  return asyncImpl(std::function<decltype(F())()>(std::forward<Func>(F)),
                   &Group);
}

instead we could have this:

template <typename Func>
auto TaskGroup::async(TaskGroup &Group, Func &&F) -> std::shared_future<decltype(F())> {
  return asyncImpl([]() {
     // ...TaskGroup logic...
     F();
     // ...TaskGroup logic... (for example, notify the local TaskGroup condition when all group tasks are done)
  }, &Group);
}

thus it makes sense to make async() and wait() members of TaskGroup and store the necessary state in the TaskGroup object itself (such a task counter -- currently stored by DenseMap<TaskGroup *, unsigned> ActiveGroups, a condition variable, etc.).
I feel the TaskGroup's logic is a subset of the generalized case we already have.

llvm/include/llvm/Support/ThreadPool.h

61

@clayborg Right now the freezing is implicit through the use of a condition_variable & a mutex, see llvm/lib/Support/ThreadPool.cpp, L124. It is a interesting question, should we make it support dynamic additions on the fly, while waiting? I would be tempted to wait for a practical use-case, but perhaps there's already one?

llvm/tools/llvm-profdata/llvm-profdata.cpp

41

Ah, good catch, thanks :-)

[llunak]

Let's first deal with the conceptual stuff, no point in dealing with the small code things as long as there's not agreement that this is the way to implement it.

In D123225#3434132, @aganea wrote:

Not sure @chandlerc will be able to review your patch, adding some people that have been working in this space recently.

I used what CODE_OWNERS.txt lists for 'Support'.

In D123225#3435887, @aganea wrote:

I'm just saying that we shouldn't be modifying ThreadPool::grow and instead implement the TaskGroup logic into a intermediate function, if that's possible.

I think that would be possible only with non-recursive use of ThreadPool. But D122975 requires creating running tasks that themselves run tasks and wait for them, which requires two things:

• Waiting just for a specific subset of tasks, otherwise a task could deadlock waiting for itself. This requires the queue-processing code to signal such state.
• Not wasting thread pool slots on waiting, otherwise they all could end up waiting for tasks that wouldn't have slots to run and deadlock. My approach handles that by letting such threads process tasks while waiting for a group. Another approach could be launching additional temporary threads, but I don't see how that would make anything simpler or better.

I don't see how either of these could be done without altering the queue-processing code in ThreadPool itself. So unless you have a good idea there, I think the most I can move to TaskGroup is syntactic sugar.

[clayborg]

llvm/include/llvm/Support/ThreadPool.h

101

That would be good. Maybe add headerdoc before the TaskGroup class. If we end up making the TaskGroup constructor take a reference to the ThreadPool, we should mention that the ThreadPool must outlive the TaskGroup as far as lifetime goes. More documentation or examples would be nice for:

• Making a TaskGroup 1 and then during work for TaskGroup 1 creating a TaskGroup 2 and then waiting on that within a worker thread
• details on if you can add to a group while work is ongoing for that group

103

This is the new code that is adding the ability to run work in the groups

192

Is this needed now that we have the TaskGroup objects? Or does this get signaled when ever all TaskGroups complete all of their work? Maybe update the documentation stating this is used for ThreadPool::wait() only for when all work is done?

llvm/lib/Support/ThreadPool.cpp

54

initialize to nullptr

98

Yeah this seems is should just be:

NotifyGroup = GroupOfTask != nullptr;

135–137

Can we detect recursive calls to wait on the same group here?

[clayborg]

llvm/include/llvm/Support/ThreadPool.h

103

Ignore my comment, I see that this type was used below where a queue was being used,.

[llunak]

llvm/include/llvm/Support/ThreadPool.h

61

As an alternative, given your current usage, the tasks could be even queued in a std::vector prior, and passed to the TaskGroup constructor.

Why? That seems unnecessary and inconsistent with how ThreadPool is used now.

61

Do we need to lock down the task group once "wait()" is called with a TaskGroup or can users keep adding things to the TaskGroup even if work is currently going on? Should we freeze the contents of a TaskGroup once we start waiting on this?

I see no reason to do it differently from ThreadPool, and ThreadPool currently says it's an error to add tasks while waiting but then does nothing about it (and AFAICT it's actually an unnecessary restriction).

89

Not without extra variables added just for detecting that. Is that something that would be done in LLVM code?

101

I'm confused about this part about documentation, because I think all of this is either documented or obvious. Did you miss those or are they not as obvious as I find them to be? It seems to me that you expect this change to be more complex than it is - it's just an ability to tag tasks with a TaskGroup pointer and then wait on tasks with a specific tag. The only somewhat complex thing here is making sure wait() called from within a task doesn't not deadlock, and that's not an API thing.

• in particular ownership model / lifetime expectation for the TaskGroup objects.

The TaskGroup objects are passed by reference, so I don't see how there could be any ownership. And TaskGroup objects are groups, so lifetime of TaskGroup objects and lifetime of groups are the same thing.

• the existing APIs that don't take group should like be updated to be documented with respect to groups (is there a concept of a "default group" represented by the nullptr?).

Tasks without a group are tasks without a group. It really matters only for wait(), which already says that it waits for all threads.

• Making a TaskGroup 1 and then during work for TaskGroup 1 creating a TaskGroup 2 and then waiting on that within a worker thread

This is done by simply doing it, there's nothing special about it from the API point of view, and wait() documentation says that this is possible. What more documentation would you need?

• details on if you can add to a group while work is ongoing for that group

The description of wait() already says that no. What other details do you need?

103

std::queue has only front() and back(), which is insufficient for checking only tasks in a specific group.

192

Yes, it is still needed. TaskGroup are dumb IDs, so they change nothing about this. This gets signalled when either all tasks in a group get finished, or when all tasks get finished (I'll add this to the description).

llvm/lib/Support/ThreadPool.cpp

54

It gets set on all paths (=only one) before it's used, and if it wouldn't, then initializing it here would prevent a warning about the mistake from -Wsometimes-uninitialized.

98

Because 'Notify' is to be done if work is done for the group or for all tasks (nullptr), while "NotifyGroup' is to be done if the work is done for the group. But I can replace the second call with 'Notify' since it's the same value.

135–137

Not without extra debug variables (see a similar question for wait()). Unless you count a guaranteed deadlock as detecting.

llvm/tools/llvm-profdata/llvm-profdata.cpp

41

Yes.

[llunak]

Small tweaks based on feedback.
Changed ThreadPool::TaskGroup to standalone ThreadPoolTaskGroup that has trivial calls forwarding to ThreadPool functions.

[mehdi_amini]

I like the change to move some of the API on the group themselves! In particular waiting in the destructor makes it "safer" to me (can't have dangling group pointers in the map in the pool)

llvm/include/llvm/Support/ThreadPool.h

38

Adding a mention of the concept of groups here would be valuable as well I think.

89

Yes it is common done, we guard such code in header within #if LLVM_ENABLE_ABI_BREAKING_CHECKS, you can grep the code base for examples (here is one: https://github.com/llvm/llvm-project/blob/main/llvm/include/llvm/Analysis/LoopInfo.h#L83-L87 )

101

I'm confused about this part about documentation, because I think all of this is either documented or obvious. Did you miss those or are they not as obvious as I find them to be?

I think there is always a natural tendency to find things "obvious" when we design some thing with a mental model and other assumptions in mind. But someone coming there new with fresh eyes won't find things as obvious.

in particular ownership model / lifetime expectation for the TaskGroup objects.
The TaskGroup objects are passed by reference, so I don't see how there could be any ownership. And TaskGroup objects are groups, so lifetime of TaskGroup objects and lifetime of groups are the same thing.

Right but that has implication on the lifetime: for example you should destroy a group while there are still tasks in flight using this group. I think this particular point isn't a concern anymore now that you call wait() in the ThreadPoolTaskGroup destructor.

Making a TaskGroup 1 and then during work for TaskGroup 1 creating a TaskGroup 2 and then waiting on that within a worker thread
This is done by simply doing it, there's nothing special about it from the API point of view, and wait() documentation says that this is possible. What more documentation would you need?

Waiting within a worker thread wasn't allowed until now: it is again more of a question of forming a mental model about the system.
Adding this concept of groups (which I welcome: it is solving a major limitation of the pool) is really making the existing model of a simple ThreadPool no longer as "simple", hence why I'm may seem overly cautious :)

Some things can be surprising as well for someone who does not really know or think about the implementation details but something like this:

group1.async([&]() {
  auto f = group2.async([&]() {
    return 1;
  });
  f.wait(); // May deadlock
})

but:

group1.async([&]() {
  auto f = group2.async([&]() {
    return 1;
  });
  group2.wait();  // yield current thread to run tasks from group2 if needed
  f.wait(); // May not deadlock
})

Not that I have a great suggestion on how to express the system and its invariant concisely, but this isn't easy to figure all of this out without reading carefully the implementation right now.
(and to be fair, even the "simple" pre-existing model isn't carefully documented in the API)

[llunak]

llvm/include/llvm/Support/ThreadPool.h

89

I see. I've added an assert for self-wait, but I actually see no good reason to enforce that there can be only one wait() for a group, so I'll instead remove that comment.

101

I think the concept is still fairly simple, as long as one doesn't forget that waiting for oneself will deadlock. But I have added some comments saying this explicitly, I hope that's enough.

[llunak]

Added asserts that wait() will not deadlock waiting for itself.
Added more documentation about deadlocks and usage from within threadpool's threads.

[mehdi_amini]

LGTM, but please wait for @aganea to have another look at it.

[clayborg]

Looks great, thanks for implementing many changes!

[MaskRay]

llvm/include/llvm/Support/ThreadPool.h

103

Prefer using

153

Use emplace_back

166

Use emplace_back

llvm/unittests/Support/ThreadPool.cpp

307

Task A runs in the first thread. It

[MaskRay]

llvm/include/llvm/Support/ThreadPool.h

61

The inline keyword can be removed.

llvm/lib/Support/ThreadPool.cpp

72

workCompletedUnlocked(WaitingForGroup) is slow when WaitingForGroup is not null. You may check the inverse of other conditions or cache the result of workCompletedUnlocked(WaitingForGroup)

llvm/unittests/Support/ThreadPool.cpp

327

What does this do? Use std::this_thread::sleep_for?

[aganea]

llvm/include/llvm/Support/ThreadPool.h

104

This is a bit confusing, we already have a llvm::TaskQueue, can we change this to something else? Just Queue maybe?

220

Would you mind please inserting a line break here, and after each other function below, just to match the style in this file?

llvm/lib/Support/ThreadPool.cpp

53

Shouldn't this be a stack since we're re-entering processTasks()? The following might not assert:

int main() {
  ThreadPool TP{hardware_concurrency(1)};
  ThreadPoolTaskGroup G1(TP);
  ThreadPoolTaskGroup G2(TP);
  G1.async([]{
    G2.wait(); // commenting this line would assert below.
    G1.wait(); // will deadlock without assert if the line above is there.
  });
  G2.async([]{
    // nop
  });
  return 0;
}

142

I think there's a bug here, it was there before but now we're using isWorkerThread() a lot more. If the ThreadPool is shutting down, thus ThreadLock is locked in ThreadPool::~ThreadPool(), then if anyone calls .wait() in a managed thread, isWorkerThread() would deadlock. Probably better replacing with a llvm::sys::RWMutex ThreadLock, and use a std::shared_lock<> here instead.

[MaskRay]

llvm/include/llvm/Support/ThreadPool.h

104

We can even omit the type alias. I think this is only used once for the member variable.

[rriddle]

Really excited for this! Thanks for taking it on.

[aganea]

llvm/lib/Support/ThreadPool.cpp

140

One more thing perhaps: tasks .wait()-ing will be "suspended" by re-entering processTasks. Any remaining tasks in the queue will be scheduled randomly over the waiting task(s), and this could easily create stack-overflows, since the scheduling is non-deterministic (so we could have several tasks waiting, piled on the top of each other). Depending on the stack size per platform, and how much stack each task eats up, this could potentially cause random crashes in production.

Probably the less intrusive approach would be to use fibers for suspended tasks. I suppose we could do that as a secondary step after this patch.

[llunak]

llvm/lib/Support/ThreadPool.cpp

140

Yes, but this is going to cost some resource no matter what, so it's just a choice of what resource it will be. And note that the recursion will be limited to the number of groups, since waiting loops are not allowed.

[llunak]

Updated according to review comments.

[aganea]

LGTM, thanks for all the changes @llunak!

llvm/lib/Support/ThreadPool.cpp

140

I think it is fine for now, let's see first if this is really a problem in practice. Worst case, it could be fixed by ensuring that the scope calling the .wait doesn't hold on too many stack variables.

[clayborg]

Looks great, really looking forward to seeing this go in!

[llunak]

The "Build Status" here lists a failure, but I cannot reproduce any test failure locally, and in the remote log I do not see any test failure, it looks like the testsuite itself failing. Is that something I should ignore, or am I missing something?

[mehdi_amini]

llvm-lit: /var/lib/buildkite-agent/builds/llvm-project/llvm/utils/lit/lit/discovery.py:247: warning: test suite 'lld-Unit' contained no tests

llvm-lit: /var/lib/buildkite-agent/builds/llvm-project/llvm/utils/lit/lit/llvm/config.py:438: note: using clang: /usr/bin/clang

llvm-lit: /var/lib/buildkite-agent/builds/llvm-project/llvm/utils/lit/lit/llvm/config.py:438: note: using ld.lld: /usr/bin/ld.lld

Traceback (most recent call last):

  File "/var/lib/buildkite-agent/builds/llvm-project/build/./bin/llvm-lit", line 59, in <module>

    main(builtin_parameters)

  File "/var/lib/buildkite-agent/builds/llvm-project/llvm/utils/lit/lit/main.py", line 111, in main

    selected_tests, discovered_tests = GoogleTest.post_process_shard_results(

  File "/var/lib/buildkite-agent/builds/llvm-project/llvm/utils/lit/lit/formats/googletest.py", line 231, in post_process_shard_results

    testsuites = json.load(f)['testsuites']

  File "/usr/lib/python3.9/json/__init__.py", line 293, in load

    return loads(fp.read(),

I'm not sure what's going on, maybe rebase and re-upload a patch?

[mehdi_amini]

It may be also safe to ignore, it's not this patch that's responsible for this.

llvm/include/llvm/Support/ThreadPool.h

 //===-- llvm/Support/ThreadPool.h - 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 defines a crude C++11 based thread pool.
 //
 //===----------------------------------------------------------------------===//
 
 #ifndef LLVM_SUPPORT_THREADPOOL_H
 #define LLVM_SUPPORT_THREADPOOL_H
 
+#include "llvm/ADT/DenseMap.h"
 #include "llvm/Config/llvm-config.h"
 #include "llvm/Support/Threading.h"
 #include "llvm/Support/thread.h"
 
 #include <future>
 
 #include <condition_variable>
+#include <deque>
 #include <functional>
 #include <memory>
 #include <mutex>
-#include <queue>
 #include <utility>
 
 namespace llvm {
 
+class ThreadPoolTaskGroup;
+
 /// A ThreadPool 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 {
 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());
 
   /// 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));
   }
 
+  /// Overload, task will be in the given task group.
+  template <typename Function, typename... Args>
+  inline auto async(ThreadPoolTaskGroup &Group, Function &&F,
+                    Args &&...ArgList) {
+    auto Task =
+        std::bind(std::forward<Function>(F), std::forward<Args>(ArgList)...);
+    return async(Group, 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)));
+    return asyncImpl(std::function<decltype(F())()>(std::forward<Func>(F)),
+                     nullptr);
+  }
+
+  template <typename Func>
+  auto async(ThreadPoolTaskGroup &Group, Func &&F)
+      -> std::shared_future<decltype(F())> {
+    return asyncImpl(std::function<decltype(F())()>(std::forward<Func>(F)),
+                     &Group);
   }
 
   /// 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();
 
+  /// Blocking wait for only all the threads in the given group to complete.
+  /// It is possible to recursively wait even inside a task, if the group is
+  /// different. There may be only one active wait() call for a given group.
+  /// It is possible to add new tasks while blocking on this call, if those
+  /// tasks are from a different group.
+  void wait(ThreadPoolTaskGroup &Group);
+
   // TODO: misleading legacy name warning!
   // Returns the 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:
+  typedef std::deque<std::pair<std::function<void()>, ThreadPoolTaskGroup *>>
+      TaskQueue;
+
   /// 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>
   static std::pair<std::function<void()>, std::future<ResTy>>
   createTaskAndFuture(std::function<ResTy()> Task) {
     std::shared_ptr<std::promise<ResTy>> Promise =
         std::make_shared<std::promise<ResTy>>();
createTaskAndFuture(std::function<void()> Task) {
     auto F = Promise->get_future();
     return {[Promise = std::move(Promise), Task]() {
               Task();
               Promise->set_value();
             },
             std::move(F)};
   }
 
-  bool workCompletedUnlocked() { return !ActiveThreads && Tasks.empty(); }
+  /// Returns true if all tasks in the given group have finished (nullptr means
+  /// all tasks regardless of their group).
+  bool workCompletedUnlocked(ThreadPoolTaskGroup *Group) const;
 
   /// 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 ResTy>
-  std::shared_future<ResTy> asyncImpl(std::function<ResTy()> Task) {
+  std::shared_future<ResTy> asyncImpl(std::function<ResTy()> Task,
+                                      ThreadPoolTaskGroup *Group) {
 
 #if LLVM_ENABLE_THREADS
     /// Wrap the Task in a std::function<void()> that sets the result of the
     /// corresponding future.
     auto R = createTaskAndFuture(Task);
 
     int requestedThreads;
     {
       // 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));
+      Tasks.push_back(std::make_pair(std::move(R.first), Group));
       requestedThreads = ActiveThreads + Tasks.size();
     }
     QueueCondition.notify_one();
     grow(requestedThreads);
     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(); });
+    Tasks.push_back(std::make_pair([Future]() { Future.get(); }, Group));
     return Future;
 #endif
   }
 
 #if LLVM_ENABLE_THREADS
   // Grow to ensure that we have at least `requested` Threads, but do not go
   // over MaxThreadCount.
   void grow(int requested);
+
+  void processTasks(ThreadPoolTaskGroup *WaitingForGroup);
 #endif
 
   /// Threads in flight
   std::vector<llvm::thread> Threads;
   /// Lock protecting access to the Threads vector.
   mutable std::mutex ThreadsLock;
 
   /// Tasks waiting for execution in the pool.
-  std::queue<std::function<void()>> Tasks;
+  TaskQueue Tasks;
 
   /// Locking and signaling for accessing the Tasks queue.
   std::mutex QueueLock;
   std::condition_variable QueueCondition;
 
-  /// Signaling for job completion
+  /// Signaling for job completion (all tasks or all tasks in a group).
   std::condition_variable CompletionCondition;
 
   /// Keep track of the number of thread actually busy
   unsigned ActiveThreads = 0;
+  /// Number of threads active for tasks in the given group (only non-zero).
+  DenseMap<ThreadPoolTaskGroup *, unsigned> ActiveGroups;
 
 #if LLVM_ENABLE_THREADS // avoids warning for unused variable
   /// Signal for the destruction of the pool, asking thread to exit.
   bool EnableFlag = true;
 #endif
 
   const ThreadPoolStrategy Strategy;
 
   /// Maximum number of threads to potentially grow this pool to.
   const unsigned MaxThreadCount;
 };
+
+/// A group of tasks to be run on a thread pool. Thread pool tasks in different
+/// groups can run on the same threadpool but can be waited for separately.
+class ThreadPoolTaskGroup {
+public:
+  /// The ThreadPool argument is the thread pool to forward calls to.
+  ThreadPoolTaskGroup(ThreadPool &Pool) : Pool(Pool) {}
+  /// Blocking destructor: will wait for all the tasks in the group to complete
+  /// by calling ThreadPool::wait().
+  ~ThreadPoolTaskGroup() { wait(); }
+  /// Calls ThreadPool::async() for this group.
+  template <typename Function, typename... Args>
+  inline auto async(Function &&F, Args &&...ArgList) {
+    return Pool.async(*this, std::forward<Function>(F),
+                      std::forward<Args>(ArgList)...);
   }
+  /// Calls ThreadPool::wait() for this group.
+  void wait() { Pool.wait(*this); }
+
+private:
+  ThreadPool &Pool;
+};
+
+} // namespace llvm
 
 #endif // LLVM_SUPPORT_THREADPOOL_H

llvm/lib/Support/ThreadPool.cpp

 #else
 #include "llvm/Support/raw_ostream.h"
 #endif
 
 using namespace llvm;
 
 #if LLVM_ENABLE_THREADS
 
+// A note on thread groups: Tasks are by default in no group (represented
+// by nullptr ThreadPoolTaskGroup pointer in the Tasks queue) and functionality
+// here normally works on all tasks regardless of their group (functions
+// in that case receive nullptr ThreadPoolTaskGroup pointer as argument).
+// A task in a group has a pointer to that ThreadPoolTaskGroup in the Tasks queue,
+// and functions called to work only on tasks from one group take that pointer.
+
 ThreadPool::ThreadPool(ThreadPoolStrategy S)
     : Strategy(S), MaxThreadCount(S.compute_thread_count()) {}
 
 void ThreadPool::grow(int requested) {
   std::unique_lock<std::mutex> LockGuard(ThreadsLock);
   if (Threads.size() >= MaxThreadCount)
     return; // Already hit the max thread pool size.
   int newThreadCount = std::min<int>(requested, MaxThreadCount);
   while (static_cast<int>(Threads.size()) < newThreadCount) {
     int ThreadID = Threads.size();
     Threads.emplace_back([this, ThreadID] {
       Strategy.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
+      processTasks(nullptr);
+    });
+  }
+}
+
+// WaitingForGroup == nullptr means all tasks regardless of their group.
+void ThreadPool::processTasks(ThreadPoolTaskGroup *WaitingForGroup) {
+    ThreadPoolTaskGroup *GroupOfTask;
+      QueueCondition.wait(LockGuard, [&] {
+        return !EnableFlag || !Tasks.empty() ||
+               (WaitingForGroup != nullptr &&
+                workCompletedUnlocked(WaitingForGroup));
+      });
+      if (WaitingForGroup != nullptr && workCompletedUnlocked(WaitingForGroup))
+        return;
 
-          // 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();
+      Task = std::move(Tasks.front().first);
+      GroupOfTask = Tasks.front().second;
+      // Need to count active threads in each group separately, ActiveThreads
+      // would never be 0 if waiting for another group inside a wait.
+      if (GroupOfTask != nullptr)
+        ++ActiveGroups[GroupOfTask]; // Increment or set to 1 if new item
+      Tasks.pop_front();
 
-        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();
-      }
-    });
-  }
+    bool NotifyGroup;
+      if (GroupOfTask != nullptr) {
+        auto A = ActiveGroups.find(GroupOfTask);
+        if (--(A->second) == 0)
+          ActiveGroups.erase(A);
+      }
+      Notify = workCompletedUnlocked(GroupOfTask);
+      NotifyGroup = GroupOfTask != nullptr && Notify;
+    // If this was a task in a group, notify also threads waiting for tasks
+    // in this function on QueueCondition, to make a recursive wait() return
+    // after the group it's been waiting for has finished.
+    if (NotifyGroup)
+      QueueCondition.notify_all();
+
+bool ThreadPool::workCompletedUnlocked(ThreadPoolTaskGroup *Group) const {
+  if (Group == nullptr)
+    return !ActiveThreads && Tasks.empty();
+  return ActiveGroups.count(Group) == 0 &&
+         !llvm::any_of(Tasks,
+                       [Group](const auto &T) { return T.second == Group; });
 }
 
 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(); });
+  CompletionCondition.wait(LockGuard,
+                           [&] { return workCompletedUnlocked(nullptr); });
+}
+
+void ThreadPool::wait(ThreadPoolTaskGroup &Group) {
+  // Wait for all threads in the group to complete.
+  if (!isWorkerThread()) {
+    std::unique_lock<std::mutex> LockGuard(QueueLock);
+    CompletionCondition.wait(LockGuard,
+                             [&] { return workCompletedUnlocked(&Group); });
+    return;
+  }
+  // Handle the case of recursive call from another task in a different group,
+  // in which case process tasks while waiting to keep the thread busy and avoid
+  // possible deadlock.
+  processTasks(&Group);
 }
 
 bool ThreadPool::isWorkerThread() const {
   std::unique_lock<std::mutex> LockGuard(ThreadsLock);
   llvm::thread::id CurrentThreadId = llvm::this_thread::get_id();
   for (const llvm::thread &Thread : Threads)
     if (CurrentThreadId == Thread.get_id())
       return true;
   return false;
 }
 
 // The destructor joins all threads, waiting for completion.
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();
+    auto Task = std::move(Tasks.front().first);
+    Tasks.pop_front();
     Task();
   }
 }
 
+void ThreadPool::wait(ThreadPoolTaskGroup &) {
+  // Simply wait for all, this works even if recursive (the running task
+  // is already removed from the queue).
+  wait();
+}
+
 bool ThreadPool::isWorkerThread() const {
   report_fatal_error("LLVM compiled without multithreading");
 }
 
 ThreadPool::~ThreadPool() { wait(); }
 
 #endif

llvm/tools/llvm-profdata/llvm-profdata.cpp

 #include "llvm/Support/InitLLVM.h"
 #include "llvm/Support/MemoryBuffer.h"
 #include "llvm/Support/Path.h"
 #include "llvm/Support/ThreadPool.h"
 #include "llvm/Support/Threading.h"
 #include "llvm/Support/WithColor.h"
 #include "llvm/Support/raw_ostream.h"
 #include <algorithm>
+#include <queue>
 
 using namespace llvm;
 
 enum ProfileFormat {
   PF_None = 0,
   PF_Text,
   PF_Compact_Binary,
   PF_Ext_Binary,

llvm/unittests/Support/ThreadPool.cpp

 #include "llvm/ADT/SmallVector.h"
 #include "llvm/ADT/Triple.h"
 #include "llvm/Support/CommandLine.h"
 #include "llvm/Support/Host.h"
 #include "llvm/Support/Program.h"
 #include "llvm/Support/TargetSelect.h"
 #include "llvm/Support/Threading.h"
 
+#include <chrono>
+
 #include "gtest/gtest.h"
 
 using namespace llvm;
 
 // Fixture for the unittests, allowing to *temporarily* disable the unittests
 // on a particular platform
 class ThreadPoolTest : public testing::Test {
   Triple Host;
   SmallVector<Triple::ArchType, 4> UnsupportedArchs;
   SmallVector<Triple::OSType, 4> UnsupportedOSs;
   SmallVector<Triple::EnvironmentType, 1> UnsupportedEnvironments;
+
 protected:
   // This is intended for platform as a temporary "XFAIL"
   bool isUnsupportedOSOrEnvironment() {
     Triple Host(Triple::normalize(sys::getProcessTriple()));
 
     if (find(UnsupportedEnvironments, Host.getEnvironment()) !=
         UnsupportedEnvironments.end())
       return true;
ThreadPoolTest() {
     //   UnsupportedArchs.push_back(Triple::x86_64);
 
     // See https://llvm.org/bugs/show_bug.cgi?id=25829
     UnsupportedArchs.push_back(Triple::ppc64le);
     UnsupportedArchs.push_back(Triple::ppc64);
   }
 
   /// Make sure this thread not progress faster than the main thread.
-  void waitForMainThread() {
-    std::unique_lock<std::mutex> LockGuard(WaitMainThreadMutex);
-    WaitMainThread.wait(LockGuard, [&] { return MainThreadReady; });
-  }
+  void waitForMainThread() { waitForPhase(1); }
 
   /// Set the readiness of the main thread.
-  void setMainThreadReady() {
+  void setMainThreadReady() { setPhase(1); }
+
+  /// Wait until given phase is set using setPhase(); first "main" phase is 1.
+  /// See also PhaseResetHelper below.
+  void waitForPhase(int Phase) {
+    std::unique_lock<std::mutex> LockGuard(CurrentPhaseMutex);
+    CurrentPhaseCondition.wait(
+        LockGuard, [&] { return CurrentPhase == Phase || CurrentPhase < 0; });
+  }
+  /// If a thread waits on another phase, the test could bail out on a failed
+  /// assertion and ThreadPool destructor would wait() on all threads, which
+  /// would deadlock on the task waiting. Create this helper to automatically
+  /// reset the phase and unblock such threads.
+  struct PhaseResetHelper {
+    PhaseResetHelper(ThreadPoolTest *test) : test(test) {}
+    ~PhaseResetHelper() { test->setPhase(-1); }
+    ThreadPoolTest *test;
+  };
+
+  /// Advance to the given phase.
+  void setPhase(int Phase) {
     {
-      std::unique_lock<std::mutex> LockGuard(WaitMainThreadMutex);
-      MainThreadReady = true;
+      std::unique_lock<std::mutex> LockGuard(CurrentPhaseMutex);
+      assert(Phase == CurrentPhase + 1 || Phase < 0);
+      CurrentPhase = Phase;
     }
-    WaitMainThread.notify_all();
+    CurrentPhaseCondition.notify_all();
   }
 
-  void SetUp() override { MainThreadReady = false; }
+  void SetUp() override { CurrentPhase = 0; }
 
   std::vector<llvm::BitVector> RunOnAllSockets(ThreadPoolStrategy S);
 
-  std::condition_variable WaitMainThread;
-  std::mutex WaitMainThreadMutex;
-  bool MainThreadReady = false;
+  std::condition_variable CurrentPhaseCondition;
+  std::mutex CurrentPhaseMutex;
+  int CurrentPhase; // -1 = error, 0 = setup, 1 = ready, 2+ = custom
 };
 
 #define CHECK_UNSUPPORTED()                                                    \
   do {                                                                         \
     if (isUnsupportedOSOrEnvironment())                                        \
       GTEST_SKIP();                                                            \
   } while (0);
 
for (size_t i = 0; i < 5; ++i) {
       });
     }
     ASSERT_EQ(0, checked_in);
     setMainThreadReady();
   }
   ASSERT_EQ(5, checked_in);
 }
 
+// Check running tasks in different groups.
+TEST_F(ThreadPoolTest, Groups) {
+  CHECK_UNSUPPORTED();
+  // Need at least two threads, as the task in group2
+  // might block a thread until all tasks in group1 finish.
+  ThreadPoolStrategy S = hardware_concurrency(2);
+  if (S.compute_thread_count() < 2)
+    return;
+  ThreadPool Pool(S);
+  PhaseResetHelper Helper(this);
+  ThreadPoolTaskGroup Group1(Pool);
+  ThreadPoolTaskGroup Group2(Pool);
+
+  // Check that waiting for an empty group is a no-op.
+  Group1.wait();
+
+  std::atomic_int checked_in1{0};
+  std::atomic_int checked_in2{0};
+
+  for (size_t i = 0; i < 5; ++i) {
+    Group1.async([this, &checked_in1] {
+      waitForMainThread();
+      ++checked_in1;
+    });
+  }
+  Group2.async([this, &checked_in2] {
+    waitForPhase(2);
+    ++checked_in2;
+  });
+  ASSERT_EQ(0, checked_in1);
+  ASSERT_EQ(0, checked_in2);
+  // Start first group and wait for it.
+  setMainThreadReady();
+  Group1.wait();
+  ASSERT_EQ(5, checked_in1);
+  // Second group has not yet finished, start it and wait for it.
+  ASSERT_EQ(0, checked_in2);
+  setPhase(2);
+  Group2.wait();
+  ASSERT_EQ(5, checked_in1);
+  ASSERT_EQ(1, checked_in2);
+}
+
+// Check recursive tasks.
+TEST_F(ThreadPoolTest, RecursiveGroups) {
+  CHECK_UNSUPPORTED();
+  ThreadPool Pool;
+  ThreadPoolTaskGroup Group(Pool);
+
+  std::atomic_int checked_in1{0};
+
+  for (size_t i = 0; i < 5; ++i) {
+    Group.async([this, &Pool, &checked_in1] {
+      waitForMainThread();
+
+      ThreadPoolTaskGroup LocalGroup(Pool);
+
+      // Check that waiting for an empty group is a no-op.
+      LocalGroup.wait();
+
+      std::atomic_int checked_in2{0};
+      for (size_t i = 0; i < 5; ++i) {
+        LocalGroup.async([&checked_in2] { ++checked_in2; });
+      }
+      LocalGroup.wait();
+      ASSERT_EQ(5, checked_in2);
+
+      ++checked_in1;
+    });
+  }
+  ASSERT_EQ(0, checked_in1);
+  setMainThreadReady();
+  Group.wait();
+  ASSERT_EQ(5, checked_in1);
+}
+
+TEST_F(ThreadPoolTest, RecursiveWaitDeadlock) {
+  CHECK_UNSUPPORTED();
+  ThreadPoolStrategy S = hardware_concurrency(2);
+  if (S.compute_thread_count() < 2)
+    return;
+  ThreadPool Pool(S);
+  PhaseResetHelper Helper(this);
+  ThreadPoolTaskGroup Group(Pool);
+
+  // Test that a thread calling wait() for a group and is waiting for more tasks
+  // returns when the last task finishes in a different thread while the waiting
+  // thread was waiting for more tasks to process while waiting.
+
+  // Task A is run in a first thread, it finishes and leaves
+  // the background thread waiting for more tasks.
+  Group.async([this] {
+    waitForMainThread();
+    setPhase(2);
+  });
+  // Task B is run in a second thread, it launches yet another
+  // task C in a different group, which will be handled by the waiting
+  // thread started above.
+  Group.async([this, &Pool] {
+    waitForPhase(2);
+    ThreadPoolTaskGroup LocalGroup(Pool);
+    LocalGroup.async([this] {
+      waitForPhase(3);
+      // Give the other thread enough time to check that there's no task
+      // to process and suspend waiting for a notification. This is indeed racy,
+      // but probably the best that can be done.
+      auto start = std::chrono::system_clock::now();
+      while (std::chrono::system_clock::now() - start <
+             std::chrono::milliseconds(10))
+        ;
+    });
+    // And task B only now will wait for the tasks in the group (=task C)
+    // to finish. This test checks that it does not deadlock. If the
+    // `NotifyGroup` handling in ThreadPool::processTasks() didn't take place,
+    // this task B would be stuck waiting for tasks to arrive.
+    setPhase(3);
+    LocalGroup.wait();
+  });
+  setMainThreadReady();
+  Group.wait();
+}
+
 #if LLVM_ENABLE_THREADS == 1
 
 // FIXME: Skip some tests below on non-Windows because multi-socket systems
 // were not fully tested on Unix yet, and llvm::get_thread_affinity_mask()
 // isn't implemented for Unix (need AffinityMask in Support/Unix/Program.inc).
 #ifdef _WIN32
 
 std::vector<llvm::BitVector>