[Mlir-commits] [mlir] [MLIR] Introduce support for early exits (PR #166688)
Mehdi Amini
llvmlistbot at llvm.org
Wed Jul 22 03:58:06 PDT 2026
https://github.com/joker-eph updated https://github.com/llvm/llvm-project/pull/166688
>From 4a20d82d2d0e4b7361c68642e0f695233570650a Mon Sep 17 00:00:00 2001
From: Mehdi Amini <joker.eph at gmail.com>
Date: Sat, 26 Apr 2025 04:51:22 -0700
Subject: [PATCH 1/2] [mlir][scf] Add token-targeted early exits for scf.loop
Summary:
- Introduce a structured region-exit control-flow model for MLIR regions.
- Model region exits with RegionExitTerminatorOpInterface. Normal exits may
target the immediate parent operation; breaking-control-flow exits request
propagation to a specific ancestor receiver. Every intermediate parent must
opt in with PropagateControlFlowBreak until the addressed receiver handles
the request. In the SCF dialect, scf.loop requests are addressed with the
control token defined by the loop body.
Core IR:
- Add RegionExitTerminatorOpInterface for terminators that identify possible
receiver operations for a region exit.
- Add PropagateControlFlowBreak for intermediate operations that forward
breaking-control-flow requests without consuming them.
- Add HasBreakingControlFlowOpInterface for operations that receive addressed
breaking-control-flow requests, plus helpers to find the receiver, collect
nested predecessors, and detect escaping breaking control flow.
- Represent propagated RegionBranchOpInterface successors with explicit target
operations instead of a parent or propagation sentinel. Add a helper to
determine which operation owns a successor input list.
SCF dialect:
- Add scf.loop, which defines a control token, accepts optional iter_args and
results, and implements the breaking-control-flow receiver interface.
- Add scf.break and scf.continue as token-consuming region-exit terminators.
scf.break requests exit from the target loop and supplies loop results;
scf.continue requests re-entry into the target loop with next iteration
values.
- Let scf.yield and scf.reduce implement RegionExitTerminatorOpInterface by
targeting their immediate parent operation.
- Verify token receivers, propagation traits on forwarding operations, and
region-branch operand/result type consistency.
- Let scf.if propagate breaks and continues through its regions. Its parser,
printer, result inference, implicit terminator handling, and
canonicalizations now account for scf.yield, scf.break, and scf.continue
terminators.
- Guard yield-specific scf.if rewrites and add a canonicalization that removes
dead code after an if whose branches both transfer breaking control flow.
Analyses and utilities:
- Teach dead code analysis, dense and sparse backward dataflow, region-branch
graph traversal, post-dominance, and region inlining to follow concrete exit
targets instead of treating the immediate parent as the final successor.
Lowering:
- Lower scf.loop in the SCF-to-CF conversion by legalizing nested regions,
collecting direct break and continue predecessors, inlining the body, turning
breaks into branches to the continuation block, turning continues into
backedges, removing the loop token argument, and wiring results through
continuation block arguments.
- Adjust scf.if lowering to branch only from scf.yield terminators, leaving
early-exit terminators for the target loop lowering.
Tests and docs:
- Document the cooperative region-exit propagation model in the LangRef.
- Add tests for valid and invalid early-exit IR, region-branch interface
printing, dominance, dead code analysis, SCCP, SCF-to-CF lowering,
SCF canonicalization, and an end-to-end early-exit integration case.
Assisted-by: Claude Code
Assisted-by: Codex
---
mlir/docs/LangRef.md | 122 +++-
mlir/include/mlir/Dialect/SCF/IR/SCF.h | 70 +++
mlir/include/mlir/Dialect/SCF/IR/SCFOps.td | 181 +++++-
mlir/include/mlir/IR/Dominance.h | 9 +
mlir/include/mlir/IR/OpBase.td | 7 +
mlir/include/mlir/IR/OpDefinition.h | 13 +
mlir/include/mlir/IR/RegionKindInterface.h | 121 ++++
mlir/include/mlir/IR/RegionKindInterface.td | 83 +++
.../mlir/Interfaces/ControlFlowInterfaces.h | 6 +
.../Analysis/DataFlow/DeadCodeAnalysis.cpp | 12 +-
mlir/lib/Analysis/DataFlow/DenseAnalysis.cpp | 9 +-
mlir/lib/Analysis/DataFlow/SparseAnalysis.cpp | 23 +-
.../SCFToControlFlow/SCFToControlFlow.cpp | 237 ++++++--
mlir/lib/Dialect/SCF/IR/SCF.cpp | 571 +++++++++++++++---
.../SCF/IR/ValueBoundsOpInterfaceImpl.cpp | 4 +-
.../BufferizableOpInterfaceImpl.cpp | 4 +-
mlir/lib/Dialect/XeGPU/Utils/XeGPUUtils.cpp | 14 +-
mlir/lib/IR/Dominance.cpp | 94 ++-
mlir/lib/IR/RegionKindInterface.cpp | 168 ++++++
mlir/lib/IR/Verifier.cpp | 1 +
mlir/lib/Interfaces/ControlFlowInterfaces.cpp | 141 ++++-
mlir/lib/Transforms/Utils/CMakeLists.txt | 3 +
mlir/lib/Transforms/Utils/InliningUtils.cpp | 6 +
.../test-dead-code-analysis-early-exit.mlir | 100 +++
mlir/test/Analysis/test-dominance.mlir | 60 ++
.../convert-early-exit-to-cfg.mlir | 135 +++++
mlir/test/Dialect/SCF/loop_canonicalize.mlir | 300 +++++++++
mlir/test/IR/early-exit-invalid.mlir | 142 +++++
mlir/test/IR/early-exit.mlir | 82 +++
.../Integration/Dialect/SCF/early_exit.mlir | 82 +++
mlir/test/Transforms/sccp-early-exit.mlir | 117 ++++
mlir/test/lib/Dialect/Test/TestOps.td | 6 +
mlir/test/lib/IR/TestDominance.cpp | 33 +
mlir/test/lib/Interfaces/CMakeLists.txt | 1 +
.../RegionBranchOpInterface/CMakeLists.txt | 9 +
.../TestRegionBranchOpInterface.cpp | 77 +++
mlir/tools/mlir-opt/CMakeLists.txt | 1 +
mlir/tools/mlir-opt/mlir-opt.cpp | 2 +
38 files changed, 2844 insertions(+), 202 deletions(-)
create mode 100644 mlir/test/Analysis/DataFlow/test-dead-code-analysis-early-exit.mlir
create mode 100644 mlir/test/Conversion/SCFToControlFlow/convert-early-exit-to-cfg.mlir
create mode 100644 mlir/test/Dialect/SCF/loop_canonicalize.mlir
create mode 100644 mlir/test/IR/early-exit-invalid.mlir
create mode 100644 mlir/test/IR/early-exit.mlir
create mode 100644 mlir/test/Integration/Dialect/SCF/early_exit.mlir
create mode 100644 mlir/test/Transforms/sccp-early-exit.mlir
create mode 100644 mlir/test/lib/Interfaces/RegionBranchOpInterface/CMakeLists.txt
create mode 100644 mlir/test/lib/Interfaces/RegionBranchOpInterface/TestRegionBranchOpInterface.cpp
diff --git a/mlir/docs/LangRef.md b/mlir/docs/LangRef.md
index fea34c093418d..6ba4277ffac9f 100644
--- a/mlir/docs/LangRef.md
+++ b/mlir/docs/LangRef.md
@@ -509,7 +509,18 @@ operation produces results and the operation returns, those results also have
well-defined values. Execution then proceeds to the next operation in the block
until the terminator operation at the end of the block is reached; the
terminator determines the next continuation, if any. The determination of the
-next instruction to execute is the 'passing of control flow'.
+next instruction to execute is the 'passing of control flow'. A nested
+[region-exit terminator](#region-exit-terminators) can also terminate the
+current region with a request for a breaking-control-flow event addressed to a
+specific ancestor operation that implements
+`HasBreakingControlFlowOpInterface`. The terminator itself does not jump
+directly to that ancestor; it transfers control back to the containing
+operation with that request. If the containing operation is not the addressed
+receiver, it must define `PropagateControlFlowBreak` and propagate the same
+request to its own parent. This repeats until the addressed receiver handles
+the request, so every intermediate operation on the path from the terminator to
+the receiving `HasBreakingControlFlowOpInterface` operation must define
+`PropagateControlFlowBreak`.
In general, when control flow is passed to an operation, MLIR does not restrict
when control flow enters or exits the regions contained in that operation.
@@ -519,25 +530,26 @@ represent possible continuations. Successors explicitly specify destination
blocks, so control flow within the region can only pass to one of the specified
successor blocks, as in a `branch` operation. When a terminator has no
successors, it may pass 1) control back to the containing operation, as in a
-`return` operation, or 2) define that control flow does not continue, as in
-`ub.unreachable`. Terminators without successors therefore do not necessarily
-imply a return to the containing operation, the specific dialect operation
-determines the terminator's semantics. Blocks (other than the entry block) that
-are not listed as a successor of any terminator operation are defined to be
-unreachable and can be removed without affecting the semantics of the
-containing operation.
+`return` operation, 2) define that control flow does not continue, as in
+`ub.unreachable`, or 3) request that the containing operation propagate breaking
+control flow outward toward a specific ancestor. This target designator is
+dialect-defined. Terminators without successors therefore do not
+necessarily imply a return to the containing operation. Blocks (other than the
+entry block) that are not listed as a successor of any terminator operation are
+defined to be unreachable and can be removed without affecting the semantics of
+the containing operation.
Although control flow always enters a region through the entry block, control
flow may exit a region through any block with an appropriate terminator. The
-standard dialect leverages this capability to define operations with
+SCF dialect for example leverages this capability to define operations with
Single-Entry-Multiple-Exit (SEME) regions, possibly flowing through different
blocks in the region and exiting through any block with a `return` operation.
-This behavior is similar to that of a function body in most programming
-languages. In addition, control flow may also not reach the end of a block or
-region, for example if a function call does not return. Such an operation
-prevents control flow from reaching later operations, but does not remove the
-structural requirement that the block must end with a terminator unless the
-enclosing operation opts out with `NoTerminator`.
+This behavior can model that of a function body in most programming languages.
+In addition, control flow may also not reach the end of a block or region, for
+example if a function call does not return. Such an operation prevents control
+flow from reaching later operations, but does not remove the structural
+requirement that the block must end with a terminator unless the enclosing
+operation opts out with `NoTerminator`.
Example:
@@ -565,6 +577,86 @@ func.func @accelerator_compute(i64, i1) -> i64 { // An SSACFG region
}
```
+#### Region-Exit Terminators
+
+A terminator that implements `RegionExitTerminatorOpInterface` terminates the
+current region and identifies the operation that may receive the region exit. A
+normal region exit can target the immediately containing operation. For
+example, `scf.yield` exits only its own immediately enclosing region and
+returns control to the parent operation.
+
+Some terminators instead request a breaking-control-flow event addressed to a
+specific ancestor. The terminator transfers control back only to its immediate
+parent operation; reaching the addressed ancestor requires each intermediate
+parent operation to collaborate by propagating the request outward. The way a
+terminator designates the ancestor is dialect-defined: it may be a builtin
+`token` operand, an integer count of region levels to exit, a named symbolic
+relationship to a parent operation, or another dialect-specific mechanism.
+
+Every intermediate operation between a breaking terminator and the addressed
+receiver must define `PropagateControlFlowBreak`. The receiver must implement
+`HasBreakingControlFlowOpInterface`. For example, a loop operation nested
+inside another loop operation body carries both traits simultaneously: it
+handles breaks whose token identifies itself, and propagates breaks whose token
+identifies an outer loop.
+
+Region-exit terminators may carry values, which are propagated to the
+target operation. For example, when breaking out of a loop that produces
+results, the terminator supplies those result values. The exact mapping between
+terminator operands and the receiving op's results is dialect-defined (the
+receiving op may also ignore operands entirely).
+
+Examples:
+
+```mlir
+// scf.yield is the standard immediate region-exit terminator.
+// It exits only its own immediately enclosing region.
+scf.if %cond {
+ scf.yield // returns control to the immediate parent of scf.if
+}
+```
+
+```mlir
+// Trait legend:
+// [H] = HasBreakingControlFlowOpInterface (receives/catches the break)
+// [P] = PropagateControlFlowBreak (passes the break request upward)
+// [H][P] = both: handles breaks whose token identifies it, and propagates
+// breaks whose token identifies an outer loop
+scf.loop token(%outer) { // [H][P]
+ scf.loop token(%inner) { // [H][P]
+ scf.if %cond1 { // [P]
+ // Requests a break of the inner loop.
+ scf.break [%inner]
+ }
+ scf.if %cond2 { // [P]
+ // Requests a break of the outer loop. The scf.if and inner scf.loop
+ // both propagate the request upward.
+ scf.break [%outer]
+ }
+ scf.if %cond3 { // [P]
+ // Requests re-entry into the inner loop.
+ scf.continue [%inner]
+ }
+ }
+}
+return
+```
+
+```mlir
+// A loop that yields a result value on early exit.
+// scf.break carries operands that become the loop's results. (scf.continue
+// would instead carry operands that become the next iteration's iter_args,
+// but this loop has none.)
+%result = scf.loop token(%loop) -> f32 { // [H]
+ scf.if %found { // [P]
+ // %value becomes the loop result.
+ scf.break [%loop] %value : f32
+ }
+ // Re-enter the loop for the next iteration (no iter_args here).
+ scf.continue [%loop]
+}
+```
+
#### Operations with Multiple Regions
An operation containing multiple regions also completely determines the
diff --git a/mlir/include/mlir/Dialect/SCF/IR/SCF.h b/mlir/include/mlir/Dialect/SCF/IR/SCF.h
index 44cbb458d94fe..8918a1e89efdf 100644
--- a/mlir/include/mlir/Dialect/SCF/IR/SCF.h
+++ b/mlir/include/mlir/Dialect/SCF/IR/SCF.h
@@ -30,6 +30,11 @@
namespace mlir {
namespace scf {
void buildTerminatedBody(OpBuilder &builder, Location loc);
+
+namespace op_impl {
+struct IfOpImplicitTerminatorType;
+struct LoopOpImplicitTerminatorType;
+} // namespace op_impl
} // namespace scf
} // namespace mlir
@@ -112,6 +117,71 @@ SmallVector<Value> replaceAndCastForOpIterArg(RewriterBase &rewriter,
OpOperand &operand,
Value replacement,
const ValueTypeCastFnTy &castFn);
+namespace op_impl {
+
+//===----------------------------------------------------------------------===//
+// ControlFlowImplicitTerminatorOperation
+//===----------------------------------------------------------------------===//
+
+/// This class provides an interface compatible with
+/// SingleBlockImplicitTerminator, but allows multiple types of potential
+/// terminators aside from just one. If a terminator isn't present, this will
+/// generate a `ImplicitOpT` operation.
+template <typename ImplicitOpT, typename... OtherTerminatorOpTs>
+struct ControlFlowImplicitTerminatorOpType {
+ /// Implementation of `classof` that supports all of the potential terminator
+ /// operations.
+ static bool classof(Operation *op) {
+ return isa<ImplicitOpT, OtherTerminatorOpTs...>(op);
+ }
+
+ //===--------------------------------------------------------------------===//
+ // Implicit Terminator Methods
+
+ /// The following methods are all used when interacting with the "implicit"
+ /// terminator.
+
+ template <typename... Args>
+ static void build(Args &&...args) {
+ ImplicitOpT::build(std::forward<Args>(args)...);
+ }
+ static constexpr StringLiteral getOperationName() {
+ return ImplicitOpT::getOperationName();
+ }
+};
+/// An implicit terminator type for `if` operations. Empty regions still receive
+/// an `scf.yield`, but explicitly terminated regions may use any terminator
+/// implementing RegionExitTerminatorOpInterface so dialect-defined exits can
+/// propagate through `scf.if`.
+struct IfOpImplicitTerminatorType {
+ static bool classof(Operation *op) {
+ return isa<RegionExitTerminatorOpInterface>(op);
+ }
+
+ template <typename... Args>
+ static void build(Args &&...args) {
+ YieldOp::build(std::forward<Args>(args)...);
+ }
+ static constexpr StringLiteral getOperationName() {
+ return YieldOp::getOperationName();
+ }
+};
+struct LoopOpImplicitTerminatorType
+ : public ControlFlowImplicitTerminatorOpType<ContinueOp, BreakOp> {
+ /// Build the implicit `scf.continue` terminator. The control token consumed
+ /// by the terminator is the loop body's entry block argument #0; the
+ /// insertion block is guaranteed to be that body when the implicit
+ /// terminator is materialized. This keeps `scf.continue`'s public builders
+ /// free of any hidden block-argument assumption.
+ static void build(OpBuilder &builder, OperationState &state) {
+ Block *block = builder.getInsertionBlock();
+ assert(block && block->getNumArguments() > 0 &&
+ isa<TokenType>(block->getArgument(0).getType()) &&
+ "expected insertion block with a token-typed loop control argument");
+ state.addOperands(block->getArgument(0));
+ }
+};
+} // namespace op_impl
/// Helper function to compute the difference between two values. This is used
/// by the loop implementations to compute the trip count.
diff --git a/mlir/include/mlir/Dialect/SCF/IR/SCFOps.td b/mlir/include/mlir/Dialect/SCF/IR/SCFOps.td
index c0d1ac501cc77..05a3e37dc8676 100644
--- a/mlir/include/mlir/Dialect/SCF/IR/SCFOps.td
+++ b/mlir/include/mlir/Dialect/SCF/IR/SCFOps.td
@@ -15,6 +15,7 @@
include "mlir/Interfaces/ControlFlowInterfaces.td"
include "mlir/Interfaces/LoopLikeInterface.td"
+include "mlir/IR/OpAsmInterface.td"
include "mlir/IR/RegionKindInterface.td"
include "mlir/Dialect/SCF/IR/DeviceMappingInterface.td"
include "mlir/Interfaces/DestinationStyleOpInterface.td"
@@ -146,6 +147,168 @@ def ExecuteRegionOp : SCF_Op<"execute_region", [
let hasVerifier = 1;
}
+//===----------------------------------------------------------------------===//
+// LoopOp
+//===----------------------------------------------------------------------===//
+
+def LoopOp : SCF_Op<"loop",[
+ AutomaticAllocationScope,
+ OpAsmOpInterface,
+ RecursiveMemoryEffects,
+ PropagateControlFlowBreak,
+ TokenProducerTrait,
+ DeclareOpInterfaceMethods<RegionBranchOpInterface,
+ ["getEntrySuccessorOperands", "getSuccessorInputs"]>,
+ SingleBlockImplicitTerminator<"op_impl::LoopOpImplicitTerminatorType">,
+ HasBreakingControlFlowOpInterface,
+ HasNestedTerminator<["ContinueOp", "BreakOp"]>
+ ]> {
+ let summary = "Loop until a break operation";
+ let description = [{
+ The `loop` operation represents a loop that executes until a `break` is reached.
+
+ The loop consists of (1) a set of loop-carried values which are initialized by
+ `initValues` and updated by each iteration of the loop, and
+ (2) a region which represents the loop body.
+
+ The loop will execute the body of the loop until a `break` is executed.
+
+ Each control path of the loop must be terminated by:
+
+ - a `continue` that yields the next iteration's value for each loop carried variable.
+ - a `break` that terminates the loop and yields the final loop results.
+ - a nested terminator requesting to propagate the control-flow to a parent.
+ }];
+
+ let arguments = (ins Variadic<AnyType>:$initValues);
+ let results = (outs Variadic<AnyType>:$resultValues);
+ let regions = (region SizedRegion<1>:$region);
+
+ let extraClassDeclaration = [{
+ /// Return the iteration values of the loop region.
+ Block::BlockArgListType getRegionIterValues() {
+ return getRegion().getArguments().drop_front();
+ }
+
+ /// Return the `index`-th region iteration value.
+ BlockArgument getRegionIterValue(unsigned index) {
+ return getRegionIterValues()[index];
+ }
+
+ /// Return the loop control token.
+ BlockArgument getControlToken() {
+ return getRegion().getArgument(0);
+ }
+
+ /// Returns the number of region arguments for loop-carried values.
+ unsigned getNumRegionIterValues() {
+ return getRegion().getNumArguments() - 1;
+ }
+
+ /// Returns the loop block body
+ Block *getBody() { return &getRegion().front(); }
+ }];
+
+ let hasCustomAssemblyFormat = 1;
+ let hasRegionVerifier = 1;
+ let hasCanonicalizer = 1;
+}
+
+
+//===----------------------------------------------------------------------===//
+// BreakOp
+//===----------------------------------------------------------------------===//
+
+def BreakOp : SCF_Op<"break", [
+ Terminator,
+ DeclareOpInterfaceMethods<RegionExitTerminatorOpInterface>,
+ DeclareOpInterfaceMethods<RegionBranchTerminatorOpInterface,
+ ["getMutableSuccessorOperands"]>,
+ ParentOneOf<["IfOp", "LoopOp"]>
+ ]> {
+ let summary = "Break from loop";
+ let description = [{
+ The `break` operation is a terminator that exits one or more nested
+ regions and terminates the `scf.loop` that defined its control token.
+
+ The `break` may yield any number of operands; their types must match the
+ result types of the target `scf.loop`.
+
+ Example — break out of the immediately enclosing loop:
+ ```mlir
+ scf.loop token(%loop) -> i32 {
+ scf.break [%loop] %result : i32
+ }
+ ```
+
+ Example — break out of a loop through an enclosing `scf.if`:
+ ```mlir
+ scf.loop token(%loop) {
+ scf.if %cond {
+ scf.break [%loop]
+ }
+ scf.continue [%loop]
+ }
+ ```
+ }];
+
+
+ let arguments = (ins Token:$targetToken, Variadic<AnyType>:$args);
+ let assemblyFormat = [{
+ ` ` `[` $targetToken `]` ($args^ `:` type($args))? attr-dict
+ }];
+ let hasVerifier = 1;
+}
+
+
+//===----------------------------------------------------------------------===//
+// ContinueOp
+//===----------------------------------------------------------------------===//
+
+def ContinueOp : SCF_Op<"continue", [
+ Terminator,
+ DeclareOpInterfaceMethods<RegionExitTerminatorOpInterface>,
+ DeclareOpInterfaceMethods<RegionBranchTerminatorOpInterface,
+ ["getMutableSuccessorOperands"]>,
+ ParentOneOf<["IfOp", "LoopOp"]>
+ ]> {
+ let summary = "Continue to next loop iteration";
+ let description = [{
+ The `continue` operation is a terminator that re-enters a `scf.loop`
+ for its next iteration. The target loop is the one that defined the control
+ token operand.
+
+ The operands of `continue` become the loop-carried values (iter_args) for
+ the next iteration; their types must match the loop's iter_arg types.
+
+ Example — continue the immediately enclosing loop:
+ ```mlir
+ scf.loop token(%loop) iter_args(%i = %init) : i32 {
+ %next = arith.addi %i, %one : i32
+ scf.continue [%loop] %next : i32
+ }
+ ```
+
+ Example — continue an outer loop from inside a nested `scf.if`:
+ ```mlir
+ scf.loop token(%outer) iter_args(%counter = %init) : i64 {
+ scf.loop token(%inner) iter_args(%inner_arg = %counter) : i64 {
+ scf.if %restart_outer {
+ scf.continue [%outer] %inner_arg : i64
+ }
+ scf.continue [%inner] %inner_arg : i64
+ }
+ scf.continue [%outer] %counter : i64
+ }
+ ```
+ }];
+
+ let arguments = (ins Token:$targetToken, Variadic<AnyType>:$args);
+ let assemblyFormat = [{
+ ` ` `[` $targetToken `]` ($args^ `:` type($args))? attr-dict
+ }];
+ let hasVerifier = 1;
+}
//===----------------------------------------------------------------------===//
// ForOp
@@ -714,8 +877,8 @@ def IfOp : SCF_Op<"if", [DeclareOpInterfaceMethods<RegionBranchOpInterface, [
"getNumRegionInvocations", "getRegionInvocationBounds",
"getEntrySuccessorRegions", "getSuccessorInputs"]>,
DeclareOpInterfaceMethods<PromotableRegionOpInterface>,
- InferTypeOpAdaptor, SingleBlockImplicitTerminator<"scf::YieldOp">,
- RecursiveMemoryEffects, RecursivelySpeculatable, NoRegionArguments]> {
+ InferTypeOpAdaptor, SingleBlockImplicitTerminator<"op_impl::IfOpImplicitTerminatorType">,
+ RecursiveMemoryEffects, RecursivelySpeculatable, NoRegionArguments, PropagateControlFlowBreak]> {
let summary = "if-then-else operation";
let description = [{
The `scf.if` operation represents an if-then-else construct for
@@ -798,9 +961,17 @@ def IfOp : SCF_Op<"if", [DeclareOpInterfaceMethods<RegionBranchOpInterface, [
: OpBuilder::atBlockEnd(body, listener);
}
Block* thenBlock();
- YieldOp thenYield();
+ /// Returns the terminator of the then block. May be scf.break,
+ /// scf.continue, or scf.yield.
+ Operation *thenTerminator() {
+ return thenBlock()->getTerminator();
+ }
Block* elseBlock();
- YieldOp elseYield();
+ /// Returns the terminator of the else block. May be scf.break,
+ /// scf.continue, or scf.yield.
+ Operation *elseTerminator() {
+ return elseBlock()->getTerminator();
+ }
}];
let hasFolder = 1;
let hasCanonicalizer = 1;
@@ -918,6 +1089,7 @@ def ParallelOp : SCF_Op<"parallel",
def ReduceOp : SCF_Op<"reduce", [
Terminator, HasParent<"ParallelOp">, RecursiveMemoryEffects,
+ DeclareOpInterfaceMethods<RegionExitTerminatorOpInterface>,
DeclareOpInterfaceMethods<PromotableRegionOpInterface>,
DeclareOpInterfaceMethods<RegionBranchTerminatorOpInterface,
["getMutableSuccessorOperands"]>]> {
@@ -1219,6 +1391,7 @@ def IndexSwitchOp : SCF_Op<"index_switch", [RecursiveMemoryEffects,
//===----------------------------------------------------------------------===//
def YieldOp : SCF_Op<"yield", [Pure, ReturnLike, Terminator,
+ DeclareOpInterfaceMethods<RegionExitTerminatorOpInterface>,
ParentOneOf<["ExecuteRegionOp", "ForOp", "IfOp", "IndexSwitchOp",
"WhileOp"]>]> {
let summary = "loop yield and termination operation";
diff --git a/mlir/include/mlir/IR/Dominance.h b/mlir/include/mlir/IR/Dominance.h
index 70924a2e9ae59..1585f1881136b 100644
--- a/mlir/include/mlir/IR/Dominance.h
+++ b/mlir/include/mlir/IR/Dominance.h
@@ -123,6 +123,11 @@ class DominanceInfoBase {
Block::iterator bIt,
bool enclosingOk = true) const;
+ /// Return true if `op` contains a nested breaking terminator that escapes
+ /// through it to an ancestor. This cache is invalidated together with
+ /// dominance info.
+ bool hasBreakingControlFlowOpsCached(Operation *op) const;
+
/// A mapping of regions to their base dominator tree and a cached
/// "hasSSADominance" bit. This map does not contain dominator trees for
/// single block CFG regions, but we do want to cache the "hasSSADominance"
@@ -131,6 +136,10 @@ class DominanceInfoBase {
///
mutable DenseMap<Region *, llvm::PointerIntPair<DomTree *, 1, bool>>
dominanceInfos;
+
+ /// Cached escaping breaking-terminator query results. This is stable for the
+ /// lifetime of the analyzed IR and cleared on dominance invalidation.
+ mutable DenseMap<Operation *, bool> breakingControlFlowOpsCache;
};
extern template class DominanceInfoBase</*IsPostDom=*/true>;
diff --git a/mlir/include/mlir/IR/OpBase.td b/mlir/include/mlir/IR/OpBase.td
index 0d0669e90c3f7..ce7c891076420 100644
--- a/mlir/include/mlir/IR/OpBase.td
+++ b/mlir/include/mlir/IR/OpBase.td
@@ -134,6 +134,13 @@ class SingleBlockImplicitTerminatorImpl<string op>
class SingleBlockImplicitTerminator<string op>
: TraitList<[SingleBlock, SingleBlockImplicitTerminatorImpl<op>]>;
+// This operation has nested regions with the supplied list of terminator
+// operations.
+class HasNestedTerminator<list<string> ops>
+ : ParamNativeOpTrait<"HasNestedTerminators", !interleave(ops, ", ")>,
+ StructuralOpTrait;
+
+
// Op's regions don't have terminator.
def NoTerminator : NativeOpTrait<"NoTerminator">, StructuralOpTrait;
diff --git a/mlir/include/mlir/IR/OpDefinition.h b/mlir/include/mlir/IR/OpDefinition.h
index bd7fa1ffd4428..a146d17433ce3 100644
--- a/mlir/include/mlir/IR/OpDefinition.h
+++ b/mlir/include/mlir/IR/OpDefinition.h
@@ -1362,6 +1362,19 @@ struct HasAncestor {
};
};
+/// This class provides a verifier for ops that are expecting to have nested
+/// predecessors.
+template <typename... NestedPredecessorOpTypes>
+struct HasNestedTerminators {
+ template <typename ConcreteType>
+ class Impl : public TraitBase<ConcreteType, Impl> {
+ public:
+ static bool acceptsTerminator(Operation *predecessor) {
+ return llvm::isa_and_nonnull<NestedPredecessorOpTypes...>(predecessor);
+ }
+ };
+};
+
/// A trait for operations that have an attribute specifying operand segments.
///
/// Certain operations can have multiple variadic operands and their size
diff --git a/mlir/include/mlir/IR/RegionKindInterface.h b/mlir/include/mlir/IR/RegionKindInterface.h
index d6d3aeeb9bd05..2ada56ba6c193 100644
--- a/mlir/include/mlir/IR/RegionKindInterface.h
+++ b/mlir/include/mlir/IR/RegionKindInterface.h
@@ -36,6 +36,26 @@ class HasOnlyGraphRegion : public TraitBase<ConcreteType, HasOnlyGraphRegion> {
static RegionKind getRegionKind(unsigned index) { return RegionKind::Graph; }
static bool hasSSADominance(unsigned index) { return false; }
};
+
+/// Indicates that this operation is transparent to breaking control flow:
+/// a terminator (e.g. scf.break / scf.continue) can propagate through
+/// this op on its way to the addressed HasBreakingControlFlowOpInterface
+/// ancestor. The op does NOT consume the break; it simply passes it upward.
+/// All intermediate ops that receive and forward the break request must carry
+/// this trait.
+template <typename ConcreteType>
+class PropagateControlFlowBreak
+ : public TraitBase<ConcreteType, PropagateControlFlowBreak> {
+public:
+ static LogicalResult verifyTrait(Operation *op) {
+ // Verify the operation has regions and can handle breaking control flow
+ if (op->getNumRegions() == 0)
+ return op->emitOpError(
+ "operation with PropagateControlFlowBreak trait must have regions");
+ return success();
+ }
+};
+
} // namespace OpTrait
/// Return "true" if the given region may have SSA dominance. This function also
@@ -49,8 +69,109 @@ bool mayHaveSSADominance(Region ®ion);
/// implement the RegionKindInterface.
bool mayBeGraphRegion(Region ®ion);
+/// Summary of breaking terminator operations nested under an op.
+struct NestedBreakingControlFlowInfo {
+ /// Breaking terminators that may target the queried op directly. These
+ /// transfer control to the queried op rather than only to an intermediate
+ /// PropagateControlFlowBreak op.
+ SmallVector<Operation *> predecessors;
+
+ /// True if any direct predecessor targets the queried op from below the
+ /// queried op's immediate region.
+ bool hasNestedPredecessors = false;
+
+ /// True if any breaking terminator under the queried op targets one of the
+ /// queried op's ancestors, escaping through the queried op.
+ bool hasBreakingControlFlowOps = false;
+};
+
+/// Collect nested breaking-control-flow information for `op` with a single walk
+/// of its nested region tree.
+NestedBreakingControlFlowInfo getNestedBreakingControlFlowInfo(Operation *op);
+
+/// Return true if `op` (which implements HasBreakingControlFlowOpInterface)
+/// contains at least one breaking terminator that directly targets it from a
+/// nested region. Such a terminator is a "nested predecessor" of `op` because
+/// control flow may re-enter or exit `op` through a request propagated from a
+/// deeply nested site rather than only through an immediately enclosing
+/// terminator.
+bool hasNestedPredecessors(Operation *op);
+
+/// Return true if `op` contains any breaking terminator that would "break
+/// through" `op` towards an outer HasBreakingControlFlowOpInterface ancestor.
+/// This is used to detect whether an op's post-dominance is disrupted by an
+/// early-exit request that propagates through it.
+bool hasBreakingControlFlowOps(Operation *op);
+
+/// Collect all breaking terminators nested inside `op` that potentially
+/// directly target `op`. These are the ops that may transfer control flow to
+/// `op` on an early exit.
+void collectAllNestedPredecessors(Operation *op,
+ SmallVector<Operation *> &predecessors);
} // namespace mlir
#include "mlir/IR/RegionKindInterface.h.inc"
+namespace mlir {
+namespace detail {
+/// Implementation helper for visitNestedBreakingControlFlowOps. Walks the
+/// regions of `op` and invokes `callback` for every breaking terminator that
+/// either targets `op` or propagates further upward through `op`.
+/// The `nestedLevel` argument passed to the callback is the 1-based depth of
+/// the terminator relative to `op`'s outermost region.
+void visitNestedBreakingControlFlowOpsImpl(
+ Operation *op,
+ function_ref<WalkResult(RegionExitTerminatorOpInterface, int nestedLevel)>
+ callback);
+} // namespace detail
+
+/// Walk all breaking terminators that are relevant to breaking control
+/// flow inside `op` (see visitNestedBreakingControlFlowOpsImpl). The callback
+/// receives the terminator op and its 1-based nesting level. Callbacks
+/// returning WalkResult support early termination via WalkResult::interrupt();
+/// void-returning callbacks always continue.
+template <typename CallbackT>
+void visitNestedBreakingControlFlowOps(Operation *op, CallbackT &&callback) {
+ using RetT =
+ decltype(callback(std::declval<Operation *>(), std::declval<int>()));
+ if constexpr (std::is_same_v<RetT, WalkResult>) {
+ detail::visitNestedBreakingControlFlowOpsImpl(op, callback);
+ } else {
+ detail::visitNestedBreakingControlFlowOpsImpl(
+ op, [&](RegionExitTerminatorOpInterface visitedOp, int nestedLevel) {
+ callback(visitedOp, nestedLevel);
+ return WalkResult::advance();
+ });
+ }
+}
+
+/// Walk all breaking terminators relevant to breaking control flow
+/// across all top-level ops in `region`.
+template <typename CallbackT>
+void visitNestedBreakingControlFlowOps(Region ®ion, CallbackT &&callback) {
+ // Pass `callback` as an lvalue: it is reused across iterations, so it must
+ // not be forwarded (which could move from it on the first iteration).
+ for (Operation &op : region.getOps())
+ visitNestedBreakingControlFlowOps(&op, callback);
+}
+
+/// Return true if the given region may contain breaking control flow — either
+/// because its parent op propagates breaks (PropagateControlFlowBreak) or
+/// because it is the body of a HasBreakingControlFlowOpInterface op. Used to
+/// decide whether post-dominance analysis must account for early-exit paths.
+inline bool mightHaveBreakingControlFlow(Region *region) {
+ Operation *parentOp = region->getParentOp();
+ return (!parentOp->isRegistered() ||
+ parentOp->hasTrait<OpTrait::PropagateControlFlowBreak>() ||
+ isa<HasBreakingControlFlowOpInterface>(parentOp));
+}
+
+/// Return the HasBreakingControlFlowOpInterface operations potentially
+/// addressed by the given terminator. Returns an empty vector for non-breaking
+/// terminators or malformed target designators.
+SmallVector<HasBreakingControlFlowOpInterface>
+findPotentialBreakTargets(Operation *terminator);
+
+} // namespace mlir
+
#endif // MLIR_IR_REGIONKINDINTERFACE_H_
diff --git a/mlir/include/mlir/IR/RegionKindInterface.td b/mlir/include/mlir/IR/RegionKindInterface.td
index 607001a89250e..25f88249e547d 100644
--- a/mlir/include/mlir/IR/RegionKindInterface.td
+++ b/mlir/include/mlir/IR/RegionKindInterface.td
@@ -61,4 +61,87 @@ def GraphRegionNoTerminator : TraitList<[
HasOnlyGraphRegion
]>;
+// Indicates that this op may propagate a breaking control-flow event from a
+// nested region upward to the addressed HasBreakingControlFlowOpInterface
+// operation. The op does NOT consume the break itself; it is merely transparent
+// to it. All ops that sit between a breaking terminator and the
+// HasBreakingControlFlowOpInterface ancestor that will ultimately receive the
+// break must carry this trait.
+def PropagateControlFlowBreak : NativeOpTrait<"PropagateControlFlowBreak">;
+
+def HasBreakingControlFlowOpInterface : OpInterface<"HasBreakingControlFlowOpInterface"> {
+ let description = [{
+ Interface for operations that act as the target of a breaking control-flow
+ event (e.g. `scf.break` or `scf.continue`). Every intermediate op must
+ carry the `PropagateControlFlowBreak` trait. The terminator may identify the
+ receiver with any dialect-defined mechanism, such as a builtin token value,
+ an attribute, or a symbolic relationship to an enclosing operation.
+ }];
+ let cppNamespace = "::mlir";
+
+ let methods = [
+ StaticInterfaceMethod<
+ /*desc=*/[{
+ Return true if this operation accepts the given terminator operation
+ as a breaking-control-flow predecessor. This can be implemented through
+ the `HasNestedTerminator<[...]>` trait, otherwise an explicit implementation
+ must be provided (e.g. `return true;` to accept all, or a type check to
+ restrict).
+ }],
+ /*retTy=*/"bool",
+ /*methodName=*/"acceptsTerminator",
+ /*args=*/(ins "Operation *":$op)
+ >,
+ InterfaceMethod<
+ /*desc=*/[{
+ Return true if this operation has at least one breaking terminator
+ nested inside it that targets this operation directly. Used to decide
+ whether post-dominance analysis must account for early-exit paths.
+ }],
+ /*retTy=*/"bool",
+ /*methodName=*/"hasNestedPredecessors",
+ /*args=*/(ins),
+ /*methodBody=*/[{}],
+ /*defaultImplementation=*/[{
+ return ::mlir::hasNestedPredecessors(this->getOperation());
+ }]
+ >
+ ];
+}
+
+// OpInterface for terminators that leave their current region and identify
+// potential receiving operations using a dialect-defined mechanism. A normal
+// region return may simply target the immediately containing operation; a
+// breaking-control-flow region exit may target an ancestor operation that
+// implements HasBreakingControlFlowOpInterface.
+def RegionExitTerminatorOpInterface : OpInterface<"RegionExitTerminatorOpInterface"> {
+ let description = [{
+ Interface for terminators that leave their current region. Implemented by
+ ops such as `scf.yield`, `scf.break`, and `scf.continue`.
+ }];
+ let cppNamespace = "::mlir";
+
+ let methods = [
+ InterfaceMethod<
+ /*desc=*/[{
+ Return the operations that may receive this region-terminating control
+ flow event. Receivers for breaking control flow must implement
+ `HasBreakingControlFlowOpInterface`; normal region exits may return the
+ immediate parent operation.
+ }],
+ /*retTy=*/"::llvm::SmallVector<::mlir::Operation *>",
+ /*methodName=*/"getPotentialTargets",
+ /*args=*/(ins),
+ /*methodBody=*/[{}]
+ >
+ ];
+
+ let verify = [{
+ static_assert(ConcreteOp::template hasTrait<OpTrait::IsTerminator>(),
+ "expected operation to be a terminator");
+ return success();
+ }];
+}
+
+
#endif // MLIR_IR_REGIONKINDINTERFACE
diff --git a/mlir/include/mlir/Interfaces/ControlFlowInterfaces.h b/mlir/include/mlir/Interfaces/ControlFlowInterfaces.h
index 48b80e2059e11..cd74de81a7b10 100644
--- a/mlir/include/mlir/Interfaces/ControlFlowInterfaces.h
+++ b/mlir/include/mlir/Interfaces/ControlFlowInterfaces.h
@@ -446,6 +446,12 @@ inline llvm::raw_ostream &operator<<(llvm::raw_ostream &os,
OpPrintingFlags().skipRegions())
<< ">";
}
+
+/// Return the RegionBranchOpInterface operation that owns `successor` and
+/// defines its successor inputs.
+RegionBranchOpInterface
+getRegionBranchSuccessorOwner(RegionSuccessor successor);
+
} // namespace mlir
#endif // MLIR_INTERFACES_CONTROLFLOWINTERFACES_H
diff --git a/mlir/lib/Analysis/DataFlow/DeadCodeAnalysis.cpp b/mlir/lib/Analysis/DataFlow/DeadCodeAnalysis.cpp
index 37aaf1f634a64..cf703785c7e08 100644
--- a/mlir/lib/Analysis/DataFlow/DeadCodeAnalysis.cpp
+++ b/mlir/lib/Analysis/DataFlow/DeadCodeAnalysis.cpp
@@ -509,11 +509,15 @@ void DeadCodeAnalysis::visitRegionTerminator(Operation *op,
if (!operands)
return;
- SmallVector<RegionSuccessor> successors;
auto terminator = dyn_cast<RegionBranchTerminatorOpInterface>(op);
if (!terminator)
return;
+
+ SmallVector<RegionSuccessor> successors;
+ // Use operand-aware resolution to prune dead successors via constant folding
+ // (e.g. a scf.while condition known to be false).
terminator.getSuccessorRegions(*operands, successors);
+
visitRegionBranchEdges(branch, op, successors);
}
@@ -540,8 +544,10 @@ void DeadCodeAnalysis::visitRegionBranchEdges(
auto *predecessors = getOrCreate<PredecessorState>(point);
propagateIfChanged(
predecessors,
- predecessors->join(predecessorOp,
- regionBranchOp.getSuccessorInputs(successor)));
+ predecessors->join(
+ predecessorOp,
+ getRegionBranchSuccessorOwner(successor).getSuccessorInputs(
+ successor)));
LDBG() << "Added region branch as predecessor for successor: " << *point;
}
}
diff --git a/mlir/lib/Analysis/DataFlow/DenseAnalysis.cpp b/mlir/lib/Analysis/DataFlow/DenseAnalysis.cpp
index a2a9018b851ea..dc30aae167ac7 100644
--- a/mlir/lib/Analysis/DataFlow/DenseAnalysis.cpp
+++ b/mlir/lib/Analysis/DataFlow/DenseAnalysis.cpp
@@ -636,6 +636,9 @@ void AbstractDenseBackwardDataFlowAnalysis::visitRegionBranchOperation(
branch.getSuccessorRegions(branchPoint, successors);
LDBG() << " Processing " << successors.size() << " successor regions";
for (const RegionSuccessor &successor : successors) {
+ RegionBranchOpInterface successorOwner =
+ getRegionBranchSuccessorOwner(successor);
+ assert(successorOwner && "expected RegionBranchOpInterface owner");
const AbstractDenseLattice *after;
if (successor.isOperation()) {
LDBG() << " Successor is operation";
@@ -643,7 +646,7 @@ void AbstractDenseBackwardDataFlowAnalysis::visitRegionBranchOperation(
getProgramPointAfter(successor.getSuccessorOp()));
} else if (successor.getSuccessor()->empty()) {
LDBG() << " Successor is empty region";
- after = getLatticeFor(point, getProgramPointAfter(branch));
+ after = getLatticeFor(point, getProgramPointAfter(successorOwner));
} else {
Region *successorRegion = successor.getSuccessor();
assert(!successorRegion->empty() && "unexpected empty successor region");
@@ -662,7 +665,7 @@ void AbstractDenseBackwardDataFlowAnalysis::visitRegionBranchOperation(
}
LDBG() << " After state: " << *after;
- visitRegionBranchControlFlowTransfer(branch, branchPoint, successor, *after,
- before);
+ visitRegionBranchControlFlowTransfer(successorOwner, branchPoint, successor,
+ *after, before);
}
}
diff --git a/mlir/lib/Analysis/DataFlow/SparseAnalysis.cpp b/mlir/lib/Analysis/DataFlow/SparseAnalysis.cpp
index 3d9a375bfb2b7..04f27c8bf622a 100644
--- a/mlir/lib/Analysis/DataFlow/SparseAnalysis.cpp
+++ b/mlir/lib/Analysis/DataFlow/SparseAnalysis.cpp
@@ -647,14 +647,23 @@ void AbstractSparseBackwardDataFlowAnalysis::
// non-contiguous in the presence of multiple successors.
BitVector unaccounted(terminator->getNumOperands(), true);
- RegionBranchSuccessorMapping mapping;
- branch.getSuccessorOperandInputMapping(mapping,
- RegionBranchPoint(terminator));
- for (const auto &[operand, inputs] : mapping) {
- for (Value input : inputs) {
- meet(getLatticeElement(operand->get()),
+ SmallVector<RegionSuccessor> successors;
+ branch.getSuccessorRegions(RegionBranchPoint(terminator), successors);
+
+ for (RegionSuccessor successor : successors) {
+ RegionBranchOpInterface successorOwner =
+ getRegionBranchSuccessorOwner(successor);
+ assert(successorOwner && "expected RegionBranchOpInterface owner");
+ OperandRange operands =
+ branch.getSuccessorOperands(RegionBranchPoint(terminator), successor);
+ ValueRange inputs = successorOwner.getSuccessorInputs(successor);
+ assert(operands.size() == inputs.size() &&
+ "expected the same number of operands and inputs");
+ MutableArrayRef<OpOperand> opOperands(operands.getBase(), operands.size());
+ for (const auto &[operand, input] : llvm::zip_equal(opOperands, inputs)) {
+ meet(getLatticeElement(operand.get()),
*getLatticeElementFor(getProgramPointAfter(terminator), input));
- unaccounted.reset(operand->getOperandNumber());
+ unaccounted.reset(operand.getOperandNumber());
}
}
diff --git a/mlir/lib/Conversion/SCFToControlFlow/SCFToControlFlow.cpp b/mlir/lib/Conversion/SCFToControlFlow/SCFToControlFlow.cpp
index 2972d79c4302f..9dc7a56152a3c 100644
--- a/mlir/lib/Conversion/SCFToControlFlow/SCFToControlFlow.cpp
+++ b/mlir/lib/Conversion/SCFToControlFlow/SCFToControlFlow.cpp
@@ -311,8 +311,81 @@ struct ForallLowering : public OpRewritePattern<mlir::scf::ForallOp> {
PatternRewriter &rewriter) const override;
};
+/// Lowers `scf.loop` to a CFG loop. The loop body is inlined; `scf.break` ops
+/// that target this loop become branches to the continuation block and
+/// `scf.continue` ops that target this loop become back-edges to the loop
+/// header.
+struct LoopOpLowering : public OpConversionPattern<LoopOp> {
+ using OpConversionPattern<LoopOp>::OpConversionPattern;
+ void initialize() { setHasBoundedRewriteRecursion(); }
+
+ LogicalResult
+ matchAndRewrite(LoopOp loopOp, OpAdaptor adaptor,
+ ConversionPatternRewriter &rewriter) const override;
+};
+
} // namespace
+static LogicalResult lowerIfOpToCFG(IfOp ifOp, PatternRewriter &rewriter) {
+ auto loc = ifOp.getLoc();
+ // Start by splitting the block containing the 'scf.if' into two parts.
+ // The part before will contain the condition, the part after will be the
+ // continuation point.
+ auto *condBlock = rewriter.getInsertionBlock();
+ auto opPosition = rewriter.getInsertionPoint();
+ auto *remainingOpsBlock = rewriter.splitBlock(condBlock, opPosition);
+ Block *continueBlock;
+ if (ifOp.getNumResults() == 0) {
+ continueBlock = remainingOpsBlock;
+ } else {
+ continueBlock =
+ rewriter.createBlock(remainingOpsBlock, ifOp.getResultTypes(),
+ SmallVector<Location>(ifOp.getNumResults(), loc));
+ cf::BranchOp::create(rewriter, loc, remainingOpsBlock);
+ }
+
+ // Move blocks from the "then" region to the region containing 'scf.if',
+ // place it before the continuation block, and branch to it.
+ auto &thenRegion = ifOp.getThenRegion();
+ if (thenRegion.empty() || thenRegion.back().empty())
+ return failure();
+ auto *thenBlock = &thenRegion.front();
+ Operation *thenTerminator = thenRegion.back().getTerminator();
+ ValueRange thenTerminatorOperands = thenTerminator->getOperands();
+ rewriter.setInsertionPointToEnd(&thenRegion.back());
+ if (isa<scf::YieldOp>(thenTerminator)) {
+ cf::BranchOp::create(rewriter, loc, continueBlock, thenTerminatorOperands);
+ rewriter.eraseOp(thenTerminator);
+ }
+ rewriter.inlineRegionBefore(thenRegion, continueBlock);
+
+ // Move blocks from the "else" region (if present) to the region containing
+ // 'scf.if', place it before the continuation block and branch to it. It
+ // will be placed after the "then" regions.
+ auto *elseBlock = continueBlock;
+ auto &elseRegion = ifOp.getElseRegion();
+ if (!elseRegion.empty() && !elseRegion.back().empty()) {
+ elseBlock = &elseRegion.front();
+ Operation *elseTerminator = elseRegion.back().getTerminator();
+ ValueRange elseTerminatorOperands = elseTerminator->getOperands();
+ rewriter.setInsertionPointToEnd(&elseRegion.back());
+ if (isa<scf::YieldOp>(elseTerminator)) {
+ cf::BranchOp::create(rewriter, loc, continueBlock,
+ elseTerminatorOperands);
+ rewriter.eraseOp(elseTerminator);
+ }
+ rewriter.inlineRegionBefore(elseRegion, continueBlock);
+ }
+
+ rewriter.setInsertionPointToEnd(condBlock);
+ cf::CondBranchOp::create(rewriter, loc, ifOp.getCondition(), thenBlock,
+ /*trueArgs=*/ArrayRef<Value>(), elseBlock,
+ /*falseArgs=*/ArrayRef<Value>());
+
+ rewriter.replaceOp(ifOp, continueBlock->getArguments());
+ return success();
+}
+
static void propagateLoopAttrs(Operation *scfOp, Operation *brOp) {
// Let the CondBranchOp carry the LLVM attributes from the ForOp, such as the
// llvm.loop_annotation attribute.
@@ -400,58 +473,7 @@ LogicalResult ForLowering::matchAndRewrite(ForOp forOp,
LogicalResult IfLowering::matchAndRewrite(IfOp ifOp,
PatternRewriter &rewriter) const {
- auto loc = ifOp.getLoc();
-
- // Start by splitting the block containing the 'scf.if' into two parts.
- // The part before will contain the condition, the part after will be the
- // continuation point.
- auto *condBlock = rewriter.getInsertionBlock();
- auto opPosition = rewriter.getInsertionPoint();
- auto *remainingOpsBlock = rewriter.splitBlock(condBlock, opPosition);
- Block *continueBlock;
- if (ifOp.getNumResults() == 0) {
- continueBlock = remainingOpsBlock;
- } else {
- continueBlock =
- rewriter.createBlock(remainingOpsBlock, ifOp.getResultTypes(),
- SmallVector<Location>(ifOp.getNumResults(), loc));
- cf::BranchOp::create(rewriter, loc, remainingOpsBlock);
- }
-
- // Move blocks from the "then" region to the region containing 'scf.if',
- // place it before the continuation block, and branch to it.
- auto &thenRegion = ifOp.getThenRegion();
- auto *thenBlock = &thenRegion.front();
- Operation *thenTerminator = thenRegion.back().getTerminator();
- ValueRange thenTerminatorOperands = thenTerminator->getOperands();
- rewriter.setInsertionPointToEnd(&thenRegion.back());
- cf::BranchOp::create(rewriter, loc, continueBlock, thenTerminatorOperands);
- rewriter.eraseOp(thenTerminator);
- rewriter.inlineRegionBefore(thenRegion, continueBlock);
-
- // Move blocks from the "else" region (if present) to the region containing
- // 'scf.if', place it before the continuation block and branch to it. It
- // will be placed after the "then" regions.
- auto *elseBlock = continueBlock;
- auto &elseRegion = ifOp.getElseRegion();
- if (!elseRegion.empty()) {
- elseBlock = &elseRegion.front();
- Operation *elseTerminator = elseRegion.back().getTerminator();
- ValueRange elseTerminatorOperands = elseTerminator->getOperands();
- rewriter.setInsertionPointToEnd(&elseRegion.back());
- cf::BranchOp::create(rewriter, loc, continueBlock, elseTerminatorOperands);
- rewriter.eraseOp(elseTerminator);
- rewriter.inlineRegionBefore(elseRegion, continueBlock);
- }
-
- rewriter.setInsertionPointToEnd(condBlock);
- cf::CondBranchOp::create(rewriter, loc, ifOp.getCondition(), thenBlock,
- /*trueArgs=*/ArrayRef<Value>(), elseBlock,
- /*falseArgs=*/ArrayRef<Value>());
-
- // Ok, we're done!
- rewriter.replaceOp(ifOp, continueBlock->getArguments());
- return success();
+ return lowerIfOpToCFG(ifOp, rewriter);
}
LogicalResult
@@ -719,10 +741,114 @@ LogicalResult ForallLowering::matchAndRewrite(ForallOp forallOp,
return scf::forallToParallelLoop(rewriter, forallOp);
}
+LogicalResult
+LoopOpLowering::matchAndRewrite(LoopOp loopOp, OpAdaptor adaptor,
+ ConversionPatternRewriter &rewriter) const {
+ {
+ OpBuilder::InsertionGuard guard(rewriter);
+ if (failed(rewriter.legalize(&loopOp.getRegion())))
+ return rewriter.notifyMatchFailure(loopOp,
+ "failed to convert nested region");
+ }
+
+ // Handle degenerate case before modifying any IR.
+ Region &bodyRegion = loopOp.getRegion();
+ if (bodyRegion.empty() || bodyRegion.front().empty()) {
+ rewriter.eraseOp(loopOp);
+ return success();
+ }
+
+ NestedBreakingControlFlowInfo nestedBreakingControlFlow =
+ getNestedBreakingControlFlowInfo(loopOp);
+ if (nestedBreakingControlFlow.hasNestedPredecessors)
+ return rewriter.notifyMatchFailure(loopOp,
+ "loop op with nested predecessors");
+
+ // Collect direct predecessors (break/continue targeting this loop).
+ SmallVector<Operation *> predecessors;
+ llvm::append_range(predecessors, nestedBreakingControlFlow.predecessors);
+
+ // Lower `scf.loop` to CFG by converting breaks/continues to branches.
+ Location loc = loopOp.getLoc();
+ BlockArgument controlToken = loopOp.getControlToken();
+ // Split the block containing loopOp into the init block and continuation.
+ Block *initBlock = rewriter.getInsertionBlock();
+ auto initPos = rewriter.getInsertionPoint();
+ Block *continueBlock = rewriter.splitBlock(initBlock, initPos);
+ continueBlock->addArguments(
+ loopOp.getResultTypes(),
+ SmallVector<Location>(loopOp.getNumResults(), loc));
+ Block *loopBody = &bodyRegion.front();
+
+ // After the split above, `initBlock` holds the ops that preceded the loop and
+ // `continueBlock` holds the ops that followed it. Position the builder at the
+ // end of `initBlock`; the branch into the loop body is emitted there below.
+ rewriter.setInsertionPoint(initBlock, initBlock->end());
+
+ // Collect the blocks from the loop body region before inlining so we can
+ // restrict the scf.yield scan to only those blocks (not the whole function).
+ SmallVector<Block *> inlinedBlocks;
+ for (Block &b : bodyRegion)
+ inlinedBlocks.push_back(&b);
+
+ // Move all blocks from the scf.loop region before continueBlock.
+ rewriter.inlineRegionBefore(bodyRegion, continueBlock);
+ // We will remember all break/continue ops to fix up after.
+ SmallVector<Operation *> toErase;
+
+ for (auto predecessor : predecessors) {
+ if (auto breakOp = dyn_cast<scf::BreakOp>(predecessor)) {
+ rewriter.setInsertionPointAfter(breakOp);
+ cf::BranchOp::create(rewriter, breakOp->getLoc(), continueBlock,
+ ValueRange{breakOp.getArgs()});
+ } else if (auto contOp = dyn_cast<scf::ContinueOp>(predecessor)) {
+ rewriter.setInsertionPointAfter(contOp);
+ cf::BranchOp::create(rewriter, contOp->getLoc(), loopBody,
+ ValueRange{contOp.getArgs()});
+ }
+ toErase.push_back(predecessor);
+ }
+
+ // Erase the old scf.break/scf.continue ops. Drop their operands first so the
+ // control token (operand #0) loses its last uses here; this lets the
+ // `controlToken.use_empty()` check below confirm no stray consumer remains
+ // before we erase the token block argument.
+ for (Operation *op : toErase) {
+ op->setOperands({});
+ rewriter.eraseOp(op);
+ }
+
+ if (!controlToken.use_empty())
+ return rewriter.notifyMatchFailure(loopOp, "loop token still has uses");
+ loopBody->eraseArgument(0);
+
+ // The loop region is now a CFG. Jump from initBlock to the loop body.
+ rewriter.setInsertionPointToEnd(initBlock);
+ cf::BranchOp::create(rewriter, loc, loopBody,
+ ValueRange{loopOp.getOperands()});
+
+ // Replace any remaining scf.yield terminators in the inlined loop blocks
+ // with branches back to the loop header (treat as implicit continue).
+ for (Block *block : inlinedBlocks) {
+ if (auto yield = dyn_cast<scf::YieldOp>(block->getTerminator())) {
+ rewriter.setInsertionPoint(yield);
+ cf::BranchOp::create(rewriter, yield.getLoc(), loopBody,
+ yield.getOperands());
+ rewriter.eraseOp(yield);
+ }
+ }
+
+ // Replace the original scf.loop op with a branch to continueBlock assigning
+ // results.
+ rewriter.replaceOp(loopOp, continueBlock->getArguments());
+ return success();
+}
+
void mlir::populateSCFToControlFlowConversionPatterns(
RewritePatternSet &patterns) {
- patterns.add<ForallLowering, ForLowering, IfLowering, ParallelLowering,
- WhileLowering, ExecuteRegionLowering, IndexSwitchLowering>(
+ patterns.add<IfLowering>(patterns.getContext(), /*benefit=*/2);
+ patterns.add<ForallLowering, ForLowering, ParallelLowering, WhileLowering,
+ ExecuteRegionLowering, IndexSwitchLowering, LoopOpLowering>(
patterns.getContext());
patterns.add<DoWhileLowering>(patterns.getContext(), /*benefit=*/2);
}
@@ -734,7 +860,8 @@ void SCFToControlFlowPass::runOnOperation() {
// Configure conversion to lower out SCF operations.
ConversionTarget target(getContext());
target.addIllegalOp<scf::ForallOp, scf::ForOp, scf::IfOp, scf::IndexSwitchOp,
- scf::ParallelOp, scf::WhileOp, scf::ExecuteRegionOp>();
+ scf::LoopOp, scf::ParallelOp, scf::WhileOp,
+ scf::ExecuteRegionOp>();
target.markUnknownOpDynamicallyLegal([](Operation *) { return true; });
ConversionConfig config;
config.allowPatternRollback = allowPatternRollback;
diff --git a/mlir/lib/Dialect/SCF/IR/SCF.cpp b/mlir/lib/Dialect/SCF/IR/SCF.cpp
index 54d28783ddc4a..aa09d59c94acc 100644
--- a/mlir/lib/Dialect/SCF/IR/SCF.cpp
+++ b/mlir/lib/Dialect/SCF/IR/SCF.cpp
@@ -132,6 +132,17 @@ std::optional<llvm::APSInt> mlir::scf::computeUbMinusLb(Value lb, Value ub,
// ExecuteRegionOp
//===----------------------------------------------------------------------===//
+/// Erase all operations after `afterOp` in the same block (not including
+/// afterOp itself). Walks backwards to avoid use-after-free.
+static void eraseOpsAfter(PatternRewriter &rewriter, Operation *afterOp) {
+ Operation *cur = &afterOp->getBlock()->back();
+ while (cur != afterOp) {
+ Operation *prev = cur->getPrevNode();
+ rewriter.eraseOp(cur);
+ cur = prev;
+ }
+}
+
///
/// (ssa-id `=`)? `execute_region` `->` function-result-type `{`
/// block+
@@ -282,6 +293,10 @@ ValueRange ExecuteRegionOp::getSuccessorInputs(RegionSuccessor successor) {
: ValueRange();
}
+SmallVector<Operation *> YieldOp::getPotentialTargets() {
+ return {getOperation()->getParentOp()};
+}
+
//===----------------------------------------------------------------------===//
// ConditionOp
//===----------------------------------------------------------------------===//
@@ -311,6 +326,316 @@ void ConditionOp::getSuccessorRegions(
regions.push_back(RegionSuccessor(whileOp.getOperation()));
}
+//===----------------------------------------------------------------------===//
+// LoopOp
+//===----------------------------------------------------------------------===//
+
+//===----------------------------------------------------------------------===//
+// Control Flow Op Utilities
+//===----------------------------------------------------------------------===//
+
+template <typename OpT>
+static ParseResult
+parseControlFlowRegion(OpAsmParser &p, Region ®ion,
+ ArrayRef<OpAsmParser::Argument> arguments = {}) {
+ if (failed(p.parseRegion(region, arguments)))
+ return failure();
+ OpT::ensureTerminator(region, p.getBuilder(),
+ p.getEncodedSourceLoc(p.getNameLoc()));
+ return success();
+}
+
+static LoopOp getLoopTargetFromToken(Value target) {
+ auto targetArg = cast<BlockArgument>(target);
+ auto loopOp = cast<LoopOp>(targetArg.getOwner()->getParentOp());
+ assert(targetArg == loopOp.getControlToken() &&
+ "expected target token to be the loop control token");
+ return loopOp;
+}
+
+static LogicalResult verifyLoopTerminatorTarget(Operation *terminator,
+ Value target) {
+ auto targetArg = dyn_cast<BlockArgument>(target);
+ if (!targetArg)
+ return terminator->emitOpError()
+ << "target token must be an entry block argument of an scf.loop";
+
+ Block *targetBlock = targetArg.getOwner();
+ auto loopOp = dyn_cast_or_null<LoopOp>(targetBlock->getParentOp());
+ if (!loopOp || targetArg != loopOp.getControlToken())
+ return terminator->emitOpError()
+ << "target token must be the control token of an scf.loop";
+
+ Operation *currentOp = terminator->getParentOp();
+ while (currentOp && currentOp != loopOp.getOperation()) {
+ if (!currentOp->mightHaveTrait<OpTrait::PropagateControlFlowBreak>())
+ return terminator->emitOpError()
+ << "target token crosses an op that does not have the "
+ "PropagateControlFlowBreak trait: "
+ << OpWithFlags(currentOp, OpPrintingFlags().skipRegions());
+ currentOp = currentOp->getParentOp();
+ }
+
+ if (!currentOp)
+ return terminator->emitOpError()
+ << "target token must be defined by an enclosing scf.loop";
+
+ if (!loopOp.acceptsTerminator(terminator))
+ return loopOp.emitOpError("does not accept terminator: ")
+ << OpWithFlags(terminator, OpPrintingFlags().skipRegions());
+
+ return success();
+}
+
+static bool terminatorPropagatesThrough(Operation *terminator, Operation *op) {
+ for (HasBreakingControlFlowOpInterface target :
+ findPotentialBreakTargets(terminator)) {
+ if (target.getOperation() != op &&
+ target.getOperation()->isProperAncestor(op))
+ return true;
+ }
+ return false;
+}
+
+static void appendPropagatedTerminatorSuccessors(
+ RegionBranchTerminatorOpInterface terminator, Operation *op,
+ SmallVectorImpl<RegionSuccessor> ®ions) {
+ for (HasBreakingControlFlowOpInterface target :
+ findPotentialBreakTargets(terminator)) {
+ Operation *targetOp = target.getOperation();
+ if (targetOp == op || !targetOp->isProperAncestor(op))
+ continue;
+ auto targetBranch = dyn_cast<RegionBranchOpInterface>(targetOp);
+ assert(targetBranch && "expected breaking target to model region branch");
+ targetBranch.getSuccessorRegions(RegionBranchPoint(terminator), regions);
+ }
+}
+
+LogicalResult BreakOp::verify() {
+ return verifyLoopTerminatorTarget(getOperation(), getTargetToken());
+}
+
+SmallVector<Operation *> BreakOp::getPotentialTargets() {
+ LoopOp loopOp = getLoopTargetFromToken(getTargetToken());
+ return {loopOp.getOperation()};
+}
+
+MutableOperandRange
+BreakOp::getMutableSuccessorOperands(RegionSuccessor point) {
+ return MutableOperandRange(getOperation(), /*start=*/1,
+ /*length=*/getOperation()->getNumOperands() - 1);
+}
+
+LogicalResult ContinueOp::verify() {
+ return verifyLoopTerminatorTarget(getOperation(), getTargetToken());
+}
+
+SmallVector<Operation *> ContinueOp::getPotentialTargets() {
+ LoopOp loopOp = getLoopTargetFromToken(getTargetToken());
+ return {loopOp.getOperation()};
+}
+
+MutableOperandRange
+ContinueOp::getMutableSuccessorOperands(RegionSuccessor point) {
+ return MutableOperandRange(getOperation(), /*start=*/1,
+ /*length=*/getOperation()->getNumOperands() - 1);
+}
+
+LogicalResult LoopOp::verifyRegions() {
+ // Check matching between the operands and the region arguments.
+ if (getRegion().empty())
+ return emitOpError("region cannot be empty");
+ if (getRegion().front().getNumArguments() != getNumOperands() + 1)
+ return emitOpError("expected the region to have one argument per "
+ "loop-carried value plus the leading control token (")
+ << getNumOperands() + 1 << " expected, but got "
+ << getRegion().front().getNumArguments() << ")";
+ if (!isa<TokenType>(getControlToken().getType()))
+ return emitOpError("first region argument must be a control token");
+ for (auto [index, argAndOperand] :
+ llvm::enumerate(llvm::zip(getRegionIterValues(), getOperands()))) {
+ auto argType = std::get<0>(argAndOperand).getType();
+ auto operandType = std::get<1>(argAndOperand).getType();
+ if (argType != operandType)
+ return emitOpError() << "types mismatch between " << index
+ << "th iter operand (" << operandType
+ << ") and defined region argument (" << argType
+ << ")";
+ }
+ return success();
+}
+
+void LoopOp::print(OpAsmPrinter &p) {
+ p << " token(" << getControlToken() << ") ";
+ bool hasIters = !getInitValues().empty();
+ bool hasReturn = !getResultTypes().empty();
+
+ if (hasIters) {
+ p << "iter_args(";
+ llvm::interleaveComma(
+ llvm::zip(getRegionIterValues(), getInitValues()), p,
+ [&](auto it) { p << std::get<0>(it) << " = " << std::get<1>(it); });
+ p << ") : ";
+ p << getInitValues().getTypes();
+ p << " ";
+ }
+ if (hasReturn) {
+ p << "-> ";
+ p << getResultTypes();
+ p << " ";
+ }
+
+ Operation *terminator = getRegion().front().getTerminator();
+ // Elide the terminator only when it is the trivial implicit terminator:
+ // a `scf.continue` that targets this loop's control token and carries no
+ // iter values. A single-operand `scf.continue` targeting an *outer* loop
+ // must be printed explicitly, otherwise the implicit terminator rebuilt on
+ // parse would silently retarget it to this loop.
+ auto continueOp = dyn_cast<ContinueOp>(terminator);
+ bool printBlockTerminators = !continueOp || !continueOp.getArgs().empty() ||
+ continueOp.getTargetToken() != getControlToken();
+ p.printRegion(getRegion(), /*printEntryBlockArgs=*/false,
+ printBlockTerminators);
+ p.printOptionalAttrDict((*this)->getAttrs());
+}
+
+ParseResult LoopOp::parse(OpAsmParser &parser, OperationState &result) {
+ SmallVector<OpAsmParser::Argument, 4> regionArgs;
+ SmallVector<OpAsmParser::Argument, 4> iterRegionArgs;
+ SmallVector<OpAsmParser::UnresolvedOperand, 4> iterOperands;
+ SmallVector<Type, 4> iterTypes;
+
+ OpAsmParser::Argument controlToken;
+ if (parser.parseKeyword("token") || parser.parseLParen() ||
+ parser.parseArgument(controlToken) || parser.parseRParen())
+ return failure();
+ controlToken.type = parser.getBuilder().getType<TokenType>();
+ regionArgs.push_back(controlToken);
+
+ if (failed(parser.parseOptionalKeyword("iter_args"))) {
+ // no iter_args, but can still have a return type
+ if (succeeded(parser.parseOptionalArrow()))
+ if (parser.parseTypeList(result.types))
+ return failure();
+ } else {
+ // iter_args are present and must have colon followed by types
+ if (parser.parseAssignmentList(iterRegionArgs, iterOperands) ||
+ parser.parseColon() || parser.parseTypeList(iterTypes))
+ return failure();
+ if (iterRegionArgs.size() != iterTypes.size())
+ return parser.emitError(parser.getCurrentLocation(),
+ "found different number of iter_args and types");
+ // check for optional result type(s)
+ if (succeeded(parser.parseOptionalArrow()))
+ if (parser.parseTypeList(result.types))
+ return failure();
+ // Set region argument types for loop body
+ for (auto [regionArg, type] : llvm::zip_equal(iterRegionArgs, iterTypes)) {
+ regionArg.type = type;
+ }
+ llvm::append_range(regionArgs, iterRegionArgs);
+ }
+
+ // Parse region and attr dict.
+ if (parseControlFlowRegion<LoopOp>(parser, *result.addRegion(), regionArgs) ||
+ parser.parseOptionalAttrDict(result.attributes))
+ return failure();
+
+ // Resolve operands.
+ if (parser.resolveOperands(iterOperands, iterTypes, parser.getNameLoc(),
+ result.operands))
+ return failure();
+
+ return success();
+}
+
+void LoopOp::getSuccessorRegions(RegionBranchPoint point,
+ SmallVectorImpl<RegionSuccessor> ®ions) {
+ if (point.isParent()) {
+ regions.push_back(RegionSuccessor(&getRegion()));
+ return;
+ }
+
+ // Otherwise, it depends on the terminator: a continue branches back to the
+ // body and a break to the parent.
+ RegionBranchTerminatorOpInterface terminator =
+ point.getTerminatorPredecessorOrNull();
+ if (terminator && terminatorPropagatesThrough(terminator, getOperation())) {
+ appendPropagatedTerminatorSuccessors(terminator, getOperation(), regions);
+ return;
+ }
+
+ if (isa<ContinueOp>(terminator)) {
+ regions.push_back(RegionSuccessor(&getRegion()));
+ return;
+ }
+ assert(isa<BreakOp>(terminator) && "expected continue or break terminator");
+
+ regions.push_back(RegionSuccessor(getOperation()));
+}
+
+OperandRange LoopOp::getEntrySuccessorOperands(RegionSuccessor successor) {
+ return getInitValues();
+}
+
+ValueRange LoopOp::getSuccessorInputs(RegionSuccessor successor) {
+ return successor.isOperation() ? ValueRange(getResults())
+ : ValueRange(getRegionIterValues());
+}
+
+namespace {
+
+/// Rewriting pattern that erases loops that have a single iteration.
+struct SimplifyTrivialLoops : public OpRewritePattern<LoopOp> {
+ using OpRewritePattern<LoopOp>::OpRewritePattern;
+
+ LogicalResult matchAndRewrite(LoopOp op,
+ PatternRewriter &rewriter) const override {
+ // Terminator must be a break.
+ auto breakOp = dyn_cast<BreakOp>(op.getBody()->getTerminator());
+ if (!breakOp)
+ return rewriter.notifyMatchFailure(op, "loop terminator isn't a break");
+ SmallVector<HasBreakingControlFlowOpInterface> targets =
+ findPotentialBreakTargets(breakOp);
+ if (targets.size() != 1 ||
+ targets.front().getOperation() != op.getOperation())
+ return rewriter.notifyMatchFailure(
+ op, "loop terminator targets another loop");
+
+ // If it has nested predecessors, it can't be trivially simplified.
+ if (hasNestedPredecessors(op))
+ return rewriter.notifyMatchFailure(op, "has nested predecessors");
+
+ // Great: it is a single iteration loop, we can simplify it.
+ Block *body = op.getBody();
+ SmallVector<Value> replacements;
+ for (Value value : breakOp.getArgs()) {
+ if (auto blockArg = dyn_cast<BlockArgument>(value);
+ blockArg && blockArg.getOwner() == body) {
+ if (blockArg.getArgNumber() == 0)
+ return rewriter.notifyMatchFailure(
+ op, "loop terminator cannot yield the control token");
+ replacements.push_back(op.getInitValues()[blockArg.getArgNumber() - 1]);
+ continue;
+ }
+ replacements.push_back(value);
+ }
+ rewriter.eraseOp(breakOp);
+ assert(op.getControlToken().use_empty() && "expected token to be unused");
+ body->eraseArgument(0);
+ rewriter.inlineBlockBefore(body, op, op.getInitValues());
+ rewriter.replaceOp(op, replacements);
+
+ return success();
+ }
+};
+} // namespace
+
+void LoopOp::getCanonicalizationPatterns(RewritePatternSet &results,
+ MLIRContext *context) {
+ results.add<SimplifyTrivialLoops>(context);
+}
+
//===----------------------------------------------------------------------===//
// ForOp
//===----------------------------------------------------------------------===//
@@ -1940,13 +2265,21 @@ IfOp::inferReturnTypes(MLIRContext *ctx, std::optional<Location> loc,
Region *r = &adaptor.getThenRegion();
if (r->empty())
return failure();
- Block &b = r->front();
- if (b.empty())
+ Block *b = &r->front();
+ if (b->empty())
return failure();
- auto yieldOp = llvm::dyn_cast<YieldOp>(b.back());
- if (!yieldOp)
- return failure();
- TypeRange types = yieldOp.getOperandTypes();
+ Operation *terminator = &b->back();
+ if (terminatorPropagatesThrough(terminator, terminator->getParentOp())) {
+ if (adaptor.getElseRegion().empty())
+ return success();
+ b = &adaptor.getElseRegion().front();
+ if (b->empty())
+ return success();
+ terminator = &b->back();
+ if (terminatorPropagatesThrough(terminator, terminator->getParentOp()))
+ return success();
+ }
+ TypeRange types = terminator->getOperandTypes();
llvm::append_range(inferredReturnTypes, types);
return success();
}
@@ -2071,7 +2404,9 @@ ParseResult IfOp::parse(OpAsmParser &parser, OperationState &result) {
}
void IfOp::print(OpAsmPrinter &p) {
- bool printBlockTerminators = false;
+ bool printBlockTerminators =
+ !isa<YieldOp>(thenBlock()->back()) ||
+ (elseBlock() && !isa<YieldOp>(elseBlock()->back()));
p << " " << getCondition();
if (!getResults().empty()) {
@@ -2101,6 +2436,16 @@ void IfOp::getSuccessorRegions(RegionBranchPoint point,
// The `then` and the `else` region branch back to the parent operation or one
// of the recursive parent operations (early exit case).
if (!point.isParent()) {
+ // Propagating breaks/continues pass through this if-op to reach an
+ // enclosing loop. Don't report parent() as a successor for them; they
+ // don't yield values to this if-op.
+ if (auto terminator = point.getTerminatorPredecessorOrNull()) {
+ if (terminatorPropagatesThrough(terminator, getOperation())) {
+ appendPropagatedTerminatorSuccessors(terminator, getOperation(),
+ regions);
+ return;
+ }
+ }
regions.push_back(RegionSuccessor(getOperation()));
return;
}
@@ -2185,9 +2530,14 @@ struct ConvertTrivialIfToSelect : public OpRewritePattern<IfOp> {
if (op->getNumResults() == 0)
return failure();
+ YieldOp thenYield = dyn_cast<YieldOp>(op.thenTerminator());
+ YieldOp elseYield = dyn_cast<YieldOp>(op.elseTerminator());
+ if (!thenYield || !elseYield)
+ return failure();
+
auto cond = op.getCondition();
- auto thenYieldArgs = op.thenYield().getOperands();
- auto elseYieldArgs = op.elseYield().getOperands();
+ auto thenYieldArgs = thenYield.getOperands();
+ auto elseYieldArgs = elseYield.getOperands();
SmallVector<Type> nonHoistable;
for (auto [trueVal, falseVal] : llvm::zip(thenYieldArgs, elseYieldArgs)) {
@@ -2231,10 +2581,12 @@ struct ConvertTrivialIfToSelect : public OpRewritePattern<IfOp> {
}
rewriter.setInsertionPointToEnd(replacement.thenBlock());
- rewriter.replaceOpWithNewOp<YieldOp>(replacement.thenYield(), trueYields);
+ rewriter.replaceOpWithNewOp<YieldOp>(replacement.thenTerminator(),
+ trueYields);
rewriter.setInsertionPointToEnd(replacement.elseBlock());
- rewriter.replaceOpWithNewOp<YieldOp>(replacement.elseYield(), falseYields);
+ rewriter.replaceOpWithNewOp<YieldOp>(replacement.elseTerminator(),
+ falseYields);
rewriter.replaceOp(op, results);
return success();
@@ -2386,36 +2738,35 @@ struct ReplaceIfYieldWithConditionOrValue : public OpRewritePattern<IfOp> {
if (op.getNumResults() == 0)
return failure();
- auto trueYield =
- cast<scf::YieldOp>(op.getThenRegion().back().getTerminator());
- auto falseYield =
- cast<scf::YieldOp>(op.getElseRegion().back().getTerminator());
+ YieldOp thenYield = dyn_cast<YieldOp>(op.thenTerminator());
+ YieldOp elseYield = dyn_cast<YieldOp>(op.elseTerminator());
+ if (!thenYield || !elseYield)
+ return failure();
rewriter.setInsertionPoint(op->getBlock(),
op.getOperation()->getIterator());
bool changed = false;
Type i1Ty = rewriter.getI1Type();
- for (auto [trueResult, falseResult, opResult] :
- llvm::zip(trueYield.getResults(), falseYield.getResults(),
- op.getResults())) {
- if (trueResult == falseResult) {
+ for (auto [thenResult, elseResult, opResult] : llvm::zip(
+ thenYield.getResults(), elseYield.getResults(), op.getResults())) {
+ if (thenResult == elseResult) {
if (!opResult.use_empty()) {
- opResult.replaceAllUsesWith(trueResult);
+ opResult.replaceAllUsesWith(thenResult);
changed = true;
}
continue;
}
- BoolAttr trueYield, falseYield;
- if (!matchPattern(trueResult, m_Constant(&trueYield)) ||
- !matchPattern(falseResult, m_Constant(&falseYield)))
+ BoolAttr thenYield, elseYield;
+ if (!matchPattern(thenResult, m_Constant(&thenYield)) ||
+ !matchPattern(elseResult, m_Constant(&elseYield)))
continue;
- bool trueVal = trueYield.getValue();
- bool falseVal = falseYield.getValue();
- if (!trueVal && falseVal) {
+ bool thenVal = thenYield.getValue();
+ bool elseVal = elseYield.getValue();
+ if (!thenVal && elseVal) {
if (!opResult.use_empty()) {
- Dialect *constDialect = trueResult.getDefiningOp()->getDialect();
+ Dialect *constDialect = thenResult.getDefiningOp()->getDialect();
Value notCond = arith::XOrIOp::create(
rewriter, op.getLoc(), op.getCondition(),
constDialect
@@ -2427,7 +2778,7 @@ struct ReplaceIfYieldWithConditionOrValue : public OpRewritePattern<IfOp> {
changed = true;
}
}
- if (trueVal && !falseVal) {
+ if (thenVal && !elseVal) {
if (!opResult.use_empty()) {
opResult.replaceAllUsesWith(op.getCondition());
changed = true;
@@ -2482,18 +2833,16 @@ struct CombineIfs : public OpRewritePattern<IfOp> {
nextThen = nextIf.thenBlock();
if (!nextIf.getElseRegion().empty())
nextElse = nextIf.elseBlock();
- }
- if (arith::XOrIOp notv =
- nextIf.getCondition().getDefiningOp<arith::XOrIOp>()) {
+ } else if (arith::XOrIOp notv =
+ nextIf.getCondition().getDefiningOp<arith::XOrIOp>()) {
if (notv.getLhs() == prevIf.getCondition() &&
matchPattern(notv.getRhs(), m_One())) {
nextElse = nextIf.thenBlock();
if (!nextIf.getElseRegion().empty())
nextThen = nextIf.elseBlock();
}
- }
- if (arith::XOrIOp notv =
- prevIf.getCondition().getDefiningOp<arith::XOrIOp>()) {
+ } else if (arith::XOrIOp notv =
+ prevIf.getCondition().getDefiningOp<arith::XOrIOp>()) {
if (notv.getLhs() == nextIf.getCondition() &&
matchPattern(notv.getRhs(), m_One())) {
nextElse = nextIf.thenBlock();
@@ -2504,14 +2853,25 @@ struct CombineIfs : public OpRewritePattern<IfOp> {
if (!nextThen && !nextElse)
return failure();
+ // Check that the terminators are all YieldOp
+ if (!isa<YieldOp>(prevIf.thenTerminator()) ||
+ (nextThen && !isa<YieldOp>(nextThen->getTerminator())))
+ return failure();
+ if (!prevIf.getElseRegion().empty() &&
+ !isa<YieldOp>(prevIf.elseTerminator()))
+ return failure();
+ if (nextElse && !nextElse->empty() &&
+ !isa<YieldOp>(nextElse->getTerminator()))
+ return failure();
SmallVector<Value> prevElseYielded;
if (!prevIf.getElseRegion().empty())
- prevElseYielded = prevIf.elseYield().getOperands();
+ prevElseYielded = prevIf.elseTerminator()->getOperands();
// Replace all uses of return values of op within nextIf with the
// corresponding yields
- for (auto it : llvm::zip(prevIf.getResults(),
- prevIf.thenYield().getOperands(), prevElseYielded))
+ for (auto it :
+ llvm::zip(prevIf.getResults(), prevIf.thenTerminator()->getOperands(),
+ prevElseYielded))
for (OpOperand &use :
llvm::make_early_inc_range(std::get<0>(it).getUses())) {
if (nextThen && nextThen->getParent()->isAncestor(
@@ -2539,16 +2899,16 @@ struct CombineIfs : public OpRewritePattern<IfOp> {
combinedIf.getThenRegion().begin());
if (nextThen) {
- YieldOp thenYield = combinedIf.thenYield();
- YieldOp thenYield2 = cast<YieldOp>(nextThen->getTerminator());
+ Operation *thenTerminator = combinedIf.thenTerminator();
+ Operation *thenTerminator2 = nextThen->getTerminator();
rewriter.mergeBlocks(nextThen, combinedIf.thenBlock());
rewriter.setInsertionPointToEnd(combinedIf.thenBlock());
- SmallVector<Value> mergedYields(thenYield.getOperands());
- llvm::append_range(mergedYields, thenYield2.getOperands());
- YieldOp::create(rewriter, thenYield2.getLoc(), mergedYields);
- rewriter.eraseOp(thenYield);
- rewriter.eraseOp(thenYield2);
+ SmallVector<Value> mergedYields(thenTerminator->getOperands());
+ llvm::append_range(mergedYields, thenTerminator2->getOperands());
+ YieldOp::create(rewriter, thenTerminator->getLoc(), mergedYields);
+ rewriter.eraseOp(thenTerminator);
+ rewriter.eraseOp(thenTerminator2);
}
rewriter.inlineRegionBefore(prevIf.getElseRegion(),
@@ -2561,18 +2921,17 @@ struct CombineIfs : public OpRewritePattern<IfOp> {
combinedIf.getElseRegion(),
combinedIf.getElseRegion().begin());
} else {
- YieldOp elseYield = combinedIf.elseYield();
- YieldOp elseYield2 = cast<YieldOp>(nextElse->getTerminator());
+ Operation *elseTerminator = combinedIf.elseTerminator();
+ Operation *elseTerminator2 = nextElse->getTerminator();
rewriter.mergeBlocks(nextElse, combinedIf.elseBlock());
-
rewriter.setInsertionPointToEnd(combinedIf.elseBlock());
- SmallVector<Value> mergedElseYields(elseYield.getOperands());
- llvm::append_range(mergedElseYields, elseYield2.getOperands());
+ SmallVector<Value> mergedElseYields(elseTerminator->getOperands());
+ llvm::append_range(mergedElseYields, elseTerminator2->getOperands());
- YieldOp::create(rewriter, elseYield2.getLoc(), mergedElseYields);
- rewriter.eraseOp(elseYield);
- rewriter.eraseOp(elseYield2);
+ YieldOp::create(rewriter, elseTerminator->getLoc(), mergedElseYields);
+ rewriter.eraseOp(elseTerminator);
+ rewriter.eraseOp(elseTerminator2);
}
}
@@ -2600,7 +2959,8 @@ struct RemoveEmptyElseBranch : public OpRewritePattern<IfOp> {
if (ifOp.getNumResults())
return failure();
Block *elseBlock = ifOp.elseBlock();
- if (!elseBlock || !llvm::hasSingleElement(*elseBlock))
+ if (!elseBlock || (!llvm::hasSingleElement(*elseBlock) ||
+ !isa<YieldOp>(elseBlock->getTerminator())))
return failure();
auto newIfOp = rewriter.cloneWithoutRegions(ifOp);
rewriter.inlineRegionBefore(ifOp.getThenRegion(), newIfOp.getThenRegion(),
@@ -2636,21 +2996,32 @@ struct CombineNestedIfs : public OpRewritePattern<IfOp> {
if (!llvm::hasSingleElement(nestedOps))
return failure();
+ auto nestedIf = dyn_cast<IfOp>(*nestedOps.begin());
+ if (!nestedIf)
+ return failure();
+
+ // Terminator must be a YieldOp
+ if (!isa<YieldOp>(op.thenTerminator()))
+ return failure();
+
// If there is an else block, it can only yield
- if (op.elseBlock() && !llvm::hasSingleElement(*op.elseBlock()))
+ if (op.elseBlock() && (!llvm::hasSingleElement(*op.elseBlock()) ||
+ !isa<YieldOp>(op.elseTerminator())))
return failure();
- auto nestedIf = dyn_cast<IfOp>(*nestedOps.begin());
- if (!nestedIf)
+ // Same for the nested if: the then and else blocks can only yield.
+ if (!isa<YieldOp>(nestedIf.thenTerminator()))
return failure();
- if (nestedIf.elseBlock() && !llvm::hasSingleElement(*nestedIf.elseBlock()))
+ if (nestedIf.elseBlock() &&
+ (!llvm::hasSingleElement(*nestedIf.elseBlock()) ||
+ !isa<YieldOp>(nestedIf.elseTerminator())))
return failure();
- SmallVector<Value> thenYield(op.thenYield().getOperands());
- SmallVector<Value> elseYield;
+ SmallVector<Value> thenTerminator(op.thenTerminator()->getOperands());
+ SmallVector<Value> elseTerminator;
if (op.elseBlock())
- llvm::append_range(elseYield, op.elseYield().getOperands());
+ llvm::append_range(elseTerminator, op.elseTerminator()->getOperands());
// A list of indices for which we should upgrade the value yielded
// in the else to a select.
@@ -2660,19 +3031,20 @@ struct CombineNestedIfs : public OpRewritePattern<IfOp> {
// only permit combining if the value yielded when the condition
// is false in the outer scf.if is the same value yielded when the
// inner scf.if condition is false.
- // Note that the array access to elseYield will not go out of bounds
- // since it must have the same length as thenYield, since they both
+ // Note that the array access to elseTerminator will not go out of bounds
+ // since it must have the same length as thenTerminator, since they both
// come from the same scf.if.
- for (const auto &tup : llvm::enumerate(thenYield)) {
+ for (const auto &tup : llvm::enumerate(thenTerminator)) {
if (tup.value().getDefiningOp() == nestedIf) {
auto nestedIdx = llvm::cast<OpResult>(tup.value()).getResultNumber();
- if (nestedIf.elseYield().getOperand(nestedIdx) !=
- elseYield[tup.index()]) {
+ if (nestedIf.elseTerminator()->getOperand(nestedIdx) !=
+ elseTerminator[tup.index()]) {
return failure();
}
// If the correctness test passes, we will yield
// corresponding value from the inner scf.if
- thenYield[tup.index()] = nestedIf.thenYield().getOperand(nestedIdx);
+ thenTerminator[tup.index()] =
+ nestedIf.thenTerminator()->getOperand(nestedIdx);
continue;
}
@@ -2704,28 +3076,72 @@ struct CombineNestedIfs : public OpRewritePattern<IfOp> {
for (auto idx : elseYieldsToUpgradeToSelect)
results[idx] =
arith::SelectOp::create(rewriter, op.getLoc(), op.getCondition(),
- thenYield[idx], elseYield[idx]);
+ thenTerminator[idx], elseTerminator[idx]);
rewriter.mergeBlocks(nestedIf.thenBlock(), newIfBlock);
rewriter.setInsertionPointToEnd(newIf.thenBlock());
- rewriter.replaceOpWithNewOp<YieldOp>(newIf.thenYield(), thenYield);
- if (!elseYield.empty()) {
+ rewriter.replaceOpWithNewOp<YieldOp>(newIf.thenTerminator(),
+ thenTerminator);
+ if (!elseTerminator.empty()) {
rewriter.createBlock(&newIf.getElseRegion());
rewriter.setInsertionPointToEnd(newIf.elseBlock());
- YieldOp::create(rewriter, loc, elseYield);
+ YieldOp::create(rewriter, loc, elseTerminator);
}
rewriter.replaceOp(op, results);
return success();
}
};
+/// Simplify if with breaking control flow in both branches.
+/// For example:
+/// scf.if %cmp {
+/// scf.break [%loop] %arg1
+/// } else {
+/// scf.continue [%loop]
+/// }
+/// print(...) // This is dead code
+/// becomes
+/// scf.if %cmp {
+/// scf.break [%loop] %arg1
+/// }
+/// scf.continue [%loop]
+struct SimplifyIfWithBreakingControlFlowInBothBranches
+ : public OpRewritePattern<IfOp> {
+ using OpRewritePattern<IfOp>::OpRewritePattern;
+ LogicalResult matchAndRewrite(IfOp op,
+ PatternRewriter &rewriter) const override {
+ if (op.getElseRegion().empty() || isa<YieldOp>(op.thenTerminator()) ||
+ isa<YieldOp>(op.elseTerminator()))
+ return failure();
+
+ // Inline the else block after the current op and erase everything after.
+ Block *block = op.elseBlock();
+
+ Operation *terminator = block->getTerminator();
+ // Inline the else block after the current op
+ rewriter.inlineBlockBefore(block, op->getNextNode());
+
+ // Erase everything that comes after the inlined terminator (dead code).
+ eraseOpsAfter(rewriter, terminator);
+
+ // The "else" region is now empty, let's clone the if op and inline the then
+ // region.
+ auto newIfOp = rewriter.cloneWithoutRegions(op);
+ rewriter.inlineRegionBefore(op.getThenRegion(), newIfOp.getThenRegion(),
+ newIfOp.getThenRegion().begin());
+ rewriter.eraseOp(op);
+
+ return success();
+ }
+};
} // namespace
void IfOp::getCanonicalizationPatterns(RewritePatternSet &results,
MLIRContext *context) {
results.add<CombineIfs, CombineNestedIfs, ConditionPropagation,
ConvertTrivialIfToSelect, RemoveEmptyElseBranch,
- ReplaceIfYieldWithConditionOrValue>(context);
+ ReplaceIfYieldWithConditionOrValue,
+ SimplifyIfWithBreakingControlFlowInBothBranches>(context);
populateRegionBranchOpInterfaceCanonicalizationPatterns(
results, IfOp::getOperationName());
populateRegionBranchOpInterfaceInliningPattern(results,
@@ -2733,14 +3149,13 @@ void IfOp::getCanonicalizationPatterns(RewritePatternSet &results,
}
Block *IfOp::thenBlock() { return &getThenRegion().back(); }
-YieldOp IfOp::thenYield() { return cast<YieldOp>(&thenBlock()->back()); }
+
Block *IfOp::elseBlock() {
Region &r = getElseRegion();
if (r.empty())
return nullptr;
return &r.back();
}
-YieldOp IfOp::elseYield() { return cast<YieldOp>(&elseBlock()->back()); }
//===----------------------------------------------------------------------===//
// ParallelOp
@@ -3184,6 +3599,10 @@ ReduceOp::getMutableSuccessorOperands(RegionSuccessor point) {
return MutableOperandRange(getOperation(), /*start=*/0, /*length=*/0);
}
+SmallVector<Operation *> ReduceOp::getPotentialTargets() {
+ return {getOperation()->getParentOp()};
+}
+
//===----------------------------------------------------------------------===//
// ReduceReturnOp
//===----------------------------------------------------------------------===//
@@ -3467,8 +3886,8 @@ struct WhileMoveIfDown : public OpRewritePattern<scf::WhileOp> {
auto it = llvm::find(ifOp->getResults(), arg);
if (it != ifOp->getResults().end()) {
size_t ifOpIdx = it.getIndex();
- Value thenValue = ifOp.thenYield()->getOperand(ifOpIdx);
- Value elseValue = ifOp.elseYield()->getOperand(ifOpIdx);
+ Value thenValue = ifOp.thenTerminator()->getOperand(ifOpIdx);
+ Value elseValue = ifOp.elseTerminator()->getOperand(ifOpIdx);
rewriter.replaceAllUsesWith(ifOp->getResults()[ifOpIdx], elseValue);
rewriter.replaceAllUsesWith(op.getAfterArguments()[idx], thenValue);
@@ -3516,7 +3935,7 @@ struct WhileMoveIfDown : public OpRewritePattern<scf::WhileOp> {
});
// Inline ifOp then region into new whileOp after region.
- rewriter.eraseOp(ifOp.thenYield());
+ rewriter.eraseOp(ifOp.thenTerminator());
rewriter.inlineBlockBefore(ifOp.thenBlock(), newWhileOp.getAfterBody(),
newWhileOp.getAfterBody()->begin());
rewriter.eraseOp(ifOp);
diff --git a/mlir/lib/Dialect/SCF/IR/ValueBoundsOpInterfaceImpl.cpp b/mlir/lib/Dialect/SCF/IR/ValueBoundsOpInterfaceImpl.cpp
index 496a7b036e65d..8a0b7f955267e 100644
--- a/mlir/lib/Dialect/SCF/IR/ValueBoundsOpInterfaceImpl.cpp
+++ b/mlir/lib/Dialect/SCF/IR/ValueBoundsOpInterfaceImpl.cpp
@@ -179,8 +179,8 @@ struct IfOpInterface
std::optional<int64_t> dim,
ValueBoundsConstraintSet &cstr) {
unsigned int resultNum = cast<OpResult>(value).getResultNumber();
- Value thenValue = ifOp.thenYield().getResults()[resultNum];
- Value elseValue = ifOp.elseYield().getResults()[resultNum];
+ Value thenValue = ifOp.thenTerminator()->getOperand(resultNum);
+ Value elseValue = ifOp.elseTerminator()->getOperand(resultNum);
auto boundsBuilder = cstr.bound(value);
if (dim)
diff --git a/mlir/lib/Dialect/SCF/Transforms/BufferizableOpInterfaceImpl.cpp b/mlir/lib/Dialect/SCF/Transforms/BufferizableOpInterfaceImpl.cpp
index f8302126a4a4e..456797b8fbec3 100644
--- a/mlir/lib/Dialect/SCF/Transforms/BufferizableOpInterfaceImpl.cpp
+++ b/mlir/lib/Dialect/SCF/Transforms/BufferizableOpInterfaceImpl.cpp
@@ -227,8 +227,8 @@ struct IfOpInterface
auto ifOp = cast<scf::IfOp>(op);
size_t resultNum = std::distance(op->getOpResults().begin(),
llvm::find(op->getOpResults(), value));
- OpOperand *thenOperand = &ifOp.thenYield()->getOpOperand(resultNum);
- OpOperand *elseOperand = &ifOp.elseYield()->getOpOperand(resultNum);
+ OpOperand *thenOperand = &ifOp.thenTerminator()->getOpOperand(resultNum);
+ OpOperand *elseOperand = &ifOp.elseTerminator()->getOpOperand(resultNum);
return {{thenOperand, BufferRelation::Equivalent, /*isDefinite=*/false},
{elseOperand, BufferRelation::Equivalent, /*isDefinite=*/false}};
}
diff --git a/mlir/lib/Dialect/XeGPU/Utils/XeGPUUtils.cpp b/mlir/lib/Dialect/XeGPU/Utils/XeGPUUtils.cpp
index 9620e21f9bfdf..306a4cbc2dcef 100644
--- a/mlir/lib/Dialect/XeGPU/Utils/XeGPUUtils.cpp
+++ b/mlir/lib/Dialect/XeGPU/Utils/XeGPUUtils.cpp
@@ -927,11 +927,17 @@ xegpu::precomputeLoopBlockArgTypes(Operation *topLevelOp,
}
if (auto ifOp = dyn_cast<scf::IfOp>(op)) {
// Each result and its then/else yield operands share one position and
- // must convert identically; derive all from the result's layout.
- scf::YieldOp thenYield = ifOp.thenYield();
- scf::YieldOp elseYield = ifOp.elseBlock() ? ifOp.elseYield() : nullptr;
+ // must convert identically; derive all from the result's layout. A
+ // branch that exits early does not yield a value to this `scf.if`, so
+ // only record operands from branches that actually end with scf.yield.
+ scf::YieldOp thenYield = dyn_cast<scf::YieldOp>(ifOp.thenTerminator());
+ scf::YieldOp elseYield;
+ if (ifOp.elseBlock())
+ elseYield = dyn_cast<scf::YieldOp>(ifOp.elseTerminator());
for (auto [idx, res] : llvm::enumerate(ifOp.getResults())) {
- SmallVector<Value> dests{res, thenYield.getOperand(idx)};
+ SmallVector<Value> dests{res};
+ if (thenYield)
+ dests.push_back(thenYield.getOperand(idx));
if (elseYield)
dests.push_back(elseYield.getOperand(idx));
recordTypes(res, dests);
diff --git a/mlir/lib/IR/Dominance.cpp b/mlir/lib/IR/Dominance.cpp
index 79fb41f2e6b30..f8820822dc01f 100644
--- a/mlir/lib/IR/Dominance.cpp
+++ b/mlir/lib/IR/Dominance.cpp
@@ -14,8 +14,12 @@
#include "mlir/IR/Dominance.h"
#include "mlir/IR/Operation.h"
#include "mlir/IR/RegionKindInterface.h"
+#include "llvm/ADT/SmallPtrSet.h"
+#include "llvm/Support/DebugLog.h"
#include "llvm/Support/GenericDomTreeConstruction.h"
+#define DEBUG_TYPE "dominance"
+
using namespace mlir;
using namespace mlir::detail;
@@ -38,6 +42,7 @@ void DominanceInfoBase<IsPostDom>::invalidate() {
for (auto entry : dominanceInfos)
delete entry.second.getPointer();
dominanceInfos.clear();
+ breakingControlFlowOpsCache.clear();
}
template <bool IsPostDom>
@@ -53,6 +58,7 @@ void DominanceInfoBase<IsPostDom>::invalidate(Region *region, bool recursive) {
region->walk([&](Region *r) { invalidate(r); });
else
invalidate(region);
+ breakingControlFlowOpsCache.clear();
}
/// Return the dom tree and "hasSSADominance" bit for the given region. The
@@ -255,6 +261,15 @@ static bool isBeforeInBlock(Block *block, Block::iterator a,
return a->isBeforeInBlock(&*b);
}
+template <bool IsPostDom>
+bool DominanceInfoBase<IsPostDom>::hasBreakingControlFlowOpsCached(
+ Operation *op) const {
+ auto [it, inserted] = breakingControlFlowOpsCache.try_emplace(op, false);
+ if (inserted)
+ it->second = hasBreakingControlFlowOps(op);
+ return it->second;
+}
+
template <bool IsPostDom>
bool DominanceInfoBase<IsPostDom>::properlyDominatesImpl(
Block *aBlock, Block::iterator aIt, Block *bBlock, Block::iterator bIt,
@@ -288,6 +303,23 @@ bool DominanceInfoBase<IsPostDom>::properlyDominatesImpl(
return true;
}
+ auto hasEscapingBreakingControlFlowInRange =
+ [&](Block *block, Block::iterator begin, Block::iterator end) {
+ assert(begin == block->end() || begin->getBlock() == block);
+ assert(end == block->end() || end->getBlock() == block);
+ for (auto it = begin; it != end; ++it) {
+ Operation *op = &*it;
+ if ((op->hasTrait<OpTrait::PropagateControlFlowBreak>() ||
+ !op->isRegistered()) &&
+ hasBreakingControlFlowOpsCached(op)) {
+ LDBG() << "Breaking control flow: "
+ << OpWithFlags(op, OpPrintingFlags().skipRegions());
+ return true;
+ }
+ }
+ return false;
+ };
+
// Ok, they are in the same region now.
if (aBlock == bBlock) {
// Dominance changes based on the region type. In a region with SSA
@@ -295,6 +327,26 @@ bool DominanceInfoBase<IsPostDom>::properlyDominatesImpl(
// regions kinds, uses and defs can come in any order inside a block.
if (!hasSSADominance(aBlock))
return true;
+
+ if (IsPostDom && aIt != aBlock->end() && bIt != bBlock->end() &&
+ mightHaveBreakingControlFlow(aBlock->getParent())) {
+ bool inRange = false;
+ for (Operation &op : *aBlock) {
+ if (&op == &*bIt)
+ inRange = true;
+ if (inRange) {
+ if ((op.hasTrait<OpTrait::PropagateControlFlowBreak>() ||
+ !op.isRegistered()) &&
+ hasBreakingControlFlowOpsCached(&op)) {
+ LDBG() << "Breaking control flow: "
+ << OpWithFlags(&op, OpPrintingFlags().skipRegions());
+ return false;
+ }
+ }
+ if (&op == &*aIt)
+ break;
+ }
+ }
if constexpr (IsPostDom) {
return isBeforeInBlock(aBlock, bIt, aIt);
} else {
@@ -303,7 +355,47 @@ bool DominanceInfoBase<IsPostDom>::properlyDominatesImpl(
}
// If the blocks are different, use DomTree to resolve the query.
- return getDomTree(aRegion).properlyDominates(aBlock, bBlock);
+ bool properlyDominates =
+ getDomTree(aRegion).properlyDominates(aBlock, bBlock);
+ if (!IsPostDom || !properlyDominates ||
+ !mightHaveBreakingControlFlow(aRegion))
+ return properlyDominates;
+
+ // A nested breaking terminator can bypass the containing block's ordinary
+ // CFG terminator, so the block post-dominator tree is not enough for
+ // operation post-dominance. Look for escaping breaks on paths from bBlock to
+ // aBlock.
+ if (bIt != bBlock->end() &&
+ hasEscapingBreakingControlFlowInRange(bBlock, bIt, bBlock->end()))
+ return false;
+
+ SmallPtrSet<Block *, 16> visited;
+ SmallVector<Block *> worklist;
+ auto enqueueSuccessor = [&](Block *successor) {
+ if (!successor || successor == aBlock || successor->getParent() != aRegion)
+ return;
+ if (visited.insert(successor).second)
+ worklist.push_back(successor);
+ };
+ for (Block *successor : bBlock->getSuccessors())
+ enqueueSuccessor(successor);
+
+ while (!worklist.empty()) {
+ Block *block = worklist.pop_back_val();
+ if (hasEscapingBreakingControlFlowInRange(block, block->begin(),
+ block->end()))
+ return false;
+ for (Block *successor : block->getSuccessors())
+ enqueueSuccessor(successor);
+ }
+
+ if (aIt != aBlock->end()) {
+ Block::iterator endIt = std::next(aIt);
+ if (hasEscapingBreakingControlFlowInRange(aBlock, aBlock->begin(), endIt))
+ return false;
+ }
+
+ return true;
}
/// Return true if the specified block is reachable from the entry block of
diff --git a/mlir/lib/IR/RegionKindInterface.cpp b/mlir/lib/IR/RegionKindInterface.cpp
index 007f4cf92dbc7..59c741ea7e11e 100644
--- a/mlir/lib/IR/RegionKindInterface.cpp
+++ b/mlir/lib/IR/RegionKindInterface.cpp
@@ -12,6 +12,12 @@
//===----------------------------------------------------------------------===//
#include "mlir/IR/RegionKindInterface.h"
+#include "mlir/IR/BuiltinTypes.h"
+#include "mlir/Support/WalkResult.h"
+
+#include "llvm/Support/DebugLog.h"
+
+#define DEBUG_TYPE "region-kind-interface"
using namespace mlir;
@@ -32,3 +38,165 @@ bool mlir::mayBeGraphRegion(Region ®ion) {
return false;
return !regionKindOp.hasSSADominance(region.getRegionNumber());
}
+
+namespace {
+// Worklist entry for walking block terminators in nested regions.
+// Tracks the current position within a region and the nesting depth.
+struct NestedOpIterator {
+ NestedOpIterator(Region *region, int nestedLevel)
+ : region(region), nestedLevel(nestedLevel) {
+ regionIt = region->begin();
+ blockIt = regionIt->end();
+ if (regionIt != region->end())
+ blockIt = regionIt->begin();
+ }
+ // Advance the iterator to the next reachable operation.
+ void advance() {
+ assert(regionIt != region->end());
+ if (blockIt == regionIt->end()) {
+ ++regionIt;
+ if (regionIt != region->end())
+ blockIt = regionIt->begin();
+ return;
+ }
+ ++blockIt;
+ if (blockIt != regionIt->end()) {
+ LDBG() << "Advancing to next op: "
+ << OpWithFlags(&*blockIt, OpPrintingFlags().skipRegions());
+ }
+ }
+
+ // The region we're iterating over.
+ Region *region;
+ // The Block currently being iterated over.
+ Region::iterator regionIt;
+ // The Operation currently being iterated over.
+ Block::iterator blockIt;
+ // The nested level of the current region relative to the starting region.
+ int nestedLevel = 0;
+};
+} // namespace
+
+/// Walk all block terminators (last operation in each block) nested within
+/// `rootOp`. The callback receives the terminator and its 1-based nesting
+/// level relative to `rootOp`. Callers filter based on breaking-control-flow
+/// properties as needed.
+static void walk(Operation *rootOp,
+ function_ref<WalkResult(Operation *, int)> callback) {
+ // Worklist of regions to visit to drive the traversal.
+ SmallVector<NestedOpIterator> worklist;
+
+ // Perform a traversal of the regions, visiting each
+ // reachable operation.
+ for (Region ®ion : rootOp->getRegions()) {
+ if (region.empty())
+ continue;
+ worklist.push_back({®ion, 1});
+ }
+ while (!worklist.empty()) {
+ NestedOpIterator &it = worklist.back();
+ if (it.regionIt == it.region->end()) {
+ // We're done with this region.
+ worklist.pop_back();
+ continue;
+ }
+ if (it.blockIt == it.regionIt->end()) {
+ // We're done with this block.
+ it.advance();
+ continue;
+ }
+ Operation *op = &*it.blockIt;
+
+ // Only call the callback if we're at the end of the block.
+ if (std::next(it.blockIt) == it.regionIt->end() &&
+ callback(op, it.nestedLevel).wasInterrupted())
+ return;
+
+ // Advance before pushing nested regions to avoid reference invalidation.
+ int currentNestedLevel = it.nestedLevel;
+ it.advance();
+
+ // Recursively visit the nested regions.
+ for (Region &nestedRegion : op->getRegions()) {
+ if (nestedRegion.empty())
+ continue;
+ worklist.push_back({&nestedRegion, currentNestedLevel + 1});
+ }
+ }
+}
+
+NestedBreakingControlFlowInfo
+mlir::getNestedBreakingControlFlowInfo(Operation *op) {
+ NestedBreakingControlFlowInfo info;
+ walk(op, [&](Operation *visitedOp, int nestedLevel) {
+ bool directlyTargetsOp = false;
+ bool escapesThroughOp = false;
+ for (HasBreakingControlFlowOpInterface target :
+ findPotentialBreakTargets(visitedOp)) {
+ if (!target)
+ continue;
+ Operation *targetOp = target.getOperation();
+ if (targetOp == op)
+ directlyTargetsOp = true;
+ else if (targetOp->isProperAncestor(op))
+ escapesThroughOp = true;
+ }
+
+ if (directlyTargetsOp) {
+ info.predecessors.push_back(visitedOp);
+ if (nestedLevel > 1)
+ info.hasNestedPredecessors = true;
+ }
+ if (escapesThroughOp)
+ info.hasBreakingControlFlowOps = true;
+
+ return WalkResult::advance();
+ });
+ return info;
+}
+
+bool mlir::hasNestedPredecessors(Operation *op) {
+ return getNestedBreakingControlFlowInfo(op).hasNestedPredecessors;
+}
+
+bool mlir::hasBreakingControlFlowOps(Operation *op) {
+ return getNestedBreakingControlFlowInfo(op).hasBreakingControlFlowOps;
+}
+
+void mlir::detail::visitNestedBreakingControlFlowOpsImpl(
+ Operation *op,
+ function_ref<WalkResult(RegionExitTerminatorOpInterface, int nestedLevel)>
+ callback) {
+ ::walk(op, [&](Operation *visitedOp, int nestedLevel) {
+ for (HasBreakingControlFlowOpInterface target :
+ findPotentialBreakTargets(visitedOp)) {
+ if (target && (target.getOperation() == op ||
+ target.getOperation()->isProperAncestor(op)))
+ return callback(cast<RegionExitTerminatorOpInterface>(visitedOp),
+ nestedLevel);
+ }
+ return WalkResult::advance();
+ });
+}
+
+void mlir::collectAllNestedPredecessors(
+ Operation *op, SmallVector<Operation *> &predecessors) {
+ llvm::append_range(predecessors,
+ getNestedBreakingControlFlowInfo(op).predecessors);
+}
+
+SmallVector<HasBreakingControlFlowOpInterface>
+mlir::findPotentialBreakTargets(Operation *terminator) {
+ auto breakingTerminator =
+ dyn_cast<RegionExitTerminatorOpInterface>(terminator);
+ if (!breakingTerminator)
+ return {};
+ SmallVector<HasBreakingControlFlowOpInterface> targets;
+ for (Operation *target : breakingTerminator.getPotentialTargets()) {
+ if (!target)
+ continue;
+ if (auto breakTarget = dyn_cast<HasBreakingControlFlowOpInterface>(target))
+ targets.push_back(breakTarget);
+ }
+ return targets;
+}
diff --git a/mlir/lib/IR/Verifier.cpp b/mlir/lib/IR/Verifier.cpp
index 11771e78d5f20..4209a36b1c560 100644
--- a/mlir/lib/IR/Verifier.cpp
+++ b/mlir/lib/IR/Verifier.cpp
@@ -30,6 +30,7 @@
#include "mlir/IR/Dialect.h"
#include "mlir/IR/Dominance.h"
#include "mlir/IR/Operation.h"
+#include "mlir/IR/OperationSupport.h"
#include "mlir/IR/RegionKindInterface.h"
#include "mlir/IR/Threading.h"
#include "llvm/ADT/PointerIntPair.h"
diff --git a/mlir/lib/Interfaces/ControlFlowInterfaces.cpp b/mlir/lib/Interfaces/ControlFlowInterfaces.cpp
index 475969a7c0d09..241e447a3e7f1 100644
--- a/mlir/lib/Interfaces/ControlFlowInterfaces.cpp
+++ b/mlir/lib/Interfaces/ControlFlowInterfaces.cpp
@@ -13,6 +13,7 @@
#include "mlir/IR/Matchers.h"
#include "mlir/IR/Operation.h"
#include "mlir/IR/PatternMatch.h"
+#include "mlir/IR/RegionKindInterface.h"
#include "mlir/Interfaces/ControlFlowInterfaces.h"
#include "llvm/ADT/EquivalenceClasses.h"
#include "llvm/Support/DebugLog.h"
@@ -76,14 +77,14 @@ detail::getBranchSuccessorArgument(const SuccessorOperands &operands,
LogicalResult
detail::verifyBranchSuccessorOperands(Operation *op, unsigned succNo,
const SuccessorOperands &operands) {
- LDBG() << "Verifying branch successor operands for successor #" << succNo
- << " in operation " << op->getName();
+ LDBG(3) << "Verifying branch successor operands for successor #" << succNo
+ << " in operation " << op->getName();
// Check the count.
unsigned operandCount = operands.size();
Block *destBB = op->getSuccessor(succNo);
- LDBG() << "Branch has " << operandCount << " operands, target block has "
- << destBB->getNumArguments() << " arguments";
+ LDBG(3) << "Branch has " << operandCount << " operands, target block has "
+ << destBB->getNumArguments() << " arguments";
if (operandCount != destBB->getNumArguments())
return op->emitError() << "branch has " << operandCount
@@ -92,22 +93,22 @@ detail::verifyBranchSuccessorOperands(Operation *op, unsigned succNo,
<< destBB->getNumArguments();
// Check the types.
- LDBG() << "Checking type compatibility for "
- << (operandCount - operands.getProducedOperandCount())
- << " forwarded operands";
+ LDBG(3) << "Checking type compatibility for "
+ << (operandCount - operands.getProducedOperandCount())
+ << " forwarded operands";
for (unsigned i = operands.getProducedOperandCount(); i != operandCount;
++i) {
Type operandType = operands[i].getType();
Type argType = destBB->getArgument(i).getType();
- LDBG() << "Checking type compatibility: operand type " << operandType
- << " vs argument type " << argType;
+ LDBG(3) << "Checking type compatibility: operand type " << operandType
+ << " vs argument type " << argType;
if (!cast<BranchOpInterface>(op).areTypesCompatible(operandType, argType))
return op->emitError() << "type mismatch for bb argument #" << i
<< " of successor #" << succNo;
}
- LDBG() << "Branch successor operand verification successful";
+ LDBG(3) << "Branch successor operand verification successful";
return success();
}
@@ -168,11 +169,17 @@ LogicalResult detail::verifyRegionBranchOpInterface(Operation *op) {
SmallVector<RegionSuccessor> successors;
regionInterface.getSuccessorRegions(branchPoint, successors);
for (const RegionSuccessor &successor : successors) {
+ RegionBranchOpInterface successorOwner =
+ getRegionBranchSuccessorOwner(successor);
+ if (!successorOwner)
+ return regionInterface->emitOpError()
+ << "has region successor not owned by a "
+ "RegionBranchOpInterface";
// Helper function that print the region branch point and the region
// successor.
auto emitRegionEdgeError = [&]() {
InFlightDiagnostic diag =
- regionInterface->emitOpError("along control flow edge from ");
+ successorOwner->emitOpError("along control flow edge from ");
if (branchPoint.isParent()) {
diag << "parent";
diag.attachNote(op->getLoc()) << "region branch point";
@@ -195,7 +202,7 @@ LogicalResult detail::verifyRegionBranchOpInterface(Operation *op) {
// Verify number of successor operands and successor inputs.
OperandRange succOperands =
regionInterface.getSuccessorOperands(branchPoint, successor);
- ValueRange succInputs = regionInterface.getSuccessorInputs(successor);
+ ValueRange succInputs = successorOwner.getSuccessorInputs(successor);
if (succOperands.size() != succInputs.size()) {
return emitRegionEdgeError()
<< ": region branch point has " << succOperands.size()
@@ -210,7 +217,7 @@ LogicalResult detail::verifyRegionBranchOpInterface(Operation *op) {
llvm::enumerate(llvm::zip(succOperandTypes, succInputTypes))) {
Type succOperandType = std::get<0>(typesIdx.value());
Type succInputType = std::get<1>(typesIdx.value());
- if (!regionInterface.areTypesCompatible(succOperandType, succInputType))
+ if (!successorOwner.areTypesCompatible(succOperandType, succInputType))
return emitRegionEdgeError()
<< ": successor operand type #" << typesIdx.index() << " "
<< succOperandType << " should match successor input type #"
@@ -227,6 +234,17 @@ LogicalResult detail::verifyRegionBranchOpInterface(Operation *op) {
/// regions.
using StopConditionFn = function_ref<bool(Region *, ArrayRef<bool> visited)>;
+/// Given a range of values, return a vector of attributes of the same size,
+/// where the i-th attribute is the constant value of the i-th value. If a
+/// value is not constant, the corresponding attribute is null.
+static SmallVector<Attribute> extractConstants(ValueRange values) {
+ return llvm::map_to_vector(values, [](Value value) {
+ Attribute attr;
+ matchPattern(value, m_Constant(&attr));
+ return attr;
+ });
+}
+
/// Traverse the region graph starting at `begin`. The traversal is interrupted
/// if `stopCondition` evaluates to "true" for a successor region. In that case,
/// this function returns "true". Otherwise, if the traversal was not
@@ -246,7 +264,6 @@ static bool traverseRegionGraph(Region *begin,
SmallVector<Region *> worklist;
auto enqueueAllSuccessors = [&](Region *region) {
LDBG() << "Enqueuing successors for region #" << region->getRegionNumber();
- SmallVector<Attribute> operandAttributes(op->getNumOperands());
for (Block &block : *region) {
if (block.empty())
continue;
@@ -255,8 +272,8 @@ static bool traverseRegionGraph(Region *begin,
if (!terminator)
continue;
SmallVector<RegionSuccessor> successors;
- operandAttributes.resize(terminator->getNumOperands());
- terminator.getSuccessorRegions(operandAttributes, successors);
+ terminator.getSuccessorRegions(
+ extractConstants(terminator->getOperands()), successors);
LDBG() << "Found " << successors.size()
<< " successors from terminator in block";
for (RegionSuccessor successor : successors) {
@@ -462,6 +479,14 @@ RegionBranchOpInterface::getNonSuccessorInputs(RegionSuccessor successor) {
return results;
}
+RegionBranchOpInterface
+mlir::getRegionBranchSuccessorOwner(RegionSuccessor successor) {
+ if (Operation *successorOp = successor.getSuccessorOp())
+ return dyn_cast<RegionBranchOpInterface>(successorOp);
+ return dyn_cast<RegionBranchOpInterface>(
+ successor.getSuccessor()->getParentOp());
+}
+
static MutableArrayRef<OpOperand> operandsToOpOperands(OperandRange &operands) {
return MutableArrayRef<OpOperand>(operands.getBase(), operands.size());
}
@@ -473,11 +498,14 @@ getSuccessorOperandInputMapping(RegionBranchOpInterface branchOp,
SmallVector<RegionSuccessor> successors;
branchOp.getSuccessorRegions(src, successors);
for (RegionSuccessor dst : successors) {
+ RegionBranchOpInterface successorOwner = getRegionBranchSuccessorOwner(dst);
+ assert(successorOwner && "expected RegionBranchOpInterface owner");
OperandRange operands = branchOp.getSuccessorOperands(src, dst);
- assert(operands.size() == branchOp.getSuccessorInputs(dst).size() &&
+ ValueRange inputs = successorOwner.getSuccessorInputs(dst);
+ assert(operands.size() == inputs.size() &&
"expected the same number of operands and inputs");
- for (const auto &[operand, input] : llvm::zip_equal(
- operandsToOpOperands(operands), branchOp.getSuccessorInputs(dst)))
+ for (const auto &[operand, input] :
+ llvm::zip_equal(operandsToOpOperands(operands), inputs))
mapping[&operand].push_back(input);
}
}
@@ -524,6 +552,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);
+ if (!terminator)
+ return;
+ branchPoints.push_back(RegionBranchPoint(terminator));
+ });
return branchPoints;
}
@@ -628,6 +712,12 @@ static bool isDefinedBefore(Operation *regionBranchOp, Value a, Value b) {
return true;
}
+static bool isSuccessorInputOwnedBy(Operation *regionBranchOp, Value value) {
+ if (Operation *definingOp = value.getDefiningOp())
+ return definingOp == regionBranchOp;
+ return value.getParentRegion()->getParentOp() == regionBranchOp;
+}
+
/// Compute all non-successor-input values that a successor input could have
/// based on the given successor input to successor operand mapping.
///
@@ -733,6 +823,8 @@ struct MakeRegionBranchOpSuccessorInputsDead : public RewritePattern {
// Try to replace the uses of each successor input one-by-one.
bool changed = false;
for (Value value : inputToOperands.keys()) {
+ if (!isSuccessorInputOwnedBy(regionBranchOp, value))
+ continue;
// Nothing to do for successor inputs that are already dead.
if (value.use_empty())
continue;
@@ -1070,17 +1162,6 @@ struct RemoveDuplicateSuccessorInputUses : public RewritePattern {
}
};
-/// Given a range of values, return a vector of attributes of the same size,
-/// where the i-th attribute is the constant value of the i-th value. If a
-/// value is not constant, the corresponding attribute is null.
-static SmallVector<Attribute> extractConstants(ValueRange values) {
- return llvm::map_to_vector(values, [](Value value) {
- Attribute attr;
- matchPattern(value, m_Constant(&attr));
- return attr;
- });
-}
-
/// Return all successor regions when branching from the given region branch
/// point. This helper functions extracts all constant operand values and
/// passes them to the `RegionBranchOpInterface`.
diff --git a/mlir/lib/Transforms/Utils/CMakeLists.txt b/mlir/lib/Transforms/Utils/CMakeLists.txt
index 3ecf3859a7f4c..cb84fb26b2408 100644
--- a/mlir/lib/Transforms/Utils/CMakeLists.txt
+++ b/mlir/lib/Transforms/Utils/CMakeLists.txt
@@ -15,6 +15,9 @@ add_mlir_library(MLIRTransformUtils
ADDITIONAL_HEADER_DIRS
${MLIR_MAIN_INCLUDE_DIR}/mlir/Transforms
+ DEPENDS
+ MLIRRegionKindInterfaceIncGen
+
LINK_LIBS PUBLIC
MLIRAnalysis
MLIRCallInterfaces
diff --git a/mlir/lib/Transforms/Utils/InliningUtils.cpp b/mlir/lib/Transforms/Utils/InliningUtils.cpp
index 73107cfc36ea9..4bfc20e7ed885 100644
--- a/mlir/lib/Transforms/Utils/InliningUtils.cpp
+++ b/mlir/lib/Transforms/Utils/InliningUtils.cpp
@@ -286,6 +286,12 @@ static LogicalResult inlineRegionImpl(
[&](BlockArgument arg) { return !mapper.contains(arg); }))
return failure();
+ // Block inlining only if breaks escape the region (propagate through the
+ // parent op toward an ancestor). Self-contained breaks that target ops
+ // within the region are fine.
+ if (hasBreakingControlFlowOps(src->getParentOp()))
+ return failure();
+
// Check that the operations within the source region are valid to inline.
Region *insertRegion = inlineBlock->getParent();
if (!interface.isLegalToInline(insertRegion, src, shouldCloneInlinedRegion,
diff --git a/mlir/test/Analysis/DataFlow/test-dead-code-analysis-early-exit.mlir b/mlir/test/Analysis/DataFlow/test-dead-code-analysis-early-exit.mlir
new file mode 100644
index 0000000000000..49338ac634462
--- /dev/null
+++ b/mlir/test/Analysis/DataFlow/test-dead-code-analysis-early-exit.mlir
@@ -0,0 +1,100 @@
+// RUN: mlir-opt -test-dead-code-analysis --split-input-file 2>&1 %s | FileCheck %s
+
+// Tests verifying that DeadCodeAnalysis correctly propagates breaking control
+// flow through PropagateControlFlowBreak ops (e.g. scf.if). The fix is in
+// visitRegionTerminator: when the immediate parent RegionBranchOpInterface
+// returns a propagating successor sentinel for a propagating break, we resolve
+// the actual HasBreakingControlFlowOp ancestor and re-dispatch through it, so
+// the correct predecessor edge is established at the loop's exit point.
+
+// -----
+
+// A loop whose only exit is scf.break propagated through scf.if.
+// The loop's exit point (op_preds) must show the break as a predecessor.
+
+// CHECK-LABEL: loop:
+// CHECK: region #0
+// CHECK: ^bb0 = live
+// CHECK: region_preds: (all) predecessors:
+// CHECK: %0 = scf.loop
+// CHECK: scf.continue
+// CHECK: op_preds: (all) predecessors:
+// CHECK: scf.break
+func.func @test_break_through_if(%cond: i1) -> i32 {
+ %result = scf.loop token(%loop) -> i32 {
+ scf.if %cond {
+ %c42 = arith.constant 42 : i32
+ scf.break [%loop] %c42 : i32
+ }
+ scf.continue [%loop]
+ } {tag = "loop"}
+ return %result : i32
+}
+
+// -----
+
+// For comparison: a loop with a direct scf.break.
+
+// CHECK-LABEL: loop:
+// CHECK: region #0
+// CHECK: ^bb0 = live
+// CHECK: region_preds: (all) predecessors:
+// CHECK: %0 = scf.loop
+// CHECK: op_preds: (all) predecessors:
+// CHECK: scf.break
+func.func @test_direct_break(%cond: i1) -> i32 {
+ %result = scf.loop token(%loop) -> i32 {
+ scf.if %cond {
+ }
+ %c42 = arith.constant 42 : i32
+ scf.break [%loop] %c42 : i32
+ } {tag = "loop"}
+ return %result : i32
+}
+
+// -----
+
+// A loop whose only exit is scf.break propagated through two nested scf.if
+// ops. The loop's op_preds must show the break. The intermediate if ops do not
+// appear in op_preds for the break — the break bypasses them without yielding
+// values to them, which is correct.
+
+// CHECK-LABEL: inner_if:
+// CHECK: region #0
+// CHECK: ^bb0 = live
+// CHECK: region_preds: (all) predecessors:
+// CHECK: scf.if {{.*}} {tag = "inner_if"}
+// CHECK: region #1
+// CHECK: op_preds: (all) predecessors:
+// CHECK: scf.if {{.*}} {tag = "inner_if"}
+
+// CHECK-LABEL: outer_if:
+// CHECK: region #0
+// CHECK: ^bb0 = live
+// CHECK: region_preds: (all) predecessors:
+// CHECK: scf.if {{.*}} {tag = "outer_if"}
+// CHECK: region #1
+// CHECK: op_preds: (all) predecessors:
+// CHECK: scf.if {{.*}} {tag = "outer_if"}
+
+// CHECK-LABEL: loop:
+// CHECK: region #0
+// CHECK: ^bb0 = live
+// CHECK: region_preds: (all) predecessors:
+// CHECK: %0 = scf.loop
+// CHECK: scf.continue
+// CHECK: op_preds: (all) predecessors:
+// CHECK: scf.break
+func.func @test_break_through_nested_ifs(%cond1: i1, %cond2: i1) -> i32 {
+ %result = scf.loop token(%loop) -> i32 {
+ scf.if %cond1 {
+ scf.if %cond2 {
+ %c99 = arith.constant 99 : i32
+ scf.break [%loop] %c99 : i32
+ } {tag = "inner_if"}
+ scf.continue [%loop]
+ } {tag = "outer_if"}
+ scf.continue [%loop]
+ } {tag = "loop"}
+ return %result : i32
+}
diff --git a/mlir/test/Analysis/test-dominance.mlir b/mlir/test/Analysis/test-dominance.mlir
index a926a8271200a..5627c9994b0e8 100644
--- a/mlir/test/Analysis/test-dominance.mlir
+++ b/mlir/test/Analysis/test-dominance.mlir
@@ -680,3 +680,63 @@ func.func @func_loop_nested_region(
// CHECK: ^{{.*}}
// CHECK: }
// CHECK: }
+
+
+// -----
+
+// CHECK-LABEL: Testing : func_loop_early_exit
+func.func @func_loop_early_exit(%cond : i1, %arg0 : index) -> index {
+ %0 = scf.loop token(%outer) -> index {
+ scf.loop token(%inner) {
+ scf.if %cond {
+ scf.break [%outer] %arg0 : index
+ }
+ "test.foo"() : () -> ()
+ scf.break [%inner] {test.print_dominance = true}
+ }
+ }
+ return %0 : index
+}
+
+// CHECK: postdominates(scf.break {{.*}} {test.print_dominance = true} {{.*}} scf.break
+// CHECK-SAME: = 0
+// CHECK: postdominates(scf.break {{.*}} {test.print_dominance = true} {{.*}} scf.if
+// CHECK-SAME: = 0
+// CHECK: postdominates(scf.break {{.*}} {test.print_dominance = true} {{.*}} "test.foo"
+// CHECK-SAME: = 1
+// CHECK: postdominates(scf.break {{.*}} {test.print_dominance = true} {{.*}} scf.break
+// CHECK-SAME: = 1
+
+// -----
+
+// CHECK-LABEL: Testing : func_loop_early_exit_cross_block
+func.func @func_loop_early_exit_cross_block(%cond : i1, %arg0 : index) -> index {
+ %0 = scf.loop token(%outer) -> index {
+ scf.loop token(%inner) {
+ test.propagate_control_flow_break {
+ ^bb0:
+ cf.cond_br %cond, ^bb1, ^bb2
+ ^bb1:
+ scf.if %cond {
+ scf.break [%outer] %arg0 : index
+ }
+ cf.br ^bb2
+ ^bb2:
+ "test.foo"() {test.print_dominance = true} : () -> ()
+ "test.return"() : () -> ()
+ }
+ scf.break [%inner]
+ }
+ }
+ return %0 : index
+}
+
+// CHECK: --- PostDominanceInfo ---
+// CHECK: postdominates("test.foo"() {test.print_dominance = true} {{.*}} cf.cond_br
+// CHECK-SAME: = 0
+// CHECK: postdominates("test.foo"() {test.print_dominance = true} {{.*}} scf.if
+// CHECK-SAME: = 0
+// CHECK: postdominates("test.foo"() {test.print_dominance = true} {{.*}} cf.br
+// CHECK-SAME: = 1
+// CHECK: postdominates("test.foo"() {test.print_dominance = true} {{.*}} "test.foo"
+// CHECK-SAME: = 1
diff --git a/mlir/test/Conversion/SCFToControlFlow/convert-early-exit-to-cfg.mlir b/mlir/test/Conversion/SCFToControlFlow/convert-early-exit-to-cfg.mlir
new file mode 100644
index 0000000000000..c5f5a05e47c46
--- /dev/null
+++ b/mlir/test/Conversion/SCFToControlFlow/convert-early-exit-to-cfg.mlir
@@ -0,0 +1,135 @@
+// RUN: mlir-opt -convert-scf-to-cf -split-input-file %s | FileCheck %s
+
+
+func.func @loop_break(%cond : i1) {
+ // CHECK: test.op1
+ "test.op1"() : () -> ()
+ // CHECK-NEXT: cf.br [[LOOP1_ENTRY:.*]]
+ // CHECK-NEXT: [[LOOP1_ENTRY]]
+ scf.loop token(%loop) {
+ // CHECK-NEXT: test.op2
+ "test.op2"() : () -> ()
+ // CHECK-NEXT: cf.cond_br %arg0, [[IF_ENTRY:.*]], [[IF_CONTINUE:.*]]
+ // CHECK-NEXT: [[IF_ENTRY]]
+ scf.if %cond {
+ "test.op3"() : () -> ()
+ scf.break [%loop] loc("break1")
+ }
+ "test.op3"() : () -> ()
+ } loc("loop1")
+ "test.op4"() : () -> ()
+ return
+}
+
+// -----
+
+// Bug regression test: IfLowering was using thenTerminator on line 461 instead
+// of elseTerminator, causing the else block's scf.yield to not be replaced with
+// a cf.branch when the then-branch has a scf.break.
+// CHECK-LABEL: func @if_break_then_yield_else
+func.func @if_break_then_yield_else(%cond : i1) {
+ // CHECK: cf.br ^[[LOOP:.*]]
+ scf.loop token(%loop) {
+ // CHECK: ^[[LOOP]]:
+ // CHECK: cf.cond_br %arg0, ^[[THEN:.*]], ^[[ELSE:.*]]
+ scf.if %cond {
+ // CHECK: ^[[THEN]]:
+ // CHECK-NEXT: cf.br ^[[EXIT:.*]]
+ scf.break [%loop]
+ } else {
+ // CHECK: ^[[ELSE]]:
+ // CHECK-NEXT: cf.br ^[[AFTER_IF:.*]]
+ }
+ // CHECK: ^[[AFTER_IF]]:
+ // CHECK-NEXT: cf.br ^[[LOOP]]
+ scf.continue [%loop]
+ }
+ // CHECK: ^[[EXIT]]:
+ // CHECK-NEXT: return
+ return
+}
+
+// -----
+
+// CHECK-LABEL: func @nested_loops_and_ifs(
+// CHECK-SAME: %[[COND1:.*]]: i1,
+// CHECK-SAME: %[[COND2:.*]]: i1
+func.func @nested_loops_and_ifs(%cond1 : i1, %cond2 : i1) {
+ // CHECK: test.op1
+ "test.op1"() : () -> ()
+ // CHECK-NEXT: cf.br ^[[OUTER_LOOP_ENTRY:.*]]
+ scf.loop token(%outer) {
+ // CHECK-NEXT: ^[[OUTER_LOOP_ENTRY]]:
+ // CHECK-NEXT: cf.cond_br %[[COND1]], ^[[IF1_THEN_BLOCK:.*]], ^[[IF1_EXIT:.*]]
+ scf.if %cond1 {
+ // CHECK-NEXT: ^[[IF1_THEN_BLOCK]]:
+ // CHECK-NEXT: test.op2
+ "test.op2"() : () -> ()
+ // CHECK-NEXT: cf.br ^[[INNER_LOOP_ENTRY:.*]]
+ scf.loop token(%inner) {
+ // CHECK-NEXT: ^[[INNER_LOOP_ENTRY]]:
+ // CHECK-NEXT: test.op3
+ "test.op3"() : () -> ()
+ // CHECK-NEXT: cf.cond_br %[[COND1]], ^[[IF2_THEN_BLOCK:.*]], ^[[IF2_EXIT:.*]]
+ scf.if %cond1 {
+ // CHECK-NEXT: ^[[IF2_THEN_BLOCK]]:
+ // CHECK-NEXT: test.op4
+ "test.op4"() : () -> ()
+ // CHECK-NEXT: cf.br ^[[INNER_LOOP_ENTRY]]
+ scf.continue [%inner] loc("continue1")
+ }
+ // CHECK-NEXT: ^[[IF2_EXIT]]:
+ // CHECK-NEXT: test.op5
+ "test.op5"() : () -> ()
+ // CHECK-NEXT: cf.cond_br %[[COND2]], ^[[IF3_THEN_BLOCK:.*]], ^[[IF3_EXIT:.*]]
+ scf.if %cond2 {
+ // CHECK-NEXT: ^[[IF3_THEN_BLOCK]]:
+ // CHECK-NEXT: test.op6
+ "test.op6"() : () -> ()
+ // CHECK-NEXT: cf.br ^[[FUNC_EXIT:.*]]
+ scf.break [%outer] loc("break2")
+ }
+ // CHECK-NEXT: ^[[IF3_EXIT]]:
+ // CHECK-NEXT: test.op7
+ "test.op7"() : () -> ()
+ // CHECK-NEXT: cf.br ^[[INNER_LOOP_ENTRY]]
+ scf.continue [%inner] loc("continue2")
+ } loc("loop3")
+ // CHECK-NEXT: ^[[AFTER_INNER_LOOP:.*]]:
+ // CHECK-NEXT: test.op8
+ "test.op8"() : () -> ()
+ // CHECK-NEXT: cf.br ^[[IF1_EXIT]]
+ } loc("if1")
+ // CHECK-NEXT: ^[[IF1_EXIT]]:
+ // CHECK-NEXT: cf.br ^[[OUTER_LOOP_ENTRY]]
+ scf.continue [%outer] loc("continue3")
+ } loc("loop2")
+ // CHECK-NEXT: ^[[FUNC_EXIT]]:
+ // CHECK-NEXT: test.op9
+ "test.op9"() : () -> ()
+ // CHECK-NEXT: return
+ return
+}
+
+// -----
+
+// CHECK-LABEL: func @loop_with_iter_args
+// CHECK-SAME: %[[INIT:.*]]: i32, %[[COND:.*]]: i1
+func.func @loop_with_iter_args(%init: i32, %cond: i1) -> i32 {
+ // CHECK: cf.br ^[[LOOP:.*]](%[[INIT]] : i32)
+ %result = scf.loop token(%loop) iter_args(%arg = %init) : i32 -> i32 {
+ // CHECK: ^[[LOOP]](%[[ARG:.*]]: i32):
+ // CHECK: cf.cond_br %[[COND]], ^[[THEN:.*]], ^[[ELSE:.*]]
+ scf.if %cond {
+ // CHECK: ^[[THEN]]:
+ // CHECK-NEXT: cf.br ^[[EXIT:.*]](%[[ARG]] : i32)
+ scf.break [%loop] %arg : i32
+ }
+ // CHECK: ^[[ELSE]]:
+ // CHECK-NEXT: cf.br ^[[LOOP]](%[[ARG]] : i32)
+ scf.continue [%loop] %arg : i32
+ }
+ // CHECK: ^[[EXIT]](%[[RES:.*]]: i32):
+ // CHECK-NEXT: return %[[RES]]
+ return %result : i32
+}
diff --git a/mlir/test/Dialect/SCF/loop_canonicalize.mlir b/mlir/test/Dialect/SCF/loop_canonicalize.mlir
new file mode 100644
index 0000000000000..90e2068290082
--- /dev/null
+++ b/mlir/test/Dialect/SCF/loop_canonicalize.mlir
@@ -0,0 +1,300 @@
+// RUN: mlir-opt %s -pass-pipeline='builtin.module(func.func(canonicalize{test-convergence}))' -split-input-file | FileCheck %s
+
+// CHECK-LABEL: func @fold_single_iteration_loop1
+func.func @fold_single_iteration_loop1(%arg0 : index) -> index {
+ // CHECK-NOT: loop
+ %0 = scf.loop token(%loop) -> index {
+ scf.break [%loop] %arg0 : index
+ }
+ return %0 : index
+}
+
+// -----
+
+// CHECK-LABEL: func @fold_single_iteration_loop_with_propagating_control_flow
+func.func @fold_single_iteration_loop_with_propagating_control_flow(%cond : i1, %arg0 : index) -> index {
+ %0 = scf.loop token(%outer) -> index {
+ scf.loop token(%inner) {
+ scf.if %cond {
+ scf.break [%outer] %arg0 : index
+ }
+ scf.break [%inner]
+ }
+ }
+ return %0 : index
+}
+
+// -----
+
+// CHECK-LABEL: func @loop_not_combine_ifs
+func.func @loop_not_combine_ifs(%arg0 : i1, %arg2: i64) {
+ // Verify that we don't combine ifs when terminator mismatches
+ scf.loop token(%loop) {
+ // CHECK: scf.if
+ %res = scf.if %arg0 -> i32 {
+ %v = "test.firstCodeTrue"() : () -> i32
+ scf.yield %v : i32
+ } else {
+ %v2 = "test.firstCodeFalse"() : () -> i32
+ scf.break [%loop]
+ }
+ // CHECK: scf.if
+ %res2 = scf.if %arg0 -> i32 {
+ %v = "test.secondCodeTrue"() : () -> i32
+ scf.yield %v : i32
+ } else {
+ %v2 = "test.secondCodeFalse"() : () -> i32
+ scf.continue [%loop]
+ }
+ }
+ return
+}
+
+// -----
+
+// CHECK-LABEL: func @loop_combine_ifs
+func.func @loop_combine_ifs(%arg0 : i1, %arg2: i64) {
+ // CombineIfs is yield-only, so breaking terminators remain separate.
+ scf.loop token(%loop) {
+ // CHECK: scf.if
+ %res = scf.if %arg0 -> i32 {
+ %v = "test.firstCodeTrue"() : () -> i32
+ scf.yield %v : i32
+ } else {
+ %v2 = "test.firstCodeFalse"() : () -> i32
+ scf.break [%loop]
+ }
+ %res2 = scf.if %arg0 -> i32 {
+ %v = "test.secondCodeTrue"() : () -> i32
+ scf.yield %v : i32
+ } else {
+ %v2 = "test.secondCodeFalse"() : () -> i32
+ scf.break [%loop]
+ }
+ }
+ return
+}
+
+// -----
+
+// CHECK-LABEL: @do_not_merge_nested_if_with_breaking_control_flow1
+func.func @do_not_merge_nested_if_with_breaking_control_flow1(%arg0: i1, %arg1: i1) {
+// The outer if then terminator isn't a yield, blocking the merge.
+// CHECK: scf.loop
+// CHECK: scf.if
+// CHECK: scf.if
+ scf.loop token(%loop) {
+ scf.if %arg0 {
+ scf.if %arg1 {
+ "test.op"() : () -> ()
+ scf.yield
+ }
+ scf.break [%loop]
+ }
+ }
+ return
+}
+
+// -----
+
+// CHECK-LABEL: @do_not_merge_nested_if_with_breaking_control_flow2
+func.func @do_not_merge_nested_if_with_breaking_control_flow2(%arg0: i1, %arg1: i1) {
+// The outer if else block terminator isn't a yield, blocking the merge.
+// CHECK: scf.loop
+// CHECK: scf.if
+// CHECK: scf.if
+// CHECK: else
+// CHECK-NEXT: scf.break
+ scf.loop token(%loop) {
+ scf.if %arg0 {
+ scf.if %arg1 {
+ "test.op"() : () -> ()
+ scf.yield
+ }
+ scf.yield
+ } else {
+ scf.break [%loop]
+ }
+ }
+ return
+}
+
+// -----
+
+// CHECK-LABEL: @do_not_merge_nested_if_with_breaking_control_flow3
+func.func @do_not_merge_nested_if_with_breaking_control_flow3(%arg0: i1, %arg1: i1) {
+// The nested if then block terminator isn't a yield, blocking the merge.
+// CHECK: scf.loop
+// CHECK: scf.if
+// CHECK: scf.if
+// CHECK: test.op
+// CHECK-NEXT: scf.break
+ scf.loop token(%loop) {
+ scf.if %arg0 {
+ scf.if %arg1 {
+ "test.op"() : () -> ()
+ scf.break [%loop]
+ }
+ }
+ }
+ return
+}
+
+// -----
+
+// CHECK-LABEL: @do_not_merge_nested_if_with_breaking_control_flow4
+func.func @do_not_merge_nested_if_with_breaking_control_flow4(%arg0: i1, %arg1: i1) {
+// The nested if else block terminator isn't a yield, blocking the merge.
+// CHECK: scf.loop
+// CHECK: scf.if
+// CHECK: scf.if
+// CHECK: else
+// CHECK-NEXT: scf.break
+ scf.loop token(%loop) {
+ scf.if %arg0 {
+ scf.if %arg1 {
+ "test.op"() : () -> ()
+ } else {
+ scf.break [%loop]
+ }
+ }
+ }
+ return
+}
+
+// -----
+
+// CHECK-LABEL: func @do_not_convert_if_to_select1
+func.func @do_not_convert_if_to_select1(%cond: i1, %arg0 : index, %arg1 : index) -> index {
+ %loop_res = scf.loop token(%loop) -> index {
+ // Inner then terminator is not a yield, blocking transform to select.
+ // CHECK: scf.if
+ %0 = scf.if %cond -> index {
+ scf.break [%loop] %arg0 : index
+ } else {
+ scf.yield %arg1 : index
+ }
+ scf.break [%loop] %0 : index
+ }
+ return %loop_res : index
+}
+
+// -----
+
+// CHECK-LABEL: func @do_not_convert_if_to_select2
+func.func @do_not_convert_if_to_select2(%cond: i1, %arg0 : index, %arg1 : index) -> index {
+ %loop_res = scf.loop token(%loop) -> index {
+ // Inner then terminator is not a yield, blocking transform to select.
+ // CHECK: scf.if
+ %0 = scf.if %cond -> index {
+ scf.yield %arg0 : index
+ } else {
+ scf.break [%loop] %arg1 : index
+ }
+ scf.break [%loop] %0 : index
+ }
+ return %loop_res : index
+}
+
+
+// -----
+
+// Verify that removing the unused results of an if with nested breaking control flow
+// operation works.
+// CHECK-LABEL: func @remove_unused_if_results1
+func.func @remove_unused_if_results1(%cond : i1, %arg0 : index) -> index {
+ // CHECK: scf.loop
+ %0 = scf.loop token(%loop) -> index {
+ // CHECK: %[[FOO:.*]]:3 = "test.foo"
+ %foo:3 = "test.foo" () : () -> (i32, i64, index)
+ // CHECK-NOT: %[[RES:.*]] = scf.if
+ // CHECK: scf.if
+ %res:3 = scf.if %cond -> (i32, i64, index) {
+ // CHECK: scf.yield
+ scf.yield %foo#0, %foo#1, %foo#2 : i32, i64, index
+ } else {
+ // CHECK: scf.break {{.*}} %[[FOO]]#2 : index
+ scf.break [%loop] %foo#2 : index
+ }
+ // CHECK: "test.op"(%[[FOO]]#1)
+ "test.op"(%res#1) : (i64) -> ()
+ }
+ return %0 : index
+}
+
+// -----
+
+// Verify that removing the unused results of an if with nested breaking control flow
+// operation works.
+// CHECK-LABEL: func @simplify_if_with_breaking_controlflow_in_both_branches
+func.func @simplify_if_with_breaking_controlflow_in_both_branches(%cond : i1, %cond2 : i1, %arg0 : index) -> index {
+ // CHECK: scf.loop
+ %0 = scf.loop token(%loop) -> index {
+ // CHECK: %[[FOO:.*]] = "test.foo"
+ %foo = "test.foo" () : () -> (index)
+ // CHECK: scf.if
+ scf.if %cond {
+ // CHECK: scf.break {{.*}} %[[FOO]] : index
+ scf.break [%loop] %foo : index
+ // CHECK-NOT: else
+ } else {
+ // CHECK: %[[BAR:.*]] = "test.bar"
+ %bar = "test.bar" () : () -> (index)
+ // CHECK: scf.if
+ scf.if %cond2 {
+ // verify that this is correctly updated when inlining the parent region.
+ // CHECK: scf.break {{.*}} %[[BAR]] : index
+ scf.break [%loop] %bar : index
+ }
+ scf.continue [%loop]
+ }
+ "test.op"() : () -> ()
+ }
+ return %0 : index
+}
+
+// -----
+
+// Verify simplification when both branches have continue.
+// CHECK-LABEL: func @simplify_if_with_continue_in_both_branches
+func.func @simplify_if_with_continue_in_both_branches(%cond : i1, %init : i32) {
+ // CHECK: scf.loop
+ scf.loop token(%loop) iter_args(%arg = %init) : i32 {
+ // CHECK: scf.if
+ scf.if %cond {
+ // CHECK: scf.continue {{.*}} %{{.*}} : i32
+ scf.continue [%loop] %arg : i32
+ } else {
+ scf.continue [%loop] %arg : i32
+ }
+ // The dead code after the if should be removed.
+ // CHECK-NOT: test.op
+ "test.op"() : () -> ()
+ scf.continue [%loop] %arg : i32
+ }
+ return
+}
+
+// -----
+
+// Verify simplification when both branches have break.
+// CHECK-LABEL: func @simplify_if_with_break_in_both_branches
+func.func @simplify_if_with_break_in_both_branches(%cond : i1) -> index {
+ // CHECK: scf.loop
+ %0 = scf.loop token(%loop) -> index {
+ %c0 = arith.constant 0 : index
+ %c1 = arith.constant 1 : index
+ // CHECK: scf.if
+ scf.if %cond {
+ // CHECK: scf.break
+ scf.break [%loop] %c0 : index
+ } else {
+ scf.break [%loop] %c1 : index
+ }
+ // Dead code after the if.
+ // CHECK-NOT: test.op
+ "test.op"() : () -> ()
+ scf.continue [%loop]
+ }
+ return %0 : index
+}
diff --git a/mlir/test/IR/early-exit-invalid.mlir b/mlir/test/IR/early-exit-invalid.mlir
new file mode 100644
index 0000000000000..509b39629ebc7
--- /dev/null
+++ b/mlir/test/IR/early-exit-invalid.mlir
@@ -0,0 +1,142 @@
+
+// RUN: mlir-opt %s --split-input-file --verify-diagnostics
+
+func.func @loop_result_mismatch(%value : f32) {
+ // expected-error @+1 {{'scf.loop' op along control flow edge from Operation scf.break to Operation scf.loop: successor operand type #0 'f32' should match successor input type #0 'i32'}}
+ %result = scf.loop token(%loop) -> i32 {
+ scf.break [%loop] %value : f32 // expected-note {{region branch point}}
+ }
+ return
+}
+
+// -----
+
+// A PropagateControlFlowBreak op that is immediately nested in a
+// RegionBranchOpInterface can contain a region terminator that targets an
+// ancestor of that RegionBranchOpInterface. The intermediate RegionBranchOp
+// must expose this as a possible propagated edge so the addressed receiver can
+// verify the break operands against its results.
+func.func @outer_loop_result_mismatch_through_if_in_inner_loop(%cond : i1,
+ %value : f32) {
+ // expected-error @+1 {{'scf.loop' op along control flow edge from Operation scf.break to Operation scf.loop: successor operand type #0 'f32' should match successor input type #0 'i32'}}
+ %result = scf.loop token(%outer) -> i32 {
+ scf.loop token(%inner) {
+ scf.if %cond {
+ scf.break [%outer] %value : f32 // expected-note {{region branch point}}
+ }
+ scf.continue [%inner]
+ }
+ scf.continue [%outer]
+ }
+ return
+}
+
+// -----
+
+func.func @loop_result_number_mismatch(%value : f32) {
+ // expected-error @+1 {{'scf.loop' op along control flow edge from Operation scf.break to Operation scf.loop: region branch point has 1 operands, but region successor needs 2 inputs}}
+ %result:2 = scf.loop token(%loop) -> f32, f32 {
+ scf.break [%loop] %value : f32 // expected-note {{region branch point}}
+ }
+ return
+}
+
+// -----
+
+func.func @loop_continue_mismatch(%init : i32, %value : f32) {
+ // expected-error @+1 {{'scf.loop' op along control flow edge from Operation scf.continue to Region #0: successor operand type #0 'f32' should match successor input type #0 'i32'}}
+ scf.loop token(%loop) iter_args(%next = %init) : i32 {
+ scf.continue [%loop] %value : f32 // expected-note {{region branch point}}
+ }
+ return
+}
+
+
+// -----
+
+func.func @loop_iterargs_mismatch(%init : i32, %value : f32) {
+ // expected-error @+2 {{'scf.loop' op along control flow edge from parent to Region #0: successor operand type #0 'i32' should match successor input type #0 'f32'}}
+ // expected-note @+1 {{region branch point}}
+ "scf.loop"(%init) ({
+ ^body(%token : token, %next : f32):
+ scf.continue [%token] %init : i32
+ }) : (i32) -> ()
+ return
+}
+
+// -----
+
+func.func @loop_iterargs_mismatch(%init : i32, %value : f32) {
+ // expected-error @+2 {{'scf.loop' op along control flow edge from parent to Region #0: region branch point has 1 operands, but region successor needs 2 inputs}}
+ // expected-note @+1 {{region branch point}}
+ "scf.loop"(%init) ({
+ ^body(%token : token, %next : i32, %next2 : f32):
+ scf.continue [%token] %init : i32
+ }) : (i32) -> ()
+ return
+}
+
+// -----
+
+// scf.for lacks PropagateControlFlowBreak, so it cannot be an intermediate
+// parent for a break targeting an enclosing loop token.
+func.func @break_through_for_missing_trait(%lb: index, %ub: index, %step: index, %cond: i1) {
+ scf.loop token(%loop) {
+ scf.for %i = %lb to %ub step %step {
+ scf.if %cond {
+ // expected-error @+1 {{target token crosses an op that does not have the PropagateControlFlowBreak trait}}
+ scf.break [%loop]
+ }
+ }
+ }
+ return
+}
+
+// -----
+
+// scf.while lacks PropagateControlFlowBreak, so break through it is rejected.
+// The scf.while verifier catches this first as the after region must terminate
+// with scf.yield.
+func.func @break_through_while(%cond: i1) {
+ %init = arith.constant true
+ scf.loop token(%loop) {
+ // expected-error @+1 {{'scf.while' op expects the 'after' region to terminate with 'scf.yield'}}
+ scf.while(%arg = %init) : (i1) -> i1 {
+ scf.condition(%arg) %arg : i1
+ } do {
+ ^bb0(%arg2: i1):
+ // expected-note @+1 {{terminator here}}
+ scf.break [%loop]
+ }
+ }
+ return
+}
+
+// -----
+
+// Verify that a break without a token target is rejected.
+func.func @break_without_token() {
+ scf.loop token(%loop) {
+ // expected-error @+1 {{'scf.break' op expected 1 or more operands}}
+ "scf.break"() : () -> ()
+ }
+ return
+}
+
+// -----
+
+// Verify that a loop whose first region argument is not a control token is
+// rejected. The break targets the outer loop so the inner loop does not need a
+// self-targeting terminator, isolating the control-token check.
+func.func @loop_non_token_control(%cond: i1) {
+ "scf.loop"() ({
+ ^bb0(%outer: token):
+ // expected-error @+1 {{'scf.loop' op first region argument must be a control token}}
+ "scf.loop"() ({
+ ^bb1(%bad: i32):
+ "scf.break"(%outer) : (token) -> ()
+ }) : () -> ()
+ "scf.continue"(%outer) : (token) -> ()
+ }) : () -> ()
+ return
+}
diff --git a/mlir/test/IR/early-exit.mlir b/mlir/test/IR/early-exit.mlir
new file mode 100644
index 0000000000000..82e7f92070e5b
--- /dev/null
+++ b/mlir/test/IR/early-exit.mlir
@@ -0,0 +1,82 @@
+// RUN: mlir-opt --print-region-branch-op-interface %s --split-input-file | FileCheck %s
+// RUN: mlir-opt %s --mlir-print-debuginfo --mlir-print-op-generic --split-input-file | mlir-opt --print-region-branch-op-interface --split-input-file | FileCheck %s
+
+
+func.func @loop_break(%cond : i1) {
+ // CHECK: Found RegionBranchOpInterface operation: scf.loop {{.*}} {...} loc("loop1")
+ // CHECK: - Successor is region #0
+ // CHECK: - Found 2 predecessor(s)
+ // CHECK: - Predecessor is scf.break {{.*}} loc("break1")
+ // CHECK: - Predecessor is scf.continue
+ scf.loop token(%loop) {
+ scf.if %cond {
+ scf.break [%loop] loc("break1")
+ }
+ } loc("loop1")
+ return
+}
+
+// -----
+
+func.func @loop_continue(%cond1 : i1, %cond2 : i1) {
+ // CHECK: Found RegionBranchOpInterface operation: scf.loop {{.*}} {...} loc("loop2")
+ // CHECK: - Successor is region #0
+ // CHECK: - Found 2 predecessor(s)
+ // CHECK: - Predecessor is scf.break {{.*}} loc("break2")
+ // CHECK: - Predecessor is scf.continue
+ scf.loop token(%outer) {
+ // CHECK: Found RegionBranchOpInterface operation: scf.loop {{.*}} {...} loc("loop3")
+ // CHECK: - Successor is region #0
+ // CHECK: - Found 2 predecessor(s)
+ // CHECK: - Predecessor is scf.continue {{.*}} loc("continue1")
+ // CHECK: - Predecessor is scf.continue
+ scf.loop token(%inner) {
+ scf.if %cond1 {
+ scf.continue [%inner] loc("continue1")
+ }
+ scf.if %cond2 {
+ scf.break [%outer] loc("break2")
+ }
+ } loc("loop3")
+ } loc("loop2")
+ return
+}
+
+// -----
+
+// CHECK-LABEL: func @loop_with_results(
+func.func @loop_with_results(%value : f32) -> f32 {
+ %result = scf.loop token(%loop) -> f32 {
+ scf.break [%loop] %value : f32
+ }
+ return %result : f32
+}
+
+// -----
+
+// CHECK-LABEL: func @loop_continue_iterargs(
+func.func @loop_continue_iterargs(%init : i32) {
+ scf.loop token(%loop) iter_args(%next = %init) : i32 {
+ scf.continue [%loop] %next : i32
+ }
+ return
+}
+
+// -----
+
+// A single-operand continue that targets an *outer* loop must be printed
+// explicitly. If the printer elided it as the trivial implicit terminator, the
+// parser would rebuild it as a continue of the inner loop, silently changing
+// the target.
+// CHECK-LABEL: func @continue_outer_from_inner
+func.func @continue_outer_from_inner() {
+ // CHECK: scf.loop token(%[[OUTER:.*]]) {
+ scf.loop token(%outer) {
+ // CHECK: scf.loop token(%[[INNER:.*]]) {
+ scf.loop token(%inner) {
+ // CHECK: scf.continue [%[[OUTER]]]
+ scf.continue [%outer]
+ }
+ }
+ return
+}
diff --git a/mlir/test/Integration/Dialect/SCF/early_exit.mlir b/mlir/test/Integration/Dialect/SCF/early_exit.mlir
new file mode 100644
index 0000000000000..518b602f790aa
--- /dev/null
+++ b/mlir/test/Integration/Dialect/SCF/early_exit.mlir
@@ -0,0 +1,82 @@
+// RUN: mlir-opt %s -convert-scf-to-cf --canonicalize --convert-cf-to-llvm --convert-to-llvm | \
+// RUN: mlir-runner -e entry -entry-point-result=void \
+// RUN: -shared-libs=%mlir_c_runner_utils | \
+// RUN: FileCheck %s
+
+
+
+// End-to-end test of all fp reduction intrinsics (not exhaustive unit tests).
+module {
+ llvm.func @entry() {
+ // Constant for the iteration space and various conditions
+ %one = llvm.mlir.constant(1 : i64) : i64
+ %two = llvm.mlir.constant(2 : i64) : i64
+ %three = llvm.mlir.constant(3 : i64) : i64
+ %four = llvm.mlir.constant(4 : i64) : i64
+ %counter_init = llvm.mlir.constant(0 : i64) : i64
+
+
+// CHECK: Outer Loop Begin with counter: 0
+// CHECK-NEXT: Inner Loop Begin, counter: 1
+// CHECK-NEXT: continue inner loop
+// CHECK-NEXT: Inner Loop Begin, counter: 2
+// CHECK-NEXT: Iteration 2, loop back to outer loop
+// CHECK-NEXT: Outer Loop Begin with counter: 2
+// CHECK-NEXT: Inner Loop Begin, counter: 3
+// CHECK-NEXT: continue inner loop
+// CHECK-NEXT: Inner Loop Begin, counter: 4
+// CHECK-NEXT: continue inner loop
+// CHECK-NEXT: Inner Loop Begin, counter: 5
+// CHECK-NEXT: Last iteration, break out of outer loop
+// CHECK-NEXT: Outer loop finished with result: 4
+
+
+ %result = scf.loop token(%outer) iter_args(%counter_out = %counter_init) : i64 -> i64 {
+ // Outer loop iteration
+ vector.print str "Outer Loop Begin with counter: "
+ vector.print %counter_out : i64
+
+ scf.loop token(%inner) iter_args(%counter = %counter_out) : i64 {
+ // %counter will go from 0 to 4
+ // %counter_update will go from 1 to 5
+ %counter_update = llvm.add %counter, %one : i64
+
+ // Inner loop iteration
+ // print from 1..5
+ vector.print str "Inner Loop Begin, counter: "
+ vector.print %counter_update : i64
+
+ // On the second iteration, print 2.3 and loop back to the outer loop.
+ %cond1 = llvm.icmp "eq" %counter_update, %two : i64
+ scf.if %cond1 {
+ vector.print str "Iteration 2, loop back to outer loop\n"
+ scf.continue [%outer] %counter_update : i64
+ }
+
+ // Exit condition when counter>4
+ %cond2 = llvm.icmp "sge" %counter, %four : i64
+ scf.if %cond2 {
+ vector.print str "Last iteration, break out of outer loop\n"
+ // return the counter from the previous iteration here (pre-update)
+ scf.break [%outer] %counter : i64
+ }
+
+ %cond3 = llvm.icmp "eq" %counter_update, %three : i64
+ scf.if %cond2 {
+ vector.print str "Iteration 3, break out of inner loop"
+ scf.break [%inner]
+ }
+ vector.print str "continue inner loop\n"
+ scf.continue [%inner] %counter_update : i64
+ }
+ vector.print str "continue outer loop\n"
+ scf.continue [%outer] %counter_out : i64
+ }
+
+// After the loop nest finishes
+ vector.print str "Outer loop finished with result: "
+ vector.print %result : i64
+
+ llvm.return
+ }
+}
diff --git a/mlir/test/Transforms/sccp-early-exit.mlir b/mlir/test/Transforms/sccp-early-exit.mlir
new file mode 100644
index 0000000000000..4a4136b94b891
--- /dev/null
+++ b/mlir/test/Transforms/sccp-early-exit.mlir
@@ -0,0 +1,117 @@
+// RUN: mlir-opt -allow-unregistered-dialect %s -pass-pipeline="builtin.module(func.func(sccp))" -split-input-file | FileCheck %s
+
+/// The inner loop's only normal exit (break 1) carries the constant 5.
+/// The early-exit path (break 3) carries -5 but bypasses the code after the
+/// inner loop entirely — it exits both loops at once.
+///
+/// Without early-exit support the analysis would conservatively join {5, -5}
+/// and mark %inner as overdefined, keeping the dead branch alive.
+/// With early-exit support SCCP sees that %inner is always 5, folds the
+/// comparison, and removes the dead scf.if.
+///
+/// Pseudocode:
+/// loop {
+/// a = loop {
+/// if (cond) { break_all -5; } // early exit from both loops
+/// break 5; // normal exit from inner loop
+/// };
+/// // here a == 5 always
+/// if (a < 0) "dead"(); // dead code
+/// break a;
+/// }
+
+// CHECK-LABEL: func @early_exit_dead_code(
+// CHECK-DAG: %[[FALSE:.*]] = arith.constant false
+// CHECK-DAG: %[[C5:.*]] = arith.constant 5 : i32
+// CHECK: scf.loop
+// CHECK: scf.loop
+// CHECK: scf.if %[[FALSE]]
+// CHECK: "test.dead_op"
+// CHECK: scf.break {{.*}} %[[C5]] : i32
+func.func @early_exit_dead_code(%cond: i1) -> i32 {
+ %c0 = arith.constant 0 : i32
+ %c5 = arith.constant 5 : i32
+ %cm5 = arith.constant -5 : i32
+ %outer = scf.loop token(%outer_token) -> i32 {
+ %inner = scf.loop token(%inner_token) -> i32 {
+ scf.if %cond {
+ scf.break [%outer_token] %cm5 : i32
+ }
+ scf.break [%inner_token] %c5 : i32
+ }
+ %is_neg = arith.cmpi slt, %inner, %c0 : i32
+ scf.if %is_neg {
+ "test.dead_op"() : () -> ()
+ }
+ scf.break [%outer_token] %inner : i32
+ }
+ return %outer : i32
+}
+
+// -----
+
+/// For contrast: the same logic emulated with scf.while and a boolean flag
+/// instead of early exit. Without break 3 to bypass the code after the inner
+/// loop, both values (-5 and 5) flow through %inner and SCCP sees it as
+/// overdefined — the comparison and dead branch survive.
+///
+/// Pseudocode:
+/// done = false
+/// while (!done) {
+/// a, should_break = while (!done) {
+/// if (cond) yield -5, true // want to break all, but can't
+/// else yield 5, false
+/// };
+/// if (!should_break) {
+/// // a ∈ {5, -5} — flag can't help the analysis narrow it
+/// if (a < 0) "not dead"();
+/// }
+/// done = true
+/// }
+
+// CHECK-LABEL: func @no_early_exit_not_foldable(
+// CHECK: scf.while
+// CHECK: scf.while
+// CHECK: scf.if
+// CHECK: arith.cmpi slt,
+// CHECK: scf.if
+// CHECK: "test.not_dead_op"
+func.func @no_early_exit_not_foldable(%cond: i1) -> i32 {
+ %c0 = arith.constant 0 : i32
+ %c5 = arith.constant 5 : i32
+ %cm5 = arith.constant -5 : i32
+ %false = arith.constant false
+ %true = arith.constant true
+
+ // Outer while: (result, done)
+ %outer, %_ = scf.while(%o_res = %c0, %o_done = %false) : (i32, i1) -> (i32, i1) {
+ %o_go = arith.xori %o_done, %true : i1
+ scf.condition(%o_go) %o_res, %o_done : i32, i1
+ } do {
+ ^bb0(%o_res: i32, %o_done: i1):
+ // Inner while: (result, done, should_break_outer)
+ %inner, %_2, %break_flag = scf.while(%i_res = %c0, %i_done = %false, %i_brk = %false)
+ : (i32, i1, i1) -> (i32, i1, i1) {
+ %i_go = arith.xori %i_done, %true : i1
+ scf.condition(%i_go) %i_res, %i_done, %i_brk : i32, i1, i1
+ } do {
+ ^bb0(%i_res: i32, %i_done: i1, %i_brk: i1):
+ // if (cond) a = -5, flag = true; else a = 5, flag = false
+ %a = arith.select %cond, %cm5, %c5 : i32
+ %brk = arith.select %cond, %true, %false : i1
+ scf.yield %a, %true, %brk : i32, i1, i1
+ }
+ // %inner ∈ {5, -5}: both values reach here (no early exit to bypass)
+ // %break_flag correlates with %inner but SCCP can't exploit that
+ %not_break = arith.xori %break_flag, %true : i1
+ scf.if %not_break {
+ // %inner ∈ {5} in theory: but the analysis can't recover this here.
+ %is_neg = arith.cmpi slt, %inner, %c0 : i32
+ scf.if %is_neg {
+ "test.not_dead_op"() : () -> ()
+ }
+ }
+ scf.yield %inner, %true : i32, i1
+ }
+ return %outer : i32
+}
diff --git a/mlir/test/lib/Dialect/Test/TestOps.td b/mlir/test/lib/Dialect/Test/TestOps.td
index c12e2abb3ab75..9d7e21cabda6f 100644
--- a/mlir/test/lib/Dialect/Test/TestOps.td
+++ b/mlir/test/lib/Dialect/Test/TestOps.td
@@ -2385,6 +2385,12 @@ def TestMergeBlocksOp : TEST_Op<"merge_blocks"> {
let results = (outs Variadic<AnyType>:$result);
}
+def TestPropagateControlFlowBreakOp
+ : TEST_Op<"propagate_control_flow_break", [PropagateControlFlowBreak]> {
+ let regions = (region AnyRegion:$body);
+ let assemblyFormat = "attr-dict-with-keyword $body";
+}
+
def TestRemappedValueRegionOp : TEST_Op<"remapped_value_region",
[SingleBlock]> {
let summary = "remapped_value_region operation";
diff --git a/mlir/test/lib/IR/TestDominance.cpp b/mlir/test/lib/IR/TestDominance.cpp
index b34149b3e2cbd..24696c8dd028c 100644
--- a/mlir/test/lib/IR/TestDominance.cpp
+++ b/mlir/test/lib/IR/TestDominance.cpp
@@ -14,6 +14,7 @@
#include "mlir/IR/Builders.h"
#include "mlir/IR/Dominance.h"
+#include "mlir/IR/Operation.h"
#include "mlir/IR/SymbolTable.h"
#include "mlir/Pass/Pass.h"
@@ -21,6 +22,14 @@ using namespace mlir;
/// Overloaded helper to call the right function based on whether we are testing
/// dominance or post-dominance.
+static bool dominatesOrPostDominates(DominanceInfo &dominanceInfo, Operation *a,
+ Operation *b) {
+ return dominanceInfo.dominates(a, b);
+}
+static bool dominatesOrPostDominates(PostDominanceInfo &dominanceInfo,
+ Operation *a, Operation *b) {
+ return dominanceInfo.postDominates(a, b);
+}
static bool dominatesOrPostDominates(DominanceInfo &dominanceInfo, Block *a,
Block *b) {
return dominanceInfo.dominates(a, b);
@@ -72,6 +81,30 @@ class DominanceTest {
template <typename DominanceT>
void printDominance(DominanceT &dominanceInfo,
bool printCommonDominatorInfo) {
+ if (printCommonDominatorInfo) {
+ operation->walk([&](Operation *op) {
+ if (!op->getDiscardableAttr("test.print_dominance"))
+ return;
+ operation->walk([&](Operation *nested) {
+ if (std::is_same<DominanceInfo, DominanceT>::value)
+ llvm::outs() << "dominates(";
+ else
+ llvm::outs() << "postdominates(";
+ bool isDominated =
+ dominatesOrPostDominates(dominanceInfo, op, nested);
+ llvm::outs() << OpWithFlags(op, OpPrintingFlags()
+ .skipRegions()
+ .enableDebugInfo()
+ .assumeVerified())
+ << ", "
+ << OpWithFlags(nested, OpPrintingFlags()
+ .skipRegions()
+ .enableDebugInfo()
+ .assumeVerified())
+ << ") = " << std::to_string(isDominated) << "\n";
+ });
+ });
+ }
DenseSet<Block *> parentVisited;
operation->walk([&](Operation *op) {
Block *block = op->getBlock();
diff --git a/mlir/test/lib/Interfaces/CMakeLists.txt b/mlir/test/lib/Interfaces/CMakeLists.txt
index 6a21ed10eec6f..3aa5097b7ed20 100644
--- a/mlir/test/lib/Interfaces/CMakeLists.txt
+++ b/mlir/test/lib/Interfaces/CMakeLists.txt
@@ -1,2 +1,3 @@
add_subdirectory(LoopLikeInterface)
+add_subdirectory(RegionBranchOpInterface)
add_subdirectory(TilingInterface)
diff --git a/mlir/test/lib/Interfaces/RegionBranchOpInterface/CMakeLists.txt b/mlir/test/lib/Interfaces/RegionBranchOpInterface/CMakeLists.txt
new file mode 100644
index 0000000000000..8e003942e41c0
--- /dev/null
+++ b/mlir/test/lib/Interfaces/RegionBranchOpInterface/CMakeLists.txt
@@ -0,0 +1,9 @@
+add_mlir_library(MLIRTestRegionBranchOpInterface
+ TestRegionBranchOpInterface.cpp
+
+ EXCLUDE_FROM_LIBMLIR
+ )
+mlir_target_link_libraries(MLIRTestRegionBranchOpInterface PUBLIC
+ MLIRControlFlowInterfaces
+ MLIRPass
+ )
diff --git a/mlir/test/lib/Interfaces/RegionBranchOpInterface/TestRegionBranchOpInterface.cpp b/mlir/test/lib/Interfaces/RegionBranchOpInterface/TestRegionBranchOpInterface.cpp
new file mode 100644
index 0000000000000..878de7d44da01
--- /dev/null
+++ b/mlir/test/lib/Interfaces/RegionBranchOpInterface/TestRegionBranchOpInterface.cpp
@@ -0,0 +1,77 @@
+//===- TestRegionBranchOpInterface.cpp - RegionBranch test pass -----------===//
+//
+// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
+// See https://llvm.org/LICENSE.txt for license information.
+// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+//
+//===----------------------------------------------------------------------===//
+
+#include "mlir/Dialect/Func/IR/FuncOps.h"
+#include "mlir/IR/BuiltinOps.h"
+#include "mlir/Interfaces/ControlFlowInterfaces.h"
+#include "mlir/Pass/Pass.h"
+#include "llvm/Support/raw_ostream.h"
+
+using namespace mlir;
+
+namespace {
+/// Test pass that prints RegionBranchOpInterface successor information and
+/// nested breaking-control-flow predecessors.
+struct PrintRegionBranchOpInterfacePass
+ : public PassWrapper<PrintRegionBranchOpInterfacePass, OperationPass<>> {
+ MLIR_DEFINE_EXPLICIT_INTERNAL_INLINE_TYPE_ID(PrintRegionBranchOpInterfacePass)
+
+ StringRef getArgument() const final {
+ return "print-region-branch-op-interface";
+ }
+ StringRef getDescription() const final {
+ return "Print control-flow edges represented by "
+ "mlir::RegionBranchOpInterface";
+ }
+
+ void runOnOperation() override {
+ Operation *op = getOperation();
+ op->walk<WalkOrder::PreOrder>([&](RegionBranchOpInterface branchOp) {
+ llvm::outs() << "Found RegionBranchOpInterface operation: "
+ << OpWithFlags(
+ branchOp,
+ OpPrintingFlags().skipRegions().enableDebugInfo())
+ << "\n";
+ SmallVector<RegionSuccessor> regions;
+ branchOp.getSuccessorRegions(RegionBranchPoint::parent(), regions);
+ for (auto &successor : regions) {
+ if (successor.isOperation()) {
+ llvm::outs() << " - Successor is operation "
+ << successor.getSuccessorOp()->getName() << "\n";
+ } else {
+ llvm::outs() << " - Successor is region #"
+ << successor.getSuccessor()->getRegionNumber() << "\n";
+ }
+ }
+ if (auto breakingControlFlowOp =
+ dyn_cast<HasBreakingControlFlowOpInterface>(
+ branchOp.getOperation())) {
+ SmallVector<Operation *> predecessors;
+ llvm::outs() << " - Collecting all nested predecessors\n";
+ collectAllNestedPredecessors(breakingControlFlowOp, predecessors);
+ llvm::outs() << " - Found " << predecessors.size()
+ << " predecessor(s)\n";
+ for (auto &predecessor : predecessors) {
+ llvm::outs() << " - Predecessor is "
+ << OpWithFlags(
+ predecessor,
+ OpPrintingFlags().skipRegions().enableDebugInfo())
+ << "\n";
+ }
+ }
+ });
+ }
+};
+
+} // namespace
+
+namespace mlir {
+void registerRegionBranchOpInterfaceTestPasses() {
+ PassRegistration<PrintRegionBranchOpInterfacePass>();
+}
+} // namespace mlir
diff --git a/mlir/tools/mlir-opt/CMakeLists.txt b/mlir/tools/mlir-opt/CMakeLists.txt
index c607ccfa80e3c..821bad3a4166a 100644
--- a/mlir/tools/mlir-opt/CMakeLists.txt
+++ b/mlir/tools/mlir-opt/CMakeLists.txt
@@ -22,6 +22,7 @@ if(MLIR_INCLUDE_TESTS)
MLIRGPUTestPasses
MLIRLinalgTestPasses
MLIRLoopLikeInterfaceTestPasses
+ MLIRTestRegionBranchOpInterface
MLIRMathTestPasses
MLIRTestMathToVCIX
MLIRMemRefTestPasses
diff --git a/mlir/tools/mlir-opt/mlir-opt.cpp b/mlir/tools/mlir-opt/mlir-opt.cpp
index 13c0934f34656..8c8813b480a99 100644
--- a/mlir/tools/mlir-opt/mlir-opt.cpp
+++ b/mlir/tools/mlir-opt/mlir-opt.cpp
@@ -38,6 +38,7 @@ void registerLazyLoadingTestPasses();
void registerLoopLikeInterfaceTestPasses();
void registerPassManagerTestPass();
void registerPrintSpirvAvailabilityPass();
+void registerRegionBranchOpInterfaceTestPasses();
void registerRegionTestPasses();
void registerPrintTosaAvailabilityPass();
void registerShapeFunctionTestPasses();
@@ -191,6 +192,7 @@ static void registerTestPasses() {
registerPassManagerTestPass();
registerPrintSpirvAvailabilityPass();
registerRegionTestPasses();
+ registerRegionBranchOpInterfaceTestPasses();
registerShapeFunctionTestPasses();
registerSideEffectTestPasses();
registerSliceAnalysisTestPass();
>From 968e49b11932bc78f5870972edb987bec35fd8ff Mon Sep 17 00:00:00 2001
From: Mehdi Amini <joker.eph at gmail.com>
Date: Thu, 11 Jun 2026 07:11:07 -0700
Subject: [PATCH 2/2] [mlir][test] Add dynamic-depth early-exit loop ops
Add Test dialect loop and terminators that model dynamic breaking control-flow
targets with an index depth. Constant depths resolve to one enclosing
test.breakable_loop; dynamic depths expose compatible reachable loop targets
through getPotentialTargets().
Filter dynamic potential targets by the terminator payload types. Runtime depths
that select an incompatible filtered-out target are undefined, while constant
depths still diagnose payload mismatches against the selected loop.
Add a test conversion pass that lowers test.breakable_loop to scf.loop and
rewrites dynamic break/continue terminators to token-targeted SCF terminators.
Dynamic depths lower through a dispatch ladder over compatible target loops
while preserving each target's original enclosing-loop depth.
Add verifier coverage for invalid depths, blocked propagation, and payload type
mismatches, RegionBranchOpInterface coverage for dynamic targets, FileCheck
coverage for the SCF conversion, and execution tests that lower through
SCF-to-CF.
Assisted-by: Codex
---
.../IR/early-exit-test-dialect-invalid.mlir | 93 +++++
mlir/test/IR/early-exit-test-dialect.mlir | 124 +++++++
.../test-convert-breakable-loop-to-scf.mlir | 139 +++++++
mlir/test/lib/Dialect/Test/TestOpDefs.cpp | 340 ++++++++++++++++++
mlir/test/lib/Dialect/Test/TestOps.td | 67 ++++
mlir/test/lib/Dialect/Test/TestPatterns.cpp | 213 +++++++++++
.../test-convert-breakable-loop-to-scf.mlir | 145 ++++++++
7 files changed, 1121 insertions(+)
create mode 100644 mlir/test/IR/early-exit-test-dialect-invalid.mlir
create mode 100644 mlir/test/IR/early-exit-test-dialect.mlir
create mode 100644 mlir/test/Transforms/test-convert-breakable-loop-to-scf.mlir
create mode 100644 mlir/test/mlir-runner/test-convert-breakable-loop-to-scf.mlir
diff --git a/mlir/test/IR/early-exit-test-dialect-invalid.mlir b/mlir/test/IR/early-exit-test-dialect-invalid.mlir
new file mode 100644
index 0000000000000..c5ca385e4aa6b
--- /dev/null
+++ b/mlir/test/IR/early-exit-test-dialect-invalid.mlir
@@ -0,0 +1,93 @@
+// RUN: mlir-opt %s --split-input-file --verify-diagnostics
+
+func.func @non_index_depth(%depth: i32) {
+ test.breakable_loop {
+ // expected-error @+1 {{operand #0 must be index}}
+ "test.dynamic_break"(%depth) : (i32) -> ()
+ }
+ return
+}
+
+// -----
+
+func.func @no_reachable_loop(%depth: index) {
+ // expected-error @+1 {{must be nested inside a reachable test.breakable_loop}}
+ test.dynamic_break %depth
+}
+
+// -----
+
+func.func @zero_depth() {
+ %c0 = arith.constant 0 : index
+ test.breakable_loop {
+ // expected-error @+1 {{depth must be positive}}
+ test.dynamic_continue %c0
+ }
+ return
+}
+
+// -----
+
+func.func @too_deep() {
+ %c2 = arith.constant 2 : index
+ test.breakable_loop {
+ // expected-error @+1 {{constant depth exceeds the number of reachable test.breakable_loop operations}}
+ test.dynamic_break %c2
+ }
+ return
+}
+
+// -----
+
+func.func @blocked_parent() {
+ %c1 = arith.constant 1 : index
+ test.breakable_loop {
+ "test.any_cond"() ({
+ // expected-error @+1 {{depth target crosses an op that does not have the PropagateControlFlowBreak trait}}
+ test.dynamic_break %c1
+ }) : () -> ()
+ test.dynamic_continue %c1
+ }
+ return
+}
+
+// -----
+
+func.func @dynamic_depth_blocked_outer(%depth: index) {
+ %c1 = arith.constant 1 : index
+ test.breakable_loop {
+ test.single_no_terminator_custom_asm_op {
+ test.breakable_loop {
+ // expected-error @+1 {{dynamic depth may target across an op that does not have the PropagateControlFlowBreak trait}}
+ test.dynamic_break %depth
+ }
+ }
+ test.dynamic_continue %c1
+ }
+ return
+}
+
+// -----
+
+// Constant depths select one concrete loop, so payload mismatches are still
+// diagnosed against that selected target.
+func.func @constant_break_payload_mismatch(%value: f32) -> i32 {
+ %c1 = arith.constant 1 : index
+ %result = test.breakable_loop -> i32 {
+ // expected-error @+1 {{payload operand #0 has type 'f32', but target results}}
+ test.dynamic_break %c1 %value : f32
+ }
+ return %result : i32
+}
+
+// -----
+
+// Dynamic depths filter out incompatible targets. This is still invalid when
+// filtering leaves no target that can accept the terminator payload.
+func.func @dynamic_break_no_compatible_target(%depth: index, %value: f32) -> i32 {
+ %result = test.breakable_loop -> i32 {
+ // expected-error @+1 {{dynamic depth has no compatible test.breakable_loop target}}
+ test.dynamic_break %depth %value : f32
+ }
+ return %result : i32
+}
diff --git a/mlir/test/IR/early-exit-test-dialect.mlir b/mlir/test/IR/early-exit-test-dialect.mlir
new file mode 100644
index 0000000000000..3ea7f0e8154b8
--- /dev/null
+++ b/mlir/test/IR/early-exit-test-dialect.mlir
@@ -0,0 +1,124 @@
+// RUN: mlir-opt %s --split-input-file | FileCheck %s --check-prefix=PRINT
+// RUN: mlir-opt --print-region-branch-op-interface %s --split-input-file | FileCheck %s --check-prefix=BRANCH
+
+// PRINT-LABEL: func.func @immediate_continue(
+func.func @immediate_continue(%depth: index) {
+ // PRINT: test.breakable_loop
+ test.breakable_loop {
+ // PRINT: test.dynamic_continue %arg0
+ test.dynamic_continue %depth
+ }
+ return
+}
+
+// -----
+
+// PRINT-LABEL: func.func @break_with_result(
+func.func @break_with_result(%depth: index, %value: i32) -> i32 {
+ // PRINT: test.breakable_loop -> i32
+ %result = test.breakable_loop -> i32 {
+ // PRINT: test.dynamic_break %arg0 %arg1 : i32
+ test.dynamic_break %depth %value : i32
+ }
+ return %result : i32
+}
+
+// -----
+
+// PRINT-LABEL: func.func @continue_iter_args(
+func.func @continue_iter_args(%depth: index, %init: i32) {
+ // PRINT: test.breakable_loop iter_args(%{{.*}} = %arg1) : i32
+ test.breakable_loop iter_args(%iter = %init) : i32 {
+ // PRINT: test.dynamic_continue %arg0 %{{.*}} : i32
+ test.dynamic_continue %depth %iter : i32
+ }
+ return
+}
+
+// -----
+
+// PRINT-LABEL: func.func @constant_depth_selects_outer(
+func.func @constant_depth_selects_outer(%outer_init: i32, %inner_init: f32) {
+ %c1 = arith.constant 1 : index
+ %c2 = arith.constant 2 : index
+ test.breakable_loop iter_args(%outer = %outer_init) : i32 {
+ test.breakable_loop iter_args(%inner = %inner_init) : f32 {
+ // PRINT: test.dynamic_continue %c2 %{{.*}} : i32
+ test.dynamic_continue %c2 %outer : i32
+ }
+ test.dynamic_continue %c1 %outer : i32
+ }
+ return
+}
+
+// -----
+
+// PRINT-LABEL: func.func @through_scf_if(
+func.func @through_scf_if(%cond: i1, %depth: index) {
+ %c1 = arith.constant 1 : index
+ // BRANCH-LABEL: Found RegionBranchOpInterface operation: test.breakable_loop {{.*}} loc("loop")
+ // BRANCH: - Successor is region #0
+ // BRANCH: - Found 2 predecessor(s)
+ // BRANCH: - Predecessor is test.dynamic_break {{.*}} loc("if_break")
+ // BRANCH: - Predecessor is test.dynamic_continue {{.*}} loc("after_if")
+ test.breakable_loop {
+ // PRINT: scf.if
+ scf.if %cond {
+ // PRINT: test.dynamic_break %arg1
+ test.dynamic_break %depth loc("if_break")
+ } loc("if")
+ test.dynamic_continue %c1 loc("after_if")
+ } loc("loop")
+ return
+}
+
+// -----
+
+// A dynamic depth can target either the immediately enclosing loop or an outer
+// compatible loop. The inner loop sees the dynamic break as a direct
+// predecessor, and the outer loop sees both the propagated dynamic break and
+// its own explicit continue.
+func.func @nested_dynamic(%depth: index) {
+ %c1 = arith.constant 1 : index
+ // BRANCH-LABEL: Found RegionBranchOpInterface operation: test.breakable_loop {{.*}} loc("outer")
+ // BRANCH: - Successor is region #0
+ // BRANCH: - Found 2 predecessor(s)
+ // BRANCH: - Predecessor is test.dynamic_break {{.*}} loc("dyn_break")
+ // BRANCH: - Predecessor is test.dynamic_continue {{.*}} loc("outer_continue")
+ test.breakable_loop {
+ // BRANCH-LABEL: Found RegionBranchOpInterface operation: test.breakable_loop {{.*}} loc("inner")
+ // BRANCH: - Successor is region #0
+ // BRANCH: - Found 1 predecessor(s)
+ // BRANCH: - Predecessor is test.dynamic_break {{.*}} loc("dyn_break")
+ test.breakable_loop {
+ test.dynamic_break %depth loc("dyn_break")
+ } loc("inner")
+ test.dynamic_continue %c1 loc("outer_continue")
+ } loc("outer")
+ return
+}
+
+// -----
+
+// The dynamic continue payload is f32. That makes the inner f32 loop a
+// potential target, but filters out the enclosing i32 loop. The outer loop
+// should only see its explicit continue predecessor.
+func.func @dynamic_continue_filters_incompatible_outer(%depth: index, %i: i32,
+ %f: f32) {
+ %c1 = arith.constant 1 : index
+ // BRANCH-LABEL: Found RegionBranchOpInterface operation: test.breakable_loop {{.*}} loc("filtered_outer")
+ // BRANCH: - Successor is region #0
+ // BRANCH: - Found 1 predecessor(s)
+ // BRANCH: - Predecessor is test.dynamic_continue {{.*}} loc("filtered_outer_continue")
+ test.breakable_loop iter_args(%outer = %i) : i32 {
+ // BRANCH-LABEL: Found RegionBranchOpInterface operation: test.breakable_loop {{.*}} loc("filtered_inner")
+ // BRANCH: - Successor is region #0
+ // BRANCH: - Found 1 predecessor(s)
+ // BRANCH: - Predecessor is test.dynamic_continue {{.*}} loc("filtered_inner_continue")
+ test.breakable_loop iter_args(%inner = %f) : f32 {
+ test.dynamic_continue %depth %inner : f32 loc("filtered_inner_continue")
+ } loc("filtered_inner")
+ test.dynamic_continue %c1 %outer : i32 loc("filtered_outer_continue")
+ } loc("filtered_outer")
+ return
+}
diff --git a/mlir/test/Transforms/test-convert-breakable-loop-to-scf.mlir b/mlir/test/Transforms/test-convert-breakable-loop-to-scf.mlir
new file mode 100644
index 0000000000000..f33530e6d91d1
--- /dev/null
+++ b/mlir/test/Transforms/test-convert-breakable-loop-to-scf.mlir
@@ -0,0 +1,139 @@
+// RUN: mlir-opt %s --test-convert-breakable-loop-to-scf --split-input-file | FileCheck %s
+
+// CHECK-LABEL: func.func @immediate_continue(
+// CHECK-SAME: %[[INIT:.*]]: i32
+func.func @immediate_continue(%init: i32) -> i32 {
+ %c1 = arith.constant 1 : index
+ // CHECK: %[[RESULT:.*]] = scf.loop token(%[[TOKEN:.*]]) iter_args(%[[ITER:.*]] = %[[INIT]]) : i32 -> i32
+ %result = test.breakable_loop iter_args(%iter = %init) : i32 -> i32 {
+ // CHECK: scf.continue [%[[TOKEN]]] %[[ITER]] : i32
+ test.dynamic_continue %c1 %iter : i32
+ }
+ // CHECK: return %[[RESULT]] : i32
+ return %result : i32
+}
+
+// -----
+
+// CHECK-LABEL: func.func @constant_outer_break(
+func.func @constant_outer_break(%init: i32) -> i32 {
+ %c1 = arith.constant 1 : index
+ %c2 = arith.constant 2 : index
+ // CHECK: %[[RESULT:.*]] = scf.loop token(%[[OUTER_TOKEN:.*]]) iter_args(%[[OUTER_ARG:.*]] =
+ %result = test.breakable_loop iter_args(%outer = %init) : i32 -> i32 {
+ // CHECK: scf.loop token(%[[INNER_TOKEN:.*]]) {
+ test.breakable_loop {
+ // CHECK: scf.break [%[[OUTER_TOKEN]]] %[[OUTER_ARG]] : i32
+ test.dynamic_break %c2 %outer : i32
+ }
+ test.dynamic_break %c1 %outer : i32
+ }
+ // CHECK: return %[[RESULT]] : i32
+ return %result : i32
+}
+
+// -----
+
+// Dynamic break over two compatible loops lowers to a dispatch ladder with one
+// arm for depth 1 (inner) and one for depth 2 (outer), plus a deterministic
+// fallback for UB depth values.
+// CHECK-LABEL: func.func @dynamic_break_dispatch(
+// CHECK-SAME: %[[CHOOSE:.*]]: i1, %[[INIT:.*]]: i32
+func.func @dynamic_break_dispatch(%choose_outer: i1, %init: i32) -> i32 {
+ %c1 = arith.constant 1 : index
+ %c2 = arith.constant 2 : index
+ // CHECK: %[[DEPTH:.*]] = arith.select %[[CHOOSE]],
+ %depth = arith.select %choose_outer, %c2, %c1 : index
+ // CHECK: %[[RESULT:.*]] = scf.loop token(%[[OUTER_TOKEN:.*]]) iter_args(%[[OUTER_ARG:.*]] =
+ %result = test.breakable_loop iter_args(%outer = %init) : i32 -> i32 {
+ // CHECK: %[[INNER_RESULT:.*]] = scf.loop token(%[[INNER_TOKEN:.*]]) iter_args(%[[INNER_ARG:.*]] =
+ %inner_result = test.breakable_loop iter_args(%inner = %outer) : i32 -> i32 {
+ // CHECK: %[[DEPTH_ONE:.*]] = arith.constant 1 : index
+ // CHECK: %[[IS_ONE:.*]] = arith.cmpi eq, %[[DEPTH]], %[[DEPTH_ONE]] : index
+ // CHECK: scf.if %[[IS_ONE]] {
+ // CHECK: scf.break [%[[INNER_TOKEN]]] %[[INNER_ARG]] : i32
+ // CHECK: %[[DEPTH_TWO:.*]] = arith.constant 2 : index
+ // CHECK: %[[IS_TWO:.*]] = arith.cmpi eq, %[[DEPTH]], %[[DEPTH_TWO]] : index
+ // CHECK: scf.if %[[IS_TWO]] {
+ // CHECK: scf.break [%[[OUTER_TOKEN]]] %[[INNER_ARG]] : i32
+ // CHECK: scf.break [%[[INNER_TOKEN]]] %[[INNER_ARG]] : i32
+ test.dynamic_break %depth %inner : i32
+ }
+ test.dynamic_break %c1 %inner_result : i32
+ }
+ // CHECK: return %[[RESULT]] : i32
+ return %result : i32
+}
+
+// -----
+
+// The same dispatch shape is required for dynamic continue, except each arm
+// targets the next iteration of the selected loop.
+// CHECK-LABEL: func.func @dynamic_continue_dispatch(
+// CHECK-SAME: %[[CHOOSE:.*]]: i1, %[[INIT:.*]]: i32
+func.func @dynamic_continue_dispatch(%choose_outer: i1, %init: i32) -> i32 {
+ %c1 = arith.constant 1 : index
+ %c2 = arith.constant 2 : index
+ // CHECK: %[[DEPTH:.*]] = arith.select %[[CHOOSE]],
+ %depth = arith.select %choose_outer, %c2, %c1 : index
+ // CHECK: %[[RESULT:.*]] = scf.loop token(%[[OUTER_TOKEN:.*]]) iter_args(%[[OUTER_ARG:.*]] =
+ %result = test.breakable_loop iter_args(%outer = %init) : i32 -> i32 {
+ // CHECK: scf.loop token(%[[INNER_TOKEN:.*]]) iter_args(%[[INNER_ARG:.*]] =
+ test.breakable_loop iter_args(%inner = %outer) : i32 {
+ // CHECK: %[[DEPTH_ONE:.*]] = arith.constant 1 : index
+ // CHECK: %[[IS_ONE:.*]] = arith.cmpi eq, %[[DEPTH]], %[[DEPTH_ONE]] : index
+ // CHECK: scf.if %[[IS_ONE]] {
+ // CHECK: scf.continue [%[[INNER_TOKEN]]] %[[INNER_ARG]] : i32
+ // CHECK: %[[DEPTH_TWO:.*]] = arith.constant 2 : index
+ // CHECK: %[[IS_TWO:.*]] = arith.cmpi eq, %[[DEPTH]], %[[DEPTH_TWO]] : index
+ // CHECK: scf.if %[[IS_TWO]] {
+ // CHECK: scf.continue [%[[OUTER_TOKEN]]] %[[INNER_ARG]] : i32
+ // CHECK: scf.continue [%[[INNER_TOKEN]]] %[[INNER_ARG]] : i32
+ test.dynamic_continue %depth %inner : i32
+ }
+ test.dynamic_break %c1 %outer : i32
+ }
+ // CHECK: return %[[RESULT]] : i32
+ return %result : i32
+}
+
+// -----
+
+// The innermost loop carries f32, so it is not compatible with the i32 break
+// payload and must be filtered out. The remaining compatible targets are still
+// addressed by their original depths: 2 for the middle loop and 3 for the outer
+// loop.
+// CHECK-LABEL: func.func @dynamic_break_dispatch_filtered_depths(
+// CHECK-SAME: %[[CHOOSE:.*]]: i1, %[[INIT:.*]]: i32, %[[F_INIT:.*]]: f32
+func.func @dynamic_break_dispatch_filtered_depths(%choose_outer: i1,
+ %init: i32, %inner_init: f32)
+ -> i32 {
+ %c1 = arith.constant 1 : index
+ %c2 = arith.constant 2 : index
+ %c3 = arith.constant 3 : index
+ // CHECK: %[[DEPTH:.*]] = arith.select %[[CHOOSE]],
+ %depth = arith.select %choose_outer, %c3, %c2 : index
+ // CHECK: %[[RESULT:.*]] = scf.loop token(%[[OUTER_TOKEN:.*]]) iter_args(%[[OUTER_ARG:.*]] =
+ %result = test.breakable_loop iter_args(%outer = %init) : i32 -> i32 {
+ // CHECK: %[[MIDDLE_RESULT:.*]] = scf.loop token(%[[MIDDLE_TOKEN:.*]]) iter_args(%[[MIDDLE_ARG:.*]] =
+ %middle_result = test.breakable_loop iter_args(%middle = %outer) : i32 -> i32 {
+ // CHECK: scf.loop token(%[[INNER_TOKEN:.*]]) iter_args(%[[INNER_ARG:.*]] =
+ test.breakable_loop iter_args(%inner = %inner_init) : f32 {
+ // CHECK: %[[DEPTH_TWO:.*]] = arith.constant 2 : index
+ // CHECK: %[[IS_TWO:.*]] = arith.cmpi eq, %[[DEPTH]], %[[DEPTH_TWO]] : index
+ // CHECK: scf.if %[[IS_TWO]] {
+ // CHECK: scf.break [%[[MIDDLE_TOKEN]]] %[[MIDDLE_ARG]] : i32
+ // CHECK: %[[DEPTH_THREE:.*]] = arith.constant 3 : index
+ // CHECK: %[[IS_THREE:.*]] = arith.cmpi eq, %[[DEPTH]], %[[DEPTH_THREE]] : index
+ // CHECK: scf.if %[[IS_THREE]] {
+ // CHECK: scf.break [%[[OUTER_TOKEN]]] %[[MIDDLE_ARG]] : i32
+ // CHECK: scf.break [%[[MIDDLE_TOKEN]]] %[[MIDDLE_ARG]] : i32
+ test.dynamic_break %depth %middle : i32
+ }
+ test.dynamic_break %c1 %middle : i32
+ }
+ test.dynamic_break %c1 %middle_result : i32
+ }
+ // CHECK: return %[[RESULT]] : i32
+ return %result : i32
+}
diff --git a/mlir/test/lib/Dialect/Test/TestOpDefs.cpp b/mlir/test/lib/Dialect/Test/TestOpDefs.cpp
index c6470da1b7852..b38e6b1455c23 100644
--- a/mlir/test/lib/Dialect/Test/TestOpDefs.cpp
+++ b/mlir/test/lib/Dialect/Test/TestOpDefs.cpp
@@ -10,6 +10,7 @@
#include "TestOps.h"
#include "mlir/Dialect/Bufferization/IR/Bufferization.h"
#include "mlir/Dialect/Tensor/IR/Tensor.h"
+#include "mlir/IR/Matchers.h"
#include "mlir/IR/Verifier.h"
#include "mlir/Interfaces/FunctionImplementation.h"
#include "mlir/Interfaces/MemorySlotInterfaces.h"
@@ -827,6 +828,345 @@ void AnyCondOp::getRegionInvocationBounds(
invocationBounds.emplace_back(1, 1);
}
+//===----------------------------------------------------------------------===//
+// TestBreakableLoopOp / TestDynamicBreakOp / TestDynamicContinueOp
+//===----------------------------------------------------------------------===//
+
+static std::optional<int64_t> getStaticDepth(Value depth) {
+ APInt value;
+ if (!matchPattern(depth, m_ConstantInt(&value)))
+ return std::nullopt;
+ return value.getSExtValue();
+}
+
+static SmallVector<TestBreakableLoopOp>
+getReachableBreakableLoops(Operation *terminator) {
+ SmallVector<TestBreakableLoopOp> loops;
+ for (Operation *currentOp = terminator->getParentOp(); currentOp;
+ currentOp = currentOp->getParentOp()) {
+ if (auto loopOp = dyn_cast<TestBreakableLoopOp>(currentOp))
+ loops.push_back(loopOp);
+ if (!currentOp->mightHaveTrait<OpTrait::PropagateControlFlowBreak>())
+ break;
+ }
+ return loops;
+}
+
+static Operation *getPropagationBlockingParent(Operation *terminator) {
+ for (Operation *currentOp = terminator->getParentOp(); currentOp;
+ currentOp = currentOp->getParentOp()) {
+ if (currentOp->mightHaveTrait<OpTrait::PropagateControlFlowBreak>())
+ continue;
+ for (Operation *parentOp = currentOp->getParentOp(); parentOp;
+ parentOp = parentOp->getParentOp())
+ if (isa<TestBreakableLoopOp>(parentOp))
+ return currentOp;
+ return nullptr;
+ }
+ return nullptr;
+}
+
+static SmallVector<TestBreakableLoopOp>
+getResolvedBreakableLoops(Operation *terminator, Value depth) {
+ SmallVector<TestBreakableLoopOp> loops =
+ getReachableBreakableLoops(terminator);
+ std::optional<int64_t> staticDepth = getStaticDepth(depth);
+ if (!staticDepth)
+ return loops;
+ if (*staticDepth <= 0 || static_cast<uint64_t>(*staticDepth) > loops.size())
+ return {};
+ return {loops[*staticDepth - 1]};
+}
+
+static OperandRange getTerminatorPayload(Operation *terminator) {
+ if (auto breakOp = dyn_cast<TestDynamicBreakOp>(terminator))
+ return breakOp.getArgs();
+ return cast<TestDynamicContinueOp>(terminator).getArgs();
+}
+
+static bool isContinueTerminator(Operation *terminator) {
+ if (isa<TestDynamicContinueOp>(terminator))
+ return true;
+ assert(isa<TestDynamicBreakOp>(terminator) &&
+ "expected test.dynamic_break or test.dynamic_continue");
+ return false;
+}
+
+static ValueRange getTargetExpectedValues(TestBreakableLoopOp target,
+ bool isContinue) {
+ return isContinue ? ValueRange(target.getRegionIterArgs())
+ : ValueRange(target.getResults());
+}
+
+static bool isTerminatorPayloadCompatible(OperandRange args,
+ ValueRange expectedValues) {
+ if (args.size() != expectedValues.size())
+ return false;
+ return llvm::all_of(llvm::zip(args, expectedValues), [](auto values) {
+ return std::get<0>(values).getType() == std::get<1>(values).getType();
+ });
+}
+
+static bool isPotentialTerminatorTarget(Operation *terminator,
+ OperandRange args,
+ TestBreakableLoopOp target,
+ bool isContinue) {
+ return target.acceptsTerminator(terminator) &&
+ isTerminatorPayloadCompatible(
+ args, getTargetExpectedValues(target, isContinue));
+}
+
+static SmallVector<Operation *>
+getPotentialBreakableLoopTargets(Operation *terminator, Value depth) {
+ SmallVector<Operation *> targets;
+ OperandRange args = getTerminatorPayload(terminator);
+ bool isContinue = isContinueTerminator(terminator);
+ for (TestBreakableLoopOp loopOp :
+ getResolvedBreakableLoops(terminator, depth)) {
+ if (!isPotentialTerminatorTarget(terminator, args, loopOp, isContinue))
+ continue;
+ targets.push_back(loopOp.getOperation());
+ }
+ return targets;
+}
+
+static void
+appendTestBreakableLoopSuccessor(TestBreakableLoopOp loopOp, bool isContinue,
+ SmallVectorImpl<RegionSuccessor> ®ions) {
+ regions.push_back(isContinue ? RegionSuccessor(&loopOp.getBody())
+ : RegionSuccessor(loopOp.getOperation()));
+}
+
+static LogicalResult verifyTerminatorPayload(Operation *terminator,
+ OperandRange args,
+ ValueRange expectedValues,
+ TestBreakableLoopOp target,
+ StringRef targetKind) {
+ if (args.size() != expectedValues.size())
+ return terminator->emitOpError()
+ << "has " << args.size() << " payload operands, but target "
+ << targetKind << " expects " << expectedValues.size();
+
+ for (auto [index, argAndExpected] :
+ llvm::enumerate(llvm::zip(args, expectedValues))) {
+ Value arg = std::get<0>(argAndExpected);
+ Value expected = std::get<1>(argAndExpected);
+ if (arg.getType() == expected.getType())
+ continue;
+ return terminator->emitOpError()
+ << "payload operand #" << index << " has type " << arg.getType()
+ << ", but target " << targetKind << " of "
+ << OpWithFlags(target, OpPrintingFlags().skipRegions())
+ << " expects " << expected.getType();
+ }
+ return success();
+}
+
+static LogicalResult verifyDynamicTerminator(Operation *terminator, Value depth,
+ OperandRange args,
+ bool isContinue) {
+ SmallVector<TestBreakableLoopOp> reachableLoops =
+ getReachableBreakableLoops(terminator);
+ Operation *blockingOp = getPropagationBlockingParent(terminator);
+ if (reachableLoops.empty()) {
+ if (blockingOp)
+ return terminator->emitOpError()
+ << "depth target crosses an op that does not have the "
+ "PropagateControlFlowBreak trait: "
+ << OpWithFlags(blockingOp, OpPrintingFlags().skipRegions());
+ return terminator->emitOpError()
+ << "must be nested inside a reachable test.breakable_loop";
+ }
+
+ SmallVector<TestBreakableLoopOp> targets = reachableLoops;
+ if (std::optional<int64_t> staticDepth = getStaticDepth(depth)) {
+ if (*staticDepth <= 0)
+ return terminator->emitOpError() << "depth must be positive";
+ if (static_cast<uint64_t>(*staticDepth) > reachableLoops.size()) {
+ if (blockingOp)
+ return terminator->emitOpError()
+ << "depth target crosses an op that does not have the "
+ "PropagateControlFlowBreak trait: "
+ << OpWithFlags(blockingOp, OpPrintingFlags().skipRegions());
+ return terminator->emitOpError()
+ << "constant depth exceeds the number of reachable "
+ "test.breakable_loop operations";
+ }
+ targets.clear();
+ targets.push_back(reachableLoops[*staticDepth - 1]);
+ } else if (blockingOp) {
+ return terminator->emitOpError()
+ << "dynamic depth may target across an op that does not have the "
+ "PropagateControlFlowBreak trait: "
+ << OpWithFlags(blockingOp, OpPrintingFlags().skipRegions());
+ }
+
+ if (!getStaticDepth(depth)) {
+ for (TestBreakableLoopOp target : targets)
+ if (isPotentialTerminatorTarget(terminator, args, target, isContinue))
+ return success();
+ return terminator->emitOpError()
+ << "dynamic depth has no compatible test.breakable_loop target";
+ }
+
+ for (TestBreakableLoopOp target : targets) {
+ if (!target.acceptsTerminator(terminator))
+ return target.emitOpError("does not accept terminator: ")
+ << OpWithFlags(terminator, OpPrintingFlags().skipRegions());
+ ValueRange expectedValues = getTargetExpectedValues(target, isContinue);
+ if (failed(verifyTerminatorPayload(terminator, args, expectedValues, target,
+ isContinue ? "loop-carried values"
+ : "results")))
+ return failure();
+ }
+ return success();
+}
+
+void TestBreakableLoopOp::print(OpAsmPrinter &p) {
+ p << " ";
+ if (!getInitArgs().empty()) {
+ p << "iter_args(";
+ llvm::interleaveComma(
+ llvm::zip(getRegionIterArgs(), getInitArgs()), p,
+ [&](auto it) { p << std::get<0>(it) << " = " << std::get<1>(it); });
+ p << ") : " << getInitArgs().getTypes() << " ";
+ }
+ if (!getResultTypes().empty())
+ p << "-> " << getResultTypes() << " ";
+ p.printRegion(getBody(), /*printEntryBlockArgs=*/false,
+ /*printBlockTerminators=*/true);
+ p.printOptionalAttrDict((*this)->getAttrs());
+}
+
+ParseResult TestBreakableLoopOp::parse(OpAsmParser &parser,
+ OperationState &result) {
+ SmallVector<OpAsmParser::Argument, 4> regionArgs;
+ SmallVector<OpAsmParser::Argument, 4> iterRegionArgs;
+ SmallVector<OpAsmParser::UnresolvedOperand, 4> iterOperands;
+ SmallVector<Type, 4> iterTypes;
+
+ if (failed(parser.parseOptionalKeyword("iter_args"))) {
+ if (succeeded(parser.parseOptionalArrow()))
+ if (parser.parseTypeList(result.types))
+ return failure();
+ } else {
+ if (parser.parseAssignmentList(iterRegionArgs, iterOperands) ||
+ parser.parseColon() || parser.parseTypeList(iterTypes))
+ return failure();
+ if (iterRegionArgs.size() != iterTypes.size())
+ return parser.emitError(parser.getCurrentLocation(),
+ "found different number of iter_args and types");
+ if (succeeded(parser.parseOptionalArrow()))
+ if (parser.parseTypeList(result.types))
+ return failure();
+ for (auto [regionArg, type] : llvm::zip_equal(iterRegionArgs, iterTypes))
+ regionArg.type = type;
+ llvm::append_range(regionArgs, iterRegionArgs);
+ }
+
+ Region *body = result.addRegion();
+ if (parser.parseRegion(*body, regionArgs) ||
+ parser.parseOptionalAttrDict(result.attributes))
+ return failure();
+
+ return parser.resolveOperands(iterOperands, iterTypes, parser.getNameLoc(),
+ result.operands);
+}
+
+LogicalResult TestBreakableLoopOp::verifyRegions() {
+ if (getBody().empty())
+ return emitOpError("region cannot be empty");
+ if (getBody().front().getNumArguments() != getNumOperands())
+ return emitOpError("expected the region to have one argument per "
+ "loop-carried value (")
+ << getNumOperands() << " expected, but got "
+ << getBody().front().getNumArguments() << ")";
+ for (auto [index, argAndOperand] :
+ llvm::enumerate(llvm::zip(getRegionIterArgs(), getOperands()))) {
+ Type argType = std::get<0>(argAndOperand).getType();
+ Type operandType = std::get<1>(argAndOperand).getType();
+ if (argType != operandType)
+ return emitOpError() << "types mismatch between " << index
+ << "th iter operand (" << operandType
+ << ") and defined region argument (" << argType
+ << ")";
+ }
+ return success();
+}
+
+void TestBreakableLoopOp::getSuccessorRegions(
+ RegionBranchPoint point, SmallVectorImpl<RegionSuccessor> ®ions) {
+ if (point.isParent()) {
+ regions.push_back(RegionSuccessor(&getBody()));
+ return;
+ }
+
+ Operation *terminator = point.getTerminatorPredecessorOrNull();
+ Value depth;
+ bool isContinue = false;
+ if (auto continueOp = dyn_cast<TestDynamicContinueOp>(terminator)) {
+ depth = continueOp.getDepth();
+ isContinue = true;
+ } else if (auto breakOp = dyn_cast<TestDynamicBreakOp>(terminator)) {
+ depth = breakOp.getDepth();
+ } else {
+ llvm_unreachable("expected test.dynamic_break or test.dynamic_continue");
+ }
+
+ bool isImmediateBranch = terminator->getParentOp() == getOperation();
+ for (Operation *target :
+ getPotentialBreakableLoopTargets(terminator, depth)) {
+ if (target == getOperation()) {
+ appendTestBreakableLoopSuccessor(*this, isContinue, regions);
+ continue;
+ }
+ if (!isImmediateBranch || !target->isProperAncestor(getOperation()))
+ continue;
+ appendTestBreakableLoopSuccessor(cast<TestBreakableLoopOp>(target),
+ isContinue, regions);
+ }
+}
+
+OperandRange
+TestBreakableLoopOp::getEntrySuccessorOperands(RegionSuccessor successor) {
+ return getInitArgs();
+}
+
+ValueRange TestBreakableLoopOp::getSuccessorInputs(RegionSuccessor successor) {
+ return successor.isOperation() ? ValueRange(getResults())
+ : ValueRange(getRegionIterArgs());
+}
+
+LogicalResult TestDynamicBreakOp::verify() {
+ return verifyDynamicTerminator(getOperation(), getDepth(), getArgs(),
+ /*isContinue=*/false);
+}
+
+SmallVector<Operation *> TestDynamicBreakOp::getPotentialTargets() {
+ return getPotentialBreakableLoopTargets(getOperation(), getDepth());
+}
+
+MutableOperandRange
+TestDynamicBreakOp::getMutableSuccessorOperands(RegionSuccessor point) {
+ return MutableOperandRange(getOperation(), /*start=*/1,
+ /*length=*/getOperation()->getNumOperands() - 1);
+}
+
+LogicalResult TestDynamicContinueOp::verify() {
+ return verifyDynamicTerminator(getOperation(), getDepth(), getArgs(),
+ /*isContinue=*/true);
+}
+
+SmallVector<Operation *> TestDynamicContinueOp::getPotentialTargets() {
+ return getPotentialBreakableLoopTargets(getOperation(), getDepth());
+}
+
+MutableOperandRange
+TestDynamicContinueOp::getMutableSuccessorOperands(RegionSuccessor point) {
+ return MutableOperandRange(getOperation(), /*start=*/1,
+ /*length=*/getOperation()->getNumOperands() - 1);
+}
+
//===----------------------------------------------------------------------===//
// SingleBlockImplicitTerminatorOp
//===----------------------------------------------------------------------===//
diff --git a/mlir/test/lib/Dialect/Test/TestOps.td b/mlir/test/lib/Dialect/Test/TestOps.td
index 9d7e21cabda6f..1439cf405d28a 100644
--- a/mlir/test/lib/Dialect/Test/TestOps.td
+++ b/mlir/test/lib/Dialect/Test/TestOps.td
@@ -2850,6 +2850,73 @@ def LoopBlockTerminatorOp : TEST_Op<"loop_block_term",
}];
}
+def TestBreakableLoopOp : TEST_Op<"breakable_loop", [
+ AutomaticAllocationScope,
+ RecursiveMemoryEffects,
+ PropagateControlFlowBreak,
+ SingleBlock,
+ DeclareOpInterfaceMethods<RegionBranchOpInterface,
+ ["getEntrySuccessorOperands", "getSuccessorInputs"]>,
+ HasBreakingControlFlowOpInterface,
+ HasNestedTerminator<["TestDynamicContinueOp", "TestDynamicBreakOp"]>
+ ]> {
+ let summary = "test loop with dynamically addressed break/continue";
+ let description = [{
+ Test dialect loop operation that models breaking control flow with
+ dynamically selected targets. Its terminators use an SSA index value to
+ count enclosing `test.breakable_loop` operations. A depth of one targets the
+ immediately enclosing test loop, a depth of two targets the next enclosing
+ test loop, and so on. A constant depth must identify a reachable target with
+ compatible payload types. A dynamic depth may evaluate to any reachable depth
+ with compatible payload types; evaluating to zero, to an incompatible target,
+ or to a depth greater than the number of reachable loops is undefined
+ behavior.
+ }];
+
+ let arguments = (ins Variadic<AnyType>:$initArgs);
+ let results = (outs Variadic<AnyType>:$results);
+ let regions = (region SizedRegion<1>:$body);
+
+ let extraClassDeclaration = [{
+ ::mlir::Block::BlockArgListType getRegionIterArgs() {
+ return getBody().getArguments();
+ }
+ }];
+
+ let hasCustomAssemblyFormat = 1;
+ let hasRegionVerifier = 1;
+}
+
+def TestDynamicBreakOp : TEST_Op<"dynamic_break", [
+ Terminator, RegionExitTerminatorOpInterface,
+ DeclareOpInterfaceMethods<RegionBranchTerminatorOpInterface,
+ ["getMutableSuccessorOperands"]>
+ ]> {
+ let summary = "dynamically targeted break from test.breakable_loop";
+ let arguments = (ins Index:$depth, Variadic<AnyType>:$args);
+ let assemblyFormat = "$depth ($args^ `:` type($args))? attr-dict";
+ let extraClassDeclaration = [{
+ ::llvm::SmallVector<::mlir::Operation *>
+ getPotentialTargets();
+ }];
+ let hasVerifier = 1;
+}
+
+def TestDynamicContinueOp : TEST_Op<"dynamic_continue", [
+ Terminator, RegionExitTerminatorOpInterface,
+ DeclareOpInterfaceMethods<RegionBranchTerminatorOpInterface,
+ ["getMutableSuccessorOperands"]>
+ ]> {
+ let summary = "dynamically targeted continue of test.breakable_loop";
+ let arguments = (ins Index:$depth, Variadic<AnyType>:$args);
+ let assemblyFormat = "$depth ($args^ `:` type($args))? attr-dict";
+ let extraClassDeclaration = [{
+ ::llvm::SmallVector<::mlir::Operation *>
+ getPotentialTargets();
+ }];
+ let hasVerifier = 1;
+}
+
def TestNoTerminatorOp : TEST_Op<"switch_with_no_break", [
NoTerminator,
DeclareOpInterfaceMethods<RegionBranchOpInterface>
diff --git a/mlir/test/lib/Dialect/Test/TestPatterns.cpp b/mlir/test/lib/Dialect/Test/TestPatterns.cpp
index 552a1a473c9fd..1ace1d1101df0 100644
--- a/mlir/test/lib/Dialect/Test/TestPatterns.cpp
+++ b/mlir/test/lib/Dialect/Test/TestPatterns.cpp
@@ -14,9 +14,11 @@
#include "mlir/Dialect/ControlFlow/Transforms/StructuralTypeConversions.h"
#include "mlir/Dialect/Func/IR/FuncOps.h"
#include "mlir/Dialect/Func/Transforms/FuncConversions.h"
+#include "mlir/Dialect/SCF/IR/SCF.h"
#include "mlir/Dialect/SCF/Transforms/Patterns.h"
#include "mlir/Dialect/Tensor/IR/Tensor.h"
#include "mlir/IR/BuiltinAttributes.h"
+#include "mlir/IR/BuiltinTypes.h"
#include "mlir/IR/Matchers.h"
#include "mlir/IR/PatternMatch.h"
#include "mlir/IR/Visitors.h"
@@ -25,6 +27,7 @@
#include "mlir/Transforms/FoldUtils.h"
#include "mlir/Transforms/GreedyPatternRewriteDriver.h"
#include "mlir/Transforms/WalkPatternRewriteDriver.h"
+#include "llvm/ADT/DenseMap.h"
#include "llvm/ADT/ScopeExit.h"
#include <cstdint>
@@ -2490,6 +2493,214 @@ struct TestFoldTypeConvertingOp
signalPassFailure();
}
};
+
+class TestConvertBreakableLoopToSCFPass
+ : public PassWrapper<TestConvertBreakableLoopToSCFPass, OperationPass<>> {
+public:
+ MLIR_DEFINE_EXPLICIT_INTERNAL_INLINE_TYPE_ID(
+ TestConvertBreakableLoopToSCFPass)
+
+ StringRef getArgument() const final {
+ return "test-convert-breakable-loop-to-scf";
+ }
+ StringRef getDescription() const final {
+ return "Convert test.breakable_loop operations to scf.loop operations";
+ }
+
+ void runOnOperation() override;
+
+ struct TargetInfo {
+ Operation *target;
+ int64_t depth;
+ };
+ using TargetMap = llvm::DenseMap<Operation *, SmallVector<TargetInfo>>;
+ using TokenMap = llvm::DenseMap<Operation *, Value>;
+
+private:
+ LogicalResult rewriteLoop(TestBreakableLoopOp loopOp,
+ const TargetMap &terminatorTargets,
+ TokenMap &loopTokens);
+ LogicalResult rewriteTerminator(Operation *terminator,
+ const TargetMap &terminatorTargets,
+ const TokenMap &loopTokens,
+ PatternRewriter &rewriter);
+};
+
+static void appendTerminatorTargets(
+ Operation *terminator,
+ SmallVectorImpl<TestConvertBreakableLoopToSCFPass::TargetInfo> &targets,
+ SmallVector<Operation *> potentialTargets) {
+ for (Operation *targetOp : potentialTargets) {
+ int64_t depth = 0;
+ for (Operation *currentOp = terminator->getParentOp(); currentOp;
+ currentOp = currentOp->getParentOp()) {
+ if (isa<TestBreakableLoopOp>(currentOp)) {
+ ++depth;
+ if (currentOp == targetOp) {
+ targets.push_back({targetOp, depth});
+ break;
+ }
+ }
+ if (!currentOp->mightHaveTrait<OpTrait::PropagateControlFlowBreak>())
+ break;
+ }
+ }
+}
+
+static bool isBreakableLoopTerminator(Operation *op) {
+ return isa<TestDynamicBreakOp, TestDynamicContinueOp>(op);
+}
+
+static OperandRange getTerminatorPayload(Operation *terminator) {
+ if (auto breakOp = dyn_cast<TestDynamicBreakOp>(terminator))
+ return breakOp.getArgs();
+ return cast<TestDynamicContinueOp>(terminator).getArgs();
+}
+
+static Value getTerminatorDepth(Operation *terminator) {
+ if (auto breakOp = dyn_cast<TestDynamicBreakOp>(terminator))
+ return breakOp.getDepth();
+ return cast<TestDynamicContinueOp>(terminator).getDepth();
+}
+
+static void createSCFTerminator(PatternRewriter &rewriter, Location loc,
+ Operation *terminator, Value targetToken,
+ ValueRange args) {
+ if (isa<TestDynamicBreakOp>(terminator)) {
+ scf::BreakOp::create(rewriter, loc, targetToken, args);
+ return;
+ }
+ scf::ContinueOp::create(rewriter, loc, targetToken, args);
+}
+
+static std::optional<Value>
+lookupToken(Operation *terminator, Operation *target,
+ const TestConvertBreakableLoopToSCFPass::TokenMap &loopTokens) {
+ auto it = loopTokens.find(target);
+ if (it != loopTokens.end())
+ return it->second;
+ terminator->emitOpError()
+ << "cannot convert because target loop has not been converted";
+ return std::nullopt;
+}
+
+LogicalResult TestConvertBreakableLoopToSCFPass::rewriteTerminator(
+ Operation *terminator, const TargetMap &terminatorTargets,
+ const TokenMap &loopTokens, PatternRewriter &rewriter) {
+ auto targetIt = terminatorTargets.find(terminator);
+ if (targetIt == terminatorTargets.end() || targetIt->second.empty())
+ return terminator->emitOpError()
+ << "cannot convert terminator without a target loop";
+
+ ArrayRef<TargetInfo> targets = targetIt->second;
+ Location loc = terminator->getLoc();
+ SmallVector<Value> args(getTerminatorPayload(terminator));
+
+ if (targets.size() == 1) {
+ std::optional<Value> token =
+ lookupToken(terminator, targets.front().target, loopTokens);
+ if (!token)
+ return failure();
+ rewriter.setInsertionPoint(terminator);
+ createSCFTerminator(rewriter, loc, terminator, *token, args);
+ rewriter.eraseOp(terminator);
+ return success();
+ }
+
+ rewriter.setInsertionPoint(terminator);
+ Value depth = getTerminatorDepth(terminator);
+ for (TargetInfo target : targets) {
+ std::optional<Value> token =
+ lookupToken(terminator, target.target, loopTokens);
+ if (!token)
+ return failure();
+
+ Value depthValue =
+ arith::ConstantIndexOp::create(rewriter, loc, target.depth);
+ Value isTarget = arith::CmpIOp::create(
+ rewriter, loc, arith::CmpIPredicate::eq, depth, depthValue);
+ auto ifOp = scf::IfOp::create(rewriter, loc, TypeRange{}, isTarget,
+ /*addThenBlock=*/true,
+ /*addElseBlock=*/false);
+ {
+ OpBuilder::InsertionGuard guard(rewriter);
+ rewriter.setInsertionPointToStart(ifOp.thenBlock());
+ createSCFTerminator(rewriter, loc, terminator, *token, args);
+ }
+ rewriter.setInsertionPointAfter(ifOp);
+ }
+
+ // Dynamic depths are verified to be able to target any listed loop. If the
+ // value falls outside that set, selects an incompatible target, or otherwise
+ // has undefined behavior, the fallback may choose any valid target. Use the
+ // first compatible target to keep the generated IR deterministic.
+ std::optional<Value> fallbackToken =
+ lookupToken(terminator, targets.front().target, loopTokens);
+ if (!fallbackToken)
+ return failure();
+ createSCFTerminator(rewriter, loc, terminator, *fallbackToken, args);
+ rewriter.eraseOp(terminator);
+ return success();
+}
+
+LogicalResult TestConvertBreakableLoopToSCFPass::rewriteLoop(
+ TestBreakableLoopOp loopOp, const TargetMap &terminatorTargets,
+ TokenMap &loopTokens) {
+ PatternRewriter rewriter(loopOp.getContext());
+ rewriter.setInsertionPoint(loopOp);
+
+ auto scfLoop = scf::LoopOp::create(
+ rewriter, loopOp.getLoc(), loopOp.getResultTypes(), loopOp.getInitArgs());
+ scfLoop->setAttrs(loopOp->getAttrs());
+ scfLoop.getRegion().takeBody(loopOp.getBody());
+ scfLoop.getBody()->insertArgument(
+ /*index=*/0u, TokenType::get(loopOp.getContext()), loopOp.getLoc());
+ loopTokens[loopOp.getOperation()] = scfLoop.getControlToken();
+
+ SmallVector<Operation *> terminators;
+ scfLoop.getRegion().walk([&](Operation *op) {
+ if (!isBreakableLoopTerminator(op))
+ return;
+ if (op->getParentOfType<TestBreakableLoopOp>())
+ return;
+ terminators.push_back(op);
+ });
+
+ for (Operation *terminator : terminators)
+ if (failed(rewriteTerminator(terminator, terminatorTargets, loopTokens,
+ rewriter)))
+ return failure();
+
+ rewriter.replaceOp(loopOp, scfLoop.getResults());
+ return success();
+}
+
+void TestConvertBreakableLoopToSCFPass::runOnOperation() {
+ SmallVector<TestBreakableLoopOp> loops;
+ TargetMap terminatorTargets;
+ getOperation()->walk<WalkOrder::PreOrder>([&](Operation *op) {
+ if (auto loopOp = dyn_cast<TestBreakableLoopOp>(op)) {
+ loops.push_back(loopOp);
+ return;
+ }
+ if (auto breakOp = dyn_cast<TestDynamicBreakOp>(op)) {
+ appendTerminatorTargets(op, terminatorTargets[op],
+ breakOp.getPotentialTargets());
+ return;
+ }
+ if (auto continueOp = dyn_cast<TestDynamicContinueOp>(op))
+ appendTerminatorTargets(op, terminatorTargets[op],
+ continueOp.getPotentialTargets());
+ });
+
+ TokenMap loopTokens;
+ for (TestBreakableLoopOp loopOp : loops) {
+ if (failed(rewriteLoop(loopOp, terminatorTargets, loopTokens))) {
+ signalPassFailure();
+ return;
+ }
+ }
+}
} // namespace
//===----------------------------------------------------------------------===//
@@ -2522,6 +2733,8 @@ void registerPatternsTestPass() {
PassRegistration<TestSelectiveReplacementPatternDriver>();
PassRegistration<TestFoldTypeConvertingOp>();
+
+ PassRegistration<TestConvertBreakableLoopToSCFPass>();
}
} // namespace test
} // namespace mlir
diff --git a/mlir/test/mlir-runner/test-convert-breakable-loop-to-scf.mlir b/mlir/test/mlir-runner/test-convert-breakable-loop-to-scf.mlir
new file mode 100644
index 0000000000000..d2a9ccfc84650
--- /dev/null
+++ b/mlir/test/mlir-runner/test-convert-breakable-loop-to-scf.mlir
@@ -0,0 +1,145 @@
+// RUN: mlir-opt %s -pass-pipeline="builtin.module(func.func(test-convert-breakable-loop-to-scf,convert-scf-to-cf,canonicalize,convert-arith-to-llvm),convert-func-to-llvm,convert-cf-to-llvm,reconcile-unrealized-casts)" -o %t.mlir
+// RUN: mlir-runner %t.mlir -e immediate_continue -entry-point-result=i32 | FileCheck %s --check-prefix=IMMEDIATE
+// RUN: mlir-runner %t.mlir -e break_with_result -entry-point-result=i32 | FileCheck %s --check-prefix=BREAK
+// RUN: mlir-runner %t.mlir -e constant_outer_break -entry-point-result=i32 | FileCheck %s --check-prefix=OUTER-BREAK
+// RUN: mlir-runner %t.mlir -e constant_outer_continue -entry-point-result=i32 | FileCheck %s --check-prefix=OUTER-CONTINUE
+// RUN: mlir-runner %t.mlir -e dynamic_break_depth_1 -entry-point-result=i32 | FileCheck %s --check-prefix=DYNAMIC-ONE
+// RUN: mlir-runner %t.mlir -e dynamic_break_depth_2 -entry-point-result=i32 | FileCheck %s --check-prefix=DYNAMIC-TWO
+// RUN: mlir-runner %t.mlir -e dynamic_filtered_depth_2 -entry-point-result=i32 | FileCheck %s --check-prefix=FILTERED-TWO
+// RUN: mlir-runner %t.mlir -e dynamic_filtered_depth_3 -entry-point-result=i32 | FileCheck %s --check-prefix=FILTERED-THREE
+// XFAIL: system-aix
+
+func.func @immediate_continue() -> i32 {
+ %depth = arith.constant 1 : index
+ %c0 = arith.constant 0 : i32
+ %c1 = arith.constant 1 : i32
+ %c4 = arith.constant 4 : i32
+ %result = test.breakable_loop iter_args(%i = %c0) : i32 -> i32 {
+ %done = arith.cmpi eq, %i, %c4 : i32
+ scf.if %done {
+ test.dynamic_break %depth %i : i32
+ }
+ %next = arith.addi %i, %c1 : i32
+ test.dynamic_continue %depth %next : i32
+ }
+ return %result : i32
+}
+// IMMEDIATE: 4
+
+func.func @break_with_result() -> i32 {
+ %depth = arith.constant 1 : index
+ %value = arith.constant 17 : i32
+ %result = test.breakable_loop -> i32 {
+ test.dynamic_break %depth %value : i32
+ }
+ return %result : i32
+}
+// BREAK: 17
+
+func.func @constant_outer_break() -> i32 {
+ %inner_depth = arith.constant 1 : index
+ %outer_depth = arith.constant 2 : index
+ %outer_value = arith.constant 42 : i32
+ %fallback = arith.constant 13 : i32
+ %result = test.breakable_loop -> i32 {
+ test.breakable_loop {
+ test.dynamic_break %outer_depth %outer_value : i32
+ }
+ test.dynamic_break %inner_depth %fallback : i32
+ }
+ return %result : i32
+}
+// OUTER-BREAK: 42
+
+func.func @constant_outer_continue() -> i32 {
+ %inner_depth = arith.constant 1 : index
+ %outer_depth = arith.constant 2 : index
+ %c0 = arith.constant 0 : i32
+ %c1 = arith.constant 1 : i32
+ %c3 = arith.constant 3 : i32
+ %result = test.breakable_loop iter_args(%i = %c0) : i32 -> i32 {
+ %done = arith.cmpi eq, %i, %c3 : i32
+ scf.if %done {
+ test.dynamic_break %inner_depth %i : i32
+ }
+ test.breakable_loop {
+ %next = arith.addi %i, %c1 : i32
+ test.dynamic_continue %outer_depth %next : i32
+ }
+ test.dynamic_break %inner_depth %i : i32
+ }
+ return %result : i32
+}
+// OUTER-CONTINUE: 3
+
+// Distinguish the two compatible dynamic break targets at runtime. Depth 1
+// breaks the inner loop, then the outer loop adds one; depth 2 skips that outer
+// post-inner add by breaking directly to the outer loop result.
+func.func @dynamic_break(%choose_outer: i1) -> i32 {
+ %inner_depth = arith.constant 1 : index
+ %outer_depth = arith.constant 2 : index
+ %depth = arith.select %choose_outer, %outer_depth, %inner_depth : index
+ %value = arith.constant 7 : i32
+ %one = arith.constant 1 : i32
+ %result = test.breakable_loop -> i32 {
+ %inner_result = test.breakable_loop -> i32 {
+ test.dynamic_break %depth %value : i32
+ }
+ %after_inner = arith.addi %inner_result, %one : i32
+ test.dynamic_break %inner_depth %after_inner : i32
+ }
+ return %result : i32
+}
+
+func.func @dynamic_break_depth_1() -> i32 {
+ %choose_outer = arith.constant false
+ %result = func.call @dynamic_break(%choose_outer) : (i1) -> i32
+ return %result : i32
+}
+// DYNAMIC-ONE: 8
+
+func.func @dynamic_break_depth_2() -> i32 {
+ %choose_outer = arith.constant true
+ %result = func.call @dynamic_break(%choose_outer) : (i1) -> i32
+ return %result : i32
+}
+// DYNAMIC-TWO: 7
+
+// The innermost f32 loop is incompatible with the i32 break payload and is
+// filtered out. Depth 2 therefore targets the middle loop and observes the
+// outer post-middle add, while depth 3 targets the outer loop directly.
+func.func @dynamic_break_filtered(%choose_outer: i1) -> i32 {
+ %inner_depth = arith.constant 1 : index
+ %middle_depth = arith.constant 2 : index
+ %outer_depth = arith.constant 3 : index
+ %depth = arith.select %choose_outer, %outer_depth, %middle_depth : index
+ %value = arith.constant 5 : i32
+ %one = arith.constant 1 : i32
+ %f0 = arith.constant 0.000000e+00 : f32
+ %result = test.breakable_loop -> i32 {
+ %middle_result = test.breakable_loop -> i32 {
+ test.breakable_loop iter_args(%inner = %f0) : f32 {
+ test.dynamic_break %depth %value : i32
+ }
+ %after_inner = arith.addi %value, %one : i32
+ test.dynamic_break %inner_depth %after_inner : i32
+ }
+ %after_middle = arith.addi %middle_result, %one : i32
+ test.dynamic_break %inner_depth %after_middle : i32
+ }
+ return %result : i32
+}
+
+func.func @dynamic_filtered_depth_2() -> i32 {
+ %choose_outer = arith.constant false
+ %result = func.call @dynamic_break_filtered(%choose_outer) : (i1) -> i32
+ return %result : i32
+}
+// FILTERED-TWO: 6
+
+func.func @dynamic_filtered_depth_3() -> i32 {
+ %choose_outer = arith.constant true
+ %result = func.call @dynamic_break_filtered(%choose_outer) : (i1) -> i32
+ return %result : i32
+}
+// FILTERED-THREE: 5
More information about the Mlir-commits
mailing list