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