[CIR][ThroughMLIR] Handle ContinueOp directly under a WhileOp (llvm#1669)

Currently we can't handle continues nested under `IfOp`, because if we
replace it with a yield, then it only breaks out of that `if`-statement,
rather than continuing the whole loop.
Perhaps that should be done by changing the whole structure of the while
loop.

Co-authored-by: Yue Huang <[email protected]>

Author: 2 people

clang/lib/CIR/Lowering/ThroughMLIR/LowerCIRLoopToSCF.cpp

@@ -66,11 +66,12 @@ class SCFWhileLoop {
   SCFWhileLoop(cir::WhileOp op, cir::WhileOp::Adaptor adaptor,
                mlir::ConversionPatternRewriter *rewriter)
       : whileOp(op), adaptor(adaptor), rewriter(rewriter) {}
-  void transferToSCFWhileOp();
+  mlir::scf::WhileOp transferToSCFWhileOp();
 
 private:
   cir::WhileOp whileOp;
   cir::WhileOp::Adaptor adaptor;
+  mlir::scf::WhileOp scfWhileOp;
   mlir::ConversionPatternRewriter *rewriter;
 };
 
@@ -356,7 +357,7 @@ void SCFLoop::transformToSCFWhileOp() {
                               scfWhileOp.getAfterBody()->end());
 }
 
-void SCFWhileLoop::transferToSCFWhileOp() {
+mlir::scf::WhileOp SCFWhileLoop::transferToSCFWhileOp() {
   auto scfWhileOp = rewriter->create<mlir::scf::WhileOp>(
       whileOp->getLoc(), whileOp->getResultTypes(), adaptor.getOperands());
   rewriter->createBlock(&scfWhileOp.getBefore());
@@ -367,6 +368,7 @@ void SCFWhileLoop::transferToSCFWhileOp() {
   rewriter->inlineBlockBefore(&whileOp.getBody().front(),
                               scfWhileOp.getAfterBody(),
                               scfWhileOp.getAfterBody()->end());
+  return scfWhileOp;
 }
 
 void SCFDoLoop::transferToSCFWhileOp() {
@@ -412,14 +414,53 @@ class CIRForOpLowering : public mlir::OpConversionPattern<cir::ForOp> {
 };
 
 class CIRWhileOpLowering : public mlir::OpConversionPattern<cir::WhileOp> {
+  void rewriteContinue(mlir::scf::WhileOp whileOp,
+                       mlir::ConversionPatternRewriter &rewriter) const {
+    // Collect all ContinueOp inside this while.
+    llvm::SmallVector<cir::ContinueOp> continues;
+    whileOp->walk([&](mlir::Operation *op) {
+      if (auto continueOp = dyn_cast<ContinueOp>(op))
+        continues.push_back(continueOp);
+    });
+
+    if (continues.empty())
+      return;
+
+    for (auto continueOp : continues) {
+      // When the break is under an IfOp, a direct replacement of `scf.yield`
+      // won't work: the yield would jump out of that IfOp instead. We might
+      // need to change the whileOp itself to achieve the same effect.
+      for (mlir::Operation *parent = continueOp->getParentOp();
+           parent != whileOp; parent = parent->getParentOp()) {
+        if (isa<mlir::scf::IfOp>(parent) || isa<cir::IfOp>(parent))
+          llvm_unreachable("NYI");
+      }
+
+      // Operations after this break has to be removed.
+      for (mlir::Operation *runner = continueOp->getNextNode(); runner;) {
+        mlir::Operation *next = runner->getNextNode();
+        runner->erase();
+        runner = next;
+      }
+
+      // Blocks after this break also has to be removed.
+      for (mlir::Block *block = continueOp->getBlock()->getNextNode(); block;) {
+        mlir::Block *next = block->getNextNode();
+        block->erase();
+        block = next;
+      }
+    }
+  }
+
 public:
   using OpConversionPattern<cir::WhileOp>::OpConversionPattern;
 
   mlir::LogicalResult
   matchAndRewrite(cir::WhileOp op, OpAdaptor adaptor,
                   mlir::ConversionPatternRewriter &rewriter) const override {
     SCFWhileLoop loop(op, adaptor, &rewriter);
-    loop.transferToSCFWhileOp();
+    auto whileOp = loop.transferToSCFWhileOp();
+    rewriteContinue(whileOp, rewriter);
     rewriter.eraseOp(op);
     return mlir::success();
   }

clang/test/CIR/Lowering/ThroughMLIR/while-with-continue.cpp

@@ -0,0 +1,28 @@
+// XFAIL: *
+// RUN: %clang_cc1 -triple x86_64-unknown-linux-gnu -fclangir -fno-clangir-direct-lowering -emit-mlir=core %s -o %t.mlir
+// RUN: FileCheck --input-file=%t.mlir %s
+ 
+void for_with_break() {
+  int i = 0;
+  while (i < 100) {
+    i++;
+    continue;
+    i++;
+  }
+  // Only the first `i++` will be emitted.
+
+  // CHECK: scf.while : () -> () {
+  // CHECK:   %[[TMP0:.+]] = memref.load %alloca[]
+  // CHECK:   %[[HUNDRED:.+]] = arith.constant 100
+  // CHECK:   %[[TMP1:.+]] = arith.cmpi slt, %[[TMP0]], %[[HUNDRED]]
+  // CHECK:   scf.condition(%[[TMP1]])
+  // CHECK: } do {
+  // CHECK:   memref.alloca_scope  {
+  // CHECK:     %[[TMP2:.+]] = memref.load %alloca[]
+  // CHECK:     %[[ONE:.+]] = arith.constant 1
+  // CHECK:     %[[TMP3:.+]] = arith.addi %[[TMP2]], %[[ONE]]
+  // CHECK:     memref.store %[[TMP3]], %alloca[]
+  // CHECK:   }
+  // CHECK:   scf.yield
+  // CHECK: }
+}