Use llvm::unique_function in the async APIs

[chandlerc]

This is needed to allow using these APIs with callable objects that transitively capture move-only constructs. These come up very widely when writing concurrent code such a std::future, std::promise, std::unique_lock, etc.

[llvmbot]

@llvm/pr-subscribers-llvm-support

Author: Chandler Carruth (chandlerc)

Changes

This is needed to allow using these APIs with callable objects that transitively capture move-only constructs. These come up very widely when writing concurrent code such a std::future, std::promise, std::unique_lock, etc.

---

Full diff: https://github.com/llvm/llvm-project/pull/166727.diff

2 Files Affected:

• (modified) llvm/include/llvm/Support/ThreadPool.h (+13-10)
• (modified) llvm/lib/Support/ThreadPool.cpp (+2-2)

diff --git a/llvm/include/llvm/Support/ThreadPool.h b/llvm/include/llvm/Support/ThreadPool.h
index c20efc7396b79..d3276a18dc2c6 100644
--- a/llvm/include/llvm/Support/ThreadPool.h
+++ b/llvm/include/llvm/Support/ThreadPool.h
@@ -14,6 +14,7 @@
 #define LLVM_SUPPORT_THREADPOOL_H
 
 #include "llvm/ADT/DenseMap.h"
+#include "llvm/ADT/FunctionExtras.h"
 #include "llvm/Config/llvm-config.h"
 #include "llvm/Support/Compiler.h"
 #include "llvm/Support/Jobserver.h"
@@ -51,7 +52,7 @@ class ThreadPoolTaskGroup;
 class LLVM_ABI ThreadPoolInterface {
   /// The actual method to enqueue a task to be defined by the concrete
   /// implementation.
-  virtual void asyncEnqueue(std::function<void()> Task,
+  virtual void asyncEnqueue(llvm::unique_function<void()> Task,
                             ThreadPoolTaskGroup *Group) = 0;
 
 public:
@@ -95,22 +96,22 @@ class LLVM_ABI ThreadPoolInterface {
   /// 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)),
-                     nullptr);
+    return asyncImpl(
+        llvm::unique_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);
+    return asyncImpl(
+        llvm::unique_function<decltype(F())()>(std::forward<Func>(F)), &Group);
   }
 
 private:
   /// 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(llvm::unique_function<ResTy()> Task,
                                       ThreadPoolTaskGroup *Group) {
     auto Future = std::async(std::launch::deferred, std::move(Task)).share();
     asyncEnqueue([Future]() { Future.wait(); }, Group);
@@ -160,7 +161,7 @@ class LLVM_ABI StdThreadPool : public ThreadPoolInterface {
 
   /// 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.
-  void asyncEnqueue(std::function<void()> Task,
+  void asyncEnqueue(llvm::unique_function<void()> Task,
                     ThreadPoolTaskGroup *Group) override {
     int requestedThreads;
     {
@@ -189,7 +190,8 @@ class LLVM_ABI StdThreadPool : public ThreadPoolInterface {
   mutable llvm::sys::RWMutex ThreadsLock;
 
   /// Tasks waiting for execution in the pool.
-  std::deque<std::pair<std::function<void()>, ThreadPoolTaskGroup *>> Tasks;
+  std::deque<std::pair<llvm::unique_function<void()>, ThreadPoolTaskGroup *>>
+      Tasks;
 
   /// Locking and signaling for accessing the Tasks queue.
   std::mutex QueueLock;
@@ -239,13 +241,14 @@ class LLVM_ABI SingleThreadExecutor : public ThreadPoolInterface {
 private:
   /// 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.
-  void asyncEnqueue(std::function<void()> Task,
+  void asyncEnqueue(llvm::unique_function<void()> Task,
                     ThreadPoolTaskGroup *Group) override {
     Tasks.emplace_back(std::make_pair(std::move(Task), Group));
   }
 
   /// Tasks waiting for execution in the pool.
-  std::deque<std::pair<std::function<void()>, ThreadPoolTaskGroup *>> Tasks;
+  std::deque<std::pair<llvm::unique_function<void()>, ThreadPoolTaskGroup *>>
+      Tasks;
 };
 
 #if LLVM_ENABLE_THREADS
diff --git a/llvm/lib/Support/ThreadPool.cpp b/llvm/lib/Support/ThreadPool.cpp
index 69602688cf3fd..4779e673cc055 100644
--- a/llvm/lib/Support/ThreadPool.cpp
+++ b/llvm/lib/Support/ThreadPool.cpp
@@ -73,7 +73,7 @@ static LLVM_THREAD_LOCAL std::vector<ThreadPoolTaskGroup *>
 // WaitingForGroup == nullptr means all tasks regardless of their group.
 void StdThreadPool::processTasks(ThreadPoolTaskGroup *WaitingForGroup) {
   while (true) {
-    std::function<void()> Task;
+    llvm::unique_function<void()> Task;
     ThreadPoolTaskGroup *GroupOfTask;
     {
       std::unique_lock<std::mutex> LockGuard(QueueLock);
@@ -189,7 +189,7 @@ void StdThreadPool::processTasksWithJobserver() {
 
     // While we hold a job slot, process tasks from the internal queue.
     while (true) {
-      std::function<void()> Task;
+      llvm::unique_function<void()> Task;
       ThreadPoolTaskGroup *GroupOfTask = nullptr;
 
       {

[dwblaikie]

Could you update llvm/unittests/Support/ThreadPool.cpp to exercise this?

[chandlerc]

Could you update llvm/unittests/Support/ThreadPool.cpp to exercise this?

Done.