[Mlir-commits] [mlir] [MLIR] Introduce support for early exits (PR #166688)

[Matthias Springer]

================
@@ -524,6 +558,62 @@ RegionBranchOpInterface::getAllRegionBranchPoints() {
         branchPoints.push_back(RegionBranchPoint(terminator));
     }
   }
+
+  Operation *op = getOperation();
+  if (!op->hasTrait<OpTrait::PropagateControlFlowBreak>() &&
+      !isa<HasBreakingControlFlowOpInterface>(op))
+    return branchPoints;
+
+  // The loop above only records terminators that are immediate branch points of
+  // this op: terminators ending blocks directly contained in one of this op's
+  // regions. Breaking control flow adds another class of branch points. A
+  // nested RegionBranchOpInterface that defines PropagateControlFlowBreak can
+  // contain a breaking terminator whose addressed receiver is this op or one of
+  // this op's ancestors. Even though the terminator is not directly contained
+  // in this op's region, it can still create a control-flow edge that leaves or
+  // propagates through this op.
+  //
+  // For example, consider:
+  //
+  //   scf.loop token(%outer) {
+  //     scf.loop token(%inner) {
+  //       scf.if %cond {
+  //         scf.break [%outer]
+  //       }
+  //       scf.continue [%inner]
+  //     }
+  //     scf.continue [%outer]
+  //   }
+  //
+  // When enumerating branch points for the inner loop, its direct body
+  // terminator is only `scf.continue [%inner]`. The `scf.break [%outer]` is
+  // hidden behind the immediately nested `scf.if`, but the inner loop must
+  // still expose a possible edge for that break: the break request propagates
+  // through `scf.if`, then through the inner loop, and is ultimately handled by
+  // the outer loop. Without adding the nested break as a branch point of the
+  // inner loop, generic RegionBranchOpInterface verification and data-flow
+  // consumers only see the continue edge and miss the propagated exit edge.
+  //
+  // This is intentionally broader than collectAllNestedPredecessors(op). That
+  // helper collects only terminators that directly target `op`, which is enough
+  // for a receiver to find incoming breaks but insufficient for an intermediate
+  // PropagateControlFlowBreak op: an escaping terminator targets an ancestor,
+  // not the intermediate op it propagates through.
+  // visitNestedBreakingControlFlowOps reports both nested terminators that
+  // target `op` and those that target an ancestor of `op`, which is exactly the
+  // set that may affect this op's RegionBranchOpInterface edges.
+  visitNestedBreakingControlFlowOps(
+      op, [&](Operation *nestedTerminator, int nestedLevel) {
+        // Immediate region terminators are already covered above. Deeper
+        // terminators may add a propagated control-flow edge through this op.
+        if (nestedLevel <= 1)
+          return;
+        auto terminator =
+            dyn_cast<RegionBranchTerminatorOpInterface>(nestedTerminator);
----------------
matthias-springer wrote:

Does a `RegionExitTerminatorOpInterface` have to implement the `RegionBranchTerminatorOpInterface` when it wants to propagate through a `RegionBranchOpInterface` op? In other words, should this be a `cast<RegionBranchTerminatorOpInterface>`? If not, we may fail to enumerate some propagating CF edges.

(Maybe this is something that the `RegionBranchOpInterface` verifier could check.)


https://github.com/llvm/llvm-project/pull/166688