[mlir:async] Fix deadlock in async runtime await-and-execute functions

`emplace???` functions running concurrently can set the ready flag and then pending awaiter will never be executed

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

1 files changed, 6 insertions, 3 deletions

diff --git a/mlir/lib/ExecutionEngine/AsyncRuntime.cpp b/mlir/lib/ExecutionEngine/AsyncRuntime.cpp
index e38ebf9..c5978eb 100644
--- a/mlir/lib/ExecutionEngine/AsyncRuntime.cpp
+++ b/mlir/lib/ExecutionEngine/AsyncRuntime.cpp
@@ -320,10 +320,11 @@ extern "C" void mlirAsyncRuntimeAwaitTokenAndExecute(AsyncToken *token,
                                                      CoroHandle handle,
                                                      CoroResume resume) {
   auto execute = [handle, resume]() { (*resume)(handle); };
+  std::unique_lock<std::mutex> lock(token->mu);
   if (token->ready) {
+    lock.unlock();
     execute();
   } else {
-    std::unique_lock<std::mutex> lock(token->mu);
     token->awaiters.push_back([execute]() { execute(); });
   }
 }
@@ -332,10 +333,11 @@ extern "C" void mlirAsyncRuntimeAwaitValueAndExecute(AsyncValue *value,
                                                      CoroHandle handle,
                                                      CoroResume resume) {
   auto execute = [handle, resume]() { (*resume)(handle); };
+  std::unique_lock<std::mutex> lock(value->mu);
   if (value->ready) {
+    lock.unlock();
     execute();
   } else {
-    std::unique_lock<std::mutex> lock(value->mu);
     value->awaiters.push_back([execute]() { execute(); });
   }
 }
@@ -344,10 +346,11 @@ extern "C" void mlirAsyncRuntimeAwaitAllInGroupAndExecute(AsyncGroup *group,
                                                           CoroHandle handle,
                                                           CoroResume resume) {
   auto execute = [handle, resume]() { (*resume)(handle); };
+  std::unique_lock<std::mutex> lock(group->mu);
   if (group->pendingTokens == 0) {
+    lock.unlock();
     execute();
   } else {
-    std::unique_lock<std::mutex> lock(group->mu);
     group->awaiters.push_back([execute]() { execute(); });
   }
 }