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