[CIR][IR] Implement cir.break operation

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

@@ -605,14 +605,13 @@ def ConditionOp : CIR_Op<"condition", [
 // YieldOp
 //===----------------------------------------------------------------------===//
 
-def YieldOpKind_BK : I32EnumAttrCase<"Break", 1, "break">;
 def YieldOpKind_FT : I32EnumAttrCase<"Fallthrough", 2, "fallthrough">;
 def YieldOpKind_NS : I32EnumAttrCase<"NoSuspend", 4, "nosuspend">;
 
 def YieldOpKind : I32EnumAttr<
     "YieldOpKind",
     "yield kind",
-    [YieldOpKind_BK, YieldOpKind_FT, YieldOpKind_NS]> {
+    [YieldOpKind_FT, YieldOpKind_NS]> {
   let cppNamespace = "::mlir::cir";
 }
 
@@ -629,8 +628,6 @@ def YieldOp : CIR_Op<"yield", [ReturnLike, Terminator,
     defined by the parent operation.
 
     Optionally, `cir.yield` can be annotated with extra kind specifiers:
-    - `break`: breaking out of the innermost `cir.switch` / `cir.loop` semantics,
-    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.
     - `nosuspend`: specific to the `ready` region inside `cir.await` op, it makes
@@ -707,9 +704,6 @@ def YieldOp : CIR_Op<"yield", [ReturnLike, Terminator,
     bool isFallthrough() {
       return !isPlain() && *getKind() == YieldOpKind::Fallthrough;
     }
-    bool isBreak() {
-      return !isPlain() && *getKind() == YieldOpKind::Break;
-    }
     bool isNoSuspend() {
       return !isPlain() && *getKind() == YieldOpKind::NoSuspend;
     }
@@ -718,6 +712,21 @@ def YieldOp : CIR_Op<"yield", [ReturnLike, Terminator,
   let hasVerifier = 1;
 }
 
+//===----------------------------------------------------------------------===//
+// BreakOp
+//===----------------------------------------------------------------------===//
+
+def BreakOp : CIR_Op<"break", [Terminator]> {
+  let summary = "C/C++ `break` statement equivalent";
+  let description = [{
+    The `cir.break` operation is used to cease the control flow to the parent
+    operation, exiting its region's control flow. It is only allowed if it is
+    within a breakable operation (loops and `switch`).
+  }];
+  let assemblyFormat = "attr-dict";
+  let hasVerifier = 1;
+}
+
 //===----------------------------------------------------------------------===//
 // ContinueOp
 //===----------------------------------------------------------------------===//

clang/lib/CIR/CodeGen/CIRGenBuilder.h

@@ -583,6 +583,11 @@ class CIRGenBuilderTy : public CIRBaseBuilderTy {
     return create<mlir::cir::ConditionOp>(condition.getLoc(), condition);
   }
 
+  /// Create a break operation.
+  mlir::cir::BreakOp createBreak(mlir::Location loc) {
+    return create<mlir::cir::BreakOp>(loc);
+  }
+
   /// Create a continue operation.
   mlir::cir::ContinueOp createContinue(mlir::Location loc) {
     return create<mlir::cir::ContinueOp>(loc);

clang/lib/CIR/CodeGen/CIRGenStmt.cpp

@@ -574,11 +574,7 @@ CIRGenFunction::buildContinueStmt(const clang::ContinueStmt &S) {
 }
 
 mlir::LogicalResult CIRGenFunction::buildBreakStmt(const clang::BreakStmt &S) {
-  builder.create<YieldOp>(
-      getLoc(S.getBreakLoc()),
-      mlir::cir::YieldOpKindAttr::get(builder.getContext(),
-                                      mlir::cir::YieldOpKind::Break),
-      mlir::ValueRange({}));
+  builder.createBreak(getLoc(S.getBreakLoc()));
   return mlir::success();
 }
 

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

@@ -226,6 +226,17 @@ void AllocaOp::build(::mlir::OpBuilder &odsBuilder,
   odsState.addTypes(addr);
 }
 
+//===----------------------------------------------------------------------===//
+// BreakOp
+//===----------------------------------------------------------------------===//
+
+LogicalResult BreakOp::verify() {
+  if (!getOperation()->getParentOfType<LoopOp>() &&
+      !getOperation()->getParentOfType<SwitchOp>())
+    return emitOpError("must be within a loop or switch");
+  return success();
+}
+
 //===----------------------------------------------------------------------===//
 // ConditionOp
 //===-----------------------------------------------------------------------===//
@@ -775,17 +786,6 @@ void TernaryOp::build(OpBuilder &builder, OperationState &result, Value cond,
 //===----------------------------------------------------------------------===//
 
 mlir::LogicalResult YieldOp::verify() {
-  auto isDominatedByLoopOrSwitch = [&](Operation *parentOp) {
-    while (!llvm::isa<cir::FuncOp>(parentOp)) {
-      if (llvm::isa<cir::SwitchOp, cir::LoopOp>(parentOp))
-        return true;
-      parentOp = parentOp->getParentOp();
-    }
-
-    emitOpError() << "shall be dominated by 'cir.loop' or 'cir.switch'";
-    return false;
-  };
-
   auto isDominatedByProperAwaitRegion = [&](Operation *parentOp,
                                             mlir::Region *currRegion) {
     while (!llvm::isa<cir::FuncOp>(parentOp)) {
@@ -814,12 +814,6 @@ mlir::LogicalResult YieldOp::verify() {
     return mlir::success();
   }
 
-  if (isBreak()) {
-    if (!isDominatedByLoopOrSwitch(getOperation()->getParentOp()))
-      return mlir::failure();
-    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

@@ -356,32 +356,6 @@ mlir::LLVM::Linkage convertLinkage(mlir::cir::GlobalLinkageKind linkage) {
   };
 }
 
-static void lowerNestedYield(mlir::cir::YieldOpKind targetKind,
-                             mlir::ConversionPatternRewriter &rewriter,
-                             mlir::Region &body, mlir::Block *dst) {
-  // top-level yields are lowered in matchAndRewrite of the parent operations
-  auto isNested = [&](mlir::Operation *op) {
-    return op->getParentRegion() != &body;
-  };
-
-  body.walk<mlir::WalkOrder::PreOrder>([&](mlir::Operation *op) {
-    if (!isNested(op))
-      return mlir::WalkResult::advance();
-
-    // don't process breaks/continues in nested loops and switches
-    if (isa<mlir::cir::LoopOp, mlir::cir::SwitchOp>(*op))
-      return mlir::WalkResult::skip();
-
-    auto yield = dyn_cast<mlir::cir::YieldOp>(*op);
-    if (yield && yield.getKind() == targetKind) {
-      rewriter.setInsertionPoint(op);
-      rewriter.replaceOpWithNewOp<mlir::cir::BrOp>(op, yield.getArgs(), dst);
-    }
-
-    return mlir::WalkResult::advance();
-  });
-}
-
 class CIRCopyOpLowering : public mlir::OpConversionPattern<mlir::cir::CopyOp> {
 public:
   using mlir::OpConversionPattern<mlir::cir::CopyOp>::OpConversionPattern;
@@ -462,7 +436,6 @@ class CIRLoopOpLowering : public mlir::OpConversionPattern<mlir::cir::LoopOp> {
     auto &bodyFrontBlock = bodyRegion.front();
     auto bodyYield =
         dyn_cast<mlir::cir::YieldOp>(bodyRegion.back().getTerminator());
-    assert(bodyYield && "unstructured while loops are NYI");
 
     // Fetch required info from the step region.
     auto &stepRegion = loopOp.getStep();
@@ -471,9 +444,6 @@ class CIRLoopOpLowering : public mlir::OpConversionPattern<mlir::cir::LoopOp> {
         dyn_cast<mlir::cir::YieldOp>(stepRegion.back().getTerminator());
     auto &stepBlock = (kind == LoopKind::For ? stepFrontBlock : condFrontBlock);
 
-    lowerNestedYield(mlir::cir::YieldOpKind::Break, rewriter, bodyRegion,
-                     continueBlock);
-
     // Lower continue statements.
     mlir::Block &dest =
         (kind != LoopKind::For ? condFrontBlock : stepFrontBlock);
@@ -483,6 +453,13 @@ class CIRLoopOpLowering : public mlir::OpConversionPattern<mlir::cir::LoopOp> {
             lowerTerminator(op, &dest, rewriter);
         });
 
+    // Lower break statements.
+    walkRegionSkipping<mlir::cir::LoopOp, mlir::cir::SwitchOp>(
+        loopOp.getBody(), [&](mlir::Operation *op) {
+          if (isa<mlir::cir::BreakOp>(op))
+            lowerTerminator(op, continueBlock, rewriter);
+        });
+
     // Move loop op region contents to current CFG.
     rewriter.inlineRegionBefore(condRegion, continueBlock);
     rewriter.inlineRegionBefore(bodyRegion, continueBlock);
@@ -500,11 +477,10 @@ class CIRLoopOpLowering : public mlir::OpConversionPattern<mlir::cir::LoopOp> {
     lowerConditionOp(conditionOp, &bodyFrontBlock, continueBlock, rewriter);
 
     // Branch from body to condition or to step on for-loop cases.
-    rewriter.setInsertionPoint(bodyYield);
-    auto bodyYieldDest = bodyYield.getKind() == mlir::cir::YieldOpKind::Break
-                             ? continueBlock
-                             : &stepBlock;
-    rewriter.replaceOpWithNewOp<mlir::cir::BrOp>(bodyYield, bodyYieldDest);
+    if (bodyYield) {
+      rewriter.setInsertionPoint(bodyYield);
+      rewriter.replaceOpWithNewOp<mlir::cir::BrOp>(bodyYield, &stepBlock);
+    }
 
     // Is a for loop: branch from step to condition.
     if (kind == LoopKind::For) {
@@ -711,10 +687,6 @@ class CIRCastOpLowering : public mlir::OpConversionPattern<mlir::cir::CastOp> {
   }
 };
 
-static bool isBreak(mlir::cir::YieldOp &op) {
-  return op.getKind() == mlir::cir::YieldOpKind::Break;
-}
-
 class CIRIfLowering : public mlir::OpConversionPattern<mlir::cir::IfOp> {
 public:
   using mlir::OpConversionPattern<mlir::cir::IfOp>::OpConversionPattern;
@@ -743,10 +715,8 @@ class CIRIfLowering : public mlir::OpConversionPattern<mlir::cir::IfOp> {
     rewriter.setInsertionPointToEnd(thenAfterBody);
     if (auto thenYieldOp =
             dyn_cast<mlir::cir::YieldOp>(thenAfterBody->getTerminator())) {
-      if (!isBreak(thenYieldOp)) // lowering of parent loop yields is
-                                 // deferred to loop lowering
-        rewriter.replaceOpWithNewOp<mlir::cir::BrOp>(
-            thenYieldOp, thenYieldOp.getArgs(), continueBlock);
+      rewriter.replaceOpWithNewOp<mlir::cir::BrOp>(
+          thenYieldOp, thenYieldOp.getArgs(), continueBlock);
     }
 
     rewriter.setInsertionPointToEnd(continueBlock);
@@ -772,10 +742,8 @@ class CIRIfLowering : public mlir::OpConversionPattern<mlir::cir::IfOp> {
       rewriter.setInsertionPointToEnd(elseAfterBody);
       if (auto elseYieldOp =
               dyn_cast<mlir::cir::YieldOp>(elseAfterBody->getTerminator())) {
-        if (!isBreak(elseYieldOp)) // lowering of parent loop yields
-                                   // is deferred to loop lowering
-          rewriter.replaceOpWithNewOp<mlir::cir::BrOp>(
-              elseYieldOp, elseYieldOp.getArgs(), continueBlock);
+        rewriter.replaceOpWithNewOp<mlir::cir::BrOp>(
+            elseYieldOp, elseYieldOp.getArgs(), continueBlock);
       }
     }
 
@@ -829,18 +797,13 @@ class CIRScopeOpLowering
     // Replace the scopeop return with a branch that jumps out of the body.
     // Stack restore before leaving the body region.
     rewriter.setInsertionPointToEnd(afterBody);
-    auto yieldOp = dyn_cast<mlir::cir::YieldOp>(afterBody->getTerminator());
-
-    if (yieldOp && !isBreak(yieldOp)) {
-      auto branchOp = rewriter.replaceOpWithNewOp<mlir::cir::BrOp>(
-          yieldOp, yieldOp.getArgs(), continueBlock);
-
-      // // Insert stack restore before jumping out of the body of the region.
-      rewriter.setInsertionPoint(branchOp);
+    if (auto yieldOp =
+            dyn_cast<mlir::cir::YieldOp>(afterBody->getTerminator())) {
+      rewriter.replaceOpWithNewOp<mlir::cir::BrOp>(yieldOp, yieldOp.getArgs(),
+                                                   continueBlock);
     }
 
-    // TODO(CIR): stackrestore?
-    // rewriter.create<mlir::LLVM::StackRestoreOp>(loc, stackSaveOp);
+    // TODO(cir): stackrestore?
 
     // Replace the op with values return from the body region.
     rewriter.replaceOp(scopeOp, continueBlock->getArguments());
@@ -1423,24 +1386,23 @@ class CIRSwitchOpLowering
           // TODO(cir): Ensure every yield instead of dealing with optional
           // values.
           assert(yieldOp.getKind().has_value() && "switch yield has no kind");
-
           switch (yieldOp.getKind().value()) {
           // Fallthrough to next case: track it for the next case to handle.
           case mlir::cir::YieldOpKind::Fallthrough:
             fallthroughYieldOp = yieldOp;
             break;
-          // Break out of switch: branch to exit block.
-          case mlir::cir::YieldOpKind::Break:
-            rewriteYieldOp(rewriter, yieldOp, exitBlock);
-            break;
           default:
             return op->emitError("invalid yield kind in case statement");
           }
         }
       }
 
-      lowerNestedYield(mlir::cir::YieldOpKind::Break, rewriter, region,
-                       exitBlock);
+      // Handle break statements.
+      walkRegionSkipping<mlir::cir::LoopOp, mlir::cir::SwitchOp>(
+          region, [&](mlir::Operation *op) {
+            if (isa<mlir::cir::BreakOp>(op))
+              lowerTerminator(op, exitBlock, rewriter);
+          });
 
       // Extract region contents before erasing the switch op.
       rewriter.inlineRegionBefore(region, exitBlock);