[Mlir-commits] [mlir] [MLIR] Introduce support for early exits (PR #166688)
Mehdi Amini
llvmlistbot at llvm.org
Tue Jun 2 05:51:32 PDT 2026
https://github.com/joker-eph updated https://github.com/llvm/llvm-project/pull/166688
>From 81ff979af48039236378b71c1793036796421732 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] [MLIR] Introduce early exit support via
scf.loop/break/continue
Add new SCF operations and core IR support for multi-level region exits:
- scf.loop: infinite loop terminated by scf.break, iterated by scf.continue
- scf.break N: exits N nested region levels, terminating a target loop
- scf.continue N: exits N nested region levels, re-entering a target loop
Each terminator carries a `num-breaking-regions` count stored in
Operation's trailing properties (gated by isBreakingControlFlowFlag).
The generic format prints/parses this as `[N]` in the successor list.
Traits and interfaces:
- PropagateControlFlowBreak: marks ops transparent to breaks (scf.if)
- HasBreakingControlFlowOpInterface: marks break receivers (scf.loop)
- RegionTerminator trait: marks terminators exiting multiple regions
- HasNestedTerminators trait: declares accepted terminator types
- ControlFlowImplicitTerminatorOpType: extends SingleBlockImplicit-
Terminator to accept multiple terminator types (yield/break/continue)
RegionSuccessor uses a Kind enum {Region, Parent, Propagating} instead
of a boolean flag, with propagating() sentinel for transparent ops.
resolveTerminatorSuccessors() resolves propagating breaks to the actual
HasBreakingControlFlowOpInterface ancestor, returning nullptr on failure.
Analysis updates:
- DeadCodeAnalysis, DenseAnalysis, SparseAnalysis resolve propagating
breaks through resolveTerminatorSuccessors
- Post-dominance accounts for breaking control flow via linear scan
- Verifier validates the break chain (intermediate PropagateControlFlow-
Break, target HasBreakingControlFlowOpInterface, acceptsTerminator)
- Inlining blocked only when breaks escape the region, not for
self-contained loops
SCF dialect changes:
- scf.if gains PropagateControlFlowBreak, accepts break/continue/yield
- IfOp: thenYield()/elseYield() removed; use thenTerminator()/
elseTerminator() returning Operation*
- ReduceOp marked ImmediateRegionTerminator
- ForOp/IfOp canonicalization patterns updated for break/continue
terminators (CombineIfs, CombineNestedIfs, ConvertTrivialIfToSelect,
ReplaceIfYieldWithConditionOrValue, RemoveStaticCondition, etc.)
- New patterns: SimplifyTrivialLoops, ForOpIterArgsFolder,
SimplifyTrivialForLoops, RemoveUnusedResults,
SimplifyIfWithBreakingControlFlowInBothBranches
- SCF-to-CF lowering: LoopOpLowering converts break/continue to
cf.branch; IfLowering adjusts break levels on inline
Also touches Operation::create (new numBreakingControlRegions param),
AsmParser (generic [N] syntax), AsmPrinter, several downstream
Operation::create callers, mlir-tblgen (num-breaking-regions format
directive, dependent trait error messages), and test infrastructure
(TestRegionBranchOpInterface pass, TestDominance updates).
Assisted-by: Claude Code
Assisted-by: Codex
---
mlir/docs/LangRef.md | 117 ++-
mlir/include/mlir/Dialect/SCF/IR/SCF.h | 57 ++
mlir/include/mlir/Dialect/SCF/IR/SCFOps.td | 191 +++-
mlir/include/mlir/IR/Diagnostics.h | 2 +-
mlir/include/mlir/IR/OpBase.td | 11 +
mlir/include/mlir/IR/OpDefinition.h | 16 +-
mlir/include/mlir/IR/Operation.h | 3 +-
mlir/include/mlir/IR/RegionKindInterface.h | 113 +++
mlir/include/mlir/IR/RegionKindInterface.td | 54 ++
.../mlir/Interfaces/ControlFlowInterfaces.h | 37 +-
.../Analysis/DataFlow/DeadCodeAnalysis.cpp | 19 +-
mlir/lib/Analysis/DataFlow/DenseAnalysis.cpp | 17 +-
mlir/lib/Analysis/DataFlow/SparseAnalysis.cpp | 13 +-
mlir/lib/AsmParser/Parser.cpp | 12 +-
.../SCFToControlFlow/SCFToControlFlow.cpp | 234 +++--
mlir/lib/Conversion/SCFToEmitC/SCFToEmitC.cpp | 2 +-
.../ShapeToStandard/ShapeToStandard.cpp | 3 +-
mlir/lib/Dialect/SCF/IR/SCF.cpp | 823 ++++++++++++++++--
.../SCF/IR/ValueBoundsOpInterfaceImpl.cpp | 4 +-
.../BufferizableOpInterfaceImpl.cpp | 4 +-
mlir/lib/IR/Diagnostics.cpp | 4 +-
mlir/lib/IR/Dominance.cpp | 29 +
mlir/lib/IR/Operation.cpp | 6 +-
mlir/lib/IR/PatternMatch.cpp | 5 +-
mlir/lib/IR/RegionKindInterface.cpp | 154 ++++
mlir/lib/IR/Verifier.cpp | 1 +
mlir/lib/Interfaces/ControlFlowInterfaces.cpp | 61 +-
mlir/lib/TableGen/Operator.cpp | 10 +-
mlir/lib/Transforms/Utils/CMakeLists.txt | 3 +
mlir/lib/Transforms/Utils/InliningUtils.cpp | 6 +
.../test-dead-code-analysis-early-exit.mlir | 101 +++
mlir/test/Analysis/test-dominance.mlir | 26 +
.../convert-early-exit-to-cfg.mlir | 135 +++
mlir/test/Dialect/SCF/loop_canonicalize.mlir | 376 ++++++++
mlir/test/IR/early-exit-invalid.mlir | 120 +++
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/IR/TestDominance.cpp | 33 +
mlir/test/lib/Interfaces/CMakeLists.txt | 1 +
.../RegionBranchOpInterface/CMakeLists.txt | 9 +
.../TestRegionBranchOpInterface.cpp | 76 ++
mlir/test/mlir-tblgen/op-error.td | 2 +-
mlir/tools/mlir-opt/CMakeLists.txt | 1 +
mlir/tools/mlir-opt/mlir-opt.cpp | 2 +
mlir/tools/mlir-tblgen/OpFormatGen.cpp | 3 -
.../FileLineColLocBreakpointManagerTest.cpp | 1 +
47 files changed, 2967 insertions(+), 211 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 0e6fb006da48b..a6f39c6747da5 100644
--- a/mlir/docs/LangRef.md
+++ b/mlir/docs/LangRef.md
@@ -296,7 +296,8 @@ generic-operation ::= string-literal `(` value-use-list? `)` successor-list
custom-operation ::= bare-id custom-operation-format
op-result-list ::= op-result (`,` op-result)* `=`
op-result ::= value-id (`:` integer-literal)?
-successor-list ::= `[` successor (`,` successor)* `]`
+successor-list ::= `[` successor-list-inner `]`
+successor-list-inner ::= successor (`,` successor)*
successor ::= caret-id (`:` block-arg-list)?
dictionary-properties ::= `<` dictionary-attribute `>`
region-list ::= `(` region (`,` region)* `)`
@@ -504,10 +505,20 @@ In MLIR, control flow semantics of a region is indicated by
regions support semantics where operations in a region 'execute sequentially'.
Before an operation executes, its operands have well-defined values. After an
operation executes, the operands have the same values and results also have
-well-defined values. After an operation executes, the next operation in the
-block executes until the operation is the terminator operation at the end of a
-block, in which case some other operation will execute. The determination of the
-next instruction to execute is the 'passing of control flow'.
+well-defined values.
+
+Usually, after an operation executes, the next operation in the block executes
+until the operation is the terminator operation at the end of a block, in which
+case the control will be transferred to another block or one of the parent
+operations. The determination of the next instruction to execute is the 'passing
+of control flow'. The control-flow can be interrupted by an operation if it
+defines the `PropagateControlFlowBreak` trait. Such an operation does not handle
+the break itself; it transparently propagates it outward to an ancestor
+operation that implements `HasBreakingControlFlowOpInterface`. The actual break
+is initiated by a nested [Region Terminator](#region-terminator). Every
+operation between the `RegionTerminator` and 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.
@@ -515,22 +526,23 @@ However, when control flow enters a region, it always begins in the first block
of the region, called the *entry* block. Terminator operations ending each block
represent control flow by explicitly specifying the successor blocks of the
block. Control flow can only pass to one of the specified successor blocks as in
-a `branch` operation, or back to the containing operation as in a `return`
-operation. Terminator operations without successors can only pass control back
-to the containing operation. Within these restrictions, the particular semantics
-of terminator operations is determined by the specific dialect operations
-involved. Blocks (other than the entry block) that are not listed as a successor
-of a terminator operation are defined to be unreachable and can be removed
-without affecting the semantics of the containing operation.
+a `branch` operation, or back to one of the enclosing parent operations, as in a
+`return` operation. Terminator operations without block successors can only pass
+control back to one of the enclosing parent operations, in this case an integer
+defines the number of parent regions to break through. Within these restrictions,
+the particular semantics of terminator operations is determined by the specific
+dialect operations involved. Blocks (other than the entry block) that are not
+listed as a successor of a 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
+LLVM 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.
+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.
Example:
@@ -558,6 +570,79 @@ func.func @accelerator_compute(i64, i1) -> i64 { // An SSACFG region
}
```
+#### Region Terminator
+
+A `RegionTerminator` is a specialization of a block terminator (the `Terminator`
+trait) that transfers control back to an ancestor operation. It can exit
+multiple nested regions in a single step, bypassing the normal
+`RegionBranchOpInterface` exit path for every intermediate level. Terminators
+that can target such an ancestor consume a builtin `token` value defined by the
+target operation. For example, `scf.loop` defines a control token as an entry
+block argument, and `scf.break` / `scf.continue` consume that token to identify
+the loop that receives the early-exit edge.
+
+Every intermediate operation between the `RegionTerminator` and the
+token-defining 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 targeting itself, and propagates breaks that
+target an outer loop through it.
+
+Region 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 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 upward unchanged)
+// [H][P] = both: handles breaks targeting it, and propagates breaks that
+// target an outer loop through it
+scf.loop %outer { // [H]
+ scf.loop %inner { // [H][P]
+ scf.if %cond1 { // [P]
+ // Breaks the inner loop.
+ scf.break %inner
+ }
+ scf.if %cond2 { // [P]
+ // Breaks the outer loop.
+ scf.break %outer
+ }
+ scf.if %cond3 { // [P]
+ // Re-enters 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 carries operands that become the next iter_args.
+%result = scf.loop %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..61d9aed4abd9b 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,58 @@ 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, which can contain:
+/// break, continue, yield.
+struct IfOpImplicitTerminatorType
+ : public ControlFlowImplicitTerminatorOpType<YieldOp, BreakOp, ContinueOp> {
+};
+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 &&
+ "expected insertion block with a loop control token");
+ 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 0b33ecb48b7f2..bf21ccc8d34cb 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,178 @@ 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 an unstructured infinite 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 dynamically 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 carried values.
+
+ As long as each loop iteration is terminated by one of these operations they may be combined with other control
+ flow operations to express different control flow patterns.
+
+ The loop operation produces one return value for each loop carried variable. The type of the `i`-th return
+ value is that of the `i`-th loop carried variable and its value is the final value of the
+ `i`-th loop carried variable.
+ }];
+
+ let arguments = (ins Variadic<AnyType>:$initValues);
+ let results = (outs Variadic<AnyType>:$resultValues);
+ let regions = (region SizedRegion<1>:$region);
+
+ let extraClassDeclaration = [{
+ /// Required by HasBreakingControlFlowOpInterface. Returns true only for
+ /// scf.break and scf.continue, which are the RegionTerminators that can
+ /// target this loop as their HasBreakingControlFlowOpInterface receiver.
+ static bool acceptsTerminator(Operation *predecessor) {
+ return isa<BreakOp, ContinueOp>(predecessor);
+ }
+
+ /// 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, RegionTerminator,
+ DeclareOpInterfaceMethods<RegionBranchTerminatorOpInterface,
+ ["getMutableSuccessorOperands"]>,
+ ParentOneOf<["IfOp", "LoopOp"]>
+ ]> {
+ let summary = "Break from loop";
+ let description = [{
+ The `break` operation is a `RegionTerminator` 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 %loop -> i32 {
+ scf.break %loop %result : i32
+ }
+ ```
+
+ Example — break out of a loop through an enclosing `scf.if`:
+ ```mlir
+ scf.loop %loop {
+ scf.if %cond {
+ scf.break %loop
+ }
+ scf.continue %loop
+ }
+ ```
+ }];
+
+
+ let arguments = (ins Token:$target, Variadic<AnyType>:$args);
+ let assemblyFormat = [{
+ $target ($args^ `:` type($args))? attr-dict
+ }];
+ let hasVerifier = 1;
+}
+
+
+//===----------------------------------------------------------------------===//
+// ContinueOp
+//===----------------------------------------------------------------------===//
+
+def ContinueOp : SCF_Op<"continue", [
+ Terminator, RegionTerminator, DeclareOpInterfaceMethods<RegionBranchTerminatorOpInterface,
+ ["getMutableSuccessorOperands"]>, ParentOneOf<["IfOp", "LoopOp"]>
+ ]> {
+ let summary = "Continue to next loop iteration";
+ let description = [{
+ The `continue` operation is a `RegionTerminator` 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 %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 %outer iter_args(%counter = %init) : i64 {
+ scf.loop %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:$target, Variadic<AnyType>:$args);
+ let assemblyFormat = [{
+ $target ($args^ `:` type($args))? attr-dict
+ }];
+ let hasVerifier = 1;
+}
//===----------------------------------------------------------------------===//
// ForOp
@@ -706,8 +879,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
@@ -790,9 +963,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;
@@ -909,7 +1090,7 @@ def ParallelOp : SCF_Op<"parallel",
//===----------------------------------------------------------------------===//
def ReduceOp : SCF_Op<"reduce", [
- Terminator, HasParent<"ParallelOp">, RecursiveMemoryEffects,
+ Terminator, ImmediateRegionTerminator, HasParent<"ParallelOp">, RecursiveMemoryEffects,
DeclareOpInterfaceMethods<PromotableRegionOpInterface>,
DeclareOpInterfaceMethods<RegionBranchTerminatorOpInterface,
["getMutableSuccessorOperands"]>]> {
diff --git a/mlir/include/mlir/IR/Diagnostics.h b/mlir/include/mlir/IR/Diagnostics.h
index 3b8fb46b06a48..b5b3c55d35c75 100644
--- a/mlir/include/mlir/IR/Diagnostics.h
+++ b/mlir/include/mlir/IR/Diagnostics.h
@@ -200,7 +200,7 @@ class Diagnostic {
/// Stream in an Operation.
Diagnostic &operator<<(Operation &op);
- Diagnostic &operator<<(OpWithFlags op);
+ Diagnostic &operator<<(const OpWithFlags &opWithFlags);
Diagnostic &operator<<(Operation *op) { return *this << *op; }
/// Append an operation with the given printing flags.
Diagnostic &appendOp(Operation &op, const OpPrintingFlags &flags);
diff --git a/mlir/include/mlir/IR/OpBase.td b/mlir/include/mlir/IR/OpBase.td
index 0d0669e90c3f7..1e0b6ffbe1502 100644
--- a/mlir/include/mlir/IR/OpBase.td
+++ b/mlir/include/mlir/IR/OpBase.td
@@ -102,6 +102,10 @@ def Terminator : NativeOpTrait<"IsTerminator">;
def TokenProducerTrait : NativeOpTrait<"TokenProducerTrait">;
// Op consumes builtin token values.
def TokenConsumerTrait : NativeOpTrait<"TokenConsumerTrait">;
+// Op is a region terminator for the immediate region only.
+def ImmediateRegionTerminator : NativeOpTrait<"RegionTerminator", [Terminator]>;
+// Op is a region terminator, potentially breaking multiple regions
+def RegionTerminator : NativeOpTrait<"RegionTerminator", [Terminator]>;
// Op can be safely normalized in the presence of MemRefs with
// non-identity maps.
def MemRefsNormalizable : NativeOpTrait<"MemRefsNormalizable">;
@@ -134,6 +138,13 @@ class SingleBlockImplicitTerminatorImpl<string op>
class SingleBlockImplicitTerminator<string op>
: TraitList<[SingleBlock, SingleBlockImplicitTerminatorImpl<op>]>;
+// This operation has nested regions with the supplied list of `RegionTerminator`
+// 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 a0a36f2bd53c4..8cb8be3ecfaa3 100644
--- a/mlir/include/mlir/IR/OpDefinition.h
+++ b/mlir/include/mlir/IR/OpDefinition.h
@@ -904,7 +904,8 @@ struct SingleBlock : public TraitBase<ConcreteType, SingleBlock> {
// Non-empty regions must contain a single basic block.
if (!region.hasOneBlock())
return op->emitOpError("expects region #")
- << i << " to have 0 or 1 blocks";
+ << i << " to have 0 or 1 blocks, found "
+ << llvm::range_size(region) << " blocks";
if (!ConcreteType::template hasTrait<NoTerminator>()) {
Block &block = region.front();
@@ -1362,6 +1363,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/Operation.h b/mlir/include/mlir/IR/Operation.h
index 793c046fbf2e5..9375426524d67 100644
--- a/mlir/include/mlir/IR/Operation.h
+++ b/mlir/include/mlir/IR/Operation.h
@@ -934,8 +934,7 @@ class alignas(8) Operation final
if (propertiesStorageSize)
return PropertyRef(
name.getOpPropertiesTypeID(),
- reinterpret_cast<void *>(const_cast<detail::OpProperties *>(
- getTrailingObjects<detail::OpProperties>())));
+ const_cast<Operation *>(this)->getRawPropertiesStorageUnsafe());
return {};
}
diff --git a/mlir/include/mlir/IR/RegionKindInterface.h b/mlir/include/mlir/IR/RegionKindInterface.h
index d6d3aeeb9bd05..d85c0de58bc3e 100644
--- a/mlir/include/mlir/IR/RegionKindInterface.h
+++ b/mlir/include/mlir/IR/RegionKindInterface.h
@@ -36,6 +36,40 @@ 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 RegionTerminator (e.g. scf.break / scf.continue) can propagate through
+/// this op on its way to the token-defining HasBreakingControlFlowOpInterface
+/// ancestor. The op does NOT consume the break; it simply passes it upward.
+/// All ops that are "skipped over" by a region terminator 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();
+ }
+};
+
+/// Indicates that this operation is a block terminator that can exit a
+/// structured region by targeting a token-defining ancestor. This trait also
+/// requires IsTerminator (enforced by verifyTrait).
+template <typename ConcreteType>
+class RegionTerminator : public TraitBase<ConcreteType, RegionTerminator> {
+public:
+ static LogicalResult verifyTrait(Operation *op) {
+ if (!op->hasTrait<OpTrait::IsTerminator>())
+ return op->emitOpError(
+ "operation with region terminator trait must be a terminator");
+ return success();
+ }
+};
+
} // namespace OpTrait
/// Return "true" if the given region may have SSA dominance. This function also
@@ -49,8 +83,87 @@ bool mayHaveSSADominance(Region ®ion);
/// implement the RegionKindInterface.
bool mayBeGraphRegion(Region ®ion);
+/// Return true if `op` (which implements HasBreakingControlFlowOpInterface)
+/// contains at least one RegionTerminator 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` from a deeply nested site rather than
+/// only through the immediately enclosing terminator.
+bool hasNestedPredecessors(Operation *op);
+
+/// Return true if `op` contains any RegionTerminator 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 path that bypasses it.
+bool hasBreakingControlFlowOps(Operation *op);
+
+/// Collect all RegionTerminator operations nested inside `op` that directly
+/// target `op`. These are the ops that will transfer control flow to `op` on
+/// an early exit.
+void collectAllNestedPredecessors(Operation *op,
+ SmallVector<Operation *> &predecessors);
+
+namespace detail {
+/// Implementation helper for visitNestedBreakingControlFlowOps. Walks the
+/// regions of `op` and invokes `callback` for every RegionTerminator 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(Operation *, int nestedLevel)> callback);
+} // namespace detail
+
+/// Walk all RegionTerminator operations 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, [&](Operation *visitedOp, int nestedLevel) {
+ callback(visitedOp, nestedLevel);
+ return WalkResult::advance();
+ });
+ }
+}
+
+/// Walk all RegionTerminator operations relevant to breaking control flow
+/// across all top-level ops in `region`.
+template <typename CallbackT>
+void visitNestedBreakingControlFlowOps(Region ®ion, CallbackT &&callback) {
+ for (Operation &op : region.getOps())
+ visitNestedBreakingControlFlowOps(&op, std::forward<CallbackT>(callback));
+}
+
} // namespace mlir
#include "mlir/IR/RegionKindInterface.h.inc"
+namespace mlir {
+
+/// 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));
+}
+
+/// For a RegionTerminator that consumes a builtin token, return the
+/// HasBreakingControlFlowOpInterface operation that defines the consumed token.
+/// Returns a null interface wrapper if the token is malformed or the target
+/// does not implement HasBreakingControlFlowOpInterface.
+HasBreakingControlFlowOpInterface findBreakTarget(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..d2042481d4d41 100644
--- a/mlir/include/mlir/IR/RegionKindInterface.td
+++ b/mlir/include/mlir/IR/RegionKindInterface.td
@@ -61,4 +61,58 @@ def GraphRegionNoTerminator : TraitList<[
HasOnlyGraphRegion
]>;
+// Indicates that this op may propagate a breaking control-flow event from a
+// nested region upward to the token-defining
+// HasBreakingControlFlowOpInterface operation. The op does NOT consume the
+// break itself; it is merely transparent to it. All ops that sit between a
+// RegionTerminator and the HasBreakingControlFlowOpInterface ancestor that will
+// ultimately receive the break must carry this trait.
+def PropagateControlFlowBreak : NativeOpTrait<"PropagateControlFlowBreak">;
+
+// OpInterface for operations that can receive a breaking control-flow event
+// originating from a RegionTerminator (scf.break / scf.continue) anywhere
+// inside their (possibly deeply nested) regions. The receiver defines the token
+// consumed by the RegionTerminator, and every op between the RegionTerminator
+// and this receiver must carry the PropagateControlFlowBreak trait.
+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`). The operation defines a builtin
+ token consumed by the `RegionTerminator`. Every intermediate op must carry
+ the `PropagateControlFlowBreak` trait.
+ }];
+ let cppNamespace = "::mlir";
+
+ let methods = [
+ StaticInterfaceMethod<
+ /*desc=*/[{
+ Return true if this operation accepts the given terminator operation
+ as a breaking-control-flow predecessor. Ops that also use
+ `HasNestedTerminator<[...]>` should delegate to the terminator list
+ check from that trait. Ops without `HasNestedTerminator` must provide
+ an explicit implementation (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 RegionTerminator 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());
+ }]
+ >
+ ];
+}
+
+
#endif // MLIR_IR_REGIONKINDINTERFACE
diff --git a/mlir/include/mlir/Interfaces/ControlFlowInterfaces.h b/mlir/include/mlir/Interfaces/ControlFlowInterfaces.h
index a76dce6f2ffc5..01a8ad8656a6b 100644
--- a/mlir/include/mlir/Interfaces/ControlFlowInterfaces.h
+++ b/mlir/include/mlir/Interfaces/ControlFlowInterfaces.h
@@ -199,22 +199,32 @@ using RegionBranchInverseSuccessorMapping =
class RegionSuccessor {
public:
/// Initialize a successor that branches to a region of the parent operation.
- RegionSuccessor(Region *region) : successor(region) {
+ RegionSuccessor(Region *region) : successor(region), kind(Kind::Region) {
assert(region && "Region must not be null");
}
/// Initialize a successor that branches after/out of the parent operation.
- static RegionSuccessor parent() { return RegionSuccessor(); }
+ static RegionSuccessor parent() { return RegionSuccessor(Kind::Parent); }
+
+ /// Sentinel: the terminator propagates through this op to an ancestor.
+ /// The op is transparent to this break and does not consume it.
+ /// Use `resolveTerminatorSuccessors` to resolve to the actual target.
+ static RegionSuccessor propagating() {
+ return RegionSuccessor(Kind::Propagating);
+ }
/// Return the given region successor. Returns nullptr if the successor is the
/// parent operation.
Region *getSuccessor() const { return successor; }
/// Return true if the successor is the parent operation.
- bool isParent() const { return successor == nullptr; }
+ bool isParent() const { return kind == Kind::Parent; }
+
+ /// Return true if this is a propagating-break sentinel.
+ bool isPropagating() const { return kind == Kind::Propagating; }
bool operator==(RegionSuccessor rhs) const {
- return successor == rhs.successor;
+ return successor == rhs.successor && kind == rhs.kind;
}
bool operator==(const Region *region) const { return successor == region; }
@@ -224,10 +234,12 @@ class RegionSuccessor {
}
private:
- /// Private constructor to encourage the use of `RegionSuccessor::parent`.
- RegionSuccessor() : successor(nullptr) {}
+ enum class Kind { Region, Parent, Propagating };
+
+ explicit RegionSuccessor(Kind kind) : successor(nullptr), kind(kind) {}
Region *successor = nullptr;
+ Kind kind = Kind::Region;
};
/// This class represents a point being branched from in the methods of the
@@ -423,11 +435,24 @@ inline llvm::raw_ostream &operator<<(llvm::raw_ostream &os,
inline llvm::raw_ostream &operator<<(llvm::raw_ostream &os,
RegionSuccessor successor) {
+ if (successor.isPropagating())
+ return os << "<propagating>";
if (successor.isParent())
return os << "<to parent>";
return os << "<to region #" << successor.getSuccessor()->getRegionNumber()
<< ">";
}
+
+/// Get successor regions for a terminator, resolving propagating breaks.
+/// When the immediate parent returns a `RegionSuccessor::propagating()`
+/// sentinel, finds and queries the actual HasBreakingControlFlowOpInterface
+/// ancestor. Returns the RegionBranchOpInterface that owns the returned
+/// successors. For non-propagating terminators, this is the terminator's
+/// immediate parent.
+RegionBranchOpInterface
+resolveTerminatorSuccessors(RegionBranchTerminatorOpInterface terminator,
+ SmallVectorImpl<RegionSuccessor> &successors);
+
} // namespace mlir
#endif // MLIR_INTERFACES_CONTROLFLOWINTERFACES_H
diff --git a/mlir/lib/Analysis/DataFlow/DeadCodeAnalysis.cpp b/mlir/lib/Analysis/DataFlow/DeadCodeAnalysis.cpp
index 38811d06ecd8c..2271ca38188d5 100644
--- a/mlir/lib/Analysis/DataFlow/DeadCodeAnalysis.cpp
+++ b/mlir/lib/Analysis/DataFlow/DeadCodeAnalysis.cpp
@@ -509,12 +509,27 @@ 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);
+
+ // For propagating breaks, the immediate parent is transparent — resolve to
+ // the actual HasBreakingControlFlowOp ancestor.
+ RegionBranchOpInterface effectiveBranch = branch;
+ if (successors.size() == 1 && successors[0].isPropagating()) {
+ successors.clear();
+ // TODO: Re-query with operand-aware getSuccessorRegions after resolution
+ // to preserve constant-folding pruning for break targets that support it.
+ effectiveBranch = resolveTerminatorSuccessors(terminator, successors);
+ if (!effectiveBranch)
+ return;
+ }
+ visitRegionBranchEdges(effectiveBranch, op, successors);
}
void DeadCodeAnalysis::visitRegionBranchEdges(
diff --git a/mlir/lib/Analysis/DataFlow/DenseAnalysis.cpp b/mlir/lib/Analysis/DataFlow/DenseAnalysis.cpp
index 22bc0b32a9bd1..c746a3e035136 100644
--- a/mlir/lib/Analysis/DataFlow/DenseAnalysis.cpp
+++ b/mlir/lib/Analysis/DataFlow/DenseAnalysis.cpp
@@ -633,13 +633,22 @@ void AbstractDenseBackwardDataFlowAnalysis::visitRegionBranchOperation(
// entry block of each possible successor region, or the next operation when
// the branch is a successor of itself.
SmallVector<RegionSuccessor> successors;
- branch.getSuccessorRegions(branchPoint, successors);
+ RegionBranchOpInterface effectiveBranch = branch;
+ if (!branchPoint.isParent()) {
+ // For terminator branch points, resolve propagating breaks.
+ auto terminator = branchPoint.getTerminatorPredecessorOrNull();
+ effectiveBranch = resolveTerminatorSuccessors(terminator, successors);
+ if (!effectiveBranch)
+ return;
+ } else {
+ branch.getSuccessorRegions(branchPoint, successors);
+ }
LDBG() << " Processing " << successors.size() << " successor regions";
for (const RegionSuccessor &successor : successors) {
const AbstractDenseLattice *after;
if (successor.isParent() || successor.getSuccessor()->empty()) {
LDBG() << " Successor is parent or empty region";
- after = getLatticeFor(point, getProgramPointAfter(branch));
+ after = getLatticeFor(point, getProgramPointAfter(effectiveBranch));
} else {
Region *successorRegion = successor.getSuccessor();
assert(!successorRegion->empty() && "unexpected empty successor region");
@@ -658,7 +667,7 @@ void AbstractDenseBackwardDataFlowAnalysis::visitRegionBranchOperation(
}
LDBG() << " After state: " << *after;
- visitRegionBranchControlFlowTransfer(branch, branchPoint, successor, *after,
- before);
+ visitRegionBranchControlFlowTransfer(effectiveBranch, branchPoint,
+ successor, *after, before);
}
}
diff --git a/mlir/lib/Analysis/DataFlow/SparseAnalysis.cpp b/mlir/lib/Analysis/DataFlow/SparseAnalysis.cpp
index 90f2a588d1ca4..104da59af1087 100644
--- a/mlir/lib/Analysis/DataFlow/SparseAnalysis.cpp
+++ b/mlir/lib/Analysis/DataFlow/SparseAnalysis.cpp
@@ -647,9 +647,18 @@ void AbstractSparseBackwardDataFlowAnalysis::
// non-contiguous in the presence of multiple successors.
BitVector unaccounted(terminator->getNumOperands(), true);
+ // For propagating breaks, the immediate parent is transparent. Resolve to the
+ // actual HasBreakingControlFlowOpInterface ancestor.
+ // Successors not needed here; only the effective branch matters.
+ SmallVector<RegionSuccessor> unusedSuccessors;
+ RegionBranchOpInterface effectiveBranch =
+ resolveTerminatorSuccessors(terminator, unusedSuccessors);
+ if (!effectiveBranch)
+ effectiveBranch = branch;
+
RegionBranchSuccessorMapping mapping;
- branch.getSuccessorOperandInputMapping(mapping,
- RegionBranchPoint(terminator));
+ effectiveBranch.getSuccessorOperandInputMapping(
+ mapping, RegionBranchPoint(terminator));
for (const auto &[operand, inputs] : mapping) {
for (Value input : inputs) {
meet(getLatticeElement(operand->get()),
diff --git a/mlir/lib/AsmParser/Parser.cpp b/mlir/lib/AsmParser/Parser.cpp
index 952d7e460c6e2..c02494959a255 100644
--- a/mlir/lib/AsmParser/Parser.cpp
+++ b/mlir/lib/AsmParser/Parser.cpp
@@ -676,7 +676,8 @@ class OperationParser : public Parser {
ParseResult parseSuccessor(Block *&dest);
/// Parse a comma-separated list of operation successors in brackets.
- ParseResult parseSuccessors(SmallVectorImpl<Block *> &destinations);
+ ParseResult parseSuccessors(SmallVectorImpl<Block *> &destinations,
+ bool parseOpeningBracket = true);
/// Parse an operation instance that is in the generic form.
Operation *parseGenericOperation();
@@ -1381,8 +1382,9 @@ ParseResult OperationParser::parseSuccessor(Block *&dest) {
/// successor-list ::= `[` successor (`,` successor )* `]`
///
ParseResult
-OperationParser::parseSuccessors(SmallVectorImpl<Block *> &destinations) {
- if (parseToken(Token::l_square, "expected '['"))
+OperationParser::parseSuccessors(SmallVectorImpl<Block *> &destinations,
+ bool parseOpeningBracket) {
+ if (parseOpeningBracket && parseToken(Token::l_square, "expected '['"))
return failure();
auto parseElt = [this, &destinations] {
@@ -1434,6 +1436,8 @@ ParseResult OperationParser::parseGenericOperationAfterOpName(
}
// Parse the successor list, if not explicitly provided.
+ if (parsedSuccessors)
+ result.addSuccessors(*parsedSuccessors);
if (!parsedSuccessors) {
if (getToken().is(Token::l_square)) {
// Check if the operation is not a known terminator.
@@ -1445,8 +1449,6 @@ ParseResult OperationParser::parseGenericOperationAfterOpName(
return failure();
result.addSuccessors(successors);
}
- } else {
- result.addSuccessors(*parsedSuccessors);
}
// Parse the properties, if not explicitly provided.
diff --git a/mlir/lib/Conversion/SCFToControlFlow/SCFToControlFlow.cpp b/mlir/lib/Conversion/SCFToControlFlow/SCFToControlFlow.cpp
index 2972d79c4302f..16140286e29db 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,11 +741,110 @@ 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();
+ }
+
+ if (hasNestedPredecessors(loopOp))
+ return rewriter.notifyMatchFailure(loopOp,
+ "loop op with nested predecessors");
+
+ // Collect direct predecessors (break/continue targeting this loop).
+ SmallVector<Operation *> predecessors;
+ collectAllNestedPredecessors(loopOp, 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();
+
+ // Create the loop entry block and move the body there.
+ rewriter.setInsertionPoint(initBlock, initBlock->end());
+ // Split out everything after loopOp into continueBlock.
+ // The block before loop is now initBlock.
+
+ // 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.
+ 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.getContext());
+ patterns.add<IfLowering>(patterns.getContext(), /*benefit=*/2);
+ patterns.add<ExecuteRegionLowering, ForallLowering, ForLowering,
+ IndexSwitchLowering, LoopOpLowering, ParallelLowering,
+ WhileLowering>(patterns.getContext());
patterns.add<DoWhileLowering>(patterns.getContext(), /*benefit=*/2);
}
@@ -734,7 +855,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/Conversion/SCFToEmitC/SCFToEmitC.cpp b/mlir/lib/Conversion/SCFToEmitC/SCFToEmitC.cpp
index d7c943ef7c4f1..d7b16651e8e3a 100644
--- a/mlir/lib/Conversion/SCFToEmitC/SCFToEmitC.cpp
+++ b/mlir/lib/Conversion/SCFToEmitC/SCFToEmitC.cpp
@@ -57,7 +57,7 @@ void mlir::registerConvertSCFToEmitCInterface(DialectRegistry ®istry) {
namespace {
-struct SCFToEmitCPass : public impl::SCFToEmitCBase<SCFToEmitCPass> {
+struct SCFToEmitCPass : public ::mlir::impl::SCFToEmitCBase<SCFToEmitCPass> {
void runOnOperation() override;
};
diff --git a/mlir/lib/Conversion/ShapeToStandard/ShapeToStandard.cpp b/mlir/lib/Conversion/ShapeToStandard/ShapeToStandard.cpp
index 0ff9fb3f628ab..da85f7c2d2eaa 100644
--- a/mlir/lib/Conversion/ShapeToStandard/ShapeToStandard.cpp
+++ b/mlir/lib/Conversion/ShapeToStandard/ShapeToStandard.cpp
@@ -681,7 +681,8 @@ namespace {
namespace {
/// Conversion pass.
class ConvertShapeToStandardPass
- : public impl::ConvertShapeToStandardPassBase<ConvertShapeToStandardPass> {
+ : public ::mlir::impl::ConvertShapeToStandardPassBase<
+ ConvertShapeToStandardPass> {
void runOnOperation() override;
};
diff --git a/mlir/lib/Dialect/SCF/IR/SCF.cpp b/mlir/lib/Dialect/SCF/IR/SCF.cpp
index 9f4f4dc9f58e6..77393b1da070a 100644
--- a/mlir/lib/Dialect/SCF/IR/SCF.cpp
+++ b/mlir/lib/Dialect/SCF/IR/SCF.cpp
@@ -132,6 +132,41 @@ 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;
+ }
+}
+
+static bool terminatorPropagatesThrough(Operation *terminator, Operation *op);
+
+/// Replaces the given op with the contents of the given single-block region,
+/// using the operands of the block terminator to replace operation results.
+static void replaceOpWithRegion(PatternRewriter &rewriter, Operation *op,
+ Region ®ion, ValueRange blockArgs = {}) {
+ assert(region.hasOneBlock() && "expected single-block region");
+ Block *block = ®ion.front();
+ Operation *terminator = block->getTerminator();
+ ValueRange results = terminator->getOperands();
+ if (auto branchTerminator =
+ dyn_cast<RegionBranchTerminatorOpInterface>(terminator))
+ results = branchTerminator.getSuccessorOperands(RegionSuccessor::parent());
+ rewriter.inlineBlockBefore(block, op, blockArgs);
+ if (terminatorPropagatesThrough(terminator, op)) {
+ // Erase `op` and every op that follows it (dead code after the early exit).
+ eraseOpsAfter(rewriter, op);
+ rewriter.eraseOp(op);
+ } else {
+ rewriter.replaceOp(op, results);
+ rewriter.eraseOp(terminator);
+ }
+}
+
///
/// (ssa-id `=`)? `execute_region` `->` function-result-type `{`
/// block+
@@ -311,6 +346,284 @@ void ConditionOp::getSuccessorRegions(
regions.push_back(RegionSuccessor::parent());
}
+//===----------------------------------------------------------------------===//
+// LoopOp
+//===----------------------------------------------------------------------===//
+
+//===----------------------------------------------------------------------===//
+// Control Flow Op Utilies
+//===----------------------------------------------------------------------===//
+
+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();
+}
+
+template <typename ImplicitTerminatorOpT, typename OpT>
+static void printControlFlowRegion(OpAsmPrinter &p, OpT op, Region ®ion) {
+ // We do not print the terminator if it is implicit and has no operands.
+ bool printBlockTerminators =
+ region.front().getTerminator()->getNumOperands() != 0 ||
+ !isa<ImplicitTerminatorOpT>(region.front().getTerminator());
+ p.printRegion(region, /*printEntryBlockArgs=*/false, printBlockTerminators);
+}
+
+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) {
+ auto breakTarget = findBreakTarget(terminator);
+ return breakTarget && breakTarget.getOperation() != op &&
+ breakTarget.getOperation()->isProperAncestor(op);
+}
+
+LogicalResult BreakOp::verify() {
+ return verifyLoopTerminatorTarget(getOperation(), getTarget());
+}
+
+MutableOperandRange
+BreakOp::getMutableSuccessorOperands(RegionSuccessor point) {
+ return MutableOperandRange(getOperation(), /*start=*/1,
+ /*length=*/getOperation()->getNumOperands() - 1);
+}
+
+LogicalResult ContinueOp::verify() {
+ return verifyLoopTerminatorTarget(getOperation(), getTarget());
+}
+
+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(
+ "mismatch in number of loop-carried values and defined values");
+ 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 << " " << 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.getTarget() != 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.parseArgument(controlToken))
+ 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())) {
+ regions.push_back(RegionSuccessor::propagating());
+ return;
+ }
+
+ if (isa<ContinueOp>(terminator)) {
+ regions.push_back(RegionSuccessor(&getRegion()));
+ return;
+ }
+ assert(isa<BreakOp>(terminator) && "expected continue or break terminator");
+
+ regions.push_back(RegionSuccessor::parent());
+}
+
+OperandRange LoopOp::getEntrySuccessorOperands(RegionSuccessor successor) {
+ return getInitValues();
+}
+
+ValueRange LoopOp::getSuccessorInputs(RegionSuccessor successor) {
+ return successor.isParent() ? 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");
+ auto target = findBreakTarget(breakOp);
+ if (!target || target.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
//===----------------------------------------------------------------------===//
@@ -944,6 +1257,194 @@ mlir::scf::replaceAndCastForOpIterArg(RewriterBase &rewriter, scf::ForOp forOp,
}
namespace {
+// Fold away ForOp iter arguments when:
+// 1) The op yields the iter arguments.
+// 2) The argument's corresponding outer region iterators (inputs) are yielded.
+// 3) The iter arguments have no use and the corresponding (operation) results
+// have no use.
+//
+// These arguments must be defined outside of the ForOp region and can just be
+// forwarded after simplifying the op inits, yields and returns.
+//
+// The implementation uses `inlineBlockBefore` to steal the content of the
+// original ForOp and avoid cloning.
+struct ForOpIterArgsFolder : public OpRewritePattern<scf::ForOp> {
+ using OpRewritePattern<scf::ForOp>::OpRewritePattern;
+
+ LogicalResult matchAndRewrite(scf::ForOp forOp,
+ PatternRewriter &rewriter) const final {
+ bool canonicalize = false;
+
+ // An internal flat vector of block transfer
+ // arguments `newBlockTransferArgs` keeps the 1-1 mapping of original to
+ // transformed block argument mappings. This plays the role of a
+ // IRMapping for the particular use case of calling into
+ // `inlineBlockBefore`.
+ int64_t numResults = forOp.getNumResults();
+ SmallVector<bool, 4> keepMask;
+ keepMask.reserve(numResults);
+ SmallVector<Value, 4> newBlockTransferArgs, newIterArgs, newYieldValues,
+ newResultValues;
+ newBlockTransferArgs.reserve(1 + numResults);
+ newBlockTransferArgs.push_back(Value()); // iv placeholder with null value
+ newIterArgs.reserve(forOp.getInitArgs().size());
+ newYieldValues.reserve(numResults);
+ newResultValues.reserve(numResults);
+ DenseMap<std::pair<Value, Value>, std::pair<Value, Value>> initYieldToArg;
+ for (auto [init, arg, result, yielded] :
+ llvm::zip(forOp.getInitArgs(), // iter from outside
+ forOp.getRegionIterArgs(), // iter inside region
+ forOp.getResults(), // op results
+ forOp.getYieldedValues() // iter yield
+ )) {
+ // Forwarded is `true` when:
+ // 1) The region `iter` argument is yielded.
+ // 2) The region `iter` argument the corresponding input is yielded.
+ // 3) The region `iter` argument has no use, and the corresponding op
+ // result has no use.
+ bool forwarded = (arg == yielded) || (init == yielded) ||
+ (arg.use_empty() && result.use_empty());
+ if (forwarded) {
+ canonicalize = true;
+ keepMask.push_back(false);
+ newBlockTransferArgs.push_back(init);
+ newResultValues.push_back(init);
+ continue;
+ }
+
+ // Check if a previous kept argument always has the same values for init
+ // and yielded values.
+ if (auto it = initYieldToArg.find({init, yielded});
+ it != initYieldToArg.end()) {
+ canonicalize = true;
+ keepMask.push_back(false);
+ auto [sameArg, sameResult] = it->second;
+ rewriter.replaceAllUsesWith(arg, sameArg);
+ rewriter.replaceAllUsesWith(result, sameResult);
+ // The replacement value doesn't matter because there are no uses.
+ newBlockTransferArgs.push_back(init);
+ newResultValues.push_back(init);
+ continue;
+ }
+
+ // This value is kept.
+ initYieldToArg.insert({{init, yielded}, {arg, result}});
+ keepMask.push_back(true);
+ newIterArgs.push_back(init);
+ newYieldValues.push_back(yielded);
+ newBlockTransferArgs.push_back(Value()); // placeholder with null value
+ newResultValues.push_back(Value()); // placeholder with null value
+ }
+
+ if (!canonicalize)
+ return failure();
+
+ scf::ForOp newForOp =
+ scf::ForOp::create(rewriter, forOp.getLoc(), forOp.getLowerBound(),
+ forOp.getUpperBound(), forOp.getStep(), newIterArgs,
+ /*bodyBuilder=*/nullptr, forOp.getUnsignedCmp());
+ newForOp->setAttrs(forOp->getAttrs());
+ Block &newBlock = newForOp.getRegion().front();
+
+ // Replace the null placeholders with newly constructed values.
+ newBlockTransferArgs[0] = newBlock.getArgument(0); // iv
+ for (unsigned idx = 0, collapsedIdx = 0, e = newResultValues.size();
+ idx != e; ++idx) {
+ Value &blockTransferArg = newBlockTransferArgs[1 + idx];
+ Value &newResultVal = newResultValues[idx];
+ assert((blockTransferArg && newResultVal) ||
+ (!blockTransferArg && !newResultVal));
+ if (!blockTransferArg) {
+ blockTransferArg = newForOp.getRegionIterArgs()[collapsedIdx];
+ newResultVal = newForOp.getResult(collapsedIdx++);
+ }
+ }
+
+ Block &oldBlock = forOp.getRegion().front();
+ assert(oldBlock.getNumArguments() == newBlockTransferArgs.size() &&
+ "unexpected argument size mismatch");
+
+ // No results case: the scf::ForOp builder already created a zero
+ // result terminator. Merge before this terminator and just get rid of the
+ // original terminator that has been merged in.
+ if (newIterArgs.empty()) {
+ auto newYieldOp = cast<scf::YieldOp>(newBlock.getTerminator());
+ rewriter.inlineBlockBefore(&oldBlock, newYieldOp, newBlockTransferArgs);
+ rewriter.eraseOp(newBlock.getTerminator()->getPrevNode());
+ rewriter.replaceOp(forOp, newResultValues);
+ return success();
+ }
+
+ // No terminator case: merge and rewrite the merged terminator.
+ auto cloneFilteredTerminator = [&](scf::YieldOp mergedTerminator) {
+ OpBuilder::InsertionGuard g(rewriter);
+ rewriter.setInsertionPoint(mergedTerminator);
+ SmallVector<Value, 4> filteredOperands;
+ filteredOperands.reserve(newResultValues.size());
+ for (unsigned idx = 0, e = keepMask.size(); idx < e; ++idx)
+ if (keepMask[idx])
+ filteredOperands.push_back(mergedTerminator.getOperand(idx));
+ scf::YieldOp::create(rewriter, mergedTerminator.getLoc(),
+ filteredOperands);
+ };
+
+ rewriter.mergeBlocks(&oldBlock, &newBlock, newBlockTransferArgs);
+ auto mergedYieldOp = cast<scf::YieldOp>(newBlock.getTerminator());
+ cloneFilteredTerminator(mergedYieldOp);
+ rewriter.eraseOp(mergedYieldOp);
+ rewriter.replaceOp(forOp, newResultValues);
+ return success();
+ }
+};
+
+/// Rewriting pattern that erases loops that are known not to iterate, replaces
+/// single-iteration loops with their bodies, and removes empty loops that
+/// iterate at least once and only return values defined outside of the loop.
+struct SimplifyTrivialForLoops : public OpRewritePattern<ForOp> {
+ using OpRewritePattern<ForOp>::OpRewritePattern;
+
+ LogicalResult matchAndRewrite(ForOp op,
+ PatternRewriter &rewriter) const override {
+ std::optional<APInt> tripCount = op.getStaticTripCount();
+ if (!tripCount.has_value())
+ return rewriter.notifyMatchFailure(op,
+ "can't compute constant trip count");
+
+ if (tripCount->isZero()) {
+ LDBG() << "SimplifyTrivialForLoops tripCount is 0 for loop "
+ << OpWithFlags(op, OpPrintingFlags().skipRegions());
+ rewriter.replaceOp(op, op.getInitArgs());
+ return success();
+ }
+
+ if (tripCount->getSExtValue() == 1) {
+ LDBG() << "SimplifyTrivialForLoops tripCount is 1 for loop "
+ << OpWithFlags(op, OpPrintingFlags().skipRegions());
+ SmallVector<Value, 4> blockArgs;
+ blockArgs.reserve(op.getInitArgs().size() + 1);
+ blockArgs.push_back(op.getLowerBound());
+ llvm::append_range(blockArgs, op.getInitArgs());
+ replaceOpWithRegion(rewriter, op, op.getRegion(), blockArgs);
+ return success();
+ }
+
+ // Now we are left with loops that have more than 1 iterations.
+ Block &block = op.getRegion().front();
+ if (!llvm::hasSingleElement(block))
+ return failure();
+ // The loop is empty and iterates at least once, if it only returns values
+ // defined outside of the loop, remove it and replace it with yield values.
+ if (llvm::any_of(op.getYieldedValues(),
+ [&](Value v) { return !op.isDefinedOutsideOfLoop(v); }))
+ return failure();
+ LDBG() << "SimplifyTrivialForLoops empty body loop allows replacement with "
+ "yield operands for loop "
+ << OpWithFlags(op, OpPrintingFlags().skipRegions());
+ rewriter.replaceOp(op, op.getYieldedValues());
+ return success();
+ }
+};
+
/// Fold scf.for iter_arg/result pairs that go through incoming/ougoing
/// a tensor.cast op pair so as to pull the tensor.cast inside the scf.for:
///
@@ -1006,7 +1507,9 @@ struct ForOpTensorCastFolder : public OpRewritePattern<ForOp> {
void ForOp::getCanonicalizationPatterns(RewritePatternSet &results,
MLIRContext *context) {
- results.add<ForOpTensorCastFolder>(context);
+ results
+ .add<ForOpIterArgsFolder, SimplifyTrivialForLoops, ForOpTensorCastFolder>(
+ context);
populateRegionBranchOpInterfaceCanonicalizationPatterns(
results, ForOp::getOperationName());
populateRegionBranchOpInterfaceInliningPattern(
@@ -1940,13 +2443,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 +2582,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 +2614,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 (isa<BreakOp, ContinueOp>(terminator) &&
+ terminatorPropagatesThrough(terminator, getOperation())) {
+ regions.push_back(RegionSuccessor::propagating());
+ return;
+ }
+ }
regions.push_back(RegionSuccessor::parent());
return;
}
@@ -2175,6 +2698,82 @@ void IfOp::getRegionInvocationBounds(
}
namespace {
+// Pattern to remove unused IfOp results.
+struct RemoveUnusedResults : public OpRewritePattern<IfOp> {
+ using OpRewritePattern<IfOp>::OpRewritePattern;
+
+ void transferBody(Block *source, Block *dest, ArrayRef<OpResult> usedResults,
+ PatternRewriter &rewriter) const {
+ // Move all operations to the destination block.
+ rewriter.mergeBlocks(source, dest);
+ // Replace the yield op by one that returns only the used values.
+ auto yieldOp = dyn_cast<scf::YieldOp>(dest->getTerminator());
+ if (!yieldOp)
+ return;
+ SmallVector<Value, 4> usedOperands;
+ llvm::transform(usedResults, std::back_inserter(usedOperands),
+ [&](OpResult result) {
+ return yieldOp.getOperand(result.getResultNumber());
+ });
+ rewriter.modifyOpInPlace(yieldOp,
+ [&]() { yieldOp->setOperands(usedOperands); });
+ }
+
+ LogicalResult matchAndRewrite(IfOp op,
+ PatternRewriter &rewriter) const override {
+ // Compute the list of used results.
+ SmallVector<OpResult, 4> usedResults;
+ llvm::copy_if(op.getResults(), std::back_inserter(usedResults),
+ [](OpResult result) { return !result.use_empty(); });
+
+ // Replace the operation if only a subset of its results have uses.
+ if (usedResults.size() == op.getNumResults())
+ return failure();
+
+ // Compute the result types of the replacement operation.
+ SmallVector<Type, 4> newTypes;
+ llvm::transform(usedResults, std::back_inserter(newTypes),
+ [](OpResult result) { return result.getType(); });
+
+ // Create a replacement operation with empty then and else regions.
+ auto newOp =
+ IfOp::create(rewriter, op.getLoc(), newTypes, op.getCondition());
+ rewriter.createBlock(&newOp.getThenRegion());
+ rewriter.createBlock(&newOp.getElseRegion());
+
+ // Move the bodies and replace the terminators (note there is a then and
+ // an else region since the operation returns results).
+ transferBody(op.getBody(0), newOp.getBody(0), usedResults, rewriter);
+ transferBody(op.getBody(1), newOp.getBody(1), usedResults, rewriter);
+
+ // Replace the operation by the new one.
+ SmallVector<Value, 4> repResults(op.getNumResults());
+ for (const auto &en : llvm::enumerate(usedResults))
+ repResults[en.value().getResultNumber()] = newOp.getResult(en.index());
+ rewriter.replaceOp(op, repResults);
+ return success();
+ }
+};
+
+struct RemoveStaticCondition : public OpRewritePattern<IfOp> {
+ using OpRewritePattern<IfOp>::OpRewritePattern;
+
+ LogicalResult matchAndRewrite(IfOp op,
+ PatternRewriter &rewriter) const override {
+ BoolAttr condition;
+ if (!matchPattern(op.getCondition(), m_Constant(&condition)))
+ return failure();
+
+ if (condition.getValue())
+ replaceOpWithRegion(rewriter, op, op.getThenRegion());
+ else if (!op.getElseRegion().empty())
+ replaceOpWithRegion(rewriter, op, op.getElseRegion());
+ else
+ rewriter.eraseOp(op);
+ return success();
+ }
+};
+
/// Hoist any yielded results whose operands are defined outside
/// the if, to a select instruction.
struct ConvertTrivialIfToSelect : public OpRewritePattern<IfOp> {
@@ -2185,9 +2784,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 +2835,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 +2992,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 +3032,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 +3087,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 +3107,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 +3153,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 +3175,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 +3213,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 +3250,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 +3285,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 +3330,73 @@ 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);
+ RemoveStaticCondition, RemoveUnusedResults,
+ ReplaceIfYieldWithConditionOrValue,
+ SimplifyIfWithBreakingControlFlowInBothBranches>(context);
populateRegionBranchOpInterfaceCanonicalizationPatterns(
results, IfOp::getOperationName());
populateRegionBranchOpInterfaceInliningPattern(results,
@@ -2733,14 +3404,12 @@ 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
@@ -3467,8 +4136,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 +4185,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 16eb9aadc06f0..160f027b8f9ea 100644
--- a/mlir/lib/Dialect/SCF/Transforms/BufferizableOpInterfaceImpl.cpp
+++ b/mlir/lib/Dialect/SCF/Transforms/BufferizableOpInterfaceImpl.cpp
@@ -235,8 +235,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/IR/Diagnostics.cpp b/mlir/lib/IR/Diagnostics.cpp
index 5caf826c84bdd..1eb901ed37690 100644
--- a/mlir/lib/IR/Diagnostics.cpp
+++ b/mlir/lib/IR/Diagnostics.cpp
@@ -138,8 +138,8 @@ Diagnostic &Diagnostic::operator<<(Operation &op) {
return appendOp(op, OpPrintingFlags());
}
-Diagnostic &Diagnostic::operator<<(OpWithFlags op) {
- return appendOp(*op.getOperation(), op.flags());
+Diagnostic &Diagnostic::operator<<(const OpWithFlags &opWithFlags) {
+ return appendOp(*opWithFlags.getOperation(), opWithFlags.flags());
}
Diagnostic &Diagnostic::appendOp(Operation &op, const OpPrintingFlags &flags) {
diff --git a/mlir/lib/IR/Dominance.cpp b/mlir/lib/IR/Dominance.cpp
index 79fb41f2e6b30..1ca9dbfb2e8a8 100644
--- a/mlir/lib/IR/Dominance.cpp
+++ b/mlir/lib/IR/Dominance.cpp
@@ -14,8 +14,11 @@
#include "mlir/IR/Dominance.h"
#include "mlir/IR/Operation.h"
#include "mlir/IR/RegionKindInterface.h"
+#include "llvm/Support/DebugLog.h"
#include "llvm/Support/GenericDomTreeConstruction.h"
+#define DEBUG_TYPE "dominance"
+
using namespace mlir;
using namespace mlir::detail;
@@ -295,6 +298,31 @@ bool DominanceInfoBase<IsPostDom>::properlyDominatesImpl(
// regions kinds, uses and defs can come in any order inside a block.
if (!hasSSADominance(aBlock))
return true;
+
+ // Any operation that propagates a control flow break invalidates the
+ // post-dominance relation. Performance note: hasBreakingControlFlowOps()
+ // walks the entire region tree of each op in [bIt, aIt]. This is acceptable
+ // because the mightHaveBreakingControlFlow guard ensures we only enter this
+ // path when breaking CF is possible, and in practice the range is small.
+ 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()) &&
+ hasBreakingControlFlowOps(&op)) {
+ LDBG() << "Breaking control flow: "
+ << OpWithFlags(&op, OpPrintingFlags().skipRegions());
+ return false;
+ }
+ }
+ if (&op == &*aIt)
+ break;
+ }
+ }
if constexpr (IsPostDom) {
return isBeforeInBlock(aBlock, bIt, aIt);
} else {
@@ -302,6 +330,7 @@ bool DominanceInfoBase<IsPostDom>::properlyDominatesImpl(
}
}
+ // TODO: this should handle breaks in the block. This is not yet implemented.
// If the blocks are different, use DomTree to resolve the query.
return getDomTree(aRegion).properlyDominates(aBlock, bBlock);
}
diff --git a/mlir/lib/IR/Operation.cpp b/mlir/lib/IR/Operation.cpp
index b7227d0802ea8..ca46a1165924d 100644
--- a/mlir/lib/IR/Operation.cpp
+++ b/mlir/lib/IR/Operation.cpp
@@ -13,6 +13,7 @@
#include "mlir/IR/Dialect.h"
#include "mlir/IR/IRMapping.h"
#include "mlir/IR/Matchers.h"
+#include "mlir/IR/OpDefinition.h"
#include "mlir/IR/OpImplementation.h"
#include "mlir/IR/OperationSupport.h"
#include "mlir/IR/PatternMatch.h"
@@ -89,7 +90,8 @@ Operation *Operation::create(Location location, OperationName name,
unsigned numSuccessors = successors.size();
unsigned numOperands = operands.size();
unsigned numResults = resultTypes.size();
- int opPropertiesAllocSize = llvm::alignTo<8>(name.getOpPropertyByteSize());
+ size_t opPropertiesByteSize = name.getOpPropertyByteSize();
+ int opPropertiesAllocSize = llvm::alignTo<8>(opPropertiesByteSize);
// If the operation is known to have no operands, don't allocate an operand
// storage.
@@ -166,7 +168,7 @@ Operation::Operation(Location location, OperationName name, unsigned numResults,
"allowUnregisteredDialects() on the MLIRContext, or use "
"-allow-unregistered-dialect with the MLIR tool used.");
#endif
- if (fullPropertiesStorageSize)
+ if (name.getOpPropertyByteSize())
name.initOpProperties(getPropertiesStorage(), properties);
}
diff --git a/mlir/lib/IR/PatternMatch.cpp b/mlir/lib/IR/PatternMatch.cpp
index cd067f2cc25b3..c0ae9af2f5083 100644
--- a/mlir/lib/IR/PatternMatch.cpp
+++ b/mlir/lib/IR/PatternMatch.cpp
@@ -10,6 +10,9 @@
#include "mlir/IR/Iterators.h"
#include "mlir/IR/RegionKindInterface.h"
#include "llvm/ADT/SmallPtrSet.h"
+#include "llvm/Support/DebugLog.h"
+
+#define DEBUG_TYPE "pattern-match"
using namespace mlir;
@@ -225,7 +228,7 @@ void RewriterBase::eraseOp(Operation *op) {
// Then erase the enclosing op.
eraseSingleOp(op);
};
-
+ LDBG() << "RewriterBase::eraseOp: " << *op;
eraseTree(op);
}
diff --git a/mlir/lib/IR/RegionKindInterface.cpp b/mlir/lib/IR/RegionKindInterface.cpp
index 007f4cf92dbc7..8ba9f3eb0ff1b 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,151 @@ 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});
+ }
+ }
+}
+
+bool mlir::hasNestedPredecessors(Operation *op) {
+ bool found = false;
+ walk(op, [&](Operation *visitedOp, int nestedLevel) {
+ auto target = findBreakTarget(visitedOp);
+ if (nestedLevel > 1 && target && target.getOperation() == op)
+ found = true;
+ return found ? WalkResult::interrupt() : WalkResult::advance();
+ });
+ return found;
+}
+
+bool mlir::hasBreakingControlFlowOps(Operation *op) {
+ bool found = false;
+ walk(op, [&](Operation *visitedOp, int nestedLevel) {
+ auto target = findBreakTarget(visitedOp);
+ if (target && target.getOperation()->isProperAncestor(op))
+ found = true;
+ return found ? WalkResult::interrupt() : WalkResult::advance();
+ });
+ return found;
+}
+
+void mlir::detail::visitNestedBreakingControlFlowOpsImpl(
+ Operation *op,
+ function_ref<WalkResult(Operation *, int nestedLevel)> callback) {
+ ::walk(op, [&](Operation *visitedOp, int nestedLevel) {
+ auto target = findBreakTarget(visitedOp);
+ if (target && (target.getOperation() == op ||
+ target.getOperation()->isProperAncestor(op)))
+ return callback(visitedOp, nestedLevel);
+ return WalkResult::advance();
+ });
+}
+
+void mlir::collectAllNestedPredecessors(
+ Operation *op, SmallVector<Operation *> &predecessors) {
+ visitNestedBreakingControlFlowOps(op,
+ [&](Operation *visitedOp, int nestedLevel) {
+ auto target = findBreakTarget(visitedOp);
+ if (target && target.getOperation() == op)
+ predecessors.push_back(visitedOp);
+ return WalkResult::advance();
+ });
+}
+
+HasBreakingControlFlowOpInterface mlir::findBreakTarget(Operation *terminator) {
+ if (!terminator->mightHaveTrait<OpTrait::RegionTerminator>() ||
+ terminator->getNumOperands() == 0)
+ return {};
+ Value target = terminator->getOperand(0);
+ if (!isa<TokenType>(target.getType()))
+ return {};
+ auto targetArg = dyn_cast<BlockArgument>(target);
+ if (!targetArg)
+ return {};
+ Block *targetBlock = targetArg.getOwner();
+ if (!targetBlock || !targetBlock->isEntryBlock())
+ return {};
+ return dyn_cast_or_null<HasBreakingControlFlowOpInterface>(
+ targetBlock->getParentOp());
+}
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 c3fb73acf5ef0..dedf8d25d2cc0 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,6 +169,9 @@ LogicalResult detail::verifyRegionBranchOpInterface(Operation *op) {
SmallVector<RegionSuccessor> successors;
regionInterface.getSuccessorRegions(branchPoint, successors);
for (const RegionSuccessor &successor : successors) {
+ // Skip propagating-break sentinels — they are resolved by the ancestor.
+ if (successor.isPropagating())
+ continue;
// Helper function that print the region branch point and the region
// successor.
auto emitRegionEdgeError = [&]() {
@@ -218,6 +222,7 @@ LogicalResult detail::verifyRegionBranchOpInterface(Operation *op) {
}
}
}
+
return success();
}
@@ -246,7 +251,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 +259,7 @@ static bool traverseRegionGraph(Region *begin,
if (!terminator)
continue;
SmallVector<RegionSuccessor> successors;
- operandAttributes.resize(terminator->getNumOperands());
- terminator.getSuccessorRegions(operandAttributes, successors);
+ resolveTerminatorSuccessors(terminator, successors);
LDBG() << "Found " << successors.size()
<< " successors from terminator in block";
for (RegionSuccessor successor : successors) {
@@ -462,6 +465,31 @@ RegionBranchOpInterface::getNonSuccessorInputs(RegionSuccessor successor) {
return results;
}
+RegionBranchOpInterface mlir::resolveTerminatorSuccessors(
+ RegionBranchTerminatorOpInterface terminator,
+ SmallVectorImpl<RegionSuccessor> &successors) {
+ auto branch = dyn_cast<RegionBranchOpInterface>(terminator->getParentOp());
+ if (!branch)
+ return nullptr;
+ branch.getSuccessorRegions(RegionBranchPoint(terminator), successors);
+
+ // Check for the propagating sentinel: the immediate parent is transparent
+ // to this break. Reroute to the actual break target ancestor.
+ if (successors.size() == 1 && successors[0].isPropagating()) {
+ successors.clear();
+ if (auto breakTarget = findBreakTarget(terminator)) {
+ if (auto actualBranch =
+ dyn_cast<RegionBranchOpInterface>(breakTarget.getOperation())) {
+ actualBranch.getSuccessorRegions(RegionBranchPoint(terminator),
+ successors);
+ return actualBranch;
+ }
+ }
+ return nullptr;
+ }
+ return branch;
+}
+
static MutableArrayRef<OpOperand> operandsToOpOperands(OperandRange &operands) {
return MutableArrayRef<OpOperand>(operands.getBase(), operands.size());
}
@@ -473,6 +501,9 @@ getSuccessorOperandInputMapping(RegionBranchOpInterface branchOp,
SmallVector<RegionSuccessor> successors;
branchOp.getSuccessorRegions(src, successors);
for (RegionSuccessor dst : successors) {
+ // Skip propagating-break sentinels — they don't map to this op's inputs.
+ if (dst.isPropagating())
+ continue;
OperandRange operands = branchOp.getSuccessorOperands(src, dst);
assert(operands.size() == branchOp.getSuccessorInputs(dst).size() &&
"expected the same number of operands and inputs");
@@ -1134,6 +1165,10 @@ computeSingleAcyclicRegionBranchPath(RegionBranchOpInterface op) {
// through the region branch op.
return {};
}
+ if (successors.front().isPropagating()) {
+ // Propagating break — can't inline through this op.
+ return {};
+ }
path.push_back(successors.front());
if (successors.front().isParent()) {
// Found path that ends with "parent".
diff --git a/mlir/lib/TableGen/Operator.cpp b/mlir/lib/TableGen/Operator.cpp
index 82dfbcbfa4d4f..57f3249c8a598 100644
--- a/mlir/lib/TableGen/Operator.cpp
+++ b/mlir/lib/TableGen/Operator.cpp
@@ -747,12 +747,10 @@ void Operator::populateOpStructure() {
auto *dependentTraits = trait->getValueAsListInit("dependentTraits");
for (auto *traitInit : *dependentTraits)
if (!traitSet.contains(traitInit))
- PrintFatalError(
- def.getLoc(),
- trait->getValueAsString("trait") + " requires " +
- cast<DefInit>(traitInit)->getDef()->getValueAsString(
- "trait") +
- " to precede it in traits list");
+ PrintFatalError(def.getLoc(),
+ trait->getName() + " requires " +
+ cast<DefInit>(traitInit)->getDef()->getName() +
+ " to precede it in traits list");
};
std::function<void(const ListInit *)> insert;
diff --git a/mlir/lib/Transforms/Utils/CMakeLists.txt b/mlir/lib/Transforms/Utils/CMakeLists.txt
index 335c2cacd2a4a..a976ddb39f923 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..b2c10d929bf36
--- /dev/null
+++ b/mlir/test/Analysis/DataFlow/test-dead-code-analysis-early-exit.mlir
@@ -0,0 +1,101 @@
+// 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 an empty successor list for a propagating break, we use
+// findBreakTarget to locate 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 %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 %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 %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..a5dbadcbe6d92 100644
--- a/mlir/test/Analysis/test-dominance.mlir
+++ b/mlir/test/Analysis/test-dominance.mlir
@@ -680,3 +680,29 @@ 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 %outer -> index {
+ scf.loop %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
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..1fd028c10ac5c
--- /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 %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 %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 %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 %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 %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..111b4ba18074a
--- /dev/null
+++ b/mlir/test/Dialect/SCF/loop_canonicalize.mlir
@@ -0,0 +1,376 @@
+// 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 %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 %outer -> index {
+ scf.loop %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 %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
+}
+
+// -----
+
+// TODO: We should combine these but we don't right now
+// CHECK-LABEL: func @loop_combine_ifs
+func.func @loop_combine_ifs(%arg0 : i1, %arg2: i64) {
+ // Verify that we don't combine ifs when terminator smatches
+ scf.loop %loop {
+ // CHECK: scf.if
+ // TODO-CHECK-NOT: 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 %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 %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 %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 %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 %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 %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
+}
+
+
+// -----
+
+// CHECK-LABEL: func @fold_constant_if_with_breaking_cf1
+func.func @fold_constant_if_with_breaking_cf1(%arg0 : index, %arg1 : index) -> index {
+ %cond = arith.constant true
+ // Infinite loop here, inner if can be simplified, the "break" is
+ // unreachable.
+ // CHECK: scf.loop
+ // CHECK-NEXT: }
+ %loop_res = scf.loop %loop -> index {
+ %0 = scf.if %cond -> index {
+ scf.yield %arg0 : index
+ } else {
+ scf.break %loop %arg1 : index
+ }
+ }
+ return %loop_res : index
+}
+
+// -----
+
+// CHECK-LABEL: func @fold_constant_if_with_breaking_cf2
+func.func @fold_constant_if_with_breaking_cf2(%arg0 : index, %arg1 : index) -> index {
+ %cond = arith.constant false
+ // Infinite loop here, inner if can be simplified, the "break" is
+ // unreachable.
+ // CHECK: scf.loop
+ // CHECK-NEXT: }
+ %loop_res = scf.loop %loop -> index {
+ %0 = scf.if %cond -> index {
+ scf.break %loop %arg1 : index
+ } else {
+ scf.yield %arg0 : index
+ }
+ }
+ return %loop_res : index
+}
+
+// -----
+
+// CHECK-LABEL: func @fold_constant_if_with_breaking_cf3
+func.func @fold_constant_if_with_breaking_cf3(%arg0 : index, %arg1 : index) -> index {
+ %cond = arith.constant true
+ // Single iteration loop here, inner if can be simplified, and then the
+ // loop itself.
+ // CHECK-NOT: scf.loop
+ %loop_res = scf.loop %loop -> index {
+ %0 = scf.if %cond -> index {
+ scf.break %loop %arg1 : index
+ } else {
+ scf.yield %arg0 : index
+ }
+ }
+ return %loop_res : index
+}
+
+// -----
+
+// CHECK-LABEL: func @fold_constant_if_with_breaking_cf4
+func.func @fold_constant_if_with_breaking_cf4(%arg0 : index, %arg1 : index) -> index {
+ %cond = arith.constant false
+ // Single iteration loop here, inner if can be simplified, and then the
+ // loop itself.
+ // CHECK-NOT: scf.loop
+ %loop_res = scf.loop %loop -> index {
+ %0 = scf.if %cond -> index {
+ scf.yield %arg0 : index
+ } else {
+ scf.break %loop %arg1 : 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 %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 %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 %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 %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..c502e69dded6a
--- /dev/null
+++ b/mlir/test/IR/early-exit-invalid.mlir
@@ -0,0 +1,120 @@
+
+// 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 parent: successor operand type #0 'f32' should match successor input type #0 'i32'}}
+ %result = scf.loop %loop -> i32 {
+ scf.break %loop %value : f32 // expected-note {{region branch point}}
+ }
+ return
+}
+
+// -----
+
+func.func @loop_result_number_mismatch(%value : f32) {
+ // expected-error @+1 {{'scf.loop' op along control flow edge from Operation scf.break to parent: region branch point has 1 operands, but region successor needs 2 inputs}}
+ %result:2 = scf.loop %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 %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 %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 %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 %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..73f72c7ca56b2
--- /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 %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 %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 %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 %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 %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 %[[OUTER:.*]] {
+ scf.loop %outer {
+ // CHECK: scf.loop %[[INNER:.*]] {
+ scf.loop %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..dbbf91e0f08c4
--- /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 %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 %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..66f4276abcde8
--- /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 %outer_token -> i32 {
+ %inner = scf.loop %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/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..a1910465d93eb
--- /dev/null
+++ b/mlir/test/lib/Interfaces/RegionBranchOpInterface/TestRegionBranchOpInterface.cpp
@@ -0,0 +1,76 @@
+//===- TestBlockInLoop.cpp - Pass to test mlir::blockIsInLoop -------------===//
+//
+// 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 {
+/// This is a test pass that tests Blocks's isInLoop method by checking if each
+/// block in a function is in a loop and outputing if it is
+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.isParent()) {
+ llvm::outs() << " - Successor is parent\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/test/mlir-tblgen/op-error.td b/mlir/test/mlir-tblgen/op-error.td
index a2eab1f08df28..658a9a951da84 100644
--- a/mlir/test/mlir-tblgen/op-error.td
+++ b/mlir/test/mlir-tblgen/op-error.td
@@ -121,6 +121,6 @@ def OpInterfaceB : OpInterface<"OpInterfaceB"> {
let dependentTraits = [OpTraitA];
}
-// ERROR13: error: OpInterfaceB::Trait requires OpTraitA to precede it in traits list
+// ERROR13: error: OpInterfaceB requires OpTraitA to precede it in traits list
def OpInterfaceWithoutDependentTrait : Op<Test_Dialect, "default_value", [OpInterfaceB]> {}
#endif
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();
diff --git a/mlir/tools/mlir-tblgen/OpFormatGen.cpp b/mlir/tools/mlir-tblgen/OpFormatGen.cpp
index cbcbc8e9bc102..bf02369a44b39 100644
--- a/mlir/tools/mlir-tblgen/OpFormatGen.cpp
+++ b/mlir/tools/mlir-tblgen/OpFormatGen.cpp
@@ -1409,7 +1409,6 @@ void OperationFormat::genParser(Operator &op, OpClass &opClass) {
auto *method = opClass.addStaticMethod("::mlir::ParseResult", "parse",
std::move(paramList));
auto &body = method->body();
-
// Generate variables to store the operands and type within the format. This
// allows for referencing these variables in the presence of optional
// groupings.
@@ -2792,7 +2791,6 @@ class OpFormatParser : public FormatParser {
FailureOr<FormatElement *> parseTypeDirective(SMLoc loc, Context context);
FailureOr<FormatElement *> parseTypeDirectiveOperand(SMLoc loc,
bool isRefChild = false);
-
//===--------------------------------------------------------------------===//
// Fields
//===--------------------------------------------------------------------===//
@@ -3438,7 +3436,6 @@ OpFormatParser::parseDirectiveImpl(SMLoc loc, FormatToken::Kind kind,
return parseTypeDirective(loc, ctx);
case FormatToken::kw_oilist:
return parseOIListDirective(loc, ctx);
-
default:
return emitError(loc, "unsupported directive kind");
}
diff --git a/mlir/unittests/Debug/FileLineColLocBreakpointManagerTest.cpp b/mlir/unittests/Debug/FileLineColLocBreakpointManagerTest.cpp
index a345aec72f583..4617a8acb7457 100644
--- a/mlir/unittests/Debug/FileLineColLocBreakpointManagerTest.cpp
+++ b/mlir/unittests/Debug/FileLineColLocBreakpointManagerTest.cpp
@@ -23,6 +23,7 @@ static Operation *createOp(MLIRContext *context, Location loc,
StringRef operationName,
unsigned int numRegions = 0) {
context->allowUnregisteredDialects();
+
return Operation::create(loc, OperationName(operationName, context), {}, {},
NamedAttrList(), PropertyRef(), {}, numRegions);
}
More information about the Mlir-commits
mailing list