AsyncRuntime.h | Explore in Territory

//===- 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 …
#else
// We are using this library
#define MLIR_ASYNC_RUNTIME_EXPORT …
#endif // mlir_async_runtime_EXPORTS
#endif // MLIR_ASYNC_RUNTIME_EXPORT
#else
// Non-windows: use visibility attributes.
#define MLIR_ASYNC_RUNTIME_EXPORT …
#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.
ValueStorage;

// 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.
CoroHandle;           // coroutine handle
CoroResume; // coroutine resume function

// Async runtime uses reference counting to manage the lifetime of async values
// (values of async types like tokens, values and groups).
RefCountedObjPtr;

// 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_
llvm/mlir/include/mlir/ExecutionEngine/AsyncRuntime.h
