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 _WIN32
20 #ifndef MLIR_ASYNCRUNTIME_EXPORT
21 #ifdef mlir_async_runtime_EXPORTS
22 // We are building this library
23 #define MLIR_ASYNCRUNTIME_EXPORT __declspec(dllexport)
24 #define MLIR_ASYNCRUNTIME_DEFINE_FUNCTIONS
25 #else
26 // We are using this library
27 #define MLIR_ASYNCRUNTIME_EXPORT __declspec(dllimport)
28 #endif // mlir_async_runtime_EXPORTS
29 #endif // MLIR_ASYNCRUNTIME_EXPORT
30 #else
31 #define MLIR_ASYNCRUNTIME_EXPORT
32 #define MLIR_ASYNCRUNTIME_DEFINE_FUNCTIONS
33 #endif // _WIN32
34 
35 //===----------------------------------------------------------------------===//
36 // Async runtime API.
37 //===----------------------------------------------------------------------===//
38 
39 // Runtime implementation of `async.token` data type.
40 typedef struct AsyncToken MLIR_AsyncToken;
41 
42 // Runtime implementation of `async.group` data type.
43 typedef struct AsyncGroup MLIR_AsyncGroup;
44 
45 // Async runtime uses LLVM coroutines to represent asynchronous tasks. Task
46 // function is a coroutine handle and a resume function that continue coroutine
47 // execution from a suspension point.
48 using CoroHandle = void *;           // coroutine handle
49 using CoroResume = void (*)(void *); // coroutine resume function
50 
51 // Async runtime uses reference counting to manage the lifetime of async values
52 // (values of async types like tokens, values and groups).
53 using RefCountedObjPtr = void *;
54 
55 // Adds references to reference counted runtime object.
56 extern "C" MLIR_ASYNCRUNTIME_EXPORT void
57     mlirAsyncRuntimeAddRef(RefCountedObjPtr, int32_t);
58 
59 // Drops references from reference counted runtime object.
60 extern "C" MLIR_ASYNCRUNTIME_EXPORT void
61     mlirAsyncRuntimeDropRef(RefCountedObjPtr, int32_t);
62 
63 // Create a new `async.token` in not-ready state.
64 extern "C" MLIR_ASYNCRUNTIME_EXPORT AsyncToken *mlirAsyncRuntimeCreateToken();
65 
66 // Create a new `async.group` in empty state.
67 extern "C" MLIR_ASYNCRUNTIME_EXPORT AsyncGroup *mlirAsyncRuntimeCreateGroup();
68 
69 extern "C" MLIR_ASYNCRUNTIME_EXPORT int64_t
70 mlirAsyncRuntimeAddTokenToGroup(AsyncToken *, AsyncGroup *);
71 
72 // Switches `async.token` to ready state and runs all awaiters.
73 extern "C" MLIR_ASYNCRUNTIME_EXPORT void
74 mlirAsyncRuntimeEmplaceToken(AsyncToken *);
75 
76 // Blocks the caller thread until the token becomes ready.
77 extern "C" MLIR_ASYNCRUNTIME_EXPORT void
78 mlirAsyncRuntimeAwaitToken(AsyncToken *);
79 
80 // Blocks the caller thread until the elements in the group become ready.
81 extern "C" MLIR_ASYNCRUNTIME_EXPORT void
82 mlirAsyncRuntimeAwaitAllInGroup(AsyncGroup *);
83 
84 // Executes the task (coro handle + resume function) in one of the threads
85 // managed by the runtime.
86 extern "C" MLIR_ASYNCRUNTIME_EXPORT void mlirAsyncRuntimeExecute(CoroHandle,
87                                                                  CoroResume);
88 
89 // Executes the task (coro handle + resume function) in one of the threads
90 // managed by the runtime after the token becomes ready.
91 extern "C" MLIR_ASYNCRUNTIME_EXPORT void
92 mlirAsyncRuntimeAwaitTokenAndExecute(AsyncToken *, CoroHandle, CoroResume);
93 
94 // Executes the task (coro handle + resume function) in one of the threads
95 // managed by the runtime after the all members of the group become ready.
96 extern "C" MLIR_ASYNCRUNTIME_EXPORT void
97 mlirAsyncRuntimeAwaitAllInGroupAndExecute(AsyncGroup *, CoroHandle, CoroResume);
98 
99 //===----------------------------------------------------------------------===//
100 // Small async runtime support library for testing.
101 //===----------------------------------------------------------------------===//
102 
103 extern "C" MLIR_ASYNCRUNTIME_EXPORT void mlirAsyncRuntimePrintCurrentThreadId();
104 
105 #endif // MLIR_EXECUTIONENGINE_ASYNCRUNTIME_H_
106