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

[ezhulenev]

Summary

Depends On D94253

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

With the shutdown hook https://reviews.llvm.org/D94253 AsyncRuntime can correctly stop all threads before the LLVM support library is unloaded.

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;
 
 // Shutdown hook registration function will be resolved to the symbol defined by
 // the mlir-cpu-runner binary when it will dlopen the async runtime library.
 using ShutdownCallback = void (*)();
 extern "C" __attribute__((weak)) void
     __mlir_cpu_runner_register_shutdown_callback(ShutdownCallback);
 

 // 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:

 }
 
 // Returns a pointer to the storage owned by the async value.
 extern "C" ValueStorage mlirAsyncRuntimeGetValueStorage(AsyncValue *value) {
   return value->storage.data();
 }
 
 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)