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 #if LLVM_ENABLE_THREADS
63   std::thread thread([handle, resume]() { (*resume)(handle); });
64   thread.detach();
65 #else
66   (*resume)(handle);
67 #endif
68 }
69 
70 extern "C" MLIR_ASYNCRUNTIME_EXPORT void
71 mlirAsyncRuntimeAwaitTokenAndExecute(AsyncToken *token, CoroHandle handle,
72                                      CoroResume resume) {
73   std::unique_lock<std::mutex> lock(token->mu);
74 
75   auto execute = [token, handle, resume]() {
76     mlirAsyncRuntimeExecute(handle, resume);
77     delete token;
78   };
79 
80   if (token->ready)
81     execute();
82   else
83     token->awaiters.push_back([execute]() { execute(); });
84 }
85 
86 //===----------------------------------------------------------------------===//
87 // Small async runtime support library for testing.
88 //===----------------------------------------------------------------------===//
89 
90 extern "C" MLIR_ASYNCRUNTIME_EXPORT void
91 mlirAsyncRuntimePrintCurrentThreadId() {
92   static thread_local std::thread::id thisId = std::this_thread::get_id();
93   std::cout << "Current thread id: " << thisId << "\n";
94 }
95 
96 #endif // MLIR_ASYNCRUNTIME_DEFINE_FUNCTIONS
97