[llvm] 06b6170 - [Support][Parallel] Change check for nested TaskGroups.

[Alexey Lapshin via llvm-commits]

Author: Alexey Lapshin
Date: 2023-05-04T11:28:39+02:00
New Revision: 06b617064a997574df409c7d846b6f6b492f5124

URL: https://github.com/llvm/llvm-project/commit/06b617064a997574df409c7d846b6f6b492f5124
DIFF: https://github.com/llvm/llvm-project/commit/06b617064a997574df409c7d846b6f6b492f5124.diff

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

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

Differential Revision: https://reviews.llvm.org/D148984

Added: 
    

Modified: 
    llvm/include/llvm/Support/Parallel.h
    llvm/lib/Support/Parallel.cpp
    llvm/unittests/Support/ParallelTest.cpp

Removed: 
    


################################################################################
diff  --git a/llvm/include/llvm/Support/Parallel.h b/llvm/include/llvm/Support/Parallel.h
index 2e5476a19d7e8..75e7e8d597c44 100644
--- a/llvm/include/llvm/Support/Parallel.h
+++ b/llvm/include/llvm/Support/Parallel.h
@@ -99,6 +99,8 @@ class TaskGroup {
   void spawn(std::function<void()> f, bool Sequential = false);
 
   void sync() const { L.sync(); }
+
+  bool isParallel() const { return Parallel; }
 };
 
 namespace detail {

diff  --git a/llvm/lib/Support/Parallel.cpp b/llvm/lib/Support/Parallel.cpp
index 95956bbe7c4de..9c04af6ba3a0a 100644
--- a/llvm/lib/Support/Parallel.cpp
+++ b/llvm/lib/Support/Parallel.cpp
@@ -99,11 +99,6 @@ class ThreadPoolExecutor : public Executor {
 
   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));
@@ -185,18 +180,17 @@ Executor *Executor::getDefaultExecutor() {
 } // 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) {

diff  --git a/llvm/unittests/Support/ParallelTest.cpp b/llvm/unittests/Support/ParallelTest.cpp
index 2384a4d2b7411..7afe6501aaa8f 100644
--- a/llvm/unittests/Support/ParallelTest.cpp
+++ b/llvm/unittests/Support/ParallelTest.cpp
@@ -12,6 +12,7 @@
 //===----------------------------------------------------------------------===//
 
 #include "llvm/Support/Parallel.h"
+#include "llvm/Support/ThreadPool.h"
 #include "gtest/gtest.h"
 #include <array>
 #include <random>
@@ -102,4 +103,76 @@ TEST(Parallel, TaskGroupSequentialFor) {
   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 
diff erent 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