[llvm-branch-commits] [mlir] bb0e621 - [mlir] AsyncRuntime: use LLVM ThreadPool to run async tasks

[Eugene Zhulenev via llvm-branch-commits]

Author: Eugene Zhulenev
Date: 2021-01-09T02:39:52-08:00
New Revision: bb0e621387f86c1e40b39fcfd9127d4d18294748

URL: https://github.com/llvm/llvm-project/commit/bb0e621387f86c1e40b39fcfd9127d4d18294748
DIFF: https://github.com/llvm/llvm-project/commit/bb0e621387f86c1e40b39fcfd9127d4d18294748.diff

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

Revert https://reviews.llvm.org/D92368 after the dynamic library unloading was fixed in https://reviews.llvm.org/D94312

Reviewed By: mehdi_amini

Differential Revision: https://reviews.llvm.org/D94346

Added: 
    

Modified: 
    mlir/lib/ExecutionEngine/AsyncRuntime.cpp

Removed: 
    


################################################################################
diff  --git a/mlir/lib/ExecutionEngine/AsyncRuntime.cpp b/mlir/lib/ExecutionEngine/AsyncRuntime.cpp
index 1e0c35afebd8..a20bd6d1e996 100644
--- a/mlir/lib/ExecutionEngine/AsyncRuntime.cpp
+++ b/mlir/lib/ExecutionEngine/AsyncRuntime.cpp
@@ -25,6 +25,7 @@
 #include <vector>
 
 #include "llvm/ADT/StringMap.h"
+#include "llvm/Support/ThreadPool.h"
 
 using namespace mlir::runtime;
 
@@ -49,6 +50,7 @@ class AsyncRuntime {
   AsyncRuntime() : numRefCountedObjects(0) {}
 
   ~AsyncRuntime() {
+    threadPool.wait(); // wait for the completion of all async tasks
     assert(getNumRefCountedObjects() == 0 &&
            "all ref counted objects must be destroyed");
   }
@@ -57,6 +59,8 @@ class AsyncRuntime {
     return numRefCountedObjects.load(std::memory_order_relaxed);
   }
 
+  llvm::ThreadPool &getThreadPool() { return threadPool; }
+
 private:
   friend class RefCounted;
 
@@ -70,6 +74,7 @@ class AsyncRuntime {
   }
 
   std::atomic<int32_t> numRefCountedObjects;
+  llvm::ThreadPool threadPool;
 };
 
 // -------------------------------------------------------------------------- //
@@ -307,7 +312,8 @@ extern "C" ValueStorage mlirAsyncRuntimeGetValueStorage(AsyncValue *value) {
 }
 
 extern "C" void mlirAsyncRuntimeExecute(CoroHandle handle, CoroResume resume) {
-  (*resume)(handle);
+  auto *runtime = getDefaultAsyncRuntime();
+  runtime->getThreadPool().async([handle, resume]() { (*resume)(handle); });
 }
 
 extern "C" void mlirAsyncRuntimeAwaitTokenAndExecute(AsyncToken *token,