[Support][Parallel] Change check for nested TaskGroups.

[avl]

Summary

This patch changes check for nested TaskGroups so that it allows
parallel execution for TaskGroups. Following pattern would not work
parallelly with current check:

std::function<void()> Fn = [&]() {
  parallel::TaskGroup tg;

  tg.spawn([&]() { });
};

ThreadPool Pool;

Pool.async(Fn);
Pool.async(Fn);

Pool.wait();

One of the TaskGroup would work sequentially as current check
verifies overall number of TaskGroup. Two not nested
TaskGroups can work parallelly but current check prevents this.

Also this patch avoids parallel mode for TaskGroup
in parallel::strategy.ThreadsRequested == 1 case.

This patch is a followup of discussion from D142318

[avl]

rebased. added setting noparallel mode for TaskGroups in case parallel::strategy.ThreadsRequested == 1.

[andrewng]

In your description, perhaps change avoid -> prevents in TaskGroups can work parallelly but current check avoid this.

Apart from the nits, LGTM.

llvm/include/llvm/Support/Parallel.h

103

Make const.

llvm/unittests/Support/ParallelTest.cpp

149

Perhaps ++Count? And the same for below...

[avl]

Thanks!

llvm/include/llvm/Support/Parallel.h

public:
   // Spawn a task, but does not wait for it to finish.
   // Tasks marked with \p Sequential will be executed
   // exactly in the order which they were spawned.
   // Note: Sequential tasks may be executed on different
   // threads, but strictly in sequential order.
   void spawn(std::function<void()> f, bool Sequential = false);
 
   void sync() const { L.sync(); }
+
+  bool isParallel() { return Parallel; }
 };
 
 namespace detail {
 
 #if LLVM_ENABLE_THREADS
 const ptrdiff_t MinParallelSize = 1024;
 
 /// Inclusive median.

llvm/lib/Support/Parallel.cpp

struct Creator {
     static void *call() { return new ThreadPoolExecutor(strategy); }
   };
   struct Deleter {
     static void call(void *Ptr) { ((ThreadPoolExecutor *)Ptr)->stop(); }
   };
 
   void add(std::function<void()> F, bool Sequential = false) override {
     {
-      if (parallel::strategy.ThreadsRequested == 1) {
-        F();
-        return;
-      }
-
       std::lock_guard<std::mutex> Lock(Mutex);
       if (Sequential)
         WorkQueueSequential.emplace_front(std::move(F));
       else
         WorkQueue.emplace_back(std::move(F));
     }
     Cond.notify_one();
   }
static ManagedStatic<ThreadPoolExecutor, ThreadPoolExecutor::Creator,
       ManagedExec;
   static std::unique_ptr<ThreadPoolExecutor> Exec(&(*ManagedExec));
   return Exec.get();
 }
 } // namespace
 } // namespace detail
 #endif
 
-static std::atomic<int> TaskGroupInstances;
-
 // Latch::sync() called by the dtor may cause one thread to block. If is a dead
 // lock if all threads in the default executor are blocked. To prevent the dead
-// lock, only allow the first TaskGroup to run tasks parallelly. In the scenario
+// lock, only allow the root 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()
+    : Parallel((parallel::strategy.ThreadsRequested != 1) &&
+               (threadIndex == UINT_MAX)) {}
 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, bool Sequential) {
 #if LLVM_ENABLE_THREADS
   if (Parallel) {
     L.inc();
     detail::Executor::getDefaultExecutor()->add(
         [&, F = std::move(F)] {

llvm/unittests/Support/ParallelTest.cpp

 //===- llvm/unittest/Support/ParallelTest.cpp -----------------------------===//
 //
 // 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
 //
 //===----------------------------------------------------------------------===//
 ///
 /// \file
 /// Parallel.h unit tests.
 ///
 //===----------------------------------------------------------------------===//
 
 #include "llvm/Support/Parallel.h"
+#include "llvm/Support/ThreadPool.h"
 #include "gtest/gtest.h"
 #include <array>
 #include <random>
 
 uint32_t array[1024 * 1024];
 
 using namespace llvm;
 
TEST(Parallel, TaskGroupSequentialFor) {
   {
     parallel::TaskGroup tg;
     for (size_t Idx = 0; Idx < 500; Idx++)
       tg.spawn([&Count, Idx]() { EXPECT_EQ(Count++, Idx); }, true);
   }
   EXPECT_EQ(Count, 500ul);
 }
 
+TEST(Parallel, NestedTaskGroup) {
+  // This test checks:
+  // 1. Root TaskGroup is in Parallel mode.
+  // 2. Nested TaskGroup is not in Parallel mode.
+  parallel::TaskGroup tg;
+
+  tg.spawn([&]() {
+    EXPECT_TRUE(tg.isParallel() || (parallel::strategy.ThreadsRequested == 1));
+  });
+
+  tg.spawn([&]() {
+    parallel::TaskGroup nestedTG;
+    EXPECT_FALSE(nestedTG.isParallel());
+
+    nestedTG.spawn([&]() {
+      // Check that root TaskGroup is in Parallel mode.
+      EXPECT_TRUE(tg.isParallel() ||
+                  (parallel::strategy.ThreadsRequested == 1));
+
+      // Check that nested TaskGroup is not in Parallel mode.
+      EXPECT_FALSE(nestedTG.isParallel());
+    });
+  });
+}
+
+#if LLVM_ENABLE_THREADS
+TEST(Parallel, ParallelNestedTaskGroup) {
+  // This test checks that it is possible to have several TaskGroups
+  // run from different threads in Parallel mode.
+  std::atomic<size_t> Count{0};
+
+  {
+    std::function<void()> Fn = [&]() {
+      parallel::TaskGroup tg;
+
+      tg.spawn([&]() {
+        // Check that root TaskGroup is in Parallel mode.
+        EXPECT_TRUE(tg.isParallel() ||
+                    (parallel::strategy.ThreadsRequested == 1));
+
+        // Check that nested TaskGroup is not in Parallel mode.
+        parallel::TaskGroup nestedTG;
+        EXPECT_FALSE(nestedTG.isParallel());
+        Count++;
+
+        nestedTG.spawn([&]() {
+          // Check that root TaskGroup is in Parallel mode.
+          EXPECT_TRUE(tg.isParallel() ||
+                      (parallel::strategy.ThreadsRequested == 1));
+
+          // Check that nested TaskGroup is not in Parallel mode.
+          EXPECT_FALSE(nestedTG.isParallel());
+          Count++;
+        });
+      });
+    };
+
+    ThreadPool Pool;
+
+    Pool.async(Fn);
+    Pool.async(Fn);
+    Pool.async(Fn);
+    Pool.async(Fn);
+    Pool.async(Fn);
+    Pool.async(Fn);
+
+    Pool.wait();
+  }
+  EXPECT_EQ(Count, 12ul);
+}
+#endif
+
 #endif