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
35namespace mlir {
36namespace runtime {
37
38//===----------------------------------------------------------------------===//
39// Async runtime API.
40//===----------------------------------------------------------------------===//
41
42// Runtime implementation of `async.token` data type.
43using AsyncToken = struct AsyncToken;
44
45// Runtime implementation of `async.group` data type.
46using AsyncGroup = struct AsyncGroup;
47
48// Runtime implementation of `async.value` data type.
49using AsyncValue = struct AsyncValue;
50
51// Async value payload stored in a memory owned by the async.value.
52using 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.
57using CoroHandle = void *; // coroutine handle
58using 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).
62using RefCountedObjPtr = void *;
63
64// Adds references to reference counted runtime object.
65extern "C" MLIR_ASYNC_RUNTIME_EXPORT void
66 mlirAsyncRuntimeAddRef(RefCountedObjPtr, int64_t);
67
68// Drops references from reference counted runtime object.
69extern "C" MLIR_ASYNC_RUNTIME_EXPORT void
70 mlirAsyncRuntimeDropRef(RefCountedObjPtr, int64_t);
71
72// Create a new `async.token` in not-ready state.
73extern "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).
79extern "C" MLIR_ASYNC_RUNTIME_EXPORT AsyncValue *
80 mlirAsyncRuntimeCreateValue(int64_t);
81
82// Create a new `async.group` in empty state.
83extern "C" MLIR_ASYNC_RUNTIME_EXPORT AsyncGroup *
84mlirAsyncRuntimeCreateGroup(int64_t size);
85
86extern "C" MLIR_ASYNC_RUNTIME_EXPORT int64_t
87mlirAsyncRuntimeAddTokenToGroup(AsyncToken *, AsyncGroup *);
88
89// Switches `async.token` to ready state and runs all awaiters.
90extern "C" MLIR_ASYNC_RUNTIME_EXPORT void
91mlirAsyncRuntimeEmplaceToken(AsyncToken *);
92
93// Switches `async.value` to ready state and runs all awaiters.
94extern "C" MLIR_ASYNC_RUNTIME_EXPORT void
95mlirAsyncRuntimeEmplaceValue(AsyncValue *);
96
97// Switches `async.token` to error state and runs all awaiters.
98extern "C" MLIR_ASYNC_RUNTIME_EXPORT void
99mlirAsyncRuntimeSetTokenError(AsyncToken *);
100
101// Switches `async.value` to error state and runs all awaiters.
102extern "C" MLIR_ASYNC_RUNTIME_EXPORT void
103mlirAsyncRuntimeSetValueError(AsyncValue *);
104
105// Returns true if token is in the error state.
106extern "C" MLIR_ASYNC_RUNTIME_EXPORT bool
107mlirAsyncRuntimeIsTokenError(AsyncToken *);
108
109// Returns true if value is in the error state.
110extern "C" MLIR_ASYNC_RUNTIME_EXPORT bool
111mlirAsyncRuntimeIsValueError(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).
115extern "C" MLIR_ASYNC_RUNTIME_EXPORT bool
116mlirAsyncRuntimeIsGroupError(AsyncGroup *);
117
118// Blocks the caller thread until the token becomes ready.
119extern "C" MLIR_ASYNC_RUNTIME_EXPORT void
120mlirAsyncRuntimeAwaitToken(AsyncToken *);
121
122// Blocks the caller thread until the value becomes ready.
123extern "C" MLIR_ASYNC_RUNTIME_EXPORT void
124mlirAsyncRuntimeAwaitValue(AsyncValue *);
125
126// Blocks the caller thread until the elements in the group become ready.
127extern "C" MLIR_ASYNC_RUNTIME_EXPORT void
128mlirAsyncRuntimeAwaitAllInGroup(AsyncGroup *);
129
130// Returns a pointer to the storage owned by the async value.
131extern "C" MLIR_ASYNC_RUNTIME_EXPORT ValueStorage
132mlirAsyncRuntimeGetValueStorage(AsyncValue *);
133
134// Executes the task (coro handle + resume function) in one of the threads
135// managed by the runtime.
136extern "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.
141extern "C" MLIR_ASYNC_RUNTIME_EXPORT void
142mlirAsyncRuntimeAwaitTokenAndExecute(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.
146extern "C" MLIR_ASYNC_RUNTIME_EXPORT void
147mlirAsyncRuntimeAwaitValueAndExecute(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.
151extern "C" MLIR_ASYNC_RUNTIME_EXPORT void
152mlirAsyncRuntimeAwaitAllInGroupAndExecute(AsyncGroup *, CoroHandle, CoroResume);
153
154// Returns the current number of available worker threads in the threadpool.
155extern "C" MLIR_ASYNC_RUNTIME_EXPORT int64_t
156mlirAsyncRuntimGetNumWorkerThreads();
157
158//===----------------------------------------------------------------------===//
159// Small async runtime support library for testing.
160//===----------------------------------------------------------------------===//
161
162extern "C" MLIR_ASYNC_RUNTIME_EXPORT void
163mlirAsyncRuntimePrintCurrentThreadId();
164
165} // namespace runtime
166} // namespace mlir
167
168#endif // MLIR_EXECUTIONENGINE_ASYNCRUNTIME_H_
169

source code of mlir/include/mlir/ExecutionEngine/AsyncRuntime.h