1 //===- AsyncRuntime.cpp - Async runtime reference implementation ----------===//
2 //
3 // Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
4 // See https://llvm.org/LICENSE.txt for license information.
5 // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
6 //
7 //===----------------------------------------------------------------------===//
8 //
9 // This file implements basic Async runtime API for supporting Async dialect
10 // to LLVM dialect lowering.
11 //
12 //===----------------------------------------------------------------------===//
13 
14 #include "mlir/ExecutionEngine/AsyncRuntime.h"
15 
16 #ifdef MLIR_ASYNCRUNTIME_DEFINE_FUNCTIONS
17 
18 #include <condition_variable>
19 #include <functional>
20 #include <iostream>
21 #include <mutex>
22 #include <thread>
23 #include <vector>
24 
25 //===----------------------------------------------------------------------===//
26 // Async runtime API.
27 //===----------------------------------------------------------------------===//
28 
29 struct AsyncToken {
30   bool ready = false;
31   std::mutex mu;
32   std::condition_variable cv;
33   std::vector<std::function<void()>> awaiters;
34 };
35 
36 // Create a new `async.token` in not-ready state.
37 extern "C" MLIR_ASYNCRUNTIME_EXPORT AsyncToken *mlirAsyncRuntimeCreateToken() {
38   AsyncToken *token = new AsyncToken;
39   return token;
40 }
41 
42 // Switches `async.token` to ready state and runs all awaiters.
43 extern "C" MLIR_ASYNCRUNTIME_EXPORT void
44 mlirAsyncRuntimeEmplaceToken(AsyncToken *token) {
45   std::unique_lock<std::mutex> lock(token->mu);
46   token->ready = true;
47   token->cv.notify_all();
48   for (auto &awaiter : token->awaiters)
49     awaiter();
50 }
51 
52 extern "C" MLIR_ASYNCRUNTIME_EXPORT void
53 mlirAsyncRuntimeAwaitToken(AsyncToken *token) {
54   std::unique_lock<std::mutex> lock(token->mu);
55   if (!token->ready)
56     token->cv.wait(lock, [token] { return token->ready; });
57   delete token;
58 }
59 
60 extern "C" MLIR_ASYNCRUNTIME_EXPORT void
61 mlirAsyncRuntimeExecute(CoroHandle handle, CoroResume resume) {
62   std::thread thread([handle, resume]() { (*resume)(handle); });
63   thread.detach();
64 }
65 
66 extern "C" MLIR_ASYNCRUNTIME_EXPORT void
67 mlirAsyncRuntimeAwaitTokenAndExecute(AsyncToken *token, CoroHandle handle,
68                                      CoroResume resume) {
69   std::unique_lock<std::mutex> lock(token->mu);
70 
71   auto execute = [token, handle, resume]() {
72     mlirAsyncRuntimeExecute(handle, resume);
73     delete token;
74   };
75 
76   if (token->ready)
77     execute();
78   else
79     token->awaiters.push_back([execute]() { execute(); });
80 }
81 
82 //===----------------------------------------------------------------------===//
83 // Small async runtime support library for testing.
84 //===----------------------------------------------------------------------===//
85 
86 extern "C" MLIR_ASYNCRUNTIME_EXPORT void
87 mlirAsyncRuntimePrintCurrentThreadId() {
88   static thread_local std::thread::id thisId = std::this_thread::get_id();
89   std::cout << "Current thread id: " << thisId << "\n";
90 }
91 
92 #endif // MLIR_ASYNCRUNTIME_DEFINE_FUNCTIONS
93