[CIR][Lowering] Fixed break/continue lowering for loops

Author: gitoleg

clang/lib/CIR/Lowering/DirectToLLVM/LowerToLLVM.cpp

@@ -241,6 +241,42 @@ class CIRLoopOpLowering : public mlir::OpConversionPattern<mlir::cir::LoopOp> {
     return mlir::success();
   }
 
+  void makeYieldIf(mlir::cir::YieldOpKind kind, mlir::cir::YieldOp &op, mlir::Block *to,
+                   mlir::ConversionPatternRewriter &rewriter) const {
+    if (op.getKind() == kind) {
+      rewriter.setInsertionPoint(op);
+      rewriter.replaceOpWithNewOp<mlir::cir::BrOp>(op, op.getArgs(), to);
+    }
+  }
+
+  void lowerNestedBreakContinue(mlir::Region &loopBody, mlir::Block *exitBlock,
+                                mlir::Block *continueBlock,
+                                mlir::ConversionPatternRewriter &rewriter) const {
+
+    auto processBreak = [&](mlir::Operation *op) {
+      if (isa<mlir::cir::LoopOp, mlir::cir::SwitchOp>(*op)) // don't process breaks in nested loops and switches
+        return mlir::WalkResult::skip();
+
+      if (auto yield = dyn_cast<mlir::cir::YieldOp>(*op))
+          makeYieldIf(mlir::cir::YieldOpKind::Break, yield, exitBlock, rewriter);
+
+      return mlir::WalkResult::advance();
+    };
+
+    auto processContinue = [&](mlir::Operation *op) {
+      if (isa<mlir::cir::LoopOp>(*op)) // don't process continues in nested loops
+        return mlir::WalkResult::skip();
+
+      if (auto yield = dyn_cast<mlir::cir::YieldOp>(*op))
+        makeYieldIf(mlir::cir::YieldOpKind::Continue, yield, continueBlock, rewriter);
+
+      return mlir::WalkResult::advance();
+    };
+
+    loopBody.walk<mlir::WalkOrder::PreOrder>(processBreak);
+    loopBody.walk<mlir::WalkOrder::PreOrder>(processContinue);
+  }
+  
   mlir::LogicalResult
   matchAndRewrite(mlir::cir::LoopOp loopOp, OpAdaptor adaptor,
                   mlir::ConversionPatternRewriter &rewriter) const override {
@@ -268,6 +304,9 @@ class CIRLoopOpLowering : public mlir::OpConversionPattern<mlir::cir::LoopOp> {
     auto &stepFrontBlock = stepRegion.front();
     auto stepYield =
         dyn_cast<mlir::cir::YieldOp>(stepRegion.back().getTerminator());
+    auto &stepBlock = (kind == LoopKind::For ? stepFrontBlock : condFrontBlock);
+
+    lowerNestedBreakContinue(bodyRegion, continueBlock, &stepBlock, rewriter);
 
     // Move loop op region contents to current CFG.
     rewriter.inlineRegionBefore(condRegion, continueBlock);
@@ -290,8 +329,7 @@ class CIRLoopOpLowering : public mlir::OpConversionPattern<mlir::cir::LoopOp> {
 
     // Branch from body to condition or to step on for-loop cases.
     rewriter.setInsertionPoint(bodyYield);
-    auto &bodyExit = (kind == LoopKind::For ? stepFrontBlock : condFrontBlock);
-    rewriter.replaceOpWithNewOp<mlir::cir::BrOp>(bodyYield, &bodyExit);
+    rewriter.replaceOpWithNewOp<mlir::cir::BrOp>(bodyYield, &stepBlock);
 
     // Is a for loop: branch from step to condition.
     if (kind == LoopKind::For) {
@@ -483,6 +521,11 @@ class CIRCastOpLowering : public mlir::OpConversionPattern<mlir::cir::CastOp> {
   }
 };
 
+static bool isLoopYield(mlir::cir::YieldOp &op) {
+  return op.getKind() == mlir::cir::YieldOpKind::Break ||
+         op.getKind() == mlir::cir::YieldOpKind::Continue;
+}
+
 class CIRIfLowering : public mlir::OpConversionPattern<mlir::cir::IfOp> {
 public:
   using mlir::OpConversionPattern<mlir::cir::IfOp>::OpConversionPattern;
@@ -511,7 +554,8 @@ class CIRIfLowering : public mlir::OpConversionPattern<mlir::cir::IfOp> {
     rewriter.setInsertionPointToEnd(thenAfterBody);
     if (auto thenYieldOp =
             dyn_cast<mlir::cir::YieldOp>(thenAfterBody->getTerminator())) {
-      rewriter.replaceOpWithNewOp<mlir::cir::BrOp>(
+      if (!isLoopYield(thenYieldOp)) // lowering of parent loop yields is deferred to loop lowering
+	rewriter.replaceOpWithNewOp<mlir::cir::BrOp>(
           thenYieldOp, thenYieldOp.getArgs(), continueBlock);
     } else if (!dyn_cast<mlir::cir::ReturnOp>(thenAfterBody->getTerminator())) {
       llvm_unreachable("what are we terminating with?");
@@ -540,7 +584,8 @@ class CIRIfLowering : public mlir::OpConversionPattern<mlir::cir::IfOp> {
       rewriter.setInsertionPointToEnd(elseAfterBody);
       if (auto elseYieldOp =
               dyn_cast<mlir::cir::YieldOp>(elseAfterBody->getTerminator())) {
-        rewriter.replaceOpWithNewOp<mlir::cir::BrOp>(
+	if (!isLoopYield(elseYieldOp)) // lowering of parent loop yields is deferred to loop lowering
+	  rewriter.replaceOpWithNewOp<mlir::cir::BrOp>(
             elseYieldOp, elseYieldOp.getArgs(), continueBlock);
       } else if (!dyn_cast<mlir::cir::ReturnOp>(
                      elseAfterBody->getTerminator())) {
@@ -1081,6 +1126,8 @@ class CIRSwitchOpLowering
         case mlir::cir::YieldOpKind::Break:
           rewriteYieldOp(rewriter, yieldOp, exitBlock);
           break;
+	case mlir::cir::YieldOpKind::Continue: // Contniue is handled only in loop lowering
+          break;
         default:
           return op->emitError("invalid yield kind in case statement");
         }