================
@@ -2110,12 +2118,428 @@ class CIRTryOpFlattening : public
mlir::OpRewritePattern<cir::TryOp> {
}
};
+static void
+collectCoReturnsOps(mlir::Region ®ion,
+ llvm::SmallVectorImpl<cir::CoReturnOp> &coReturnOps) {
+ region.walk(
+ [&](cir::CoReturnOp coReturnOp) { coReturnOps.push_back(coReturnOp); });
+}
+
+// Flatten `cir.coroutine`. This runs once every nested op (awaits, cir.scope,
+// cleanup scopes, ...) has been flattened. The initialSuspend, body and
+// finalSuspend regions become the body of a cir.cleanup.scope whose cleanup is
+// the destroy region; the exit region follows the cleanup scope:
+//
+// cir.cleanup.scope {
+// initialSuspend ---> body ---> finalSuspend
+// | ^
+// +-------------+ cir.co_return
+// } cleanup all {
+// destroy cir.coro.intrinsic.free
+// }
+// exit cir.coro.intrinsic.end, cir.return
+//
+// Leaving the cleanup scope body always runs the destroy region. In
+// particular, the cir.yield emitted for an await's destroy case leaves through
+// it. A suspended coroutine (cir.coro.exit) must keep its frame, so it
+// branches straight to the exit region.
+class CIRCoroutineFlattening : public mlir::OpRewritePattern<cir::CoroutineOp>
{
+ void collectCoroExitsOps(
+ cir::CoroutineOp coroOp,
+ llvm::SmallVectorImpl<cir::CoroExitOp> &coroExitsOps) const {
+ // cir.coro.exit is produced by flattening a cir.await, which can only
+ // appear in the initialSuspend, body and finalSuspend regions.
+ auto collectCoroExit = [&](cir::CoroExitOp op) {
+ coroExitsOps.push_back(op);
+ };
+ coroOp.getInitialSuspend().walk(collectCoroExit);
+ coroOp.getBody().walk(collectCoroExit);
+ coroOp.getFinalSuspend().walk(collectCoroExit);
+ }
+
+ // Cleanup scope body: initialSuspend -> body -> finalSuspend. The three
+ // regions are moved, in this order, in front of the placeholder block the
+ // cleanup scope was created with, and the placeholder is erased once they
+ // are in place.
+ void inlineCleanUpScopeBody(mlir::PatternRewriter &rewriter,
+ cir::CoroutineOp coroOp,
+ cir::CleanupScopeOp scope) const {
+ mlir::Block *bodyPlaceholder = &scope.getBodyRegion().front();
+ // Take the entry blocks before the regions are moved.
+ mlir::Block *beforeBody = &coroOp.getBody().front();
+ mlir::Block *beforeFinal = &coroOp.getFinalSuspend().front();
+
+ inlineInitialSuspend(rewriter, coroOp, bodyPlaceholder, beforeBody);
+ inlineBody(rewriter, coroOp, bodyPlaceholder, beforeFinal);
+ inlineFinalSuspend(rewriter, coroOp, bodyPlaceholder);
+ rewriter.eraseBlock(bodyPlaceholder);
+ }
+
+ // initialSuspend ends in a cir.yield once the init await is done. Replace it
+ // with a branch to the body.
+ void inlineInitialSuspend(mlir::PatternRewriter &rewriter,
+ cir::CoroutineOp coroOp, mlir::Block *insertBefore,
+ mlir::Block *beforeBody) const {
+ mlir::Region &initialRegion = coroOp.getInitialSuspend();
+ auto initialYieldOp =
+ cast<cir::YieldOp>(initialRegion.back().getTerminator());
+ rewriter.setInsertionPoint(initialYieldOp);
+ rewriter.replaceOpWithNewOp<cir::BrOp>(
+ initialYieldOp, initialYieldOp.getArgs(), beforeBody);
+ rewriter.inlineRegionBefore(initialRegion, insertBefore);
+ }
+
+ // The body either falls off the end (cir.yield) or ends in a
+ // cir.co_return. Both, like every other cir.co_return in the body, branch
+ // to the final suspend.
+ void inlineBody(mlir::PatternRewriter &rewriter, cir::CoroutineOp coroOp,
+ mlir::Block *insertBefore, mlir::Block *beforeFinal) const {
+ mlir::Region &bodyRegion = coroOp.getBody();
+
+ llvm::SmallVector<cir::CoReturnOp> coReturnsOps;
+ collectCoReturnsOps(bodyRegion, coReturnsOps);
+
+ mlir::Operation *bodyTerminator = bodyRegion.back().getTerminator();
+ if (auto bodyYield = dyn_cast<cir::YieldOp>(bodyTerminator)) {
+ rewriter.setInsertionPoint(bodyYield);
+ rewriter.replaceOpWithNewOp<cir::BrOp>(bodyYield, mlir::ValueRange{},
+ beforeFinal);
+ }
+
+ for (auto &coReturn : coReturnsOps) {
+ rewriter.setInsertionPoint(coReturn);
+ rewriter.replaceOpWithNewOp<cir::BrOp>(coReturn, mlir::ValueRange{},
+ beforeFinal);
+ }
+
+ rewriter.inlineRegionBefore(bodyRegion, insertBefore);
+ }
+
+ // finalSuspend keeps its cir.yield, which now ends the cleanup scope body:
+ // once the final await completes, the coroutine leaves the cleanup scope
+ // normally and runs the destroy region.
+ void inlineFinalSuspend(mlir::PatternRewriter &rewriter,
+ cir::CoroutineOp coroOp,
+ mlir::Block *insertBefore) const {
+ rewriter.inlineRegionBefore(coroOp.getFinalSuspend(), insertBefore);
+ }
+
+ // Cleanup scope cleanup region: the destroy region.
+ void inlineCleanUpScopeCleanup(mlir::PatternRewriter &rewriter,
+ cir::CoroutineOp coroOp,
+ cir::CleanupScopeOp scope) const {
+ mlir::Region &destroyRegion = coroOp.getDestroy();
+ mlir::Block *cleanupPlaceholder = &scope.getCleanupRegion().front();
+ rewriter.inlineRegionBefore(destroyRegion, cleanupPlaceholder);
+ rewriter.eraseBlock(cleanupPlaceholder);
+ }
+
+ // A suspended coroutine keeps its frame, so each cir.coro.exit branches
+ // straight to the exit region without running the destroy region.
+ void lowerCoroExits(mlir::PatternRewriter &rewriter,
+ llvm::ArrayRef<cir::CoroExitOp> coroExitsOps,
+ mlir::Block *beforeExit) const {
+ for (cir::CoroExitOp coroExit : coroExitsOps) {
+ // TODO(cir): once the GRO region is supported, this may need to branch
+ // there instead of to the exit region.
+ assert(!cir::MissingFeatures::coroutineGetRetObjManager());
+ rewriter.setInsertionPoint(coroExit);
+ rewriter.replaceOpWithNewOp<cir::BrOp>(coroExit, mlir::ValueRange{},
+ beforeExit);
+ }
+ }
+
+ // Exit region: inline it right after the cleanup scope, which falls through
+ // to it.
+ void inlineExitRegion(mlir::PatternRewriter &rewriter,
+ cir::CoroutineOp coroOp, cir::CleanupScopeOp scope,
+ mlir::Block *remainingOpsBlock) const {
+ mlir::Region &exitRegion = coroOp.getExit();
+ mlir::Block *beforeExit = &exitRegion.front();
+
+ rewriter.setInsertionPointAfter(scope);
+ mlir::Block *afterCleanup =
+ rewriter.splitBlock(remainingOpsBlock, rewriter.getInsertionPoint());
+ rewriter.inlineRegionBefore(exitRegion, afterCleanup);
+
+ rewriter.setInsertionPointAfter(scope);
+ cir::BrOp::create(rewriter, scope.getLoc(), mlir::ValueRange(),
beforeExit);
+ }
+
+public:
+ using OpRewritePattern<cir::CoroutineOp>::OpRewritePattern;
+
+ mlir::LogicalResult
+ matchAndRewrite(cir::CoroutineOp coroOp,
+ mlir::PatternRewriter &rewriter) const override {
+
+ assert(!cir::MissingFeatures::coroutineUnwindRegion());
+ for (mlir::Region ®ion : coroOp->getRegions())
+ if (hasNestedOpsToFlatten(region))
+ return mlir::failure();
+
+ // Collect the cir.coro.exit ops before the regions are moved. Each one
+ // becomes a branch to the entry of the exit region.
+ llvm::SmallVector<cir::CoroExitOp> coroExitsOps;
+ collectCoroExitsOps(coroOp, coroExitsOps);
+
+ mlir::Location loc = coroOp.getLoc();
+ mlir::Block *coroBlock = rewriter.getInsertionBlock();
+ // Split the current block before the CoroutineOp to create the inlining
+ // point.
+ mlir::Block *remainingOpsBlock =
+ rewriter.splitBlock(coroBlock, rewriter.getInsertionPoint());
+
+ rewriter.setInsertionPointToEnd(coroBlock);
+ cir::BrOp::create(rewriter, loc, mlir::ValueRange(), remainingOpsBlock);
+
+ // coroOp now lives in remainingOpsBlock. Build the cleanup scope in front
+ // of it.
+ //
+ // The coroutine becomes a cir.cleanup.scope so that its frame is freed
+ // (the destroy region, as its cleanup region) on every way out of the
+ // coroutine: completing the final suspend, an await's destroy case and,
+ // because the kind is All, an exception escaping the coroutine. The
+ // regular cleanup scope flattening then routes each of those exits
+ // through the destroy region. A suspended coroutine must keep its frame,
+ // so cir.coro.exit skips the cleanup scope (see lowerCoroExits).
+ rewriter.setInsertionPoint(coroOp);
+
+ auto emitYield = [](mlir::OpBuilder &b, mlir::Location l) {
+ cir::YieldOp::create(b, l);
+ };
+
+ auto scope = cir::CleanupScopeOp::create(rewriter, loc, CleanupKind::All,
+ emitYield, emitYield);
+
+ inlineCleanUpScopeBody(rewriter, coroOp, scope);
+ inlineCleanUpScopeCleanup(rewriter, coroOp, scope);
+
+ mlir::Block *beforeExit = &coroOp.getExit().front();
+ lowerCoroExits(rewriter, coroExitsOps, beforeExit);
+ inlineExitRegion(rewriter, coroOp, scope, remainingOpsBlock);
+
+ auto func = coroOp->getParentOfType<cir::FuncOp>();
+ // The coroutine regions have been flattened, so the original coroutine
+ // attribute no longer verifies. Preserve the information that this
function
+ // originated from a coroutine for LLVM lowering.
+ func.setCoroutine(false);
+ func.setFlattenedCoroutine(true);
+
+ rewriter.eraseOp(coroOp);
+ return mlir::success();
+ }
+};
+
+class CIRAwaitOpFlattening : public mlir::OpRewritePattern<cir::AwaitOp> {
+ // The ready region ends in a cir.condition that holds the result of
+ // `await_ready()`. Replace it with a cir.brcond: if the awaiter is already
+ // ready, skip the suspension and go straight to the resume region
+ // otherwise go to the suspend region.
+ void flattenReadyRegion(mlir::PatternRewriter &rewriter, cir::AwaitOp
awaitOp,
+ mlir::Block *resumeEntry, mlir::Block *suspendEntry,
+ mlir::Block *remainingOpsBlock) const {
+ mlir::Region &readyRegion = awaitOp.getReady();
+ auto conditionOp =
+ cast<cir::ConditionOp>(readyRegion.back().getTerminator());
+ rewriter.setInsertionPoint(conditionOp);
+ rewriter.replaceOpWithNewOp<cir::BrCondOp>(
+ conditionOp, conditionOp.getCondition(), resumeEntry, suspendEntry);
+ rewriter.inlineRegionBefore(readyRegion, remainingOpsBlock);
+ }
+
+ // The suspend region runs `await_suspend()` and ends in a
+ // cir.coro.suspend_point. Emit cir.coro.intrinsic.save and
+ // cir.coro.intrinsic.suspend around it, and replace the suspend point with a
+ // cir.switch.flat on the suspend result, which decides whether the coroutine
+ // resumes, is destroyed or stays suspended.
+ void flattenSuspendRegion(mlir::PatternRewriter &rewriter,
+ cir::AwaitOp awaitOp, mlir::Block *resumeEntry,
+ mlir::Block *remainingOpsBlock) const {
+ mlir::Region &suspendRegion = awaitOp.getSuspend();
+ // TODO(cir): support an await_suspend() that returns bool (it can veto
+ // the suspension and resume right away) by flattening it into a brcond.
+ auto suspendPoint =
+ cast<cir::CoroSuspendPoint>(suspendRegion.back().getTerminator());
+ cir::CoroSuspendOp coroSuspendIntri =
+ emitCoroSuspend(rewriter, awaitOp, suspendPoint);
+
+ // Inline the suspend region before creating the destroy and suspended
+ // blocks, so they are placed right after it, in front of the continuation.
+ rewriter.inlineRegionBefore(suspendRegion, remainingOpsBlock);
+
+ mlir::Location loc = awaitOp.getLoc();
+ mlir::Block *cleanupBlock =
+ createDestroyBlock(rewriter, loc, remainingOpsBlock);
+ mlir::Block *retBlock =
+ createSuspendedBlock(rewriter, loc, remainingOpsBlock);
+
+ // 0 resumes, 1 destroys and the default (-1) stays suspended.
+ llvm::SmallVector<mlir::APInt> caseValues{mlir::APInt(8, 0),
+ mlir::APInt(8, 1)};
+ llvm::SmallVector<mlir::Block *> caseDestinations{resumeEntry,
+ cleanupBlock};
+ llvm::SmallVector<mlir::ValueRange> caseOperands{caseValues.size()};
+
+ // Dispatch to the appropriate coroutine path based on the result of
+ // `cir.coro.intrinsic.suspend`: resume, destroy, or suspend.
+ rewriter.setInsertionPoint(suspendPoint);
+ rewriter.replaceOpWithNewOp<cir::SwitchFlatOp>(
+ suspendPoint, coroSuspendIntri.getResult(), retBlock,
+ mlir::ValueRange{}, caseValues, caseDestinations, caseOperands);
+ }
+
+ // Emit cir.coro.intrinsic.save at the start of the suspend region and
+ // cir.coro.intrinsic.suspend right before its cir.coro.suspend_point.
+ cir::CoroSuspendOp emitCoroSuspend(mlir::PatternRewriter &rewriter,
+ cir::AwaitOp awaitOp,
+ cir::CoroSuspendPoint suspendPoint) const
{
+ mlir::Location loc = awaitOp.getLoc();
+
+ // cir.coro.intrinsic.suspend takes the token produced by
+ // cir.coro.intrinsic.save as its first argument, so emit the save first.
+ rewriter.setInsertionPointToStart(&awaitOp.getSuspend().front());
+
+ // The save goes first in the suspend region, before await_suspend()
+ // runs, so the coroutine is already considered suspended (and can be
+ // resumed) from that point.
+ auto voidPtrTy = cir::PointerType::get(cir::VoidType::get(getContext()));
+ auto nullPtr = cir::ConstantOp::create(
+ rewriter, loc,
+ cir::ConstPtrAttr::get(voidPtrTy, rewriter.getI64IntegerAttr(0)));
+ auto coroSave = cir::CoroSaveOp::create(rewriter, loc, nullPtr);
+ rewriter.setInsertionPoint(suspendPoint);
+
+ // The second argument to `cir.coro.intrinsic.suspend` indicates whether
+ // this is the final suspend point. Coroutines suspended at a final suspend
+ // point are considered done (`cir.coro.intrinsic.done` returns true) and
+ // may only be destroyed; resuming them is undefined behavior.
+ bool isFinalSuspend = awaitOp.getKind() == cir::AwaitKind::Final;
+ auto isFinalCoroSuspend = cir::ConstantOp::create(
+ rewriter, loc, cir::BoolAttr::get(getContext(), isFinalSuspend));
+
+ // cir.coro.intrinsic.suspend returns:
+ // -1 : coroutine suspended
+ // 0 : coroutine resumed
+ // 1 : coroutine destroyed
+ return cir::CoroSuspendOp::create(
+ rewriter, loc,
+ mlir::ValueRange{coroSave.getResult(), isFinalCoroSuspend});
+ }
+
+ // Destroy case. Classic codegen uses EmitBranchThroughCleanup() so every
+ // enclosing cleanup runs before the frame is freed. Here the destroy block
+ // is a single cir.yield, meant to leave through the coroutine cleanup scope
+ // (cir.coro.intrinsic.free) and then reach the exit region.
+ //
+ // FIXME: This only works when the await is directly in the coroutine
----------------
bcardosolopes wrote:
Until this is fixed, it's silent wrong code for very ordinary coroutines: a
`co_await` in a loop, or with a local that has a destructor, keeps running the
body after a destroy. The new tests (`silly_task_with_dtor`, `await_in_loop`,
...) check that output in. Please emit an NYI error when the await is nested in
anything other than the coroutine's own cleanup scope, and drop or `XFAIL`
those checks, rather than landing the miscompile. Is the dedicated terminator
much more work than that?
https://github.com/llvm/llvm-project/pull/203802
_______________________________________________
cfe-commits mailing list
[email protected]
https://lists.llvm.org/cgi-bin/mailman/listinfo/cfe-commits
