1 | //===- AsyncRuntime.h - 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 declares basic Async runtime API for supporting Async dialect |
10 | // to LLVM dialect lowering. |
11 | // |
12 | //===----------------------------------------------------------------------===// |
13 | |
14 | #ifndef MLIR_EXECUTIONENGINE_ASYNCRUNTIME_H_ |
15 | #define MLIR_EXECUTIONENGINE_ASYNCRUNTIME_H_ |
16 | |
17 | #include <cstddef> |
18 | #include <stdint.h> |
19 | |
20 | #ifdef _WIN32 |
21 | #ifndef MLIR_ASYNC_RUNTIME_EXPORT |
22 | #ifdef mlir_async_runtime_EXPORTS |
23 | // We are building this library |
24 | #define MLIR_ASYNC_RUNTIME_EXPORT __declspec(dllexport) |
25 | #else |
26 | // We are using this library |
27 | #define MLIR_ASYNC_RUNTIME_EXPORT __declspec(dllimport) |
28 | #endif // mlir_async_runtime_EXPORTS |
29 | #endif // MLIR_ASYNC_RUNTIME_EXPORT |
30 | #else |
31 | // Non-windows: use visibility attributes. |
32 | #define MLIR_ASYNC_RUNTIME_EXPORT __attribute__((visibility("default"))) |
33 | #endif // _WIN32 |
34 | |
35 | namespace mlir { |
36 | namespace runtime { |
37 | |
38 | //===----------------------------------------------------------------------===// |
39 | // Async runtime API. |
40 | //===----------------------------------------------------------------------===// |
41 | |
42 | // Runtime implementation of `async.token` data type. |
43 | using AsyncToken = struct AsyncToken; |
44 | |
45 | // Runtime implementation of `async.group` data type. |
46 | using AsyncGroup = struct AsyncGroup; |
47 | |
48 | // Runtime implementation of `async.value` data type. |
49 | using AsyncValue = struct AsyncValue; |
50 | |
51 | // Async value payload stored in a memory owned by the async.value. |
52 | using ValueStorage = std::byte *; |
53 | |
54 | // Async runtime uses LLVM coroutines to represent asynchronous tasks. Task |
55 | // function is a coroutine handle and a resume function that continue coroutine |
56 | // execution from a suspension point. |
57 | using CoroHandle = void *; // coroutine handle |
58 | using CoroResume = void (*)(void *); // coroutine resume function |
59 | |
60 | // Async runtime uses reference counting to manage the lifetime of async values |
61 | // (values of async types like tokens, values and groups). |
62 | using RefCountedObjPtr = void *; |
63 | |
64 | // Adds references to reference counted runtime object. |
65 | extern "C" MLIR_ASYNC_RUNTIME_EXPORT void |
66 | mlirAsyncRuntimeAddRef(RefCountedObjPtr, int64_t); |
67 | |
68 | // Drops references from reference counted runtime object. |
69 | extern "C" MLIR_ASYNC_RUNTIME_EXPORT void |
70 | mlirAsyncRuntimeDropRef(RefCountedObjPtr, int64_t); |
71 | |
72 | // Create a new `async.token` in not-ready state. |
73 | extern "C" MLIR_ASYNC_RUNTIME_EXPORT AsyncToken *mlirAsyncRuntimeCreateToken(); |
74 | |
75 | // Create a new `async.value` in not-ready state. Size parameter specifies the |
76 | // number of bytes that will be allocated for the async value storage. Storage |
77 | // is owned by the `async.value` and deallocated when the async value is |
78 | // destructed (reference count drops to zero). |
79 | extern "C" MLIR_ASYNC_RUNTIME_EXPORT AsyncValue * |
80 | mlirAsyncRuntimeCreateValue(int64_t); |
81 | |
82 | // Create a new `async.group` in empty state. |
83 | extern "C" MLIR_ASYNC_RUNTIME_EXPORT AsyncGroup * |
84 | mlirAsyncRuntimeCreateGroup(int64_t size); |
85 | |
86 | extern "C" MLIR_ASYNC_RUNTIME_EXPORT int64_t |
87 | mlirAsyncRuntimeAddTokenToGroup(AsyncToken *, AsyncGroup *); |
88 | |
89 | // Switches `async.token` to ready state and runs all awaiters. |
90 | extern "C" MLIR_ASYNC_RUNTIME_EXPORT void |
91 | mlirAsyncRuntimeEmplaceToken(AsyncToken *); |
92 | |
93 | // Switches `async.value` to ready state and runs all awaiters. |
94 | extern "C" MLIR_ASYNC_RUNTIME_EXPORT void |
95 | mlirAsyncRuntimeEmplaceValue(AsyncValue *); |
96 | |
97 | // Switches `async.token` to error state and runs all awaiters. |
98 | extern "C" MLIR_ASYNC_RUNTIME_EXPORT void |
99 | mlirAsyncRuntimeSetTokenError(AsyncToken *); |
100 | |
101 | // Switches `async.value` to error state and runs all awaiters. |
102 | extern "C" MLIR_ASYNC_RUNTIME_EXPORT void |
103 | mlirAsyncRuntimeSetValueError(AsyncValue *); |
104 | |
105 | // Returns true if token is in the error state. |
106 | extern "C" MLIR_ASYNC_RUNTIME_EXPORT bool |
107 | mlirAsyncRuntimeIsTokenError(AsyncToken *); |
108 | |
109 | // Returns true if value is in the error state. |
110 | extern "C" MLIR_ASYNC_RUNTIME_EXPORT bool |
111 | mlirAsyncRuntimeIsValueError(AsyncValue *); |
112 | |
113 | // Returns true if group is in the error state (any of the tokens or values |
114 | // added to the group are in the error state). |
115 | extern "C" MLIR_ASYNC_RUNTIME_EXPORT bool |
116 | mlirAsyncRuntimeIsGroupError(AsyncGroup *); |
117 | |
118 | // Blocks the caller thread until the token becomes ready. |
119 | extern "C" MLIR_ASYNC_RUNTIME_EXPORT void |
120 | mlirAsyncRuntimeAwaitToken(AsyncToken *); |
121 | |
122 | // Blocks the caller thread until the value becomes ready. |
123 | extern "C" MLIR_ASYNC_RUNTIME_EXPORT void |
124 | mlirAsyncRuntimeAwaitValue(AsyncValue *); |
125 | |
126 | // Blocks the caller thread until the elements in the group become ready. |
127 | extern "C" MLIR_ASYNC_RUNTIME_EXPORT void |
128 | mlirAsyncRuntimeAwaitAllInGroup(AsyncGroup *); |
129 | |
130 | // Returns a pointer to the storage owned by the async value. |
131 | extern "C" MLIR_ASYNC_RUNTIME_EXPORT ValueStorage |
132 | mlirAsyncRuntimeGetValueStorage(AsyncValue *); |
133 | |
134 | // Executes the task (coro handle + resume function) in one of the threads |
135 | // managed by the runtime. |
136 | extern "C" MLIR_ASYNC_RUNTIME_EXPORT void mlirAsyncRuntimeExecute(CoroHandle, |
137 | CoroResume); |
138 | |
139 | // Executes the task (coro handle + resume function) in one of the threads |
140 | // managed by the runtime after the token becomes ready. |
141 | extern "C" MLIR_ASYNC_RUNTIME_EXPORT void |
142 | mlirAsyncRuntimeAwaitTokenAndExecute(AsyncToken *, CoroHandle, CoroResume); |
143 | |
144 | // Executes the task (coro handle + resume function) in one of the threads |
145 | // managed by the runtime after the value becomes ready. |
146 | extern "C" MLIR_ASYNC_RUNTIME_EXPORT void |
147 | mlirAsyncRuntimeAwaitValueAndExecute(AsyncValue *, CoroHandle, CoroResume); |
148 | |
149 | // Executes the task (coro handle + resume function) in one of the threads |
150 | // managed by the runtime after the all members of the group become ready. |
151 | extern "C" MLIR_ASYNC_RUNTIME_EXPORT void |
152 | mlirAsyncRuntimeAwaitAllInGroupAndExecute(AsyncGroup *, CoroHandle, CoroResume); |
153 | |
154 | // Returns the current number of available worker threads in the threadpool. |
155 | extern "C" MLIR_ASYNC_RUNTIME_EXPORT int64_t |
156 | mlirAsyncRuntimGetNumWorkerThreads(); |
157 | |
158 | //===----------------------------------------------------------------------===// |
159 | // Small async runtime support library for testing. |
160 | //===----------------------------------------------------------------------===// |
161 | |
162 | extern "C" MLIR_ASYNC_RUNTIME_EXPORT void |
163 | mlirAsyncRuntimePrintCurrentThreadId(); |
164 | |
165 | } // namespace runtime |
166 | } // namespace mlir |
167 | |
168 | #endif // MLIR_EXECUTIONENGINE_ASYNCRUNTIME_H_ |
169 | |