| //===- AsyncRuntime.h - Async runtime reference implementation ------------===// |
| // |
| // Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions. |
| // See https://llvm.org/LICENSE.txt for license information. |
| // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception |
| // |
| //===----------------------------------------------------------------------===// |
| // |
| // This file declares basic Async runtime API for supporting Async dialect |
| // to LLVM dialect lowering. |
| // |
| //===----------------------------------------------------------------------===// |
| |
| #ifndef MLIR_EXECUTIONENGINE_ASYNCRUNTIME_H_ |
| #define MLIR_EXECUTIONENGINE_ASYNCRUNTIME_H_ |
| |
| #include <cstddef> |
| #include <stdint.h> |
| |
| #ifdef _WIN32 |
| #ifndef MLIR_ASYNC_RUNTIME_EXPORT |
| #ifdef mlir_async_runtime_EXPORTS |
| // We are building this library |
| #define MLIR_ASYNC_RUNTIME_EXPORT __declspec(dllexport) |
| #else |
| // We are using this library |
| #define MLIR_ASYNC_RUNTIME_EXPORT __declspec(dllimport) |
| #endif // mlir_async_runtime_EXPORTS |
| #endif // MLIR_ASYNC_RUNTIME_EXPORT |
| #else |
| // Non-windows: use visibility attributes. |
| #define MLIR_ASYNC_RUNTIME_EXPORT __attribute__((visibility("default"))) |
| #endif // _WIN32 |
| |
| namespace mlir { |
| namespace runtime { |
| |
| //===----------------------------------------------------------------------===// |
| // Async runtime API. |
| //===----------------------------------------------------------------------===// |
| |
| // Runtime implementation of `async.token` data type. |
| using AsyncToken = struct AsyncToken; |
| |
| // Runtime implementation of `async.group` data type. |
| using AsyncGroup = struct AsyncGroup; |
| |
| // Runtime implementation of `async.value` data type. |
| using AsyncValue = struct AsyncValue; |
| |
| // Async value payload stored in a memory owned by the async.value. |
| using ValueStorage = std::byte *; |
| |
| // Async runtime uses LLVM coroutines to represent asynchronous tasks. Task |
| // function is a coroutine handle and a resume function that continue coroutine |
| // execution from a suspension point. |
| using CoroHandle = void *; // coroutine handle |
| using CoroResume = void (*)(void *); // coroutine resume function |
| |
| // Async runtime uses reference counting to manage the lifetime of async values |
| // (values of async types like tokens, values and groups). |
| using RefCountedObjPtr = void *; |
| |
| // Adds references to reference counted runtime object. |
| extern "C" MLIR_ASYNC_RUNTIME_EXPORT void |
| mlirAsyncRuntimeAddRef(RefCountedObjPtr, int64_t); |
| |
| // Drops references from reference counted runtime object. |
| extern "C" MLIR_ASYNC_RUNTIME_EXPORT void |
| mlirAsyncRuntimeDropRef(RefCountedObjPtr, int64_t); |
| |
| // Create a new `async.token` in not-ready state. |
| extern "C" MLIR_ASYNC_RUNTIME_EXPORT AsyncToken *mlirAsyncRuntimeCreateToken(); |
| |
| // Create a new `async.value` in not-ready state. Size parameter specifies the |
| // number of bytes that will be allocated for the async value storage. Storage |
| // is owned by the `async.value` and deallocated when the async value is |
| // destructed (reference count drops to zero). |
| extern "C" MLIR_ASYNC_RUNTIME_EXPORT AsyncValue * |
| mlirAsyncRuntimeCreateValue(int64_t); |
| |
| // Create a new `async.group` in empty state. |
| extern "C" MLIR_ASYNC_RUNTIME_EXPORT AsyncGroup * |
| mlirAsyncRuntimeCreateGroup(int64_t size); |
| |
| extern "C" MLIR_ASYNC_RUNTIME_EXPORT int64_t |
| mlirAsyncRuntimeAddTokenToGroup(AsyncToken *, AsyncGroup *); |
| |
| // Switches `async.token` to ready state and runs all awaiters. |
| extern "C" MLIR_ASYNC_RUNTIME_EXPORT void |
| mlirAsyncRuntimeEmplaceToken(AsyncToken *); |
| |
| // Switches `async.value` to ready state and runs all awaiters. |
| extern "C" MLIR_ASYNC_RUNTIME_EXPORT void |
| mlirAsyncRuntimeEmplaceValue(AsyncValue *); |
| |
| // Switches `async.token` to error state and runs all awaiters. |
| extern "C" MLIR_ASYNC_RUNTIME_EXPORT void |
| mlirAsyncRuntimeSetTokenError(AsyncToken *); |
| |
| // Switches `async.value` to error state and runs all awaiters. |
| extern "C" MLIR_ASYNC_RUNTIME_EXPORT void |
| mlirAsyncRuntimeSetValueError(AsyncValue *); |
| |
| // Returns true if token is in the error state. |
| extern "C" MLIR_ASYNC_RUNTIME_EXPORT bool |
| mlirAsyncRuntimeIsTokenError(AsyncToken *); |
| |
| // Returns true if value is in the error state. |
| extern "C" MLIR_ASYNC_RUNTIME_EXPORT bool |
| mlirAsyncRuntimeIsValueError(AsyncValue *); |
| |
| // Returns true if group is in the error state (any of the tokens or values |
| // added to the group are in the error state). |
| extern "C" MLIR_ASYNC_RUNTIME_EXPORT bool |
| mlirAsyncRuntimeIsGroupError(AsyncGroup *); |
| |
| // Blocks the caller thread until the token becomes ready. |
| extern "C" MLIR_ASYNC_RUNTIME_EXPORT void |
| mlirAsyncRuntimeAwaitToken(AsyncToken *); |
| |
| // Blocks the caller thread until the value becomes ready. |
| extern "C" MLIR_ASYNC_RUNTIME_EXPORT void |
| mlirAsyncRuntimeAwaitValue(AsyncValue *); |
| |
| // Blocks the caller thread until the elements in the group become ready. |
| extern "C" MLIR_ASYNC_RUNTIME_EXPORT void |
| mlirAsyncRuntimeAwaitAllInGroup(AsyncGroup *); |
| |
| // Returns a pointer to the storage owned by the async value. |
| extern "C" MLIR_ASYNC_RUNTIME_EXPORT ValueStorage |
| mlirAsyncRuntimeGetValueStorage(AsyncValue *); |
| |
| // Executes the task (coro handle + resume function) in one of the threads |
| // managed by the runtime. |
| extern "C" MLIR_ASYNC_RUNTIME_EXPORT void mlirAsyncRuntimeExecute(CoroHandle, |
| CoroResume); |
| |
| // Executes the task (coro handle + resume function) in one of the threads |
| // managed by the runtime after the token becomes ready. |
| extern "C" MLIR_ASYNC_RUNTIME_EXPORT void |
| mlirAsyncRuntimeAwaitTokenAndExecute(AsyncToken *, CoroHandle, CoroResume); |
| |
| // Executes the task (coro handle + resume function) in one of the threads |
| // managed by the runtime after the value becomes ready. |
| extern "C" MLIR_ASYNC_RUNTIME_EXPORT void |
| mlirAsyncRuntimeAwaitValueAndExecute(AsyncValue *, CoroHandle, CoroResume); |
| |
| // Executes the task (coro handle + resume function) in one of the threads |
| // managed by the runtime after the all members of the group become ready. |
| extern "C" MLIR_ASYNC_RUNTIME_EXPORT void |
| mlirAsyncRuntimeAwaitAllInGroupAndExecute(AsyncGroup *, CoroHandle, CoroResume); |
| |
| // Returns the current number of available worker threads in the threadpool. |
| extern "C" MLIR_ASYNC_RUNTIME_EXPORT int64_t |
| mlirAsyncRuntimGetNumWorkerThreads(); |
| |
| //===----------------------------------------------------------------------===// |
| // Small async runtime support library for testing. |
| //===----------------------------------------------------------------------===// |
| |
| extern "C" MLIR_ASYNC_RUNTIME_EXPORT void |
| mlirAsyncRuntimePrintCurrentThreadId(); |
| |
| } // namespace runtime |
| } // namespace mlir |
| |
| #endif // MLIR_EXECUTIONENGINE_ASYNCRUNTIME_H_ |