[mlir] AsyncRuntime: fix concurrency bugs + fix exports in methods definitions

[ezhulenev]

Summary

1. Move ThreadPool ownership to the runtime, and wait for the async tasks completion in the destructor.
2. Remove MLIR_ASYNCRUNTIME_EXPORT from method definitions because they are unnecessary in .cpp files, as only function declarations need to be exported, not their definitions.
3. Fix concurrency bugs in group emplace and potential use-after-free in token emplace.

Tested internally 10k runs in async.mlir and async-group.mlir.

Fixed: https://bugs.llvm.org/show_bug.cgi?id=48267

[mehdi_amini]

Why the removal for all MLIR_ASYNCRUNTIME_EXPORT? (in general please explain the motivation in commit description: the "why" instead of just the "what").

[ezhulenev]

In D91988#2412353, @mehdi_amini wrote:

Why the removal for all MLIR_ASYNCRUNTIME_EXPORT? (in general please explain the motivation in commit description: the "why" instead of just the "what").

That just got out of sync with https://github.com/llvm/llvm-project/commit/6fd9e59e1b3ac15430eccd4011e063605f58039c and reference counting change. Added a motivation.

[mehdi_amini]

OK, in general please keep changes that are conceptually separated in different revision / commits, each with they description.

LGTM.

[rriddle]

mlir/lib/ExecutionEngine/AsyncRuntime.cpp

56

Extra S here and below.

[ezhulenev]

Remove unnecessary addRef/dropRef in async await

[ezhulenev]

Fix pre-merge checks warnings

[ezhulenev]

Rebase

mlir/lib/ExecutionEngine/AsyncRuntime.cpp

~AsyncRuntime() {
     assert(getNumRefCountedObjects() == 0 &&
            "all ref counted objects must be destroyed");
   }
 
   int32_t getNumRefCountedObjects() {
     return numRefCountedObjects.load(std::memory_order_relaxed);
   }
 
+#if LLVM_ENABLE_THREADSS
+  llvm::ThreadPool &getThreadPool() { return threadPool; }
+#endif
+
 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;
+
+#if LLVM_ENABLE_THREADSS
+  llvm::ThreadPool threadPool;
+#endif
 };
 
 // Returns the default per-process instance of an async runtime.
 AsyncRuntime *getDefaultAsyncRuntimeInstance() {
   static auto runtime = std::make_unique<AsyncRuntime>();
   return runtime.get();
 }
 
struct AsyncGroup : public RefCounted {
   // Internal state below guarded by a mutex.
   std::mutex mu;
   std::condition_variable cv;
 
   std::vector<std::function<void()>> awaiters;
 };
 
 // Adds references to reference counted runtime object.
-extern "C" MLIR_ASYNCRUNTIME_EXPORT void
-mlirAsyncRuntimeAddRef(RefCountedObjPtr ptr, int32_t count) {
+extern "C" void mlirAsyncRuntimeAddRef(RefCountedObjPtr ptr, int32_t count) {
   RefCounted *refCounted = static_cast<RefCounted *>(ptr);
   refCounted->addRef(count);
 }
 
 // Drops references from reference counted runtime object.
-extern "C" MLIR_ASYNCRUNTIME_EXPORT void
-mlirAsyncRuntimeDropRef(RefCountedObjPtr ptr, int32_t count) {
+extern "C" void mlirAsyncRuntimeDropRef(RefCountedObjPtr ptr, int32_t count) {
   RefCounted *refCounted = static_cast<RefCounted *>(ptr);
   refCounted->dropRef(count);
 }
 
 // Create a new `async.token` in not-ready state.
 extern "C" AsyncToken *mlirAsyncRuntimeCreateToken() {
   AsyncToken *token = new AsyncToken(getDefaultAsyncRuntimeInstance());
   return token;
 }
 
 // Create a new `async.group` in empty state.
-extern "C" MLIR_ASYNCRUNTIME_EXPORT AsyncGroup *mlirAsyncRuntimeCreateGroup() {
+extern "C" AsyncGroup *mlirAsyncRuntimeCreateGroup() {
   AsyncGroup *group = new AsyncGroup(getDefaultAsyncRuntimeInstance());
   return group;
 }
 
-extern "C" MLIR_ASYNCRUNTIME_EXPORT int64_t
-mlirAsyncRuntimeAddTokenToGroup(AsyncToken *token, AsyncGroup *group) {
+extern "C" int64_t mlirAsyncRuntimeAddTokenToGroup(AsyncToken *token,
+                                                   AsyncGroup *group) {
   std::unique_lock<std::mutex> lockToken(token->mu);
   std::unique_lock<std::mutex> lockGroup(group->mu);
 
   // Get the rank of the token inside the group before we drop the reference.
   int rank = group->rank.fetch_add(1);
   group->pendingTokens.fetch_add(1);
 
-  auto onTokenReady = [group, token](bool dropRef) {
+  auto onTokenReady = [group]() {
     // Run all group awaiters if it was the last token in the group.
     if (group->pendingTokens.fetch_sub(1) == 1) {
       group->cv.notify_all();
       for (auto &awaiter : group->awaiters)
         awaiter();
     }
-
-    // We no longer need the token or the group, drop references on them.
-    if (dropRef) {
-      group->dropRef();
-      token->dropRef();
-    }
   };
 
   if (token->ready) {
-    onTokenReady(false);
+    // Update group pending tokens immediately and maybe run awaiters.
+    onTokenReady();
+
   } else {
+    // Update group pending tokens when token will become ready. Because this
+    // will happen asynchronously we must ensure that `group` and `token` are
+    // alive until then, and re-ackquire the lock.
     group->addRef();
     token->addRef();
-    token->awaiters.push_back([onTokenReady]() { onTokenReady(true); });
+
+    token->awaiters.push_back([group, token, onTokenReady]() {
+      // Make sure that `dropRef` does not destroy the mutex owned by the lock.
+      {
+        std::unique_lock<std::mutex> lockGroup(group->mu);
+        onTokenReady();
+      }
+      group->dropRef();
+      token->dropRef();
+    });
   }
 
   return rank;
 }
 
 // Switches `async.token` to ready state and runs all awaiters.
 extern "C" void mlirAsyncRuntimeEmplaceToken(AsyncToken *token) {
+  // Make sure that `dropRef` does not destroy the mutex owned by the lock.
+  {
     std::unique_lock<std::mutex> lock(token->mu);
     token->ready = true;
     token->cv.notify_all();
     for (auto &awaiter : token->awaiters)
       awaiter();
+  }
 
   // Async tokens created with a ref count `2` to keep token alive until the
   // async task completes. Drop this reference explicitly when token emplaced.
   token->dropRef();
 }
 
 extern "C" void mlirAsyncRuntimeAwaitToken(AsyncToken *token) {
   std::unique_lock<std::mutex> lock(token->mu);
   if (!token->ready)
     token->cv.wait(lock, [token] { return token->ready; });
 }
 
-extern "C" MLIR_ASYNCRUNTIME_EXPORT void
-mlirAsyncRuntimeAwaitAllInGroup(AsyncGroup *group) {
+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) {
-#if LLVM_ENABLE_THREADS
-  static llvm::ThreadPool *threadPool = new llvm::ThreadPool();
-  threadPool->async([handle, resume]() { (*resume)(handle); });
+#if LLVM_ENABLE_THREADSS
+  auto *runtime = getDefaultAsyncRuntimeInstance();
+  runtime->getThreadPool().async([handle, resume]() { (*resume)(handle); });
 #else
   (*resume)(handle);
 #endif
 }
 
 extern "C" void mlirAsyncRuntimeAwaitTokenAndExecute(AsyncToken *token,
                                                      CoroHandle handle,
                                                      CoroResume resume) {
   std::unique_lock<std::mutex> lock(token->mu);
 
   auto execute = [handle, resume, token](bool dropRef) {
     if (dropRef)
       token->dropRef();
-    mlirAsyncRuntimeExecute(handle, resume);
+    (*resume)(handle);
   };
 
   if (token->ready) {
     execute(false);
   } else {
     token->addRef();
     token->awaiters.push_back([execute]() { execute(true); });
   }
 }
 
-extern "C" MLIR_ASYNCRUNTIME_EXPORT void
-mlirAsyncRuntimeAwaitAllInGroupAndExecute(AsyncGroup *group, CoroHandle handle,
+extern "C" void mlirAsyncRuntimeAwaitAllInGroupAndExecute(AsyncGroup *group,
+                                                          CoroHandle handle,
                                                           CoroResume resume) {
   std::unique_lock<std::mutex> lock(group->mu);
 
   auto execute = [handle, resume, group](bool dropRef) {
     if (dropRef)
       group->dropRef();
-    mlirAsyncRuntimeExecute(handle, resume);
+    (*resume)(handle);
   };
 
   if (group->pendingTokens == 0) {
     execute(false);
   } else {
     group->addRef();
     group->awaiters.push_back([execute]() { execute(true); });
   }