[mlir] AsyncRuntime: use LLVM ThreadPool to run async tasks

[ezhulenev]

Summary

Revert https://reviews.llvm.org/D92368

Can't reproduce the failures anymore. The original problem was the LLVM thread pool library was unloaded before the async runtime, but seems to be working now.

[mehdi_amini]

I'd still like to see the global ctor/dtor removed, this is incredibly fragile and known to cause issues. The solution of having a well-known interface to initialize/deinitialize mlir-cpu-runner plugins seems more robust (and portable) to me.

[ezhulenev]

Judging by the build failures this still does not work. Will work on a "mlir-cpu-runner dylib unloading" in a separate PR.

mlir/lib/ExecutionEngine/AsyncRuntime.cpp

 #include <cassert>
 #include <condition_variable>
 #include <functional>
 #include <iostream>
 #include <mutex>
 #include <thread>
 #include <vector>
 
+#include "llvm/Support/ThreadPool.h"
+
 using namespace mlir::runtime;
 
 //===----------------------------------------------------------------------===//
 // Async runtime API.
 //===----------------------------------------------------------------------===//
 
 namespace mlir {
 namespace runtime {
 namespace {
 
 // Forward declare class defined below.
 class RefCounted;
 
 // -------------------------------------------------------------------------- //
 // AsyncRuntime orchestrates all async operations and Async runtime API is built
 // on top of the default runtime instance.
 // -------------------------------------------------------------------------- //
 
 class AsyncRuntime {
 public:
   AsyncRuntime() : numRefCountedObjects(0) {}
 
   ~AsyncRuntime() {
+    threadPool.wait(); // wait for the completion of all async tasks
     assert(getNumRefCountedObjects() == 0 &&
            "all ref counted objects must be destroyed");
   }
 
   int32_t getNumRefCountedObjects() {
     return numRefCountedObjects.load(std::memory_order_relaxed);
   }
 
+  llvm::ThreadPool &getThreadPool() { return threadPool; }
+
 private:
   friend class RefCounted;
 
   // Count the total number of reference counted objects in this instance
   // of an AsyncRuntime. For debugging purposes only.
   void addNumRefCountedObjects() {
     numRefCountedObjects.fetch_add(1, std::memory_order_relaxed);
   }
   void dropNumRefCountedObjects() {
     numRefCountedObjects.fetch_sub(1, std::memory_order_relaxed);
   }
 
   std::atomic<int32_t> numRefCountedObjects;
+
+  llvm::ThreadPool threadPool;
 };
 
 // -------------------------------------------------------------------------- //
 // A base class for all reference counted objects created by the async runtime.
 // -------------------------------------------------------------------------- //
 
 class RefCounted {
 public:

 
 extern "C" void mlirAsyncRuntimeAwaitAllInGroup(AsyncGroup *group) {
   std::unique_lock<std::mutex> lock(group->mu);
   if (group->pendingTokens != 0)
     group->cv.wait(lock, [group] { return group->pendingTokens == 0; });
 }
 
 extern "C" void mlirAsyncRuntimeExecute(CoroHandle handle, CoroResume resume) {
-  (*resume)(handle);
+  auto *runtime = getDefaultAsyncRuntimeInstance();
+  runtime->getThreadPool().async([handle, resume]() { (*resume)(handle); });
 }
 
 extern "C" void mlirAsyncRuntimeAwaitTokenAndExecute(AsyncToken *token,
                                                      CoroHandle handle,
                                                      CoroResume resume) {
   std::unique_lock<std::mutex> lock(token->mu);
   auto execute = [handle, resume]() { (*resume)(handle); };
   if (token->ready)