[lld] [llvm] pr/taskgroup opt2 (PR #189293)
via llvm-commits
llvm-commits at lists.llvm.org
Sun Mar 29 17:05:40 PDT 2026
llvmbot wrote:
<!--LLVM PR SUMMARY COMMENT-->
@llvm/pr-subscribers-lld
Author: Fangrui Song (MaskRay)
<details>
<summary>Changes</summary>
- **[Support] Optimize parallel `TaskGroup`**
- **[Support] Support nested parallel TaskGroup via work-stealing**
---
Full diff: https://github.com/llvm/llvm-project/pull/189293.diff
4 Files Affected:
- (modified) lld/test/ELF/merge-piece-oob.s (+2-2)
- (modified) llvm/include/llvm/Support/Parallel.h (+10-15)
- (modified) llvm/lib/Support/Parallel.cpp (+48-28)
- (modified) llvm/unittests/Support/ParallelTest.cpp (+16-59)
``````````diff
diff --git a/lld/test/ELF/merge-piece-oob.s b/lld/test/ELF/merge-piece-oob.s
index 1ff34768a4d13..f798c9eba5a0e 100644
--- a/lld/test/ELF/merge-piece-oob.s
+++ b/lld/test/ELF/merge-piece-oob.s
@@ -3,8 +3,8 @@
## Non-section symbols and offset <= section_size are accepted, matching GNU ld.
# RUN: llvm-mc %s -o %t.o -filetype=obj -triple=x86_64
-# RUN: not ld.lld %t.o -o /dev/null -shared 2>&1 | FileCheck %s -DPREFIX=error --implicit-check-not=error:
-# RUN: ld.lld %t.o -o /dev/null -shared --noinhibit-exec 2>&1 | FileCheck %s -DPREFIX=warning --implicit-check-not=warning:
+# RUN: not ld.lld --threads=1 %t.o -o /dev/null -shared 2>&1 | FileCheck %s -DPREFIX=error --implicit-check-not=error:
+# RUN: ld.lld --threads=1 %t.o -o /dev/null -shared --noinhibit-exec 2>&1 | FileCheck %s -DPREFIX=warning --implicit-check-not=warning:
## .foo is 8 bytes with entsize=8 (1 piece). .foo+8 (offset==size) is accepted.
# CHECK: [[PREFIX]]: {{.*}}:(.foo): offset 0x9 is outside the section
diff --git a/llvm/include/llvm/Support/Parallel.h b/llvm/include/llvm/Support/Parallel.h
index b0c9e8f29f970..14ddee720a87e 100644
--- a/llvm/include/llvm/Support/Parallel.h
+++ b/llvm/include/llvm/Support/Parallel.h
@@ -17,6 +17,7 @@
#include "llvm/Support/Threading.h"
#include <algorithm>
+#include <atomic>
#include <condition_variable>
#include <functional>
#include <mutex>
@@ -58,31 +59,29 @@ inline size_t getThreadCount() { return 1; }
namespace detail {
class Latch {
- uint32_t Count;
+ std::atomic<uint32_t> Count;
mutable std::mutex Mutex;
mutable std::condition_variable Cond;
public:
explicit Latch(uint32_t Count = 0) : Count(Count) {}
- ~Latch() {
- // Ensure at least that sync() was called.
- assert(Count == 0);
- }
+ ~Latch() { assert(Count.load(std::memory_order_relaxed) == 0); }
- void inc() {
- std::lock_guard<std::mutex> lock(Mutex);
- ++Count;
- }
+ void inc() { Count.fetch_add(1, std::memory_order_relaxed); }
+ // dec() must hold Mutex so that sync() cannot observe Count==0 and
+ // destroy the Latch while dec() is still running.
void dec() {
std::lock_guard<std::mutex> lock(Mutex);
- if (--Count == 0)
+ if (Count.fetch_sub(1, std::memory_order_acq_rel) == 1)
Cond.notify_all();
}
+ uint32_t getCount() const { return Count.load(std::memory_order_acquire); }
+
void sync() const {
std::unique_lock<std::mutex> lock(Mutex);
- Cond.wait(lock, [&] { return Count == 0; });
+ Cond.wait(lock, [&] { return Count.load(std::memory_order_relaxed) == 0; });
}
};
} // namespace detail
@@ -96,12 +95,8 @@ class TaskGroup {
LLVM_ABI ~TaskGroup();
// 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.
LLVM_ABI void spawn(std::function<void()> f);
- void sync() const { L.sync(); }
-
bool isParallel() const { return Parallel; }
};
diff --git a/llvm/lib/Support/Parallel.cpp b/llvm/lib/Support/Parallel.cpp
index 2348b788d3ea7..1459bcb2d868f 100644
--- a/llvm/lib/Support/Parallel.cpp
+++ b/llvm/lib/Support/Parallel.cpp
@@ -96,14 +96,30 @@ class ThreadPoolExecutor {
static void call(void *Ptr) { ((ThreadPoolExecutor *)Ptr)->stop(); }
};
- void add(std::function<void()> F) {
+ void add(std::function<void()> F, parallel::detail::Latch &L) {
{
std::lock_guard<std::mutex> Lock(Mutex);
- WorkStack.push_back(std::move(F));
+ WorkStack.push_back({std::move(F), std::ref(L)});
}
Cond.notify_one();
}
+ // Execute tasks from the work queue until the latch reaches zero.
+ // Used by nested TaskGroups (on worker threads) to prevent deadlock:
+ // instead of blocking in sync(), actively help drain the queue.
+ void helpSync(const parallel::detail::Latch &L) {
+ while (L.getCount() != 0) {
+ std::unique_lock<std::mutex> Lock(Mutex);
+ if (WorkStack.empty())
+ return;
+ auto Item = std::move(WorkStack.back());
+ WorkStack.pop_back();
+ Lock.unlock();
+ Item.F();
+ Item.L.get().dec();
+ }
+ }
+
size_t getThreadCount() const { return ThreadCount; }
private:
@@ -141,34 +157,39 @@ class ThreadPoolExecutor {
[&] { TheJobserver->release(std::move(Slot)); });
while (true) {
- std::function<void()> Task;
- {
- std::unique_lock<std::mutex> Lock(Mutex);
- Cond.wait(Lock, [&] { return Stop || !WorkStack.empty(); });
- if (Stop && WorkStack.empty())
- return;
- if (WorkStack.empty())
- break;
- Task = std::move(WorkStack.back());
- WorkStack.pop_back();
- }
- Task();
+ std::unique_lock<std::mutex> Lock(Mutex);
+ Cond.wait(Lock, [&] { return Stop || !WorkStack.empty(); });
+ if (Stop && WorkStack.empty())
+ return;
+ if (WorkStack.empty())
+ break;
+ auto Item = std::move(WorkStack.back());
+ WorkStack.pop_back();
+ Lock.unlock();
+ Item.F();
+ Item.L.get().dec();
}
} else {
std::unique_lock<std::mutex> Lock(Mutex);
Cond.wait(Lock, [&] { return Stop || !WorkStack.empty(); });
if (Stop)
break;
- auto Task = std::move(WorkStack.back());
+ auto Item = std::move(WorkStack.back());
WorkStack.pop_back();
Lock.unlock();
- Task();
+ Item.F();
+ Item.L.get().dec();
}
}
}
+ struct WorkItem {
+ std::function<void()> F;
+ std::reference_wrapper<parallel::detail::Latch> L;
+ };
+
std::atomic<bool> Stop{false};
- std::vector<std::function<void()>> WorkStack;
+ std::vector<WorkItem> WorkStack;
std::mutex Mutex;
std::condition_variable Cond;
std::promise<void> ThreadsCreated;
@@ -205,22 +226,24 @@ size_t parallel::getThreadCount() {
}
#endif
-// 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 root TaskGroup to run tasks parallelly. In the scenario
-// of nested parallel_for_each(), only the outermost one runs parallelly.
+// Nested TaskGroups on worker threads use helpSync() which executes tasks from
+// the queue while waiting, preventing deadlock while enabling nested
+// parallelism.
TaskGroup::TaskGroup()
: Parallel(
#if LLVM_ENABLE_THREADS
- strategy.ThreadsRequested != 1 && threadIndex == UINT_MAX
+ strategy.ThreadsRequested != 1
#else
false
#endif
) {
}
TaskGroup::~TaskGroup() {
- // We must ensure that all the workloads have finished before decrementing the
- // instances count.
+#if LLVM_ENABLE_THREADS
+ // In a nested TaskGroup (threadIndex != -1u), actively help drain the queue.
+ if (Parallel && threadIndex != UINT_MAX)
+ getDefaultExecutor()->helpSync(L);
+#endif
L.sync();
}
@@ -228,10 +251,7 @@ void TaskGroup::spawn(std::function<void()> F) {
#if LLVM_ENABLE_THREADS
if (Parallel) {
L.inc();
- getDefaultExecutor()->add([&, F = std::move(F)] {
- F();
- L.dec();
- });
+ getDefaultExecutor()->add(std::move(F), L);
return;
}
#endif
diff --git a/llvm/unittests/Support/ParallelTest.cpp b/llvm/unittests/Support/ParallelTest.cpp
index c7ecc4eff6c29..ad833419cf1bb 100644
--- a/llvm/unittests/Support/ParallelTest.cpp
+++ b/llvm/unittests/Support/ParallelTest.cpp
@@ -95,73 +95,30 @@ TEST(Parallel, ForEachError) {
#if LLVM_ENABLE_THREADS
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));
- });
+ 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());
- });
+ EXPECT_TRUE(nestedTG.isParallel() ||
+ (parallel::strategy.ThreadsRequested == 1));
});
}
-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;
- });
+// Verify nested parallelFor doesn't deadlock. This is a simplified version of
+// the pattern from https://reviews.llvm.org/D61115 that originally motivated
+// serializing nested TaskGroups. With work-stealing in helpSync(), nested
+// parallelism now works without deadlock.
+TEST(Parallel, NestedParallelFor) {
+ std::atomic<uint32_t> count{0};
+ parallelFor(0, 8, [&](size_t i) {
+ parallelFor(0, 8, [&](size_t j) {
+ parallelFor(0, 8, [&](size_t k) {
+ count.fetch_add(1, std::memory_order_relaxed);
});
- };
-
- DefaultThreadPool 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);
+ });
+ });
+ EXPECT_EQ(count.load(), 512u);
}
#endif
``````````
</details>
https://github.com/llvm/llvm-project/pull/189293
More information about the llvm-commits
mailing list