[ThreadPool] Support returning futures with results.

[fhahn]

Summary

This patch adjusts ThreadPool::async to return futures that wrap
the result type of the passed in callable.

To do so, ThreadPool::asyncImpl first constructs a
std::packaged_task<ResTy()> object, then extracts
the future and converts the std::packaged_task<ResTy()>
to std::packaged_task<void()> before submitting it to the
queue.

[Meinersbur]

The pre-merge check has found additional locations where template arguments need to be added. Did you consider using using void as default template argument, or make ThreadPool a typedef of "ThreadPoolWithResult<void>"?

However, I don't think that the ThreadPool itself needs to be a template, but just the async implementation such that one can enqueue tasks with different return types to the same thread pool:

template <typename RetTy>
inline std::shared_future<RetTy> async(std::function<RetTy()> F);

AFAIU, std::packaged_task<> already handles return values itself.

Edit: I missed that packaged_task itself a template of the return type. However, after calling get_future, it it itself packaged into a std::function<void()> object that is pushed to the Tasks queue, i.e. does not depend on the result type anymore.

[fhahn]

In D114183#3141408, @Meinersbur wrote:

The pre-merge check has found additional locations where template arguments need to be added. Did you consider using using void as default template argument, or make ThreadPool a typedef of "ThreadPoolWithResult<void>"?

Thanks for the suggestion! I updated ThreadPool -> ThreadPoolWithResult and add using ThreadPool = ThreadPoolWithResult<void>, to avoid adjusting all users.

However, I don't think that the ThreadPool itself needs to be a template, but just the async implementation such that one can enqueue tasks with different return types to the same thread pool:

template <typename RetTy>
inline std::shared_future<RetTy> async(std::function<RetTy()> F);

AFAIU, std::packaged_task<> already handles return values itself.

Edit: I missed that packaged_task itself a template of the return type. However, after calling get_future, it it itself packaged into a std::function<void()> object that is pushed to the Tasks queue, i.e. does not depend on the result type anymore.

I tried using void with both TaskTy and PackagedTaskTy, but both resulted in build errors when sharing the Future.

[fhahn]

Adjust ThreadPool forward declaration in MLIRContext.h

[Meinersbur]

In D114183#3142329, @fhahn wrote:

I tried using void with both TaskTy and PackagedTaskTy, but both resulted in build errors when sharing the Future.

What were thoise build errors?

Implementing a non-templated ThreadPool should be possible, as demonstrated here.

[fhahn]

In D114183#3143037, @Meinersbur wrote:

In D114183#3142329, @fhahn wrote:

I tried using void with both TaskTy and PackagedTaskTy, but both resulted in build errors when sharing the Future.

What were thoise build errors?

Variants of

llvm/include/llvm/Support/ThreadPool.h:184:12: error: no viable conversion from returned value of type 'shared_future<void>' to function return type 'shared_future<llvm::SmallString<0>>'
    return Future.share();
           ^~~~~~~~~~~~~~

Implementing a non-templated ThreadPool should be possible, as demonstrated here.

I think the issue is that LLVM's ThreadPool uses std::packaged_task to store the task and extract the future, and the type used for the future depends on the type for the packaged_task, which in turn depends on the type of Task.

Are you suggesting adjusting the way we package stuff in LLVM's ThreadPool so the packaged task types do not depend on the result type of the packaged function? I think that should be an feasible alternative.

[mehdi_amini]

This seems like a bit of a "niche" extension: it makes it possible to return future but only for a statically predefined type. Is there a way to type-erase this in the threadpool a bit better and manager this at "enqueue time" instead so that each new task can return its own future?

Also isn't this feature something that can likely be layered on top of the existing ThreadPool without changing it?
For example by wrapping the submitted task into another object that manages the promise/future.

[fhahn]

In D114183#3143742, @mehdi_amini wrote:

This seems like a bit of a "niche" extension: it makes it possible to return future but only for a statically predefined type. Is there a way to type-erase this in the threadpool a bit better and manager this at "enqueue time" instead so that each new task can return its own future?

Also isn't this feature something that can likely be layered on top of the existing ThreadPool without changing it?
For example by wrapping the submitted task into another object that manages the promise/future.

I think that's what Michael was asking/suggestion in the previous comment :)

My main goal is to use futures with non-void results, I'm happy to go with whatever implementation people think is most useful. Introducing a predefined type was a mostly mechanical change, but as I said if changing the packaging is preferred, I can adjust the code.

[mehdi_amini]

In D114183#3143852, @fhahn wrote:

In D114183#3143742, @mehdi_amini wrote:

This seems like a bit of a "niche" extension: it makes it possible to return future but only for a statically predefined type. Is there a way to type-erase this in the threadpool a bit better and manager this at "enqueue time" instead so that each new task can return its own future?

Also isn't this feature something that can likely be layered on top of the existing ThreadPool without changing it?
For example by wrapping the submitted task into another object that manages the promise/future.

I think that's what Michael was asking/suggestion in the previous comment :)

Ah, I missed this! I think skimming at the link Michael provided, the ScheduleAndGetFuture is what I had in mind.

[Meinersbur]

In D114183#3143705, @fhahn wrote:

Variants of

llvm/include/llvm/Support/ThreadPool.h:184:12: error: no viable conversion from returned value of type 'shared_future<void>' to function return type 'shared_future<llvm::SmallString<0>>'
    return Future.share();
           ^~~~~~~~~~~~~~

This seems to be a mismatch of asyncImpl's future return type and the return type of the function passed to asyncImpl. Future should be of type shared_future<llvm::SmallString<0>> but somehow it is shared_future<void>. That is, both should be the same template parameter, like the template<typename RetTy> std::shared_future<RetTy> async(std::function<RetTy()> F); suggested before. But this is requirements is only with the asyncImpl defintion.

More crucially, the line

Tasks.push(std::move(PackagedTask));

wraps the PackagedTask into a std::function<void()> which does not depend on the return type anymore. That is, we don't need ThreadPool::Tasks to be a template of the return type, i.e. it can store tasks of any return type (I think the actually result value is stored as a closure parameter of the std::function<void()> object, the std::future can be used to retrieve it), the ThreadPool itself does not need to store the future which was only returned by asyncImpl for the caller to remember.

[fhahn]

Updated so there's no need for adding a template argument to ThreadPool.

In D114183#3144443, @Meinersbur wrote:

More crucially, the line

Tasks.push(std::move(PackagedTask));

wraps the PackagedTask into a std::function<void()> which does not depend on the return type anymore. That is, we don't need ThreadPool::Tasks to be a template of the return type, i.e. it can store tasks of any return type (I think the actually result value is stored as a closure parameter of the std::function<void()> object, the std::future can be used to retrieve it), the ThreadPool itself does not need to store the future which was only returned by asyncImpl for the caller to remember.

Thanks! It looks like it is possible to first construct a std::packaged_task<RetTy()>, extract the future and then move the object into a std::packaged_task<void()>.

It required a few changes, but it looks like it is working now without the extra template argument. It also seems async also needs to convert it's argument to a std::function object before calling asyncImpl, so asyncImpl can determine the correct template argument types for std::function.

[mehdi_amini]

You have some build failure on Windows: https://buildkite.com/llvm-project/premerge-checks/builds/66414#450a0afb-7b1b-4cbb-8654-6de57c1f4e51

[fhahn]

attempt to work around windows build failure.

[fhahn]

another stab at trying to figure out what's going on with the windows build

[fhahn]

Replace uses of packaged_tasks with manually creating promise and callable to set result, to work around MSVC build failures.

[fhahn]

I think this should be ready for another look now. Windows tests are passing now and the failure on Debian is unrelated to the patch I think.

Unfortunately I wasn't able to find a way to convert std::packaged_task<ResTy()> to std::packaged_task<void()> with MSVC. To work around that, I updated the code to manually create a promise and wrap the task function into another callable that calls the function and updates the promise.

[DanielMcIntosh-IBM]

It seems like you may have missed some issues with const vs. non-const types in the case when LLVM_ENABLE_THREADS is false.
I think it should be pretty easy to fix, but I don't know ThreadPool well enough to quickly address it myself.
This is the error:

llvm/include/llvm/Support/ThreadPool.h:127:29: error: 'this' argument to member function 'get' has type 'const std::future<void>', but function is not marked const
    Tasks.push([Future]() { Future.get(); });
                            ^~~~~~
llvm/include/llvm/Support/ThreadPool.h:60:12: note: in instantiation of function template specialization 'llvm::ThreadPool::asyncImpl<void>' requested here
    return asyncImpl(std::function<decltype(F())()>(std::forward<Func>(F)));
           ^
llvm/include/llvm/Support/ThreadPool.h:54:12: note: in instantiation of function template specialization 'llvm::ThreadPool::async<std::__bind<(lambda at llvm/lib/CodeGen/ParallelCG.cpp:80:15), llvm::SmallString<0>>>' requested here
    return async(std::move(Task));
           ^
future:1239:10: note: 'get' declared here
    void get();

Do you mind submitting a fix?

[thakis]

In D114183#3151653, @DanielMcIntosh-IBM wrote:

It seems like you may have missed some issues with const vs. non-const types in the case when LLVM_ENABLE_THREADS is false.
I think it should be pretty easy to fix, but I don't know ThreadPool well enough to quickly address it myself.
This is the error:

llvm/include/llvm/Support/ThreadPool.h:127:29: error: 'this' argument to member function 'get' has type 'const std::future<void>', but function is not marked const
    Tasks.push([Future]() { Future.get(); });
                            ^~~~~~
llvm/include/llvm/Support/ThreadPool.h:60:12: note: in instantiation of function template specialization 'llvm::ThreadPool::asyncImpl<void>' requested here
    return asyncImpl(std::function<decltype(F())()>(std::forward<Func>(F)));
           ^
llvm/include/llvm/Support/ThreadPool.h:54:12: note: in instantiation of function template specialization 'llvm::ThreadPool::async<std::__bind<(lambda at llvm/lib/CodeGen/ParallelCG.cpp:80:15), llvm::SmallString<0>>>' requested here
    return async(std::move(Task));
           ^
future:1239:10: note: 'get' declared here
    void get();

Do you mind submitting a fix?

Your revert breaks normal building: http://45.33.8.238/linux/61593/step_4.txt

PTAL.

[ronlieb]

Both of our amdgpu buildbots are broken by this commit, when I revert it locally my build will pass.
https://lab.llvm.org/buildbot/#/builders/193
https://lab.llvm.org/buildbot/#/builders/165

https://lab.llvm.org/buildbot/#/builders/165/builds/10841
https://lab.llvm.org/buildbot/#/builders/193/builds/2281

[DanielMcIntosh-IBM]

Your revert breaks normal building: http://45.33.8.238/linux/61593/step_4.txt

PTAL.

Ooops, sorry. Working on it

[fhahn]

In D114183#3154540, @ronlieb wrote:

Both of our amdgpu buildbots are broken by this commit, when I revert it locally my build will pass.
https://lab.llvm.org/buildbot/#/builders/193
https://lab.llvm.org/buildbot/#/builders/165

https://lab.llvm.org/buildbot/#/builders/165/builds/10841
https://lab.llvm.org/buildbot/#/builders/193/builds/2281

I recommitted the patch. for now and also submitted a fix for the build failure, which was caused by the return type of the future having been changed from auto to std::future bc41653a1f289a63a1472b7364437515c9659f77. Everything should build again now.

[Meinersbur]

Could you close the review again if done?

[fhahn]

In D114183#3158894, @Meinersbur wrote:

Could you close the review again if done?

I'm not sure how unfortunately :( I can't see an option to close at the moment.

[mehdi_amini]

In D114183#3159449, @fhahn wrote:

In D114183#3158894, @Meinersbur wrote:

Could you close the review again if done?

I'm not sure how unfortunately :( I can't see an option to close at the moment.

You can only "Abandon Revision" in the list of actions at the bottom.

[mehdi_amini]

Actually the "Close Revision" was hidden because of the "requested changes" from Daniel, you should see it now

[Meinersbur]

I try to make it disappear from my Phabricator dashboard. Seems the only way is to resign as a reviewer :-(

llvm/include/llvm/Support/ThreadPool.h

 
 /// A ThreadPool for asynchronous parallel execution on a defined number of
 /// threads.
 ///
 /// The pool keeps a vector of threads alive, waiting on a condition variable
 /// for some work to become available.
 class ThreadPool {
 public:
-  using TaskTy = std::function<void()>;
-  using PackagedTaskTy = std::packaged_task<void()>;
-
   /// Construct a pool using the hardware strategy \p S for mapping hardware
   /// execution resources (threads, cores, CPUs)
   /// Defaults to using the maximum execution resources in the system, but
   /// accounting for the affinity mask.
   ThreadPool(ThreadPoolStrategy S = hardware_concurrency());
 
   /// Blocking destructor: the pool will wait for all the threads to complete.
   ~ThreadPool();
 
   /// 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 Function, typename... Args>
-  inline std::shared_future<void> async(Function &&F, Args &&... ArgList) {
+  inline auto async(Function &&F, Args &&...ArgList) {
     auto Task =
         std::bind(std::forward<Function>(F), std::forward<Args>(ArgList)...);
-    return asyncImpl(std::move(Task));
+    return async(std::move(Task));
   }
 
   /// 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 Function>
-  inline std::shared_future<void> async(Function &&F) {
-    return asyncImpl(std::forward<Function>(F));
+  template <typename Func>
+  auto async(Func &&F) -> std::shared_future<decltype(F())> {
+    return asyncImpl(std::function<decltype(F())()>(std::forward<Func>(F)));
   }
 
   /// Blocking wait for all the threads to complete and the queue to be empty.
   /// It is an error to try to add new tasks while blocking on this call.
   void wait();
 
   unsigned getThreadCount() const { return ThreadCount; }
 
   /// Returns true if the current thread is a worker thread of this thread pool.
   bool isWorkerThread() const;
 
 private:
+  /// Helpers to create a promise and a callable wrapper of \p Task that sets
+  /// the result of the promise. Returns the callable and a future to access the
+  /// result.
+  template <typename ResTy>
+  static std::pair<std::function<void()>, std::future<ResTy>>
+  createTaskAndFuture(std::function<ResTy()> Task) {
+    std::shared_ptr<std::promise<ResTy>> Promise =
+        std::make_shared<std::promise<ResTy>>();
+    auto F = Promise->get_future();
+    return {
+        [Promise = std::move(Promise), Task]() { Promise->set_value(Task()); },
+        std::move(F)};
+  }
+  static std::pair<std::function<void()>, std::future<void>>
+  createTaskAndFuture(std::function<void()> Task) {
+    std::shared_ptr<std::promise<void>> Promise =
+        std::make_shared<std::promise<void>>();
+    auto F = Promise->get_future();
+    return {[Promise = std::move(Promise), Task]() {
+              Task();
+              Promise->set_value();
+            },
+            std::move(F)};
+  }
+
   bool workCompletedUnlocked() { return !ActiveThreads && Tasks.empty(); }
 
   /// 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.
-  std::shared_future<void> asyncImpl(TaskTy F);
+  template <typename ResTy>
+  std::shared_future<ResTy> asyncImpl(std::function<ResTy()> Task) {
+
+#if LLVM_ENABLE_THREADS
+    /// Wrap the Task in a std::function<void()> that sets the result of the
+    /// corresponding future.
+    auto R = createTaskAndFuture(Task);
+
+    {
+      // Lock the queue and push the new task
+      std::unique_lock<std::mutex> LockGuard(QueueLock);
+
+      // Don't allow enqueueing after disabling the pool
+      assert(EnableFlag && "Queuing a thread during ThreadPool destruction");
+      Tasks.push(std::move(R.first));
+    }
+    QueueCondition.notify_one();
+    return R.second.share();
+
+#else // LLVM_ENABLE_THREADS Disabled
+
+    // Get a Future with launch::deferred execution using std::async
+    auto Future = std::async(std::launch::deferred, std::move(Task)).share();
+    // Wrap the future so that both ThreadPool::wait() can operate and the
+    // returned future can be sync'ed on.
+    Tasks.push([Future]() { Future.get(); });
+    return Future;
+#endif
+  }
 
   /// Threads in flight
   std::vector<llvm::thread> Threads;
 
   /// Tasks waiting for execution in the pool.
-  std::queue<PackagedTaskTy> Tasks;
+  std::queue<std::function<void()>> Tasks;
 
   /// Locking and signaling for accessing the Tasks queue.
   std::mutex QueueLock;
   std::condition_variable QueueCondition;
 
   /// Signaling for job completion
   std::condition_variable CompletionCondition;
 

llvm/lib/Support/ThreadPool.cpp

ThreadPool::ThreadPool(ThreadPoolStrategy S)
     : ThreadCount(S.compute_thread_count()) {
   // Create ThreadCount threads that will loop forever, wait on QueueCondition
   // for tasks to be queued or the Pool to be destroyed.
   Threads.reserve(ThreadCount);
   for (unsigned ThreadID = 0; ThreadID < ThreadCount; ++ThreadID) {
     Threads.emplace_back([S, ThreadID, this] {
       S.apply_thread_strategy(ThreadID);
       while (true) {
-        PackagedTaskTy Task;
+        std::function<void()> Task;
         {
           std::unique_lock<std::mutex> LockGuard(QueueLock);
           // Wait for tasks to be pushed in the queue
           QueueCondition.wait(LockGuard,
                               [&] { return !EnableFlag || !Tasks.empty(); });
           // Exit condition
           if (!EnableFlag && Tasks.empty())
             return;

 bool ThreadPool::isWorkerThread() const {
   llvm::thread::id CurrentThreadId = llvm::this_thread::get_id();
   for (const llvm::thread &Thread : Threads)
     if (CurrentThreadId == Thread.get_id())
       return true;
   return false;
 }
 
-std::shared_future<void> ThreadPool::asyncImpl(TaskTy Task) {
-  /// Wrap the Task in a packaged_task to return a future object.
-  PackagedTaskTy PackagedTask(std::move(Task));
-  auto Future = PackagedTask.get_future();
-  {
-    // Lock the queue and push the new task
-    std::unique_lock<std::mutex> LockGuard(QueueLock);
-
-    // Don't allow enqueueing after disabling the pool
-    assert(EnableFlag && "Queuing a thread during ThreadPool destruction");
-
-    Tasks.push(std::move(PackagedTask));
-  }
-  QueueCondition.notify_one();
-  return Future.share();
-}
-
 // The destructor joins all threads, waiting for completion.
 ThreadPool::~ThreadPool() {
   {
     std::unique_lock<std::mutex> LockGuard(QueueLock);
     EnableFlag = false;
   }
   QueueCondition.notify_all();
   for (auto &Worker : Threads)
void ThreadPool::wait() {
   // Sequential implementation running the tasks
   while (!Tasks.empty()) {
     auto Task = std::move(Tasks.front());
     Tasks.pop();
     Task();
   }
 }
 
-std::shared_future<void> ThreadPool::asyncImpl(TaskTy Task) {
-  // Get a Future with launch::deferred execution using std::async
-  auto Future = std::async(std::launch::deferred, std::move(Task)).share();
-  // Wrap the future so that both ThreadPool::wait() can operate and the
-  // returned future can be sync'ed on.
-  PackagedTaskTy PackagedTask([Future]() { Future.get(); });
-  Tasks.push(std::move(PackagedTask));
-  return Future;
-}
-
 ThreadPool::~ThreadPool() { wait(); }
 
 #endif

llvm/unittests/Support/ThreadPool.cpp

TEST_F(ThreadPoolTest, GetFuture) {
   // Force the future using get()
   Pool.async([&i] { ++i; }).get();
   ASSERT_NE(2, i.load());
   setMainThreadReady();
   Pool.wait();
   ASSERT_EQ(2, i.load());
 }
 
+TEST_F(ThreadPoolTest, GetFutureWithResult) {
+  CHECK_UNSUPPORTED();
+  ThreadPool Pool(hardware_concurrency(2));
+  auto F1 = Pool.async([] { return 1; });
+  auto F2 = Pool.async([] { return 2; });
+
+  setMainThreadReady();
+  Pool.wait();
+  ASSERT_EQ(1, F1.get());
+  ASSERT_EQ(2, F2.get());
+}
+
+TEST_F(ThreadPoolTest, GetFutureWithResultAndArgs) {
+  CHECK_UNSUPPORTED();
+  ThreadPool Pool(hardware_concurrency(2));
+  auto Fn = [](int x) { return x; };
+  auto F1 = Pool.async(Fn, 1);
+  auto F2 = Pool.async(Fn, 2);
+
+  setMainThreadReady();
+  Pool.wait();
+  ASSERT_EQ(1, F1.get());
+  ASSERT_EQ(2, F2.get());
+}
+
 TEST_F(ThreadPoolTest, PoolDestruction) {
   CHECK_UNSUPPORTED();
   // Test that we are waiting on destruction
   std::atomic_int checked_in{0};
   {
     ThreadPool Pool;
     for (size_t i = 0; i < 5; ++i) {
       Pool.async([this, &checked_in] {