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" 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" void mlirAsyncRuntimeEmplaceToken(AsyncToken *token) {
44   std::unique_lock<std::mutex> lock(token->mu);
45   token->ready = true;
46   token->cv.notify_all();
47   for (auto &awaiter : token->awaiters)
48     awaiter();
49 }
50 
51 extern "C" void mlirAsyncRuntimeAwaitToken(AsyncToken *token) {
52   std::unique_lock<std::mutex> lock(token->mu);
53   if (!token->ready)
54     token->cv.wait(lock, [token] { return token->ready; });
55   delete token;
56 }
57 
58 extern "C" void mlirAsyncRuntimeExecute(CoroHandle handle, CoroResume resume) {
59 #if LLVM_ENABLE_THREADS
60   std::thread thread([handle, resume]() { (*resume)(handle); });
61   thread.detach();
62 #else
63   (*resume)(handle);
64 #endif
65 }
66 
67 extern "C" void mlirAsyncRuntimeAwaitTokenAndExecute(AsyncToken *token,
68                                                      CoroHandle handle,
69                                                      CoroResume resume) {
70   std::unique_lock<std::mutex> lock(token->mu);
71 
72   auto execute = [token, handle, resume]() {
73     mlirAsyncRuntimeExecute(handle, resume);
74     delete token;
75   };
76 
77   if (token->ready)
78     execute();
79   else
80     token->awaiters.push_back([execute]() { execute(); });
81 }
82 
83 //===----------------------------------------------------------------------===//
84 // Small async runtime support library for testing.
85 //===----------------------------------------------------------------------===//
86 
87 extern "C" void 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