================
@@ -2110,12 +2118,339 @@ 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`. It runs once every nested op (awaits, scopes,
+// cleanups, ...) has been flattened. The initialSuspend, body and finalSuspend
+// regions become the body of a cleanup scope whose cleanup is the destroy
+// region; the exit region follows the scope:
+//
+// cir.cleanup.scope {
+// initialSuspend ---> body ---> finalSuspend
+// | ^
+// +-------------+ cir.co_return
+// } cleanup all {
+// destroy coro.free
+// }
+// exit coro.end + cir.return
+//
+// Leaving the 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);
+ }
+
+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.
+ 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);
+ mlir::Block *bodyPlaceholder = &scope.getBodyRegion().front();
+
+ // Scope body: initialSuspend -> body -> finalSuspend.
+ mlir::Region &initialRegion = coroOp.getInitialSuspend();
+ mlir::Block *afterInitial = &initialRegion.back();
+ rewriter.inlineRegionBefore(initialRegion, bodyPlaceholder);
+
+ mlir::Region &bodyRegion = coroOp.getBody();
+ mlir::Block *beforeBody = &bodyRegion.front();
+ mlir::Block *afterBody = &bodyRegion.back();
+
+ rewriter.setInsertionPointToEnd(afterInitial);
+ auto initialYieldOp = cast<cir::YieldOp>(afterInitial->getTerminator());
+ rewriter.replaceOpWithNewOp<cir::BrOp>(
+ initialYieldOp, initialYieldOp.getArgs(), beforeBody);
+
+ llvm::SmallVector<cir::CoReturnOp> coReturnsOps;
+ collectCoReturnsOps(bodyRegion, coReturnsOps);
+
+ mlir::Region &finalRegion = coroOp.getFinalSuspend();
+ mlir::Block *beforeFinal = &finalRegion.front();
+
+ rewriter.inlineRegionBefore(bodyRegion, bodyPlaceholder);
+ rewriter.setInsertionPointToEnd(afterBody);
+
+ // 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.
+ mlir::Operation *bodyTerminator = afterBody->getTerminator();
+ if (auto bodyYield = dyn_cast<cir::YieldOp>(bodyTerminator))
+ 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(finalRegion, bodyPlaceholder);
+ rewriter.eraseBlock(bodyPlaceholder);
+
+ // Scope cleanup: the destroy region.
+ mlir::Region &destroyRegion = coroOp.getDestroy();
+ mlir::Block *cleanupPlaceholder = &scope.getCleanupRegion().front();
+ rewriter.inlineRegionBefore(destroyRegion, cleanupPlaceholder);
+ rewriter.eraseBlock(cleanupPlaceholder);
+
+ mlir::Region &exitRegion = coroOp.getExit();
+ mlir::Block *beforeExit = &exitRegion.front();
+
+ for (auto &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::coroutineGroManager());
+ rewriter.setInsertionPoint(coroExit);
+ rewriter.replaceOpWithNewOp<cir::BrOp>(coroExit, mlir::ValueRange{},
+ beforeExit);
+ }
+
+ // Exit region: inline it right after the scope, which falls through to it.
+ rewriter.setInsertionPointAfter(scope);
+ mlir::Block *afterCleanup =
+ rewriter.splitBlock(remainingOpsBlock, rewriter.getInsertionPoint());
+ rewriter.inlineRegionBefore(exitRegion, afterCleanup);
+
+ rewriter.setInsertionPointAfter(scope);
+ cir::BrOp::create(rewriter, loc, mlir::ValueRange(), beforeExit);
+
+ 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> {
+public:
+ using OpRewritePattern<cir::AwaitOp>::OpRewritePattern;
+
+ // Flatten `cir.await` into the following control flow:
+ //
+ // awaitBlock
+ // |
+ // v
+ // ready await_ready()
+ // true / \ false
+ // | v
+ // | suspend coro.save, await_suspend(),
+ // | / | \ coro.suspend
+ // | 0/ 1| \-1 (default)
+ // v v v v
+ // resume destroy suspended
+ // | cir.yield cir.coro.exit
+ // v
+ // continuation ops after the await
+ //
+ // `destroy` is meant to leave through the enclosing cleanups to the
+ // coroutine's destroy region (see the FIXME below), and `suspended` goes
+ // straight to its exit region (see CIRCoroutineFlattening).
+ mlir::LogicalResult
+ matchAndRewrite(cir::AwaitOp awaitOp,
+ mlir::PatternRewriter &rewriter) const override {
+
+ for (mlir::Region ®ion : awaitOp->getRegions())
----------------
erichkeane wrote:
This function has the same problem. Break up the responsibilities/actions best
you can so we can read/digest this better.
https://github.com/llvm/llvm-project/pull/203802
_______________________________________________
cfe-commits mailing list
[email protected]
https://lists.llvm.org/cgi-bin/mailman/listinfo/cfe-commits
