[CIR][IR] Implement cir.continue operation

Detaches the representation of the C/C++ `continue` statement into a
separate operation. This simplifies mostly lowering and verifications
related to `continue` statements, as well as the definition and lowering
of the `cir.yield` operation.

A few checks regarding region terminators were also removed from the
lowering stage, since they are already enforced by MLIR.

ghstack-source-id: 1810a48ada88fe7ef5638b0758a2298d9cfbdb8b
Pull Request resolved: llvm/clangir#394

Author: sitio-couto

clang/include/clang/CIR/Dialect/IR/CIROps.td

@@ -607,13 +607,12 @@ def ConditionOp : CIR_Op<"condition", [
 
 def YieldOpKind_BK : I32EnumAttrCase<"Break", 1, "break">;
 def YieldOpKind_FT : I32EnumAttrCase<"Fallthrough", 2, "fallthrough">;
-def YieldOpKind_CE : I32EnumAttrCase<"Continue", 3, "continue">;
 def YieldOpKind_NS : I32EnumAttrCase<"NoSuspend", 4, "nosuspend">;
 
 def YieldOpKind : I32EnumAttr<
     "YieldOpKind",
     "yield kind",
-    [YieldOpKind_BK, YieldOpKind_FT, YieldOpKind_CE, YieldOpKind_NS]> {
+    [YieldOpKind_BK, YieldOpKind_FT, YieldOpKind_NS]> {
   let cppNamespace = "::mlir::cir";
 }
 
@@ -634,8 +633,6 @@ def YieldOp : CIR_Op<"yield", [ReturnLike, Terminator,
     cannot be used if not dominated by these parent operations.
     - `fallthrough`: execution falls to the next region in `cir.switch` case list.
     Only available inside `cir.switch` regions.
-    - `continue`: only allowed under `cir.loop`, continue execution to the next
-    loop step.
     - `nosuspend`: specific to the `ready` region inside `cir.await` op, it makes
     control-flow to be transfered back to the parent, preventing suspension.
 
@@ -657,11 +654,6 @@ def YieldOp : CIR_Op<"yield", [ReturnLike, Terminator,
       }, ...
     ]
 
-    cir.loop (cond : {...}, step : {...}) {
-      ...
-      cir.yield continue
-    }
-
     cir.await(init, ready : {
       // Call std::suspend_always::await_ready
       %18 = cir.call @_ZNSt14suspend_always11await_readyEv(...)
@@ -718,9 +710,6 @@ def YieldOp : CIR_Op<"yield", [ReturnLike, Terminator,
     bool isBreak() {
       return !isPlain() && *getKind() == YieldOpKind::Break;
     }
-    bool isContinue() {
-      return !isPlain() && *getKind() == YieldOpKind::Continue;
-    }
     bool isNoSuspend() {
       return !isPlain() && *getKind() == YieldOpKind::NoSuspend;
     }
@@ -729,6 +718,20 @@ def YieldOp : CIR_Op<"yield", [ReturnLike, Terminator,
   let hasVerifier = 1;
 }
 
+//===----------------------------------------------------------------------===//
+// ContinueOp
+//===----------------------------------------------------------------------===//
+
+def ContinueOp : CIR_Op<"continue", [Terminator]> {
+  let summary = "C/C++ `continue` statement equivalent";
+  let description = [{
+    The `cir.continue` operation is used to continue execution to the next
+    iteration of a loop. It is only allowed within `cir.loop` regions.
+  }];
+  let assemblyFormat = "attr-dict";
+  let hasVerifier = 1;
+}
+
 //===----------------------------------------------------------------------===//
 // ScopeOp
 //===----------------------------------------------------------------------===//
@@ -1166,7 +1169,7 @@ def LoopOp : CIR_Op<"loop",
     `cir.loop` represents C/C++ loop forms. It defines 3 blocks:
     - `cond`: region can contain multiple blocks, terminated by regular
     `cir.yield` when control should yield back to the parent, and
-    `cir.yield continue` when execution continues to another region.
+    `cir.continue` when execution continues to the next region.
     The region destination depends on the loop form specified.
     - `step`: region with one block, containing code to compute the
     loop step, must be terminated with `cir.yield`.
@@ -1181,7 +1184,7 @@ def LoopOp : CIR_Op<"loop",
       //  i = i + 1;
       // }
       cir.loop while(cond :  {
-        cir.yield continue
+        cir.continue
       }, step :  {
         cir.yield
       })  {

clang/lib/CIR/CodeGen/CIRGenBuilder.h

@@ -585,6 +585,11 @@ class CIRGenBuilderTy : public CIRBaseBuilderTy {
     return create<mlir::cir::ConditionOp>(condition.getLoc(), condition);
   }
 
+  /// Create a continue operation.
+  mlir::cir::ContinueOp createContinue(mlir::Location loc) {
+    return create<mlir::cir::ContinueOp>(loc);
+  }
+
   mlir::cir::MemCpyOp createMemCpy(mlir::Location loc, mlir::Value dst,
                                    mlir::Value src, mlir::Value len) {
     return create<mlir::cir::MemCpyOp>(loc, dst, src, len);

clang/lib/CIR/CodeGen/CIRGenStmt.cpp

@@ -301,7 +301,6 @@ mlir::LogicalResult CIRGenFunction::buildSimpleStmt(const Stmt *S,
     return buildGotoStmt(cast<GotoStmt>(*S));
   case Stmt::ContinueStmtClass:
     return buildContinueStmt(cast<ContinueStmt>(*S));
-
   case Stmt::NullStmtClass:
     break;
 
@@ -572,11 +571,7 @@ mlir::LogicalResult CIRGenFunction::buildLabel(const LabelDecl *D) {
 
 mlir::LogicalResult
 CIRGenFunction::buildContinueStmt(const clang::ContinueStmt &S) {
-  builder.create<YieldOp>(
-      getLoc(S.getContinueLoc()),
-      mlir::cir::YieldOpKindAttr::get(builder.getContext(),
-                                      mlir::cir::YieldOpKind::Continue),
-      mlir::ValueRange({}));
+  builder.createContinue(getLoc(S.getContinueLoc()));
   return mlir::success();
 }
 

clang/lib/CIR/Dialect/IR/CIRDialect.cpp

@@ -341,6 +341,16 @@ static void printConstantValue(OpAsmPrinter &p, cir::ConstantOp op,
 
 OpFoldResult ConstantOp::fold(FoldAdaptor /*adaptor*/) { return getValue(); }
 
+//===----------------------------------------------------------------------===//
+// ContinueOp
+//===----------------------------------------------------------------------===//
+
+LogicalResult ContinueOp::verify() {
+  if (!this->getOperation()->getParentOfType<LoopOp>())
+    return emitOpError("must be within a loop");
+  return success();
+}
+
 //===----------------------------------------------------------------------===//
 // CastOp
 //===----------------------------------------------------------------------===//
@@ -802,15 +812,6 @@ mlir::LogicalResult YieldOp::verify() {
     return false;
   };
 
-  auto isDominatedByLoop = [](Operation *parentOp) {
-    while (!llvm::isa<cir::FuncOp>(parentOp)) {
-      if (llvm::isa<cir::LoopOp>(parentOp))
-        return true;
-      parentOp = parentOp->getParentOp();
-    }
-    return false;
-  };
-
   if (isNoSuspend()) {
     if (!isDominatedByProperAwaitRegion(getOperation()->getParentOp(),
                                         getOperation()->getParentRegion()))
@@ -824,12 +825,6 @@ mlir::LogicalResult YieldOp::verify() {
     return mlir::success();
   }
 
-  if (isContinue()) {
-    if (!isDominatedByLoop(getOperation()->getParentOp()))
-      return emitOpError() << "shall be dominated by 'cir.loop'";
-    return mlir::success();
-  }
-
   if (isFallthrough()) {
     if (!llvm::isa<SwitchOp>(getOperation()->getParentOp()))
       return emitOpError() << "fallthrough only expected within 'cir.switch'";

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

@@ -29,10 +29,12 @@
 #include "mlir/IR/BuiltinDialect.h"
 #include "mlir/IR/BuiltinOps.h"
 #include "mlir/IR/BuiltinTypes.h"
+#include "mlir/IR/OpDefinition.h"
 #include "mlir/IR/Operation.h"
 #include "mlir/IR/Types.h"
 #include "mlir/IR/Value.h"
 #include "mlir/IR/ValueRange.h"
+#include "mlir/IR/Visitors.h"
 #include "mlir/Interfaces/DataLayoutInterfaces.h"
 #include "mlir/Pass/Pass.h"
 #include "mlir/Pass/PassManager.h"
@@ -65,6 +67,36 @@ using namespace llvm;
 namespace cir {
 namespace direct {
 
+//===----------------------------------------------------------------------===//
+// Helper Methods
+//===----------------------------------------------------------------------===//
+
+namespace {
+
+/// Lowers operations with the terminator trait that have a single successor.
+void lowerTerminator(mlir::Operation *op, mlir::Block *dest,
+                     mlir::ConversionPatternRewriter &rewriter) {
+  assert(op->hasTrait<mlir::OpTrait::IsTerminator>() && "not a terminator");
+  mlir::OpBuilder::InsertionGuard guard(rewriter);
+  rewriter.setInsertionPoint(op);
+  rewriter.replaceOpWithNewOp<mlir::cir::BrOp>(op, dest);
+}
+
+/// Walks a region while skipping operations of type `Ops`. This ensures the
+/// callback is not applied to said operations and its children.
+template <typename... Ops>
+void walkRegionSkipping(mlir::Region &region,
+                        mlir::function_ref<void(mlir::Operation *)> callback) {
+  region.walk<mlir::WalkOrder::PreOrder>([&](mlir::Operation *op) {
+    if (isa<Ops...>(op))
+      return mlir::WalkResult::skip();
+    callback(op);
+    return mlir::WalkResult::advance();
+  });
+}
+
+} // namespace
+
 //===----------------------------------------------------------------------===//
 // Visitors for Lowering CIR Const Attributes
 //===----------------------------------------------------------------------===//
@@ -447,8 +479,15 @@ class CIRLoopOpLowering : public mlir::OpConversionPattern<mlir::cir::LoopOp> {
 
     lowerNestedYield(mlir::cir::YieldOpKind::Break, rewriter, bodyRegion,
                      continueBlock);
-    lowerNestedYield(mlir::cir::YieldOpKind::Continue, rewriter, bodyRegion,
-                     &stepBlock);
+
+    // Lower continue statements.
+    mlir::Block &dest =
+        (kind != LoopKind::For ? condFrontBlock : stepFrontBlock);
+    walkRegionSkipping<mlir::cir::LoopOp>(
+        loopOp.getBody(), [&](mlir::Operation *op) {
+          if (isa<mlir::cir::ContinueOp>(op))
+            lowerTerminator(op, &dest, rewriter);
+        });
 
     // Move loop op region contents to current CFG.
     rewriter.inlineRegionBefore(condRegion, continueBlock);
@@ -678,9 +717,8 @@ class CIRCastOpLowering : public mlir::OpConversionPattern<mlir::cir::CastOp> {
   }
 };
 
-static bool isBreakOrContinue(mlir::cir::YieldOp &op) {
-  return op.getKind() == mlir::cir::YieldOpKind::Break ||
-         op.getKind() == mlir::cir::YieldOpKind::Continue;
+static bool isBreak(mlir::cir::YieldOp &op) {
+  return op.getKind() == mlir::cir::YieldOpKind::Break;
 }
 
 class CIRIfLowering : public mlir::OpConversionPattern<mlir::cir::IfOp> {
@@ -711,12 +749,10 @@ class CIRIfLowering : public mlir::OpConversionPattern<mlir::cir::IfOp> {
     rewriter.setInsertionPointToEnd(thenAfterBody);
     if (auto thenYieldOp =
             mlir::dyn_cast<mlir::cir::YieldOp>(thenAfterBody->getTerminator())) {
-      if (!isBreakOrContinue(thenYieldOp)) // lowering of parent loop yields is
-                                           // deferred to loop lowering
+      if (!isBreak(thenYieldOp)) // lowering of parent loop yields is
+                                 // deferred to loop lowering
         rewriter.replaceOpWithNewOp<mlir::cir::BrOp>(
             thenYieldOp, thenYieldOp.getArgs(), continueBlock);
-    } else if (!mlir::dyn_cast<mlir::cir::ReturnOp>(thenAfterBody->getTerminator())) {
-      llvm_unreachable("what are we terminating with?");
     }
 
     rewriter.setInsertionPointToEnd(continueBlock);
@@ -742,12 +778,10 @@ class CIRIfLowering : public mlir::OpConversionPattern<mlir::cir::IfOp> {
       rewriter.setInsertionPointToEnd(elseAfterBody);
       if (auto elseYieldOp =
               mlir::dyn_cast<mlir::cir::YieldOp>(elseAfterBody->getTerminator())) {
-        if (!isBreakOrContinue(elseYieldOp)) // lowering of parent loop yields
-                                             // is deferred to loop lowering
+        if (!isBreak(elseYieldOp)) // lowering of parent loop yields
+                                   // is deferred to loop lowering
           rewriter.replaceOpWithNewOp<mlir::cir::BrOp>(
               elseYieldOp, elseYieldOp.getArgs(), continueBlock);
-      } else if (!mlir::dyn_cast<mlir::cir::ReturnOp>(elseAfterBody->getTerminator())) {
-        llvm_unreachable("what are we terminating with?");
       }
     }
 
@@ -803,7 +837,7 @@ class CIRScopeOpLowering
     rewriter.setInsertionPointToEnd(afterBody);
     auto yieldOp = mlir::dyn_cast<mlir::cir::YieldOp>(afterBody->getTerminator());
 
-    if (yieldOp && !isBreakOrContinue(yieldOp)) {
+    if (yieldOp && !isBreak(yieldOp)) {
       auto branchOp = rewriter.replaceOpWithNewOp<mlir::cir::BrOp>(
           yieldOp, yieldOp.getArgs(), continueBlock);
 
@@ -1404,9 +1438,6 @@ class CIRSwitchOpLowering
           case mlir::cir::YieldOpKind::Break:
             rewriteYieldOp(rewriter, yieldOp, exitBlock);
             break;
-          case mlir::cir::YieldOpKind::Continue: // Continue is handled only in
-                                                 // loop lowering
-            break;
           default:
             return op->emitError("invalid yield kind in case statement");
           }

clang/test/CIR/CodeGen/loop.cpp

@@ -180,7 +180,7 @@ void l4() {
 // CHECK-NEXT:     %11 = cir.const(#cir.int<10> : !s32i) : !s32i
 // CHECK-NEXT:     %12 = cir.cmp(lt, %10, %11) : !s32i, !cir.bool
 // CHECK-NEXT:     cir.if %12 {
-// CHECK-NEXT:       cir.yield continue
+// CHECK-NEXT:       cir.continue
 // CHECK-NEXT:     }
 // CHECK-NEXT:   }
 

clang/test/CIR/IR/invalid.cir

@@ -80,7 +80,7 @@ cir.func @yieldbreak() {
 cir.func @yieldcontinue() {
   %0 = cir.const(#true) : !cir.bool
   cir.if %0 {
-    cir.yield continue // expected-error {{shall be dominated by 'cir.loop'}}
+    cir.continue // expected-error {{op must be within a loop}}
   }
   cir.return
 }

clang/test/CIR/IR/loop.cir

@@ -52,7 +52,7 @@ cir.func @l0() {
       cir.store %6, %0 : !u32i, cir.ptr <!u32i>
       %7 = cir.const(#true) : !cir.bool
       cir.if %7 {
-        cir.yield continue
+        cir.continue
       }
       cir.yield
     }
@@ -118,7 +118,7 @@ cir.func @l0() {
 // CHECK-NEXT:     cir.store %6, %0 : !u32i, cir.ptr <!u32i>
 // CHECK-NEXT:     %7 = cir.const(#true) : !cir.bool
 // CHECK-NEXT:     cir.if %7 {
-// CHECK-NEXT:       cir.yield continue
+// CHECK-NEXT:       cir.continue
 // CHECK-NEXT:     }
 // CHECK-NEXT:     cir.yield
 // CHECK-NEXT:   }

clang/test/CIR/Lowering/loops-with-continue.cir

@@ -27,7 +27,7 @@ module {
             %4 = cir.cmp(eq, %2, %3) : !s32i, !s32i
             %5 = cir.cast(int_to_bool, %4 : !s32i), !cir.bool
             cir.if %5 {
-              cir.yield continue
+              cir.continue
             }
           }
         }
@@ -107,7 +107,7 @@ module {
                   %6 = cir.cmp(eq, %4, %5) : !s32i, !s32i
                   %7 = cir.cast(int_to_bool, %6 : !s32i), !cir.bool
                   cir.if %7 {
-                    cir.yield continue
+                    cir.continue
                   }
                 }
               }
@@ -189,7 +189,7 @@ cir.func @testWhile() {
           %6 = cir.cmp(eq, %4, %5) : !s32i, !s32i
           %7 = cir.cast(int_to_bool, %6 : !s32i), !cir.bool
           cir.if %7 {
-            cir.yield continue
+            cir.continue
           }
         }
         cir.yield
@@ -243,7 +243,7 @@ cir.func @testWhile() {
           %6 = cir.cmp(eq, %4, %5) : !s32i, !s32i
           %7 = cir.cast(int_to_bool, %6 : !s32i), !cir.bool
           cir.if %7 {
-            cir.yield continue
+            cir.continue
           }
         }
         cir.yield