[Support] Attempt to fix deadlock in ThreadGroup

This is an attempt to fix the situation described by https://reviews.llvm.org/D104207#2826290 and PR41508.
See sequence of operations leading to the bug in https://reviews.llvm.org/D104207#3004689

We ensure that the Latch is completely "free" before decrementing the number of TaskGroupInstances.

Differential revision: https://reviews.llvm.org/D109914

(cherry picked from commit 7b25fa8c7a151e94be46ed8f0a56bf4e2af1c104)

2 files changed, 10 insertions, 2 deletions

diff --git a/llvm/include/llvm/Support/Parallel.h b/llvm/include/llvm/Support/Parallel.h
index 28d171d45256..5c3b26d5754c 100644
--- a/llvm/include/llvm/Support/Parallel.h
+++ b/llvm/include/llvm/Support/Parallel.h
@@ -40,7 +40,10 @@ class Latch {
 
 public:
   explicit Latch(uint32_t Count = 0) : Count(Count) {}
-  ~Latch() { sync(); }
+  ~Latch() {
+    // Ensure at least that sync() was called.
+    assert(Count == 0);
+  }
 
   void inc() {
     std::lock_guard<std::mutex> lock(Mutex);
diff --git a/llvm/lib/Support/Parallel.cpp b/llvm/lib/Support/Parallel.cpp
index 9a2e1003da5a..71e3a1362f7e 100644
--- a/llvm/lib/Support/Parallel.cpp
+++ b/llvm/lib/Support/Parallel.cpp
@@ -151,7 +151,12 @@ static std::atomic<int> TaskGroupInstances;
 // lock, only allow the first TaskGroup to run tasks parallelly. In the scenario
 // of nested parallel_for_each(), only the outermost one runs parallelly.
 TaskGroup::TaskGroup() : Parallel(TaskGroupInstances++ == 0) {}
-TaskGroup::~TaskGroup() { --TaskGroupInstances; }
+TaskGroup::~TaskGroup() {
+  // We must ensure that all the workloads have finished before decrementing the
+  // instances count.
+  L.sync();
+  --TaskGroupInstances;
+}
 
 void TaskGroup::spawn(std::function<void()> F) {
   if (Parallel) {