================
@@ -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
----------------
bcardosolopes wrote:
#227527 is approved and adds a `cond` operand to `cir.coro.suspend_point` for
bool `await_suspend`. This ignores it, so a vetoed suspend would still suspend.
It's a small change: when `getCond()` is set, branch to `resumeEntry` on false
before the `cir.coro.intrinsic.suspend`, which is what classic does with
`CreateCondBr(SuspendRet, RealSuspendBlock, ReadyBlock)`. Please handle it
here, or at least error out on a `cond` operand, depending on which PR lands
first.
https://github.com/llvm/llvm-project/pull/203802
_______________________________________________
cfe-commits mailing list
[email protected]
https://lists.llvm.org/cgi-bin/mailman/listinfo/cfe-commits
