From bb0e621387f86c1e40b39fcfd9127d4d18294748 Mon Sep 17 00:00:00 2001 From: Eugene Zhulenev Date: Fri, 8 Jan 2021 15:23:59 -0800 Subject: [PATCH] [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 --- mlir/lib/ExecutionEngine/AsyncRuntime.cpp | 8 +++++++- 1 file changed, 7 insertions(+), 1 deletion(-) diff --git a/mlir/lib/ExecutionEngine/AsyncRuntime.cpp b/mlir/lib/ExecutionEngine/AsyncRuntime.cpp index 1e0c35a..a20bd6d 100644 --- a/mlir/lib/ExecutionEngine/AsyncRuntime.cpp +++ b/mlir/lib/ExecutionEngine/AsyncRuntime.cpp @@ -25,6 +25,7 @@ #include #include "llvm/ADT/StringMap.h" +#include "llvm/Support/ThreadPool.h" using namespace mlir::runtime; @@ -49,6 +50,7 @@ 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"); } @@ -57,6 +59,8 @@ public: return numRefCountedObjects.load(std::memory_order_relaxed); } + llvm::ThreadPool &getThreadPool() { return threadPool; } + private: friend class RefCounted; @@ -70,6 +74,7 @@ private: } std::atomic 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, -- 2.7.4