1//===- AsyncToAsyncRuntime.cpp - Lower from Async to Async Runtime --------===//
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 implements lowering from high level async operations to async.coro
10// and async.runtime operations.
11//
12//===----------------------------------------------------------------------===//
13
14#include <utility>
15
16#include "mlir/Dialect/Async/Passes.h"
17
18#include "PassDetail.h"
19#include "mlir/Conversion/SCFToControlFlow/SCFToControlFlow.h"
20#include "mlir/Dialect/Arith/IR/Arith.h"
21#include "mlir/Dialect/Async/IR/Async.h"
22#include "mlir/Dialect/ControlFlow/IR/ControlFlowOps.h"
23#include "mlir/Dialect/Func/IR/FuncOps.h"
24#include "mlir/Dialect/SCF/IR/SCF.h"
25#include "mlir/IR/IRMapping.h"
26#include "mlir/IR/ImplicitLocOpBuilder.h"
27#include "mlir/IR/PatternMatch.h"
28#include "mlir/Transforms/DialectConversion.h"
29#include "mlir/Transforms/RegionUtils.h"
30#include "llvm/ADT/SetVector.h"
31#include "llvm/Support/Debug.h"
32#include <optional>
33
34namespace mlir {
35#define GEN_PASS_DEF_ASYNCTOASYNCRUNTIMEPASS
36#define GEN_PASS_DEF_ASYNCFUNCTOASYNCRUNTIMEPASS
37#include "mlir/Dialect/Async/Passes.h.inc"
38} // namespace mlir
39
40using namespace mlir;
41using namespace mlir::async;
42
43#define DEBUG_TYPE "async-to-async-runtime"
44// Prefix for functions outlined from `async.execute` op regions.
45static constexpr const char kAsyncFnPrefix[] = "async_execute_fn";
46
47namespace {
48
49class AsyncToAsyncRuntimePass
50 : public impl::AsyncToAsyncRuntimePassBase<AsyncToAsyncRuntimePass> {
51public:
52 AsyncToAsyncRuntimePass() = default;
53 void runOnOperation() override;
54};
55
56} // namespace
57
58namespace {
59
60class AsyncFuncToAsyncRuntimePass
61 : public impl::AsyncFuncToAsyncRuntimePassBase<
62 AsyncFuncToAsyncRuntimePass> {
63public:
64 AsyncFuncToAsyncRuntimePass() = default;
65 void runOnOperation() override;
66};
67
68} // namespace
69
70/// Function targeted for coroutine transformation has two additional blocks at
71/// the end: coroutine cleanup and coroutine suspension.
72///
73/// async.await op lowering additionaly creates a resume block for each
74/// operation to enable non-blocking waiting via coroutine suspension.
75namespace {
76struct CoroMachinery {
77 func::FuncOp func;
78
79 // Async function returns an optional token, followed by some async values
80 //
81 // async.func @foo() -> !async.value<T> {
82 // %cst = arith.constant 42.0 : T
83 // return %cst: T
84 // }
85 // Async execute region returns a completion token, and an async value for
86 // each yielded value.
87 //
88 // %token, %result = async.execute -> !async.value<T> {
89 // %0 = arith.constant ... : T
90 // async.yield %0 : T
91 // }
92 std::optional<Value> asyncToken; // returned completion token
93 llvm::SmallVector<Value, 4> returnValues; // returned async values
94
95 Value coroHandle; // coroutine handle (!async.coro.getHandle value)
96 Block *entry; // coroutine entry block
97 std::optional<Block *> setError; // set returned values to error state
98 Block *cleanup; // coroutine cleanup block
99
100 // Coroutine cleanup block for destroy after the coroutine is resumed,
101 // e.g. async.coro.suspend state, [suspend], [resume], [destroy]
102 //
103 // This cleanup block is a duplicate of the cleanup block followed by the
104 // resume block. The purpose of having a duplicate cleanup block for destroy
105 // is to make the CFG clear so that the control flow analysis won't confuse.
106 //
107 // The overall structure of the lowered CFG can be the following,
108 //
109 // Entry (calling async.coro.suspend)
110 // | \
111 // Resume Destroy (duplicate of Cleanup)
112 // | |
113 // Cleanup |
114 // | /
115 // End (ends the corontine)
116 //
117 // If there is resume-specific cleanup logic, it can go into the Cleanup
118 // block but not the destroy block. Otherwise, it can fail block dominance
119 // check.
120 Block *cleanupForDestroy;
121 Block *suspend; // coroutine suspension block
122};
123} // namespace
124
125using FuncCoroMapPtr =
126 std::shared_ptr<llvm::DenseMap<func::FuncOp, CoroMachinery>>;
127
128/// Utility to partially update the regular function CFG to the coroutine CFG
129/// compatible with LLVM coroutines switched-resume lowering using
130/// `async.runtime.*` and `async.coro.*` operations. Adds a new entry block
131/// that branches into preexisting entry block. Also inserts trailing blocks.
132///
133/// The result types of the passed `func` start with an optional `async.token`
134/// and be continued with some number of `async.value`s.
135///
136/// See LLVM coroutines documentation: https://llvm.org/docs/Coroutines.html
137///
138/// - `entry` block sets up the coroutine.
139/// - `set_error` block sets completion token and async values state to error.
140/// - `cleanup` block cleans up the coroutine state.
141/// - `suspend block after the @llvm.coro.end() defines what value will be
142/// returned to the initial caller of a coroutine. Everything before the
143/// @llvm.coro.end() will be executed at every suspension point.
144///
145/// Coroutine structure (only the important bits):
146///
147/// func @some_fn(<function-arguments>) -> (!async.token, !async.value<T>)
148/// {
149/// ^entry(<function-arguments>):
150/// %token = <async token> : !async.token // create async runtime token
151/// %value = <async value> : !async.value<T> // create async value
152/// %id = async.coro.getId // create a coroutine id
153/// %hdl = async.coro.begin %id // create a coroutine handle
154/// cf.br ^preexisting_entry_block
155///
156/// /* preexisting blocks modified to branch to the cleanup block */
157///
158/// ^set_error: // this block created lazily only if needed (see code below)
159/// async.runtime.set_error %token : !async.token
160/// async.runtime.set_error %value : !async.value<T>
161/// cf.br ^cleanup
162///
163/// ^cleanup:
164/// async.coro.free %hdl // delete the coroutine state
165/// cf.br ^suspend
166///
167/// ^suspend:
168/// async.coro.end %hdl // marks the end of a coroutine
169/// return %token, %value : !async.token, !async.value<T>
170/// }
171///
172static CoroMachinery setupCoroMachinery(func::FuncOp func) {
173 assert(!func.getBlocks().empty() && "Function must have an entry block");
174
175 MLIRContext *ctx = func.getContext();
176 Block *entryBlock = &func.getBlocks().front();
177 Block *originalEntryBlock =
178 entryBlock->splitBlock(splitBefore: entryBlock->getOperations().begin());
179 auto builder = ImplicitLocOpBuilder::atBlockBegin(loc: func->getLoc(), block: entryBlock);
180
181 // ------------------------------------------------------------------------ //
182 // Allocate async token/values that we will return from a ramp function.
183 // ------------------------------------------------------------------------ //
184
185 // We treat TokenType as state update marker to represent side-effects of
186 // async computations
187 bool isStateful = isa<TokenType>(func.getResultTypes().front());
188
189 std::optional<Value> retToken;
190 if (isStateful)
191 retToken.emplace(builder.create<RuntimeCreateOp>(TokenType::get(ctx)));
192
193 llvm::SmallVector<Value, 4> retValues;
194 ArrayRef<Type> resValueTypes =
195 isStateful ? func.getResultTypes().drop_front() : func.getResultTypes();
196 for (auto resType : resValueTypes)
197 retValues.emplace_back(
198 builder.create<RuntimeCreateOp>(resType).getResult());
199
200 // ------------------------------------------------------------------------ //
201 // Initialize coroutine: get coroutine id and coroutine handle.
202 // ------------------------------------------------------------------------ //
203 auto coroIdOp = builder.create<CoroIdOp>(CoroIdType::get(ctx));
204 auto coroHdlOp =
205 builder.create<CoroBeginOp>(CoroHandleType::get(ctx), coroIdOp.getId());
206 builder.create<cf::BranchOp>(originalEntryBlock);
207
208 Block *cleanupBlock = func.addBlock();
209 Block *cleanupBlockForDestroy = func.addBlock();
210 Block *suspendBlock = func.addBlock();
211
212 // ------------------------------------------------------------------------ //
213 // Coroutine cleanup blocks: deallocate coroutine frame, free the memory.
214 // ------------------------------------------------------------------------ //
215 auto buildCleanupBlock = [&](Block *cb) {
216 builder.setInsertionPointToStart(cb);
217 builder.create<CoroFreeOp>(coroIdOp.getId(), coroHdlOp.getHandle());
218
219 // Branch into the suspend block.
220 builder.create<cf::BranchOp>(suspendBlock);
221 };
222 buildCleanupBlock(cleanupBlock);
223 buildCleanupBlock(cleanupBlockForDestroy);
224
225 // ------------------------------------------------------------------------ //
226 // Coroutine suspend block: mark the end of a coroutine and return allocated
227 // async token.
228 // ------------------------------------------------------------------------ //
229 builder.setInsertionPointToStart(suspendBlock);
230
231 // Mark the end of a coroutine: async.coro.end
232 builder.create<CoroEndOp>(coroHdlOp.getHandle());
233
234 // Return created optional `async.token` and `async.values` from the suspend
235 // block. This will be the return value of a coroutine ramp function.
236 SmallVector<Value, 4> ret;
237 if (retToken)
238 ret.push_back(Elt: *retToken);
239 llvm::append_range(C&: ret, R&: retValues);
240 builder.create<func::ReturnOp>(ret);
241
242 // `async.await` op lowering will create resume blocks for async
243 // continuations, and will conditionally branch to cleanup or suspend blocks.
244
245 // The switch-resumed API based coroutine should be marked with
246 // presplitcoroutine attribute to mark the function as a coroutine.
247 func->setAttr("passthrough", builder.getArrayAttr(
248 StringAttr::get(ctx, "presplitcoroutine")));
249
250 CoroMachinery machinery;
251 machinery.func = func;
252 machinery.asyncToken = retToken;
253 machinery.returnValues = retValues;
254 machinery.coroHandle = coroHdlOp.getHandle();
255 machinery.entry = entryBlock;
256 machinery.setError = std::nullopt; // created lazily only if needed
257 machinery.cleanup = cleanupBlock;
258 machinery.cleanupForDestroy = cleanupBlockForDestroy;
259 machinery.suspend = suspendBlock;
260 return machinery;
261}
262
263// Lazily creates `set_error` block only if it is required for lowering to the
264// runtime operations (see for example lowering of assert operation).
265static Block *setupSetErrorBlock(CoroMachinery &coro) {
266 if (coro.setError)
267 return *coro.setError;
268
269 coro.setError = coro.func.addBlock();
270 (*coro.setError)->moveBefore(block: coro.cleanup);
271
272 auto builder =
273 ImplicitLocOpBuilder::atBlockBegin(loc: coro.func->getLoc(), block: *coro.setError);
274
275 // Coroutine set_error block: set error on token and all returned values.
276 if (coro.asyncToken)
277 builder.create<RuntimeSetErrorOp>(*coro.asyncToken);
278
279 for (Value retValue : coro.returnValues)
280 builder.create<RuntimeSetErrorOp>(retValue);
281
282 // Branch into the cleanup block.
283 builder.create<cf::BranchOp>(coro.cleanup);
284
285 return *coro.setError;
286}
287
288//===----------------------------------------------------------------------===//
289// async.execute op outlining to the coroutine functions.
290//===----------------------------------------------------------------------===//
291
292/// Outline the body region attached to the `async.execute` op into a standalone
293/// function.
294///
295/// Note that this is not reversible transformation.
296static std::pair<func::FuncOp, CoroMachinery>
297outlineExecuteOp(SymbolTable &symbolTable, ExecuteOp execute) {
298 ModuleOp module = execute->getParentOfType<ModuleOp>();
299
300 MLIRContext *ctx = module.getContext();
301 Location loc = execute.getLoc();
302
303 // Make sure that all constants will be inside the outlined async function to
304 // reduce the number of function arguments.
305 cloneConstantsIntoTheRegion(execute.getBodyRegion());
306
307 // Collect all outlined function inputs.
308 SetVector<mlir::Value> functionInputs(llvm::from_range,
309 execute.getDependencies());
310 functionInputs.insert_range(execute.getBodyOperands());
311 getUsedValuesDefinedAbove(execute.getBodyRegion(), functionInputs);
312
313 // Collect types for the outlined function inputs and outputs.
314 auto typesRange = llvm::map_range(
315 C&: functionInputs, F: [](Value value) { return value.getType(); });
316 SmallVector<Type, 4> inputTypes(typesRange.begin(), typesRange.end());
317 auto outputTypes = execute.getResultTypes();
318
319 auto funcType = FunctionType::get(ctx, inputTypes, outputTypes);
320 auto funcAttrs = ArrayRef<NamedAttribute>();
321
322 // TODO: Derive outlined function name from the parent FuncOp (support
323 // multiple nested async.execute operations).
324 func::FuncOp func =
325 func::FuncOp::create(loc, kAsyncFnPrefix, funcType, funcAttrs);
326 symbolTable.insert(symbol: func);
327
328 SymbolTable::setSymbolVisibility(symbol: func, vis: SymbolTable::Visibility::Private);
329 auto builder = ImplicitLocOpBuilder::atBlockBegin(loc, block: func.addEntryBlock());
330
331 // Prepare for coroutine conversion by creating the body of the function.
332 {
333 size_t numDependencies = execute.getDependencies().size();
334 size_t numOperands = execute.getBodyOperands().size();
335
336 // Await on all dependencies before starting to execute the body region.
337 for (size_t i = 0; i < numDependencies; ++i)
338 builder.create<AwaitOp>(func.getArgument(i));
339
340 // Await on all async value operands and unwrap the payload.
341 SmallVector<Value, 4> unwrappedOperands(numOperands);
342 for (size_t i = 0; i < numOperands; ++i) {
343 Value operand = func.getArgument(numDependencies + i);
344 unwrappedOperands[i] = builder.create<AwaitOp>(loc, operand).getResult();
345 }
346
347 // Map from function inputs defined above the execute op to the function
348 // arguments.
349 IRMapping valueMapping;
350 valueMapping.map(functionInputs, func.getArguments());
351 valueMapping.map(execute.getBodyRegion().getArguments(), unwrappedOperands);
352
353 // Clone all operations from the execute operation body into the outlined
354 // function body.
355 for (Operation &op : execute.getBodyRegion().getOps())
356 builder.clone(op, valueMapping);
357 }
358
359 // Adding entry/cleanup/suspend blocks.
360 CoroMachinery coro = setupCoroMachinery(func);
361
362 // Suspend async function at the end of an entry block, and resume it using
363 // Async resume operation (execution will be resumed in a thread managed by
364 // the async runtime).
365 {
366 cf::BranchOp branch = cast<cf::BranchOp>(coro.entry->getTerminator());
367 builder.setInsertionPointToEnd(coro.entry);
368
369 // Save the coroutine state: async.coro.save
370 auto coroSaveOp =
371 builder.create<CoroSaveOp>(CoroStateType::get(ctx), coro.coroHandle);
372
373 // Pass coroutine to the runtime to be resumed on a runtime managed
374 // thread.
375 builder.create<RuntimeResumeOp>(coro.coroHandle);
376
377 // Add async.coro.suspend as a suspended block terminator.
378 builder.create<CoroSuspendOp>(coroSaveOp.getState(), coro.suspend,
379 branch.getDest(), coro.cleanupForDestroy);
380
381 branch.erase();
382 }
383
384 // Replace the original `async.execute` with a call to outlined function.
385 {
386 ImplicitLocOpBuilder callBuilder(loc, execute);
387 auto callOutlinedFunc = callBuilder.create<func::CallOp>(
388 func.getName(), execute.getResultTypes(), functionInputs.getArrayRef());
389 execute.replaceAllUsesWith(callOutlinedFunc.getResults());
390 execute.erase();
391 }
392
393 return {func, coro};
394}
395
396//===----------------------------------------------------------------------===//
397// Convert async.create_group operation to async.runtime.create_group
398//===----------------------------------------------------------------------===//
399
400namespace {
401class CreateGroupOpLowering : public OpConversionPattern<CreateGroupOp> {
402public:
403 using OpConversionPattern::OpConversionPattern;
404
405 LogicalResult
406 matchAndRewrite(CreateGroupOp op, OpAdaptor adaptor,
407 ConversionPatternRewriter &rewriter) const override {
408 rewriter.replaceOpWithNewOp<RuntimeCreateGroupOp>(
409 op, GroupType::get(op->getContext()), adaptor.getOperands());
410 return success();
411 }
412};
413} // namespace
414
415//===----------------------------------------------------------------------===//
416// Convert async.add_to_group operation to async.runtime.add_to_group.
417//===----------------------------------------------------------------------===//
418
419namespace {
420class AddToGroupOpLowering : public OpConversionPattern<AddToGroupOp> {
421public:
422 using OpConversionPattern::OpConversionPattern;
423
424 LogicalResult
425 matchAndRewrite(AddToGroupOp op, OpAdaptor adaptor,
426 ConversionPatternRewriter &rewriter) const override {
427 rewriter.replaceOpWithNewOp<RuntimeAddToGroupOp>(
428 op, rewriter.getIndexType(), adaptor.getOperands());
429 return success();
430 }
431};
432} // namespace
433
434//===----------------------------------------------------------------------===//
435// Convert async.func, async.return and async.call operations to non-blocking
436// operations based on llvm coroutine
437//===----------------------------------------------------------------------===//
438
439namespace {
440
441//===----------------------------------------------------------------------===//
442// Convert async.func operation to func.func
443//===----------------------------------------------------------------------===//
444
445class AsyncFuncOpLowering : public OpConversionPattern<async::FuncOp> {
446public:
447 AsyncFuncOpLowering(MLIRContext *ctx, FuncCoroMapPtr coros)
448 : OpConversionPattern<async::FuncOp>(ctx), coros(std::move(coros)) {}
449
450 LogicalResult
451 matchAndRewrite(async::FuncOp op, OpAdaptor adaptor,
452 ConversionPatternRewriter &rewriter) const override {
453 Location loc = op->getLoc();
454
455 auto newFuncOp =
456 rewriter.create<func::FuncOp>(loc, op.getName(), op.getFunctionType());
457
458 SymbolTable::setSymbolVisibility(symbol: newFuncOp,
459 vis: SymbolTable::getSymbolVisibility(symbol: op));
460 // Copy over all attributes other than the name.
461 for (const auto &namedAttr : op->getAttrs()) {
462 if (namedAttr.getName() != SymbolTable::getSymbolAttrName())
463 newFuncOp->setAttr(namedAttr.getName(), namedAttr.getValue());
464 }
465
466 rewriter.inlineRegionBefore(op.getBody(), newFuncOp.getBody(),
467 newFuncOp.end());
468
469 CoroMachinery coro = setupCoroMachinery(newFuncOp);
470 (*coros)[newFuncOp] = coro;
471 // no initial suspend, we should hot-start
472
473 rewriter.eraseOp(op: op);
474 return success();
475 }
476
477private:
478 FuncCoroMapPtr coros;
479};
480
481//===----------------------------------------------------------------------===//
482// Convert async.call operation to func.call
483//===----------------------------------------------------------------------===//
484
485class AsyncCallOpLowering : public OpConversionPattern<async::CallOp> {
486public:
487 AsyncCallOpLowering(MLIRContext *ctx)
488 : OpConversionPattern<async::CallOp>(ctx) {}
489
490 LogicalResult
491 matchAndRewrite(async::CallOp op, OpAdaptor adaptor,
492 ConversionPatternRewriter &rewriter) const override {
493 rewriter.replaceOpWithNewOp<func::CallOp>(
494 op, op.getCallee(), op.getResultTypes(), op.getOperands());
495 return success();
496 }
497};
498
499//===----------------------------------------------------------------------===//
500// Convert async.return operation to async.runtime operations.
501//===----------------------------------------------------------------------===//
502
503class AsyncReturnOpLowering : public OpConversionPattern<async::ReturnOp> {
504public:
505 AsyncReturnOpLowering(MLIRContext *ctx, FuncCoroMapPtr coros)
506 : OpConversionPattern<async::ReturnOp>(ctx), coros(std::move(coros)) {}
507
508 LogicalResult
509 matchAndRewrite(async::ReturnOp op, OpAdaptor adaptor,
510 ConversionPatternRewriter &rewriter) const override {
511 auto func = op->template getParentOfType<func::FuncOp>();
512 auto funcCoro = coros->find(func);
513 if (funcCoro == coros->end())
514 return rewriter.notifyMatchFailure(
515 op, "operation is not inside the async coroutine function");
516
517 Location loc = op->getLoc();
518 const CoroMachinery &coro = funcCoro->getSecond();
519 rewriter.setInsertionPointAfter(op);
520
521 // Store return values into the async values storage and switch async
522 // values state to available.
523 for (auto tuple : llvm::zip(adaptor.getOperands(), coro.returnValues)) {
524 Value returnValue = std::get<0>(tuple);
525 Value asyncValue = std::get<1>(tuple);
526 rewriter.create<RuntimeStoreOp>(loc, returnValue, asyncValue);
527 rewriter.create<RuntimeSetAvailableOp>(loc, asyncValue);
528 }
529
530 if (coro.asyncToken)
531 // Switch the coroutine completion token to available state.
532 rewriter.create<RuntimeSetAvailableOp>(loc, *coro.asyncToken);
533
534 rewriter.eraseOp(op: op);
535 rewriter.create<cf::BranchOp>(loc, coro.cleanup);
536 return success();
537 }
538
539private:
540 FuncCoroMapPtr coros;
541};
542} // namespace
543
544//===----------------------------------------------------------------------===//
545// Convert async.await and async.await_all operations to the async.runtime.await
546// or async.runtime.await_and_resume operations.
547//===----------------------------------------------------------------------===//
548
549namespace {
550template <typename AwaitType, typename AwaitableType>
551class AwaitOpLoweringBase : public OpConversionPattern<AwaitType> {
552 using AwaitAdaptor = typename AwaitType::Adaptor;
553
554public:
555 AwaitOpLoweringBase(MLIRContext *ctx, FuncCoroMapPtr coros,
556 bool shouldLowerBlockingWait)
557 : OpConversionPattern<AwaitType>(ctx), coros(std::move(coros)),
558 shouldLowerBlockingWait(shouldLowerBlockingWait) {}
559
560 LogicalResult
561 matchAndRewrite(AwaitType op, typename AwaitType::Adaptor adaptor,
562 ConversionPatternRewriter &rewriter) const override {
563 // We can only await on one the `AwaitableType` (for `await` it can be
564 // a `token` or a `value`, for `await_all` it must be a `group`).
565 if (!isa<AwaitableType>(op.getOperand().getType()))
566 return rewriter.notifyMatchFailure(op, "unsupported awaitable type");
567
568 // Check if await operation is inside the coroutine function.
569 auto func = op->template getParentOfType<func::FuncOp>();
570 auto funcCoro = coros->find(func);
571 const bool isInCoroutine = funcCoro != coros->end();
572
573 Location loc = op->getLoc();
574 Value operand = adaptor.getOperand();
575
576 Type i1 = rewriter.getI1Type();
577
578 // Delay lowering to block wait in case await op is inside async.execute
579 if (!isInCoroutine && !shouldLowerBlockingWait)
580 return failure();
581
582 // Inside regular functions we use the blocking wait operation to wait for
583 // the async object (token, value or group) to become available.
584 if (!isInCoroutine) {
585 ImplicitLocOpBuilder builder(loc, rewriter);
586 builder.create<RuntimeAwaitOp>(loc, operand);
587
588 // Assert that the awaited operands is not in the error state.
589 Value isError = builder.create<RuntimeIsErrorOp>(i1, operand);
590 Value notError = builder.create<arith::XOrIOp>(
591 isError, builder.create<arith::ConstantOp>(
592 loc, i1, builder.getIntegerAttr(i1, 1)));
593
594 builder.create<cf::AssertOp>(notError,
595 "Awaited async operand is in error state");
596 }
597
598 // Inside the coroutine we convert await operation into coroutine suspension
599 // point, and resume execution asynchronously.
600 if (isInCoroutine) {
601 CoroMachinery &coro = funcCoro->getSecond();
602 Block *suspended = op->getBlock();
603
604 ImplicitLocOpBuilder builder(loc, rewriter);
605 MLIRContext *ctx = op->getContext();
606
607 // Save the coroutine state and resume on a runtime managed thread when
608 // the operand becomes available.
609 auto coroSaveOp =
610 builder.create<CoroSaveOp>(CoroStateType::get(ctx), coro.coroHandle);
611 builder.create<RuntimeAwaitAndResumeOp>(operand, coro.coroHandle);
612
613 // Split the entry block before the await operation.
614 Block *resume = rewriter.splitBlock(block: suspended, before: Block::iterator(op));
615
616 // Add async.coro.suspend as a suspended block terminator.
617 builder.setInsertionPointToEnd(suspended);
618 builder.create<CoroSuspendOp>(coroSaveOp.getState(), coro.suspend, resume,
619 coro.cleanupForDestroy);
620
621 // Split the resume block into error checking and continuation.
622 Block *continuation = rewriter.splitBlock(block: resume, before: Block::iterator(op));
623
624 // Check if the awaited value is in the error state.
625 builder.setInsertionPointToStart(resume);
626 auto isError = builder.create<RuntimeIsErrorOp>(loc, i1, operand);
627 builder.create<cf::CondBranchOp>(isError,
628 /*trueDest=*/setupSetErrorBlock(coro),
629 /*trueArgs=*/ArrayRef<Value>(),
630 /*falseDest=*/continuation,
631 /*falseArgs=*/ArrayRef<Value>());
632
633 // Make sure that replacement value will be constructed in the
634 // continuation block.
635 rewriter.setInsertionPointToStart(continuation);
636 }
637
638 // Erase or replace the await operation with the new value.
639 if (Value replaceWith = getReplacementValue(op, operand, rewriter))
640 rewriter.replaceOp(op, replaceWith);
641 else
642 rewriter.eraseOp(op);
643
644 return success();
645 }
646
647 virtual Value getReplacementValue(AwaitType op, Value operand,
648 ConversionPatternRewriter &rewriter) const {
649 return Value();
650 }
651
652private:
653 FuncCoroMapPtr coros;
654 bool shouldLowerBlockingWait;
655};
656
657/// Lowering for `async.await` with a token operand.
658class AwaitTokenOpLowering : public AwaitOpLoweringBase<AwaitOp, TokenType> {
659 using Base = AwaitOpLoweringBase<AwaitOp, TokenType>;
660
661public:
662 using Base::Base;
663};
664
665/// Lowering for `async.await` with a value operand.
666class AwaitValueOpLowering : public AwaitOpLoweringBase<AwaitOp, ValueType> {
667 using Base = AwaitOpLoweringBase<AwaitOp, ValueType>;
668
669public:
670 using Base::Base;
671
672 Value
673 getReplacementValue(AwaitOp op, Value operand,
674 ConversionPatternRewriter &rewriter) const override {
675 // Load from the async value storage.
676 auto valueType = cast<ValueType>(operand.getType()).getValueType();
677 return rewriter.create<RuntimeLoadOp>(op->getLoc(), valueType, operand);
678 }
679};
680
681/// Lowering for `async.await_all` operation.
682class AwaitAllOpLowering : public AwaitOpLoweringBase<AwaitAllOp, GroupType> {
683 using Base = AwaitOpLoweringBase<AwaitAllOp, GroupType>;
684
685public:
686 using Base::Base;
687};
688
689} // namespace
690
691//===----------------------------------------------------------------------===//
692// Convert async.yield operation to async.runtime operations.
693//===----------------------------------------------------------------------===//
694
695class YieldOpLowering : public OpConversionPattern<async::YieldOp> {
696public:
697 YieldOpLowering(MLIRContext *ctx, FuncCoroMapPtr coros)
698 : OpConversionPattern<async::YieldOp>(ctx), coros(std::move(coros)) {}
699
700 LogicalResult
701 matchAndRewrite(async::YieldOp op, OpAdaptor adaptor,
702 ConversionPatternRewriter &rewriter) const override {
703 // Check if yield operation is inside the async coroutine function.
704 auto func = op->template getParentOfType<func::FuncOp>();
705 auto funcCoro = coros->find(func);
706 if (funcCoro == coros->end())
707 return rewriter.notifyMatchFailure(
708 op, "operation is not inside the async coroutine function");
709
710 Location loc = op->getLoc();
711 const CoroMachinery &coro = funcCoro->getSecond();
712
713 // Store yielded values into the async values storage and switch async
714 // values state to available.
715 for (auto tuple : llvm::zip(adaptor.getOperands(), coro.returnValues)) {
716 Value yieldValue = std::get<0>(tuple);
717 Value asyncValue = std::get<1>(tuple);
718 rewriter.create<RuntimeStoreOp>(loc, yieldValue, asyncValue);
719 rewriter.create<RuntimeSetAvailableOp>(loc, asyncValue);
720 }
721
722 if (coro.asyncToken)
723 // Switch the coroutine completion token to available state.
724 rewriter.create<RuntimeSetAvailableOp>(loc, *coro.asyncToken);
725
726 rewriter.eraseOp(op: op);
727 rewriter.create<cf::BranchOp>(loc, coro.cleanup);
728
729 return success();
730 }
731
732private:
733 FuncCoroMapPtr coros;
734};
735
736//===----------------------------------------------------------------------===//
737// Convert cf.assert operation to cf.cond_br into `set_error` block.
738//===----------------------------------------------------------------------===//
739
740class AssertOpLowering : public OpConversionPattern<cf::AssertOp> {
741public:
742 AssertOpLowering(MLIRContext *ctx, FuncCoroMapPtr coros)
743 : OpConversionPattern<cf::AssertOp>(ctx), coros(std::move(coros)) {}
744
745 LogicalResult
746 matchAndRewrite(cf::AssertOp op, OpAdaptor adaptor,
747 ConversionPatternRewriter &rewriter) const override {
748 // Check if assert operation is inside the async coroutine function.
749 auto func = op->template getParentOfType<func::FuncOp>();
750 auto funcCoro = coros->find(func);
751 if (funcCoro == coros->end())
752 return rewriter.notifyMatchFailure(
753 op, "operation is not inside the async coroutine function");
754
755 Location loc = op->getLoc();
756 CoroMachinery &coro = funcCoro->getSecond();
757
758 Block *cont = rewriter.splitBlock(block: op->getBlock(), before: Block::iterator(op));
759 rewriter.setInsertionPointToEnd(cont->getPrevNode());
760 rewriter.create<cf::CondBranchOp>(loc, adaptor.getArg(),
761 /*trueDest=*/cont,
762 /*trueArgs=*/ArrayRef<Value>(),
763 /*falseDest=*/setupSetErrorBlock(coro),
764 /*falseArgs=*/ArrayRef<Value>());
765 rewriter.eraseOp(op: op);
766
767 return success();
768 }
769
770private:
771 FuncCoroMapPtr coros;
772};
773
774//===----------------------------------------------------------------------===//
775void AsyncToAsyncRuntimePass::runOnOperation() {
776 ModuleOp module = getOperation();
777 SymbolTable symbolTable(module);
778
779 // Functions with coroutine CFG setups, which are results of outlining
780 // `async.execute` body regions
781 FuncCoroMapPtr coros =
782 std::make_shared<llvm::DenseMap<func::FuncOp, CoroMachinery>>();
783
784 module.walk([&](ExecuteOp execute) {
785 coros->insert(outlineExecuteOp(symbolTable, execute));
786 });
787
788 LLVM_DEBUG({
789 llvm::dbgs() << "Outlined " << coros->size()
790 << " functions built from async.execute operations\n";
791 });
792
793 // Returns true if operation is inside the coroutine.
794 auto isInCoroutine = [&](Operation *op) -> bool {
795 auto parentFunc = op->getParentOfType<func::FuncOp>();
796 return coros->contains(parentFunc);
797 };
798
799 // Lower async operations to async.runtime operations.
800 MLIRContext *ctx = module->getContext();
801 RewritePatternSet asyncPatterns(ctx);
802
803 // Conversion to async runtime augments original CFG with the coroutine CFG,
804 // and we have to make sure that structured control flow operations with async
805 // operations in nested regions will be converted to branch-based control flow
806 // before we add the coroutine basic blocks.
807 populateSCFToControlFlowConversionPatterns(patterns&: asyncPatterns);
808
809 // Async lowering does not use type converter because it must preserve all
810 // types for async.runtime operations.
811 asyncPatterns.add<CreateGroupOpLowering, AddToGroupOpLowering>(arg&: ctx);
812
813 asyncPatterns
814 .add<AwaitTokenOpLowering, AwaitValueOpLowering, AwaitAllOpLowering>(
815 ctx, coros, /*should_lower_blocking_wait=*/true);
816
817 // Lower assertions to conditional branches into error blocks.
818 asyncPatterns.add<YieldOpLowering, AssertOpLowering>(ctx, coros);
819
820 // All high level async operations must be lowered to the runtime operations.
821 ConversionTarget runtimeTarget(*ctx);
822 runtimeTarget.addLegalDialect<AsyncDialect, func::FuncDialect>();
823 runtimeTarget.addIllegalOp<CreateGroupOp, AddToGroupOp>();
824 runtimeTarget.addIllegalOp<ExecuteOp, AwaitOp, AwaitAllOp, async::YieldOp>();
825
826 // Decide if structured control flow has to be lowered to branch-based CFG.
827 runtimeTarget.addDynamicallyLegalDialect<scf::SCFDialect>(callback: [&](Operation *op) {
828 auto walkResult = op->walk(callback: [&](Operation *nested) {
829 bool isAsync = isa<async::AsyncDialect>(Val: nested->getDialect());
830 return isAsync && isInCoroutine(nested) ? WalkResult::interrupt()
831 : WalkResult::advance();
832 });
833 return !walkResult.wasInterrupted();
834 });
835 runtimeTarget.addLegalOp<cf::AssertOp, arith::XOrIOp, arith::ConstantOp,
836 func::ConstantOp, cf::BranchOp, cf::CondBranchOp>();
837
838 // Assertions must be converted to runtime errors inside async functions.
839 runtimeTarget.addDynamicallyLegalOp<cf::AssertOp>(
840 [&](cf::AssertOp op) -> bool {
841 auto func = op->getParentOfType<func::FuncOp>();
842 return !coros->contains(func);
843 });
844
845 if (failed(applyPartialConversion(module, runtimeTarget,
846 std::move(asyncPatterns)))) {
847 signalPassFailure();
848 return;
849 }
850}
851
852//===----------------------------------------------------------------------===//
853void mlir::populateAsyncFuncToAsyncRuntimeConversionPatterns(
854 RewritePatternSet &patterns, ConversionTarget &target) {
855 // Functions with coroutine CFG setups, which are results of converting
856 // async.func.
857 FuncCoroMapPtr coros =
858 std::make_shared<llvm::DenseMap<func::FuncOp, CoroMachinery>>();
859 MLIRContext *ctx = patterns.getContext();
860 // Lower async.func to func.func with coroutine cfg.
861 patterns.add<AsyncCallOpLowering>(arg&: ctx);
862 patterns.add<AsyncFuncOpLowering, AsyncReturnOpLowering>(ctx, coros);
863
864 patterns.add<AwaitTokenOpLowering, AwaitValueOpLowering, AwaitAllOpLowering>(
865 ctx, coros, /*should_lower_blocking_wait=*/false);
866 patterns.add<YieldOpLowering, AssertOpLowering>(ctx, coros);
867
868 target.addDynamicallyLegalOp<AwaitOp, AwaitAllOp, YieldOp, cf::AssertOp>(
869 [coros](Operation *op) {
870 auto exec = op->getParentOfType<ExecuteOp>();
871 auto func = op->getParentOfType<func::FuncOp>();
872 return exec || !coros->contains(func);
873 });
874}
875
876void AsyncFuncToAsyncRuntimePass::runOnOperation() {
877 ModuleOp module = getOperation();
878
879 // Lower async operations to async.runtime operations.
880 MLIRContext *ctx = module->getContext();
881 RewritePatternSet asyncPatterns(ctx);
882 ConversionTarget runtimeTarget(*ctx);
883
884 // Lower async.func to func.func with coroutine cfg.
885 populateAsyncFuncToAsyncRuntimeConversionPatterns(patterns&: asyncPatterns,
886 target&: runtimeTarget);
887
888 runtimeTarget.addLegalDialect<AsyncDialect, func::FuncDialect>();
889 runtimeTarget.addIllegalOp<async::FuncOp, async::CallOp, async::ReturnOp>();
890
891 runtimeTarget.addLegalOp<arith::XOrIOp, arith::ConstantOp, func::ConstantOp,
892 cf::BranchOp, cf::CondBranchOp>();
893
894 if (failed(applyPartialConversion(module, runtimeTarget,
895 std::move(asyncPatterns)))) {
896 signalPassFailure();
897 return;
898 }
899}
900

source code of mlir/lib/Dialect/Async/Transforms/AsyncToAsyncRuntime.cpp