AsyncToAsyncRuntime.cpp File Reference

#include "PassDetail.h"
#include "mlir/Conversion/SCFToControlFlow/SCFToControlFlow.h"
#include "mlir/Dialect/Arithmetic/IR/Arithmetic.h"
#include "mlir/Dialect/Async/IR/Async.h"
#include "mlir/Dialect/Async/Passes.h"
#include "mlir/Dialect/ControlFlow/IR/ControlFlowOps.h"
#include "mlir/Dialect/Func/IR/FuncOps.h"
#include "mlir/Dialect/SCF/SCF.h"
#include "mlir/IR/BlockAndValueMapping.h"
#include "mlir/IR/ImplicitLocOpBuilder.h"
#include "mlir/IR/PatternMatch.h"
#include "mlir/Transforms/DialectConversion.h"
#include "mlir/Transforms/RegionUtils.h"
#include "llvm/ADT/SetVector.h"
#include "llvm/Support/Debug.h"

Include dependency graph for AsyncToAsyncRuntime.cpp:

Go to the source code of this file.

Classes
class	YieldOpLowering
class	AssertOpLowering

Macros
#define	DEBUG_TYPE   "async-to-async-runtime"

Functions
value !async value< T > ***static CoroMachinery	setupCoroMachinery (func::FuncOp func)
static Block *	setupSetErrorBlock (CoroMachinery &coro)
static std::pair< func::FuncOp, CoroMachinery >	outlineExecuteOp (SymbolTable &symbolTable, ExecuteOp execute)
	Outline the body region attached to the async.execute op into a standalone function. More...
static CoroMachinery	rewriteFuncAsCoroutine (func::FuncOp func)
	Rewrite a func as a coroutine by: 1) Wrapping the results into async.value. More...
static void	rewriteCallsiteForCoroutine (func::CallOp oldCall, func::FuncOp func)
	Rewrites a call into a function that has been rewritten as a coroutine. More...
static bool	isAllowedToBlock (func::FuncOp func)
static LogicalResult	funcsToCoroutines (ModuleOp module, llvm::DenseMap< func::FuncOp, CoroMachinery > &outlinedFunctions)

Variables
static constexpr const char	kAsyncFnPrefix [] = "async_execute_fn"
**	set_error
	Utility to partially update the regular function CFG to the coroutine CFG compatible with LLVM coroutines switched-resume lowering using `async.runtime. More...
value	__pad0__

Macro Definition Documentation

DEBUG_TYPE

#define DEBUG_TYPE   "async-to-async-runtime"

Definition at line 33 of file AsyncToAsyncRuntime.cpp.

Function Documentation

funcsToCoroutines()

static LogicalResult funcsToCoroutines ( ModuleOp  module, llvm::DenseMap< func::FuncOp, CoroMachinery > &  outlinedFunctions  )

static

Definition at line 662 of file AsyncToAsyncRuntime.cpp.

References mlir::RewritePatternSet::add(), mlir::WalkResult::advance(), mlir::applyPartialConversion(), mlir::failed(), mlir::failure(), mlir::Operation::getBlock(), mlir::Operation::getDialect(), mlir::SymbolUserMap::getUsers(), mlir::WalkResult::interrupt(), isAllowedToBlock(), mlir::Operation::isBeforeInBlock(), outlineExecuteOp(), mlir::populateSCFToControlFlowConversionPatterns(), rewriteCallsiteForCoroutine(), rewriteFuncAsCoroutine(), and mlir::success().

isAllowedToBlock()

static bool isAllowedToBlock ( func::FuncOp  func )

static

Definition at line 658 of file AsyncToAsyncRuntime.cpp.

Referenced by funcsToCoroutines().

outlineExecuteOp()

static std::pair<func::FuncOp, CoroMachinery> outlineExecuteOp ( SymbolTable &  symbolTable, ExecuteOp  execute  )

static

Outline the body region attached to the async.execute op into a standalone function.

Note that this is not reversible transformation.

Definition at line 241 of file AsyncToAsyncRuntime.cpp.

References mlir::ImplicitLocOpBuilder::atBlockBegin(), mlir::Type::cast(), mlir::async::cloneConstantsIntoTheRegion(), mlir::ImplicitLocOpBuilder::create(), mlir::OpBuilder::create(), mlir::ConversionPatternRewriter::eraseOp(), mlir::Builder::getI1Type(), mlir::Builder::getIndexType(), mlir::OpBuilder::getListener(), mlir::Value::getType(), mlir::getUsedValuesDefinedAbove(), mlir::SymbolTable::insert(), kAsyncFnPrefix, mlir::BlockAndValueMapping::map(), mlir::ConversionPatternRewriter::notifyMatchFailure(), mlir::OpConversionPattern< SourceOp >::OpConversionPattern(), mlir::SymbolTable::Private, mlir::ConversionPatternRewriter::replaceOp(), mlir::RewriterBase::replaceOpWithNewOp(), mlir::OpBuilder::setInsertionPointToStart(), mlir::SymbolTable::setSymbolVisibility(), setupCoroMachinery(), setupSetErrorBlock(), mlir::ConversionPatternRewriter::splitBlock(), mlir::success(), and value.

Referenced by funcsToCoroutines().

rewriteCallsiteForCoroutine()

static void rewriteCallsiteForCoroutine ( func::CallOp  oldCall, func::FuncOp  func  )

static

Rewrites a call into a function that has been rewritten as a coroutine.

The invocation of this function is safe only when call ops are traversed in reverse order of how they appear in a single block. See funcsToCoroutines.

Definition at line 638 of file AsyncToAsyncRuntime.cpp.

References mlir::ImplicitLocOpBuilder::create().

Referenced by funcsToCoroutines().

rewriteFuncAsCoroutine()

static CoroMachinery rewriteFuncAsCoroutine ( func::FuncOp  func )

static

Rewrite a func as a coroutine by: 1) Wrapping the results into async.value.

2) Prepending the results with async.token. 3) Setting up coroutine blocks. 4) Rewriting return ops as yield op and branch op into the suspend block.

Definition at line 613 of file AsyncToAsyncRuntime.cpp.

References mlir::ImplicitLocOpBuilder::create(), and setupCoroMachinery().

Referenced by funcsToCoroutines().

setupCoroMachinery()

value !async value<T>* * * static CoroMachinery setupCoroMachinery ( func::FuncOp  func )

static

Definition at line 127 of file AsyncToAsyncRuntime.cpp.

References mlir::ImplicitLocOpBuilder::atBlockBegin(), mlir::Operation::create(), mlir::ImplicitLocOpBuilder::create(), mlir::Block::front(), mlir::Block::getOperations(), mlir::Operation::setAttr(), and mlir::Block::splitBlock().

Referenced by outlineExecuteOp(), and rewriteFuncAsCoroutine().

setupSetErrorBlock()

static Block* setupSetErrorBlock ( CoroMachinery &  coro )

static

Definition at line 215 of file AsyncToAsyncRuntime.cpp.

References mlir::ImplicitLocOpBuilder::atBlockBegin(), and mlir::ImplicitLocOpBuilder::create().

Referenced by AssertOpLowering::matchAndRewrite(), and outlineExecuteOp().

Variable Documentation

__pad0__

value __pad0__

Definition at line 124 of file AsyncToAsyncRuntime.cpp.

kAsyncFnPrefix

constexpr const char kAsyncFnPrefix[] = "async_execute_fn"

static

Definition at line 35 of file AsyncToAsyncRuntime.cpp.

Referenced by outlineExecuteOp().

set_error

* * set_error

Utility to partially update the regular function CFG to the coroutine CFG compatible with LLVM coroutines switched-resume lowering using `async.runtime.

*andasync.coro.*` operations. Adds a new entry block that branches into preexisting entry block. Also inserts trailing blocks.

The result types of the passed func must start with an async.token and be continued with some number of async.values.

The func given to this function needs to have been preprocessed to have either branch or yield ops as terminators. Branches to the cleanup block are inserted after each yield.

See LLVM coroutines documentation: https://llvm.org/docs/Coroutines.html

• entry block sets up the coroutine.
• set_error block sets completion token and async values state to error.
• cleanup block cleans up the coroutine state.
• `suspend block after the .coro.end() defines what value will be returned to the initial caller of a coroutine. Everything before the .coro.end() will be executed at every suspension point.

Coroutine structure (only the important bits):

func (<function-arguments>) -> (!async.token, !async.value<T>) { ^entry(<function-arguments>): token = <async token>=""> : !async.token // create async runtime token value = <async value>=""> : !async.value<T> // create async value id = async.coro.id // create a coroutine id hdl = async.coro.begin id // create a coroutine handle cf.br ^preexisting_entry_block

/* preexisting blocks modified to branch to the cleanup block

Definition at line 124 of file AsyncToAsyncRuntime.cpp.