125f80e16SEugene Zhulenev //===- AsyncToAsyncRuntime.cpp - Lower from Async to Async Runtime --------===// 225f80e16SEugene Zhulenev // 325f80e16SEugene Zhulenev // Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions. 425f80e16SEugene Zhulenev // See https://llvm.org/LICENSE.txt for license information. 525f80e16SEugene Zhulenev // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception 625f80e16SEugene Zhulenev // 725f80e16SEugene Zhulenev //===----------------------------------------------------------------------===// 825f80e16SEugene Zhulenev // 925f80e16SEugene Zhulenev // This file implements lowering from high level async operations to async.coro 1025f80e16SEugene Zhulenev // and async.runtime operations. 1125f80e16SEugene Zhulenev // 1225f80e16SEugene Zhulenev //===----------------------------------------------------------------------===// 1325f80e16SEugene Zhulenev 1425f80e16SEugene Zhulenev #include "PassDetail.h" 15*de7a4e53SEugene Zhulenev #include "mlir/Conversion/SCFToStandard/SCFToStandard.h" 1625f80e16SEugene Zhulenev #include "mlir/Dialect/Async/IR/Async.h" 1725f80e16SEugene Zhulenev #include "mlir/Dialect/Async/Passes.h" 18*de7a4e53SEugene Zhulenev #include "mlir/Dialect/SCF/SCF.h" 1925f80e16SEugene Zhulenev #include "mlir/Dialect/StandardOps/IR/Ops.h" 2025f80e16SEugene Zhulenev #include "mlir/IR/BlockAndValueMapping.h" 2125f80e16SEugene Zhulenev #include "mlir/IR/ImplicitLocOpBuilder.h" 2225f80e16SEugene Zhulenev #include "mlir/IR/PatternMatch.h" 2325f80e16SEugene Zhulenev #include "mlir/Transforms/DialectConversion.h" 2425f80e16SEugene Zhulenev #include "mlir/Transforms/RegionUtils.h" 2525f80e16SEugene Zhulenev #include "llvm/ADT/SetVector.h" 26297a5b7cSNico Weber #include "llvm/Support/Debug.h" 2725f80e16SEugene Zhulenev 2825f80e16SEugene Zhulenev using namespace mlir; 2925f80e16SEugene Zhulenev using namespace mlir::async; 3025f80e16SEugene Zhulenev 3125f80e16SEugene Zhulenev #define DEBUG_TYPE "async-to-async-runtime" 3225f80e16SEugene Zhulenev // Prefix for functions outlined from `async.execute` op regions. 3325f80e16SEugene Zhulenev static constexpr const char kAsyncFnPrefix[] = "async_execute_fn"; 3425f80e16SEugene Zhulenev 3525f80e16SEugene Zhulenev namespace { 3625f80e16SEugene Zhulenev 3725f80e16SEugene Zhulenev class AsyncToAsyncRuntimePass 3825f80e16SEugene Zhulenev : public AsyncToAsyncRuntimeBase<AsyncToAsyncRuntimePass> { 3925f80e16SEugene Zhulenev public: 4025f80e16SEugene Zhulenev AsyncToAsyncRuntimePass() = default; 4125f80e16SEugene Zhulenev void runOnOperation() override; 4225f80e16SEugene Zhulenev }; 4325f80e16SEugene Zhulenev 4425f80e16SEugene Zhulenev } // namespace 4525f80e16SEugene Zhulenev 4625f80e16SEugene Zhulenev //===----------------------------------------------------------------------===// 4725f80e16SEugene Zhulenev // async.execute op outlining to the coroutine functions. 4825f80e16SEugene Zhulenev //===----------------------------------------------------------------------===// 4925f80e16SEugene Zhulenev 5025f80e16SEugene Zhulenev /// Function targeted for coroutine transformation has two additional blocks at 5125f80e16SEugene Zhulenev /// the end: coroutine cleanup and coroutine suspension. 5225f80e16SEugene Zhulenev /// 5325f80e16SEugene Zhulenev /// async.await op lowering additionaly creates a resume block for each 5425f80e16SEugene Zhulenev /// operation to enable non-blocking waiting via coroutine suspension. 5525f80e16SEugene Zhulenev namespace { 5625f80e16SEugene Zhulenev struct CoroMachinery { 5739957aa4SEugene Zhulenev FuncOp func; 5839957aa4SEugene Zhulenev 5925f80e16SEugene Zhulenev // Async execute region returns a completion token, and an async value for 6025f80e16SEugene Zhulenev // each yielded value. 6125f80e16SEugene Zhulenev // 6225f80e16SEugene Zhulenev // %token, %result = async.execute -> !async.value<T> { 6325f80e16SEugene Zhulenev // %0 = constant ... : T 6425f80e16SEugene Zhulenev // async.yield %0 : T 6525f80e16SEugene Zhulenev // } 6625f80e16SEugene Zhulenev Value asyncToken; // token representing completion of the async region 6725f80e16SEugene Zhulenev llvm::SmallVector<Value, 4> returnValues; // returned async values 6825f80e16SEugene Zhulenev 6925f80e16SEugene Zhulenev Value coroHandle; // coroutine handle (!async.coro.handle value) 7039957aa4SEugene Zhulenev Block *setError; // switch completion token and all values to error state 7125f80e16SEugene Zhulenev Block *cleanup; // coroutine cleanup block 7225f80e16SEugene Zhulenev Block *suspend; // coroutine suspension block 7325f80e16SEugene Zhulenev }; 7425f80e16SEugene Zhulenev } // namespace 7525f80e16SEugene Zhulenev 7625f80e16SEugene Zhulenev /// Builds an coroutine template compatible with LLVM coroutines switched-resume 7725f80e16SEugene Zhulenev /// lowering using `async.runtime.*` and `async.coro.*` operations. 7825f80e16SEugene Zhulenev /// 7925f80e16SEugene Zhulenev /// See LLVM coroutines documentation: https://llvm.org/docs/Coroutines.html 8025f80e16SEugene Zhulenev /// 8125f80e16SEugene Zhulenev /// - `entry` block sets up the coroutine. 8239957aa4SEugene Zhulenev /// - `set_error` block sets completion token and async values state to error. 8325f80e16SEugene Zhulenev /// - `cleanup` block cleans up the coroutine state. 8425f80e16SEugene Zhulenev /// - `suspend block after the @llvm.coro.end() defines what value will be 8525f80e16SEugene Zhulenev /// returned to the initial caller of a coroutine. Everything before the 8625f80e16SEugene Zhulenev /// @llvm.coro.end() will be executed at every suspension point. 8725f80e16SEugene Zhulenev /// 8825f80e16SEugene Zhulenev /// Coroutine structure (only the important bits): 8925f80e16SEugene Zhulenev /// 9025f80e16SEugene Zhulenev /// func @async_execute_fn(<function-arguments>) 9125f80e16SEugene Zhulenev /// -> (!async.token, !async.value<T>) 9225f80e16SEugene Zhulenev /// { 9325f80e16SEugene Zhulenev /// ^entry(<function-arguments>): 9425f80e16SEugene Zhulenev /// %token = <async token> : !async.token // create async runtime token 9525f80e16SEugene Zhulenev /// %value = <async value> : !async.value<T> // create async value 9625f80e16SEugene Zhulenev /// %id = async.coro.id // create a coroutine id 9725f80e16SEugene Zhulenev /// %hdl = async.coro.begin %id // create a coroutine handle 9825f80e16SEugene Zhulenev /// br ^cleanup 9925f80e16SEugene Zhulenev /// 10039957aa4SEugene Zhulenev /// ^set_error: // this block created lazily only if needed (see code below) 10139957aa4SEugene Zhulenev /// async.runtime.set_error %token : !async.token 10239957aa4SEugene Zhulenev /// async.runtime.set_error %value : !async.value<T> 10339957aa4SEugene Zhulenev /// br ^cleanup 10439957aa4SEugene Zhulenev /// 10525f80e16SEugene Zhulenev /// ^cleanup: 10625f80e16SEugene Zhulenev /// async.coro.free %hdl // delete the coroutine state 10725f80e16SEugene Zhulenev /// br ^suspend 10825f80e16SEugene Zhulenev /// 10925f80e16SEugene Zhulenev /// ^suspend: 11025f80e16SEugene Zhulenev /// async.coro.end %hdl // marks the end of a coroutine 11125f80e16SEugene Zhulenev /// return %token, %value : !async.token, !async.value<T> 11225f80e16SEugene Zhulenev /// } 11325f80e16SEugene Zhulenev /// 11425f80e16SEugene Zhulenev /// The actual code for the async.execute operation body region will be inserted 11525f80e16SEugene Zhulenev /// before the entry block terminator. 11625f80e16SEugene Zhulenev /// 11725f80e16SEugene Zhulenev /// 11825f80e16SEugene Zhulenev static CoroMachinery setupCoroMachinery(FuncOp func) { 11925f80e16SEugene Zhulenev assert(func.getBody().empty() && "Function must have empty body"); 12025f80e16SEugene Zhulenev 12125f80e16SEugene Zhulenev MLIRContext *ctx = func.getContext(); 12225f80e16SEugene Zhulenev Block *entryBlock = func.addEntryBlock(); 12325f80e16SEugene Zhulenev 12425f80e16SEugene Zhulenev auto builder = ImplicitLocOpBuilder::atBlockBegin(func->getLoc(), entryBlock); 12525f80e16SEugene Zhulenev 12625f80e16SEugene Zhulenev // ------------------------------------------------------------------------ // 12725f80e16SEugene Zhulenev // Allocate async token/values that we will return from a ramp function. 12825f80e16SEugene Zhulenev // ------------------------------------------------------------------------ // 12925f80e16SEugene Zhulenev auto retToken = builder.create<RuntimeCreateOp>(TokenType::get(ctx)).result(); 13025f80e16SEugene Zhulenev 13125f80e16SEugene Zhulenev llvm::SmallVector<Value, 4> retValues; 13225f80e16SEugene Zhulenev for (auto resType : func.getCallableResults().drop_front()) 13325f80e16SEugene Zhulenev retValues.emplace_back(builder.create<RuntimeCreateOp>(resType).result()); 13425f80e16SEugene Zhulenev 13525f80e16SEugene Zhulenev // ------------------------------------------------------------------------ // 13625f80e16SEugene Zhulenev // Initialize coroutine: get coroutine id and coroutine handle. 13725f80e16SEugene Zhulenev // ------------------------------------------------------------------------ // 13825f80e16SEugene Zhulenev auto coroIdOp = builder.create<CoroIdOp>(CoroIdType::get(ctx)); 13925f80e16SEugene Zhulenev auto coroHdlOp = 14025f80e16SEugene Zhulenev builder.create<CoroBeginOp>(CoroHandleType::get(ctx), coroIdOp.id()); 14125f80e16SEugene Zhulenev 14225f80e16SEugene Zhulenev Block *cleanupBlock = func.addBlock(); 14325f80e16SEugene Zhulenev Block *suspendBlock = func.addBlock(); 14425f80e16SEugene Zhulenev 14525f80e16SEugene Zhulenev // ------------------------------------------------------------------------ // 14625f80e16SEugene Zhulenev // Coroutine cleanup block: deallocate coroutine frame, free the memory. 14725f80e16SEugene Zhulenev // ------------------------------------------------------------------------ // 14825f80e16SEugene Zhulenev builder.setInsertionPointToStart(cleanupBlock); 14925f80e16SEugene Zhulenev builder.create<CoroFreeOp>(coroIdOp.id(), coroHdlOp.handle()); 15025f80e16SEugene Zhulenev 15125f80e16SEugene Zhulenev // Branch into the suspend block. 15225f80e16SEugene Zhulenev builder.create<BranchOp>(suspendBlock); 15325f80e16SEugene Zhulenev 15425f80e16SEugene Zhulenev // ------------------------------------------------------------------------ // 15525f80e16SEugene Zhulenev // Coroutine suspend block: mark the end of a coroutine and return allocated 15625f80e16SEugene Zhulenev // async token. 15725f80e16SEugene Zhulenev // ------------------------------------------------------------------------ // 15825f80e16SEugene Zhulenev builder.setInsertionPointToStart(suspendBlock); 15925f80e16SEugene Zhulenev 16025f80e16SEugene Zhulenev // Mark the end of a coroutine: async.coro.end 16125f80e16SEugene Zhulenev builder.create<CoroEndOp>(coroHdlOp.handle()); 16225f80e16SEugene Zhulenev 16325f80e16SEugene Zhulenev // Return created `async.token` and `async.values` from the suspend block. 16425f80e16SEugene Zhulenev // This will be the return value of a coroutine ramp function. 16525f80e16SEugene Zhulenev SmallVector<Value, 4> ret{retToken}; 16625f80e16SEugene Zhulenev ret.insert(ret.end(), retValues.begin(), retValues.end()); 16725f80e16SEugene Zhulenev builder.create<ReturnOp>(ret); 16825f80e16SEugene Zhulenev 16925f80e16SEugene Zhulenev // Branch from the entry block to the cleanup block to create a valid CFG. 17025f80e16SEugene Zhulenev builder.setInsertionPointToEnd(entryBlock); 17125f80e16SEugene Zhulenev builder.create<BranchOp>(cleanupBlock); 17225f80e16SEugene Zhulenev 17325f80e16SEugene Zhulenev // `async.await` op lowering will create resume blocks for async 17425f80e16SEugene Zhulenev // continuations, and will conditionally branch to cleanup or suspend blocks. 17525f80e16SEugene Zhulenev 17625f80e16SEugene Zhulenev CoroMachinery machinery; 17739957aa4SEugene Zhulenev machinery.func = func; 17825f80e16SEugene Zhulenev machinery.asyncToken = retToken; 17925f80e16SEugene Zhulenev machinery.returnValues = retValues; 18025f80e16SEugene Zhulenev machinery.coroHandle = coroHdlOp.handle(); 18139957aa4SEugene Zhulenev machinery.setError = nullptr; // created lazily only if needed 18225f80e16SEugene Zhulenev machinery.cleanup = cleanupBlock; 18325f80e16SEugene Zhulenev machinery.suspend = suspendBlock; 18425f80e16SEugene Zhulenev return machinery; 18525f80e16SEugene Zhulenev } 18625f80e16SEugene Zhulenev 18739957aa4SEugene Zhulenev // Lazily creates `set_error` block only if it is required for lowering to the 18839957aa4SEugene Zhulenev // runtime operations (see for example lowering of assert operation). 18939957aa4SEugene Zhulenev static Block *setupSetErrorBlock(CoroMachinery &coro) { 19039957aa4SEugene Zhulenev if (coro.setError) 19139957aa4SEugene Zhulenev return coro.setError; 19239957aa4SEugene Zhulenev 19339957aa4SEugene Zhulenev coro.setError = coro.func.addBlock(); 19439957aa4SEugene Zhulenev coro.setError->moveBefore(coro.cleanup); 19539957aa4SEugene Zhulenev 19639957aa4SEugene Zhulenev auto builder = 19739957aa4SEugene Zhulenev ImplicitLocOpBuilder::atBlockBegin(coro.func->getLoc(), coro.setError); 19839957aa4SEugene Zhulenev 19939957aa4SEugene Zhulenev // Coroutine set_error block: set error on token and all returned values. 20039957aa4SEugene Zhulenev builder.create<RuntimeSetErrorOp>(coro.asyncToken); 20139957aa4SEugene Zhulenev for (Value retValue : coro.returnValues) 20239957aa4SEugene Zhulenev builder.create<RuntimeSetErrorOp>(retValue); 20339957aa4SEugene Zhulenev 20439957aa4SEugene Zhulenev // Branch into the cleanup block. 20539957aa4SEugene Zhulenev builder.create<BranchOp>(coro.cleanup); 20639957aa4SEugene Zhulenev 20739957aa4SEugene Zhulenev return coro.setError; 20839957aa4SEugene Zhulenev } 20939957aa4SEugene Zhulenev 21025f80e16SEugene Zhulenev /// Outline the body region attached to the `async.execute` op into a standalone 21125f80e16SEugene Zhulenev /// function. 21225f80e16SEugene Zhulenev /// 21325f80e16SEugene Zhulenev /// Note that this is not reversible transformation. 21425f80e16SEugene Zhulenev static std::pair<FuncOp, CoroMachinery> 21525f80e16SEugene Zhulenev outlineExecuteOp(SymbolTable &symbolTable, ExecuteOp execute) { 21625f80e16SEugene Zhulenev ModuleOp module = execute->getParentOfType<ModuleOp>(); 21725f80e16SEugene Zhulenev 21825f80e16SEugene Zhulenev MLIRContext *ctx = module.getContext(); 21925f80e16SEugene Zhulenev Location loc = execute.getLoc(); 22025f80e16SEugene Zhulenev 22125f80e16SEugene Zhulenev // Collect all outlined function inputs. 2224efb7754SRiver Riddle SetVector<mlir::Value> functionInputs(execute.dependencies().begin(), 22325f80e16SEugene Zhulenev execute.dependencies().end()); 22425f80e16SEugene Zhulenev functionInputs.insert(execute.operands().begin(), execute.operands().end()); 22525f80e16SEugene Zhulenev getUsedValuesDefinedAbove(execute.body(), functionInputs); 22625f80e16SEugene Zhulenev 22725f80e16SEugene Zhulenev // Collect types for the outlined function inputs and outputs. 22825f80e16SEugene Zhulenev auto typesRange = llvm::map_range( 22925f80e16SEugene Zhulenev functionInputs, [](Value value) { return value.getType(); }); 23025f80e16SEugene Zhulenev SmallVector<Type, 4> inputTypes(typesRange.begin(), typesRange.end()); 23125f80e16SEugene Zhulenev auto outputTypes = execute.getResultTypes(); 23225f80e16SEugene Zhulenev 23325f80e16SEugene Zhulenev auto funcType = FunctionType::get(ctx, inputTypes, outputTypes); 23425f80e16SEugene Zhulenev auto funcAttrs = ArrayRef<NamedAttribute>(); 23525f80e16SEugene Zhulenev 23625f80e16SEugene Zhulenev // TODO: Derive outlined function name from the parent FuncOp (support 23725f80e16SEugene Zhulenev // multiple nested async.execute operations). 23825f80e16SEugene Zhulenev FuncOp func = FuncOp::create(loc, kAsyncFnPrefix, funcType, funcAttrs); 239973ddb7dSMehdi Amini symbolTable.insert(func); 24025f80e16SEugene Zhulenev 24125f80e16SEugene Zhulenev SymbolTable::setSymbolVisibility(func, SymbolTable::Visibility::Private); 24225f80e16SEugene Zhulenev 24325f80e16SEugene Zhulenev // Prepare a function for coroutine lowering by adding entry/cleanup/suspend 24425f80e16SEugene Zhulenev // blocks, adding async.coro operations and setting up control flow. 24525f80e16SEugene Zhulenev CoroMachinery coro = setupCoroMachinery(func); 24625f80e16SEugene Zhulenev 24725f80e16SEugene Zhulenev // Suspend async function at the end of an entry block, and resume it using 24825f80e16SEugene Zhulenev // Async resume operation (execution will be resumed in a thread managed by 24925f80e16SEugene Zhulenev // the async runtime). 25025f80e16SEugene Zhulenev Block *entryBlock = &func.getBlocks().front(); 25125f80e16SEugene Zhulenev auto builder = ImplicitLocOpBuilder::atBlockTerminator(loc, entryBlock); 25225f80e16SEugene Zhulenev 25325f80e16SEugene Zhulenev // Save the coroutine state: async.coro.save 25425f80e16SEugene Zhulenev auto coroSaveOp = 25525f80e16SEugene Zhulenev builder.create<CoroSaveOp>(CoroStateType::get(ctx), coro.coroHandle); 25625f80e16SEugene Zhulenev 25725f80e16SEugene Zhulenev // Pass coroutine to the runtime to be resumed on a runtime managed thread. 25825f80e16SEugene Zhulenev builder.create<RuntimeResumeOp>(coro.coroHandle); 25925f80e16SEugene Zhulenev 26025f80e16SEugene Zhulenev // Split the entry block before the terminator (branch to suspend block). 26125f80e16SEugene Zhulenev auto *terminatorOp = entryBlock->getTerminator(); 26225f80e16SEugene Zhulenev Block *suspended = terminatorOp->getBlock(); 26325f80e16SEugene Zhulenev Block *resume = suspended->splitBlock(terminatorOp); 26425f80e16SEugene Zhulenev 26525f80e16SEugene Zhulenev // Add async.coro.suspend as a suspended block terminator. 26625f80e16SEugene Zhulenev builder.setInsertionPointToEnd(suspended); 26725f80e16SEugene Zhulenev builder.create<CoroSuspendOp>(coroSaveOp.state(), coro.suspend, resume, 26825f80e16SEugene Zhulenev coro.cleanup); 26925f80e16SEugene Zhulenev 27025f80e16SEugene Zhulenev size_t numDependencies = execute.dependencies().size(); 27125f80e16SEugene Zhulenev size_t numOperands = execute.operands().size(); 27225f80e16SEugene Zhulenev 27325f80e16SEugene Zhulenev // Await on all dependencies before starting to execute the body region. 27425f80e16SEugene Zhulenev builder.setInsertionPointToStart(resume); 27525f80e16SEugene Zhulenev for (size_t i = 0; i < numDependencies; ++i) 27625f80e16SEugene Zhulenev builder.create<AwaitOp>(func.getArgument(i)); 27725f80e16SEugene Zhulenev 27825f80e16SEugene Zhulenev // Await on all async value operands and unwrap the payload. 27925f80e16SEugene Zhulenev SmallVector<Value, 4> unwrappedOperands(numOperands); 28025f80e16SEugene Zhulenev for (size_t i = 0; i < numOperands; ++i) { 28125f80e16SEugene Zhulenev Value operand = func.getArgument(numDependencies + i); 28225f80e16SEugene Zhulenev unwrappedOperands[i] = builder.create<AwaitOp>(loc, operand).result(); 28325f80e16SEugene Zhulenev } 28425f80e16SEugene Zhulenev 28525f80e16SEugene Zhulenev // Map from function inputs defined above the execute op to the function 28625f80e16SEugene Zhulenev // arguments. 28725f80e16SEugene Zhulenev BlockAndValueMapping valueMapping; 28825f80e16SEugene Zhulenev valueMapping.map(functionInputs, func.getArguments()); 28925f80e16SEugene Zhulenev valueMapping.map(execute.body().getArguments(), unwrappedOperands); 29025f80e16SEugene Zhulenev 29125f80e16SEugene Zhulenev // Clone all operations from the execute operation body into the outlined 29225f80e16SEugene Zhulenev // function body. 29325f80e16SEugene Zhulenev for (Operation &op : execute.body().getOps()) 29425f80e16SEugene Zhulenev builder.clone(op, valueMapping); 29525f80e16SEugene Zhulenev 29625f80e16SEugene Zhulenev // Replace the original `async.execute` with a call to outlined function. 29725f80e16SEugene Zhulenev ImplicitLocOpBuilder callBuilder(loc, execute); 29825f80e16SEugene Zhulenev auto callOutlinedFunc = callBuilder.create<CallOp>( 29925f80e16SEugene Zhulenev func.getName(), execute.getResultTypes(), functionInputs.getArrayRef()); 30025f80e16SEugene Zhulenev execute.replaceAllUsesWith(callOutlinedFunc.getResults()); 30125f80e16SEugene Zhulenev execute.erase(); 30225f80e16SEugene Zhulenev 30325f80e16SEugene Zhulenev return {func, coro}; 30425f80e16SEugene Zhulenev } 30525f80e16SEugene Zhulenev 30625f80e16SEugene Zhulenev //===----------------------------------------------------------------------===// 307d43b2360SEugene Zhulenev // Convert async.create_group operation to async.runtime.create_group 30825f80e16SEugene Zhulenev //===----------------------------------------------------------------------===// 30925f80e16SEugene Zhulenev 31025f80e16SEugene Zhulenev namespace { 31125f80e16SEugene Zhulenev class CreateGroupOpLowering : public OpConversionPattern<CreateGroupOp> { 31225f80e16SEugene Zhulenev public: 31325f80e16SEugene Zhulenev using OpConversionPattern::OpConversionPattern; 31425f80e16SEugene Zhulenev 31525f80e16SEugene Zhulenev LogicalResult 31625f80e16SEugene Zhulenev matchAndRewrite(CreateGroupOp op, ArrayRef<Value> operands, 31725f80e16SEugene Zhulenev ConversionPatternRewriter &rewriter) const override { 318d43b2360SEugene Zhulenev rewriter.replaceOpWithNewOp<RuntimeCreateGroupOp>( 319d43b2360SEugene Zhulenev op, GroupType::get(op->getContext()), operands); 32025f80e16SEugene Zhulenev return success(); 32125f80e16SEugene Zhulenev } 32225f80e16SEugene Zhulenev }; 32325f80e16SEugene Zhulenev } // namespace 32425f80e16SEugene Zhulenev 32525f80e16SEugene Zhulenev //===----------------------------------------------------------------------===// 32625f80e16SEugene Zhulenev // Convert async.add_to_group operation to async.runtime.add_to_group. 32725f80e16SEugene Zhulenev //===----------------------------------------------------------------------===// 32825f80e16SEugene Zhulenev 32925f80e16SEugene Zhulenev namespace { 33025f80e16SEugene Zhulenev class AddToGroupOpLowering : public OpConversionPattern<AddToGroupOp> { 33125f80e16SEugene Zhulenev public: 33225f80e16SEugene Zhulenev using OpConversionPattern::OpConversionPattern; 33325f80e16SEugene Zhulenev 33425f80e16SEugene Zhulenev LogicalResult 33525f80e16SEugene Zhulenev matchAndRewrite(AddToGroupOp op, ArrayRef<Value> operands, 33625f80e16SEugene Zhulenev ConversionPatternRewriter &rewriter) const override { 33725f80e16SEugene Zhulenev rewriter.replaceOpWithNewOp<RuntimeAddToGroupOp>( 33825f80e16SEugene Zhulenev op, rewriter.getIndexType(), operands); 33925f80e16SEugene Zhulenev return success(); 34025f80e16SEugene Zhulenev } 34125f80e16SEugene Zhulenev }; 34225f80e16SEugene Zhulenev } // namespace 34325f80e16SEugene Zhulenev 34425f80e16SEugene Zhulenev //===----------------------------------------------------------------------===// 34525f80e16SEugene Zhulenev // Convert async.await and async.await_all operations to the async.runtime.await 34625f80e16SEugene Zhulenev // or async.runtime.await_and_resume operations. 34725f80e16SEugene Zhulenev //===----------------------------------------------------------------------===// 34825f80e16SEugene Zhulenev 34925f80e16SEugene Zhulenev namespace { 35025f80e16SEugene Zhulenev template <typename AwaitType, typename AwaitableType> 35125f80e16SEugene Zhulenev class AwaitOpLoweringBase : public OpConversionPattern<AwaitType> { 35225f80e16SEugene Zhulenev using AwaitAdaptor = typename AwaitType::Adaptor; 35325f80e16SEugene Zhulenev 35425f80e16SEugene Zhulenev public: 35539957aa4SEugene Zhulenev AwaitOpLoweringBase(MLIRContext *ctx, 35639957aa4SEugene Zhulenev llvm::DenseMap<FuncOp, CoroMachinery> &outlinedFunctions) 35725f80e16SEugene Zhulenev : OpConversionPattern<AwaitType>(ctx), 35825f80e16SEugene Zhulenev outlinedFunctions(outlinedFunctions) {} 35925f80e16SEugene Zhulenev 36025f80e16SEugene Zhulenev LogicalResult 36125f80e16SEugene Zhulenev matchAndRewrite(AwaitType op, ArrayRef<Value> operands, 36225f80e16SEugene Zhulenev ConversionPatternRewriter &rewriter) const override { 36325f80e16SEugene Zhulenev // We can only await on one the `AwaitableType` (for `await` it can be 36425f80e16SEugene Zhulenev // a `token` or a `value`, for `await_all` it must be a `group`). 36525f80e16SEugene Zhulenev if (!op.operand().getType().template isa<AwaitableType>()) 36625f80e16SEugene Zhulenev return rewriter.notifyMatchFailure(op, "unsupported awaitable type"); 36725f80e16SEugene Zhulenev 36825f80e16SEugene Zhulenev // Check if await operation is inside the outlined coroutine function. 36925f80e16SEugene Zhulenev auto func = op->template getParentOfType<FuncOp>(); 37025f80e16SEugene Zhulenev auto outlined = outlinedFunctions.find(func); 37125f80e16SEugene Zhulenev const bool isInCoroutine = outlined != outlinedFunctions.end(); 37225f80e16SEugene Zhulenev 37325f80e16SEugene Zhulenev Location loc = op->getLoc(); 37425f80e16SEugene Zhulenev Value operand = AwaitAdaptor(operands).operand(); 37525f80e16SEugene Zhulenev 37625f80e16SEugene Zhulenev // Inside regular functions we use the blocking wait operation to wait for 37725f80e16SEugene Zhulenev // the async object (token, value or group) to become available. 37825f80e16SEugene Zhulenev if (!isInCoroutine) 37925f80e16SEugene Zhulenev rewriter.create<RuntimeAwaitOp>(loc, operand); 38025f80e16SEugene Zhulenev 38125f80e16SEugene Zhulenev // Inside the coroutine we convert await operation into coroutine suspension 38225f80e16SEugene Zhulenev // point, and resume execution asynchronously. 38325f80e16SEugene Zhulenev if (isInCoroutine) { 38439957aa4SEugene Zhulenev CoroMachinery &coro = outlined->getSecond(); 38525f80e16SEugene Zhulenev Block *suspended = op->getBlock(); 38625f80e16SEugene Zhulenev 38725f80e16SEugene Zhulenev ImplicitLocOpBuilder builder(loc, op, rewriter.getListener()); 38825f80e16SEugene Zhulenev MLIRContext *ctx = op->getContext(); 38925f80e16SEugene Zhulenev 39025f80e16SEugene Zhulenev // Save the coroutine state and resume on a runtime managed thread when 39125f80e16SEugene Zhulenev // the operand becomes available. 39225f80e16SEugene Zhulenev auto coroSaveOp = 39325f80e16SEugene Zhulenev builder.create<CoroSaveOp>(CoroStateType::get(ctx), coro.coroHandle); 39425f80e16SEugene Zhulenev builder.create<RuntimeAwaitAndResumeOp>(operand, coro.coroHandle); 39525f80e16SEugene Zhulenev 39625f80e16SEugene Zhulenev // Split the entry block before the await operation. 39725f80e16SEugene Zhulenev Block *resume = rewriter.splitBlock(suspended, Block::iterator(op)); 39825f80e16SEugene Zhulenev 39925f80e16SEugene Zhulenev // Add async.coro.suspend as a suspended block terminator. 40025f80e16SEugene Zhulenev builder.setInsertionPointToEnd(suspended); 40125f80e16SEugene Zhulenev builder.create<CoroSuspendOp>(coroSaveOp.state(), coro.suspend, resume, 40225f80e16SEugene Zhulenev coro.cleanup); 40325f80e16SEugene Zhulenev 40439957aa4SEugene Zhulenev // Split the resume block into error checking and continuation. 40539957aa4SEugene Zhulenev Block *continuation = rewriter.splitBlock(resume, Block::iterator(op)); 40639957aa4SEugene Zhulenev 40739957aa4SEugene Zhulenev // Check if the awaited value is in the error state. 40839957aa4SEugene Zhulenev builder.setInsertionPointToStart(resume); 409d8c84d2aSEugene Zhulenev auto isError = 410d8c84d2aSEugene Zhulenev builder.create<RuntimeIsErrorOp>(loc, rewriter.getI1Type(), operand); 41139957aa4SEugene Zhulenev builder.create<CondBranchOp>(isError, 41239957aa4SEugene Zhulenev /*trueDest=*/setupSetErrorBlock(coro), 41339957aa4SEugene Zhulenev /*trueArgs=*/ArrayRef<Value>(), 41439957aa4SEugene Zhulenev /*falseDest=*/continuation, 41539957aa4SEugene Zhulenev /*falseArgs=*/ArrayRef<Value>()); 41639957aa4SEugene Zhulenev 41739957aa4SEugene Zhulenev // Make sure that replacement value will be constructed in the 41839957aa4SEugene Zhulenev // continuation block. 41939957aa4SEugene Zhulenev rewriter.setInsertionPointToStart(continuation); 42039957aa4SEugene Zhulenev } 42125f80e16SEugene Zhulenev 42225f80e16SEugene Zhulenev // Erase or replace the await operation with the new value. 42325f80e16SEugene Zhulenev if (Value replaceWith = getReplacementValue(op, operand, rewriter)) 42425f80e16SEugene Zhulenev rewriter.replaceOp(op, replaceWith); 42525f80e16SEugene Zhulenev else 42625f80e16SEugene Zhulenev rewriter.eraseOp(op); 42725f80e16SEugene Zhulenev 42825f80e16SEugene Zhulenev return success(); 42925f80e16SEugene Zhulenev } 43025f80e16SEugene Zhulenev 43125f80e16SEugene Zhulenev virtual Value getReplacementValue(AwaitType op, Value operand, 43225f80e16SEugene Zhulenev ConversionPatternRewriter &rewriter) const { 43325f80e16SEugene Zhulenev return Value(); 43425f80e16SEugene Zhulenev } 43525f80e16SEugene Zhulenev 43625f80e16SEugene Zhulenev private: 43739957aa4SEugene Zhulenev llvm::DenseMap<FuncOp, CoroMachinery> &outlinedFunctions; 43825f80e16SEugene Zhulenev }; 43925f80e16SEugene Zhulenev 44025f80e16SEugene Zhulenev /// Lowering for `async.await` with a token operand. 44125f80e16SEugene Zhulenev class AwaitTokenOpLowering : public AwaitOpLoweringBase<AwaitOp, TokenType> { 44225f80e16SEugene Zhulenev using Base = AwaitOpLoweringBase<AwaitOp, TokenType>; 44325f80e16SEugene Zhulenev 44425f80e16SEugene Zhulenev public: 44525f80e16SEugene Zhulenev using Base::Base; 44625f80e16SEugene Zhulenev }; 44725f80e16SEugene Zhulenev 44825f80e16SEugene Zhulenev /// Lowering for `async.await` with a value operand. 44925f80e16SEugene Zhulenev class AwaitValueOpLowering : public AwaitOpLoweringBase<AwaitOp, ValueType> { 45025f80e16SEugene Zhulenev using Base = AwaitOpLoweringBase<AwaitOp, ValueType>; 45125f80e16SEugene Zhulenev 45225f80e16SEugene Zhulenev public: 45325f80e16SEugene Zhulenev using Base::Base; 45425f80e16SEugene Zhulenev 45525f80e16SEugene Zhulenev Value 45625f80e16SEugene Zhulenev getReplacementValue(AwaitOp op, Value operand, 45725f80e16SEugene Zhulenev ConversionPatternRewriter &rewriter) const override { 45825f80e16SEugene Zhulenev // Load from the async value storage. 45925f80e16SEugene Zhulenev auto valueType = operand.getType().cast<ValueType>().getValueType(); 46025f80e16SEugene Zhulenev return rewriter.create<RuntimeLoadOp>(op->getLoc(), valueType, operand); 46125f80e16SEugene Zhulenev } 46225f80e16SEugene Zhulenev }; 46325f80e16SEugene Zhulenev 46425f80e16SEugene Zhulenev /// Lowering for `async.await_all` operation. 46525f80e16SEugene Zhulenev class AwaitAllOpLowering : public AwaitOpLoweringBase<AwaitAllOp, GroupType> { 46625f80e16SEugene Zhulenev using Base = AwaitOpLoweringBase<AwaitAllOp, GroupType>; 46725f80e16SEugene Zhulenev 46825f80e16SEugene Zhulenev public: 46925f80e16SEugene Zhulenev using Base::Base; 47025f80e16SEugene Zhulenev }; 47125f80e16SEugene Zhulenev 47225f80e16SEugene Zhulenev } // namespace 47325f80e16SEugene Zhulenev 47425f80e16SEugene Zhulenev //===----------------------------------------------------------------------===// 47525f80e16SEugene Zhulenev // Convert async.yield operation to async.runtime operations. 47625f80e16SEugene Zhulenev //===----------------------------------------------------------------------===// 47725f80e16SEugene Zhulenev 47825f80e16SEugene Zhulenev class YieldOpLowering : public OpConversionPattern<async::YieldOp> { 47925f80e16SEugene Zhulenev public: 48025f80e16SEugene Zhulenev YieldOpLowering( 48125f80e16SEugene Zhulenev MLIRContext *ctx, 48225f80e16SEugene Zhulenev const llvm::DenseMap<FuncOp, CoroMachinery> &outlinedFunctions) 48325f80e16SEugene Zhulenev : OpConversionPattern<async::YieldOp>(ctx), 48425f80e16SEugene Zhulenev outlinedFunctions(outlinedFunctions) {} 48525f80e16SEugene Zhulenev 48625f80e16SEugene Zhulenev LogicalResult 48725f80e16SEugene Zhulenev matchAndRewrite(async::YieldOp op, ArrayRef<Value> operands, 48825f80e16SEugene Zhulenev ConversionPatternRewriter &rewriter) const override { 48939957aa4SEugene Zhulenev // Check if yield operation is inside the async coroutine function. 49025f80e16SEugene Zhulenev auto func = op->template getParentOfType<FuncOp>(); 49125f80e16SEugene Zhulenev auto outlined = outlinedFunctions.find(func); 49225f80e16SEugene Zhulenev if (outlined == outlinedFunctions.end()) 49325f80e16SEugene Zhulenev return rewriter.notifyMatchFailure( 49439957aa4SEugene Zhulenev op, "operation is not inside the async coroutine function"); 49525f80e16SEugene Zhulenev 49625f80e16SEugene Zhulenev Location loc = op->getLoc(); 49725f80e16SEugene Zhulenev const CoroMachinery &coro = outlined->getSecond(); 49825f80e16SEugene Zhulenev 49925f80e16SEugene Zhulenev // Store yielded values into the async values storage and switch async 50025f80e16SEugene Zhulenev // values state to available. 50125f80e16SEugene Zhulenev for (auto tuple : llvm::zip(operands, coro.returnValues)) { 50225f80e16SEugene Zhulenev Value yieldValue = std::get<0>(tuple); 50325f80e16SEugene Zhulenev Value asyncValue = std::get<1>(tuple); 50425f80e16SEugene Zhulenev rewriter.create<RuntimeStoreOp>(loc, yieldValue, asyncValue); 50525f80e16SEugene Zhulenev rewriter.create<RuntimeSetAvailableOp>(loc, asyncValue); 50625f80e16SEugene Zhulenev } 50725f80e16SEugene Zhulenev 50825f80e16SEugene Zhulenev // Switch the coroutine completion token to available state. 50925f80e16SEugene Zhulenev rewriter.replaceOpWithNewOp<RuntimeSetAvailableOp>(op, coro.asyncToken); 51025f80e16SEugene Zhulenev 51125f80e16SEugene Zhulenev return success(); 51225f80e16SEugene Zhulenev } 51325f80e16SEugene Zhulenev 51425f80e16SEugene Zhulenev private: 51525f80e16SEugene Zhulenev const llvm::DenseMap<FuncOp, CoroMachinery> &outlinedFunctions; 51625f80e16SEugene Zhulenev }; 51725f80e16SEugene Zhulenev 51825f80e16SEugene Zhulenev //===----------------------------------------------------------------------===// 51939957aa4SEugene Zhulenev // Convert std.assert operation to cond_br into `set_error` block. 52039957aa4SEugene Zhulenev //===----------------------------------------------------------------------===// 52139957aa4SEugene Zhulenev 52239957aa4SEugene Zhulenev class AssertOpLowering : public OpConversionPattern<AssertOp> { 52339957aa4SEugene Zhulenev public: 52439957aa4SEugene Zhulenev AssertOpLowering(MLIRContext *ctx, 52539957aa4SEugene Zhulenev llvm::DenseMap<FuncOp, CoroMachinery> &outlinedFunctions) 52639957aa4SEugene Zhulenev : OpConversionPattern<AssertOp>(ctx), 52739957aa4SEugene Zhulenev outlinedFunctions(outlinedFunctions) {} 52839957aa4SEugene Zhulenev 52939957aa4SEugene Zhulenev LogicalResult 53039957aa4SEugene Zhulenev matchAndRewrite(AssertOp op, ArrayRef<Value> operands, 53139957aa4SEugene Zhulenev ConversionPatternRewriter &rewriter) const override { 53239957aa4SEugene Zhulenev // Check if assert operation is inside the async coroutine function. 53339957aa4SEugene Zhulenev auto func = op->template getParentOfType<FuncOp>(); 53439957aa4SEugene Zhulenev auto outlined = outlinedFunctions.find(func); 53539957aa4SEugene Zhulenev if (outlined == outlinedFunctions.end()) 53639957aa4SEugene Zhulenev return rewriter.notifyMatchFailure( 53739957aa4SEugene Zhulenev op, "operation is not inside the async coroutine function"); 53839957aa4SEugene Zhulenev 53939957aa4SEugene Zhulenev Location loc = op->getLoc(); 54039957aa4SEugene Zhulenev CoroMachinery &coro = outlined->getSecond(); 54139957aa4SEugene Zhulenev 54239957aa4SEugene Zhulenev Block *cont = rewriter.splitBlock(op->getBlock(), Block::iterator(op)); 54339957aa4SEugene Zhulenev rewriter.setInsertionPointToEnd(cont->getPrevNode()); 54439957aa4SEugene Zhulenev rewriter.create<CondBranchOp>(loc, AssertOpAdaptor(operands).arg(), 54539957aa4SEugene Zhulenev /*trueDest=*/cont, 54639957aa4SEugene Zhulenev /*trueArgs=*/ArrayRef<Value>(), 54739957aa4SEugene Zhulenev /*falseDest=*/setupSetErrorBlock(coro), 54839957aa4SEugene Zhulenev /*falseArgs=*/ArrayRef<Value>()); 54939957aa4SEugene Zhulenev rewriter.eraseOp(op); 55039957aa4SEugene Zhulenev 55139957aa4SEugene Zhulenev return success(); 55239957aa4SEugene Zhulenev } 55339957aa4SEugene Zhulenev 55439957aa4SEugene Zhulenev private: 55539957aa4SEugene Zhulenev llvm::DenseMap<FuncOp, CoroMachinery> &outlinedFunctions; 55639957aa4SEugene Zhulenev }; 55739957aa4SEugene Zhulenev 55839957aa4SEugene Zhulenev //===----------------------------------------------------------------------===// 55925f80e16SEugene Zhulenev 56025f80e16SEugene Zhulenev void AsyncToAsyncRuntimePass::runOnOperation() { 56125f80e16SEugene Zhulenev ModuleOp module = getOperation(); 56225f80e16SEugene Zhulenev SymbolTable symbolTable(module); 56325f80e16SEugene Zhulenev 56425f80e16SEugene Zhulenev // Outline all `async.execute` body regions into async functions (coroutines). 56525f80e16SEugene Zhulenev llvm::DenseMap<FuncOp, CoroMachinery> outlinedFunctions; 56625f80e16SEugene Zhulenev 56725f80e16SEugene Zhulenev module.walk([&](ExecuteOp execute) { 56825f80e16SEugene Zhulenev outlinedFunctions.insert(outlineExecuteOp(symbolTable, execute)); 56925f80e16SEugene Zhulenev }); 57025f80e16SEugene Zhulenev 57125f80e16SEugene Zhulenev LLVM_DEBUG({ 57225f80e16SEugene Zhulenev llvm::dbgs() << "Outlined " << outlinedFunctions.size() 57325f80e16SEugene Zhulenev << " functions built from async.execute operations\n"; 57425f80e16SEugene Zhulenev }); 57525f80e16SEugene Zhulenev 576*de7a4e53SEugene Zhulenev // Returns true if operation is inside the coroutine. 577*de7a4e53SEugene Zhulenev auto isInCoroutine = [&](Operation *op) -> bool { 578*de7a4e53SEugene Zhulenev auto parentFunc = op->getParentOfType<FuncOp>(); 579*de7a4e53SEugene Zhulenev return outlinedFunctions.find(parentFunc) != outlinedFunctions.end(); 580*de7a4e53SEugene Zhulenev }; 581*de7a4e53SEugene Zhulenev 58225f80e16SEugene Zhulenev // Lower async operations to async.runtime operations. 58325f80e16SEugene Zhulenev MLIRContext *ctx = module->getContext(); 584dc4e913bSChris Lattner RewritePatternSet asyncPatterns(ctx); 58525f80e16SEugene Zhulenev 586*de7a4e53SEugene Zhulenev // Conversion to async runtime augments original CFG with the coroutine CFG, 587*de7a4e53SEugene Zhulenev // and we have to make sure that structured control flow operations with async 588*de7a4e53SEugene Zhulenev // operations in nested regions will be converted to branch-based control flow 589*de7a4e53SEugene Zhulenev // before we add the coroutine basic blocks. 590*de7a4e53SEugene Zhulenev populateLoopToStdConversionPatterns(asyncPatterns); 591*de7a4e53SEugene Zhulenev 59225f80e16SEugene Zhulenev // Async lowering does not use type converter because it must preserve all 59325f80e16SEugene Zhulenev // types for async.runtime operations. 594dc4e913bSChris Lattner asyncPatterns.add<CreateGroupOpLowering, AddToGroupOpLowering>(ctx); 595dc4e913bSChris Lattner asyncPatterns.add<AwaitTokenOpLowering, AwaitValueOpLowering, 59625f80e16SEugene Zhulenev AwaitAllOpLowering, YieldOpLowering>(ctx, 59725f80e16SEugene Zhulenev outlinedFunctions); 59825f80e16SEugene Zhulenev 59939957aa4SEugene Zhulenev // Lower assertions to conditional branches into error blocks. 60039957aa4SEugene Zhulenev asyncPatterns.add<AssertOpLowering>(ctx, outlinedFunctions); 60139957aa4SEugene Zhulenev 60225f80e16SEugene Zhulenev // All high level async operations must be lowered to the runtime operations. 60325f80e16SEugene Zhulenev ConversionTarget runtimeTarget(*ctx); 60425f80e16SEugene Zhulenev runtimeTarget.addLegalDialect<AsyncDialect>(); 60525f80e16SEugene Zhulenev runtimeTarget.addIllegalOp<CreateGroupOp, AddToGroupOp>(); 60625f80e16SEugene Zhulenev runtimeTarget.addIllegalOp<ExecuteOp, AwaitOp, AwaitAllOp, async::YieldOp>(); 60725f80e16SEugene Zhulenev 608*de7a4e53SEugene Zhulenev // Decide if structured control flow has to be lowered to branch-based CFG. 609*de7a4e53SEugene Zhulenev runtimeTarget.addDynamicallyLegalDialect<scf::SCFDialect>([&](Operation *op) { 610*de7a4e53SEugene Zhulenev auto walkResult = op->walk([&](Operation *nested) { 611*de7a4e53SEugene Zhulenev bool isAsync = isa<async::AsyncDialect>(nested->getDialect()); 612*de7a4e53SEugene Zhulenev return isAsync && isInCoroutine(nested) ? WalkResult::interrupt() 613*de7a4e53SEugene Zhulenev : WalkResult::advance(); 614*de7a4e53SEugene Zhulenev }); 615*de7a4e53SEugene Zhulenev return !walkResult.wasInterrupted(); 616*de7a4e53SEugene Zhulenev }); 617*de7a4e53SEugene Zhulenev runtimeTarget.addLegalOp<BranchOp, CondBranchOp>(); 618*de7a4e53SEugene Zhulenev 6198f23fac4SEugene Zhulenev // Assertions must be converted to runtime errors inside async functions. 6208f23fac4SEugene Zhulenev runtimeTarget.addDynamicallyLegalOp<AssertOp>([&](AssertOp op) -> bool { 6218f23fac4SEugene Zhulenev auto func = op->getParentOfType<FuncOp>(); 6228f23fac4SEugene Zhulenev return outlinedFunctions.find(func) == outlinedFunctions.end(); 6238f23fac4SEugene Zhulenev }); 62439957aa4SEugene Zhulenev 62525f80e16SEugene Zhulenev if (failed(applyPartialConversion(module, runtimeTarget, 62625f80e16SEugene Zhulenev std::move(asyncPatterns)))) { 62725f80e16SEugene Zhulenev signalPassFailure(); 62825f80e16SEugene Zhulenev return; 62925f80e16SEugene Zhulenev } 63025f80e16SEugene Zhulenev } 63125f80e16SEugene Zhulenev 63225f80e16SEugene Zhulenev std::unique_ptr<OperationPass<ModuleOp>> mlir::createAsyncToAsyncRuntimePass() { 63325f80e16SEugene Zhulenev return std::make_unique<AsyncToAsyncRuntimePass>(); 63425f80e16SEugene Zhulenev } 635