[CIR][ThroughMLIR] Lower For to While when it contains break/continue (#1716)

We lower `cir::ForOp` into `cir::WhileOp` (rather than `scf::WhileOp`)
when it contains break and continue. This is to reuse the rewriting
functions already implemented for while loops.

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

Author: terapines-osc-cir

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

@@ -14,7 +14,6 @@
 #include "mlir/Dialect/MemRef/IR/MemRef.h"
 #include "mlir/Dialect/SCF/IR/SCF.h"
 #include "mlir/IR/Builders.h"
-#include "mlir/IR/BuiltinOps.h"
 #include "mlir/IR/Location.h"
 #include "mlir/IR/ValueRange.h"
 #include "mlir/Pass/PassManager.h"
@@ -40,6 +39,7 @@ class SCFLoop {
   mlir::Value getLowerBound() { return lowerBound; }
   mlir::Value getUpperBound() { return upperBound; }
   bool isCanonical() { return canonical; }
+  bool hasBreakOrContinue() { return hasBreakContinue; }
 
   // Returns true if successfully finds both step and induction variable.
   mlir::LogicalResult findStepAndIV();
@@ -50,13 +50,15 @@ class SCFLoop {
   mlir::Value plusConstant(mlir::Value v, mlir::Location loc, int addend);
   void transferToSCFForOp();
   void transformToSCFWhileOp();
+  void transformToCIRWhileOp(); // TODO
 
 private:
   cir::ForOp forOp;
   cir::CmpOp cmpOp;
   mlir::Value ivAddr, lowerBound = nullptr, upperBound = nullptr;
   mlir::ConversionPatternRewriter *rewriter;
   int64_t step = 0;
+  bool hasBreakContinue = false;
   bool canonical = true;
 };
 
@@ -251,6 +253,16 @@ mlir::Value SCFLoop::findIVInitValue() {
 }
 
 void SCFLoop::analysis() {
+  // Check whether this ForOp contains break or continue.
+  forOp.walk([&](mlir::Operation *op) {
+    if (isa<BreakOp, ContinueOp>(op))
+      hasBreakContinue = true;
+  });
+  if (hasBreakContinue) {
+    canonical = false;
+    return;
+  }
+
   canonical = mlir::succeeded(findStepAndIV());
   if (!canonical)
     return;
@@ -356,6 +368,30 @@ void SCFLoop::transformToSCFWhileOp() {
                               scfWhileOp.getAfterBody()->end());
 }
 
+void SCFLoop::transformToCIRWhileOp() {
+  auto cirWhileOp = rewriter->create<cir::WhileOp>(
+      forOp->getLoc(), forOp->getResultTypes(), mlir::ValueRange());
+  rewriter->createBlock(&cirWhileOp.getCond());
+  rewriter->createBlock(&cirWhileOp.getBody());
+
+  mlir::Block &condFront = cirWhileOp.getCond().front();
+  rewriter->inlineBlockBefore(&forOp.getCond().front(), &condFront,
+                              condFront.end());
+
+  mlir::Block &bodyFront = cirWhileOp.getBody().front();
+  rewriter->inlineBlockBefore(&forOp.getBody().front(), &bodyFront,
+                              bodyFront.end());
+
+  // The last operation of `bodyFront` must be a terminator.
+  // We need to place the step just before it.
+  rewriter->inlineBlockBefore(&forOp.getStep().front(), &bodyFront,
+                              --bodyFront.end());
+  // This also introduces another terminator before the one of the body.
+  // We need to erase it.
+  auto &stepTerminator = *--bodyFront.end();
+  rewriter->eraseOp(&stepTerminator);
+}
+
 mlir::scf::WhileOp SCFWhileLoop::transferToSCFWhileOp() {
   auto scfWhileOp = rewriter->create<mlir::scf::WhileOp>(
       whileOp->getLoc(), whileOp->getResultTypes(), adaptor.getOperands());
@@ -400,6 +436,14 @@ class CIRForOpLowering : public mlir::OpConversionPattern<cir::ForOp> {
                   mlir::ConversionPatternRewriter &rewriter) const override {
     SCFLoop loop(op, &rewriter);
     loop.analysis();
+    // Breaks and continues are handled in lowering of cir::WhileOp.
+    // We can reuse the code by transforming this ForOp into WhileOp.
+    if (loop.hasBreakOrContinue()) {
+      loop.transformToCIRWhileOp();
+      rewriter.eraseOp(op);
+      return mlir::success();
+    }
+
     if (!loop.isCanonical()) {
       loop.transformToSCFWhileOp();
       rewriter.eraseOp(op);

clang/test/CIR/Lowering/ThroughMLIR/for_with_continue.cpp

@@ -0,0 +1,19 @@
+// 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_continue() {
+  for (int i = 0; i < 100; i++)
+    continue;
+
+  // CHECK: scf.while : () -> () {
+  // CHECK:   %[[IV:.+]] = memref.load %alloca[]
+  // CHECK:   %[[CMP:.+]] = arith.cmpi slt, %[[IV]], %c100_i32
+  // CHECK:   scf.condition(%[[CMP]])
+  // CHECK: } do {
+  // CHECK:   %[[IV2:.+]] = memref.load %alloca[]
+  // CHECK:   %[[ONE:.+]] = arith.constant 1
+  // CHECK:   %[[CMP2:.+]] = arith.addi %[[IV2]], %[[ONE]]
+  // CHECK:   memref.store %[[CMP2]], %alloca[]
+  // CHECK:   scf.yield
+  // CHECK: }
+}