[Mlir-commits] [llvm] [mlir] [MLIR] Introduce support for early exits (PR #166688)
Mehdi Amini
llvmlistbot at llvm.org
Sat Oct 31 07:42:07 PDT 2026
https://github.com/joker-eph updated https://github.com/llvm/llvm-project/pull/166688
>From c980f895628ac7a9e790f4f2038d74b76f2d2a47 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][scf] Add token-targeted early exits for scf.loop
Introduce IR-level interfaces for region-exit terminators, operations that
receive breaking control flow, and operations that propagate it. Represent
region-branch successors with explicit destination owners and document the
propagation and potential-target contracts.
Add scf.loop with a control token, optional loop-carried arguments, and
independent result types. Add token-targeted scf.break and scf.continue, and
let scf.if propagate exits to enclosing receivers. Verify token users and
target ancestry, and handle early exits in parsing, printing, result
inference, bufferization, and value bounds.
Adapt dataflow, post-dominance, and inlining to ancestor-directed exits. Limit
nested branch-point enumeration to receivers and prune traversal at
propagation boundaries. Restrict canonicalization to owned successor inputs,
stop transitive forwarding at ancestor inputs, and reject inlining paths that
escape the matched operation.
Lower loops and exit terminators to CFG branches with rollback-safe operand
and block-signature changes. Account for deferred terminator erasure when
lowering nested loops.
Add test-dialect loops with dynamic-depth exits and a conversion to
token-targeted SCF. Dynamic target sets retain every compatible reachable
target; incompatible runtime depth choices are undefined, while constant
depths diagnose invalid targets and payloads.
Add verifier, canonicalization, analysis, lowering, callback-builder,
bufferization, value-bounds, round-trip, and execution coverage. Document the
model and update CMake and Bazel interface generation rules.
Validation: check-mlir passed with 4,738 tests passing, 171 unsupported, and
one expected failure. Nested verification time decreased by 18-21% on the
generated benchmark inputs. The Bazel build was not run.
Assisted-by: Codex
Assisted-by: Claude Code
---
mlir/docs/LangRef.md | 126 ++++-
mlir/include/mlir/Dialect/SCF/IR/SCF.h | 44 ++
mlir/include/mlir/Dialect/SCF/IR/SCFOps.td | 207 ++++++-
mlir/include/mlir/IR/CMakeLists.txt | 1 +
mlir/include/mlir/IR/Dominance.h | 9 +
mlir/include/mlir/IR/OpBase.td | 7 +
mlir/include/mlir/IR/OpDefinition.h | 13 +
mlir/include/mlir/IR/RegionExitInterfaces.h | 124 ++++
mlir/include/mlir/IR/RegionExitInterfaces.td | 95 ++++
.../mlir/Interfaces/ControlFlowInterfaces.h | 12 +-
.../mlir/Interfaces/ControlFlowInterfaces.td | 26 +-
.../Analysis/DataFlow/DeadCodeAnalysis.cpp | 12 +-
mlir/lib/Analysis/DataFlow/DenseAnalysis.cpp | 9 +-
mlir/lib/Analysis/DataFlow/SparseAnalysis.cpp | 23 +-
.../SCFToControlFlow/SCFToControlFlow.cpp | 210 +++++--
mlir/lib/Dialect/SCF/IR/SCF.cpp | 528 ++++++++++++++++--
.../SCF/IR/ValueBoundsOpInterfaceImpl.cpp | 18 +-
.../BufferizableOpInterfaceImpl.cpp | 57 +-
mlir/lib/Dialect/XeGPU/Utils/XeGPUUtils.cpp | 14 +-
mlir/lib/IR/CMakeLists.txt | 2 +
mlir/lib/IR/Dominance.cpp | 85 ++-
mlir/lib/IR/RegionExitInterfaces.cpp | 171 ++++++
mlir/lib/IR/Verifier.cpp | 1 +
mlir/lib/Interfaces/ControlFlowInterfaces.cpp | 80 ++-
mlir/lib/Transforms/Utils/CMakeLists.txt | 1 +
mlir/lib/Transforms/Utils/InliningUtils.cpp | 8 +
.../DataFlow/early-exit-backward.mlir | 46 ++
.../test-dead-code-analysis-early-exit.mlir | 96 ++++
mlir/test/Analysis/test-dominance.mlir | 60 ++
.../convert-early-exit-to-cfg.mlir | 196 +++++++
.../Dialect/SCF/early-exit-bufferize.mlir | 43 ++
mlir/test/Dialect/SCF/early-exit-invalid.mlir | 221 ++++++++
mlir/test/Dialect/SCF/early-exit.mlir | 125 +++++
.../test/Dialect/SCF/if-callback-builder.mlir | 51 ++
mlir/test/Dialect/SCF/loop_canonicalize.mlir | 496 ++++++++++++++++
.../SCF/value-bounds-op-interface-impl.mlir | 58 ++
.../IR/early-exit-test-dialect-invalid.mlir | 142 +++++
mlir/test/IR/early-exit-test-dialect.mlir | 124 ++++
.../Integration/Dialect/SCF/early_exit.mlir | 84 +++
mlir/test/Transforms/sccp-early-exit.mlir | 117 ++++
.../test-convert-breakable-loop-to-scf.mlir | 139 +++++
mlir/test/lib/Dialect/SCF/TestSCFUtils.cpp | 35 ++
mlir/test/lib/Dialect/Test/TestDialect.h | 3 +
mlir/test/lib/Dialect/Test/TestOpDefs.cpp | 347 ++++++++++++
mlir/test/lib/Dialect/Test/TestOps.h | 1 +
mlir/test/lib/Dialect/Test/TestOps.td | 74 +++
mlir/test/lib/Dialect/Test/TestPatterns.cpp | 208 +++++++
mlir/test/lib/IR/TestDominance.cpp | 30 +
mlir/test/lib/Interfaces/CMakeLists.txt | 1 +
.../RegionBranchOpInterface/CMakeLists.txt | 9 +
.../TestRegionBranchOpInterface.cpp | 78 +++
.../test-convert-breakable-loop-to-scf.mlir | 145 +++++
mlir/tools/mlir-opt/CMakeLists.txt | 1 +
mlir/tools/mlir-opt/mlir-opt.cpp | 2 +
.../llvm-project-overlay/mlir/BUILD.bazel | 3 +
55 files changed, 4616 insertions(+), 202 deletions(-)
create mode 100644 mlir/include/mlir/IR/RegionExitInterfaces.h
create mode 100644 mlir/include/mlir/IR/RegionExitInterfaces.td
create mode 100644 mlir/lib/IR/RegionExitInterfaces.cpp
create mode 100644 mlir/test/Analysis/DataFlow/early-exit-backward.mlir
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/early-exit-bufferize.mlir
create mode 100644 mlir/test/Dialect/SCF/early-exit-invalid.mlir
create mode 100644 mlir/test/Dialect/SCF/early-exit.mlir
create mode 100644 mlir/test/Dialect/SCF/if-callback-builder.mlir
create mode 100644 mlir/test/Dialect/SCF/loop_canonicalize.mlir
create mode 100644 mlir/test/IR/early-exit-test-dialect-invalid.mlir
create mode 100644 mlir/test/IR/early-exit-test-dialect.mlir
create mode 100644 mlir/test/Integration/Dialect/SCF/early_exit.mlir
create mode 100644 mlir/test/Transforms/sccp-early-exit.mlir
create mode 100644 mlir/test/Transforms/test-convert-breakable-loop-to-scf.mlir
create mode 100644 mlir/test/lib/Interfaces/RegionBranchOpInterface/CMakeLists.txt
create mode 100644 mlir/test/lib/Interfaces/RegionBranchOpInterface/TestRegionBranchOpInterface.cpp
create mode 100644 mlir/test/mlir-runner/test-convert-breakable-loop-to-scf.mlir
diff --git a/mlir/docs/LangRef.md b/mlir/docs/LangRef.md
index fea34c093418d3..326b3af366fbb6 100644
--- a/mlir/docs/LangRef.md
+++ b/mlir/docs/LangRef.md
@@ -509,7 +509,18 @@ operation produces results and the operation returns, those results also have
well-defined values. Execution then proceeds to the next operation in the block
until the terminator operation at the end of the block is reached; the
terminator determines the next continuation, if any. The determination of the
-next instruction to execute is the 'passing of control flow'.
+next instruction to execute is the 'passing of control flow'. A nested
+[region-exit terminator](#region-exit-terminators) can also terminate the
+current region with a request for a breaking-control-flow event addressed to a
+specific ancestor operation that implements
+`HasBreakingControlFlowOpInterface`. The terminator itself does not jump
+directly to that ancestor; it transfers control back to the containing
+operation with that request. If the containing operation is not the addressed
+receiver, it must define `PropagateControlFlowBreak` and propagate the same
+request to its own parent. This repeats until the addressed receiver handles
+the request, so every intermediate operation on the path from the terminator to
+the receiving `HasBreakingControlFlowOpInterface` operation must define
+`PropagateControlFlowBreak`.
In general, when control flow is passed to an operation, MLIR does not restrict
when control flow enters or exits the regions contained in that operation.
@@ -519,25 +530,26 @@ represent possible continuations. Successors explicitly specify destination
blocks, so control flow within the region can only pass to one of the specified
successor blocks, as in a `branch` operation. When a terminator has no
successors, it may pass 1) control back to the containing operation, as in a
-`return` operation, or 2) define that control flow does not continue, as in
-`ub.unreachable`. Terminators without successors therefore do not necessarily
-imply a return to the containing operation, the specific dialect operation
-determines the terminator's semantics. Blocks (other than the entry block) that
-are not listed as a successor of any terminator operation are defined to be
-unreachable and can be removed without affecting the semantics of the
-containing operation.
+`return` operation, 2) define that control flow does not continue, as in
+`ub.unreachable`, or 3) request that the containing operation propagate breaking
+control flow outward toward a specific ancestor. The target is
+dialect-defined. Terminators without successors therefore do not
+necessarily imply a return to the containing operation. Blocks (other than the
+entry block) that are not listed as a successor of any terminator operation are
+defined to be unreachable and can be removed without affecting the semantics of
+the containing operation.
Although control flow always enters a region through the entry block, control
flow may exit a region through any block with an appropriate terminator. The
-standard dialect leverages this capability to define operations with
+`func` dialect for example uses this capability to define operations with
Single-Entry-Multiple-Exit (SEME) regions, possibly flowing through different
blocks in the region and exiting through any block with a `return` operation.
-This behavior is similar to that of a function body in most programming
-languages. In addition, control flow may also not reach the end of a block or
-region, for example if a function call does not return. Such an operation
-prevents control flow from reaching later operations, but does not remove the
-structural requirement that the block must end with a terminator unless the
-enclosing operation opts out with `NoTerminator`.
+This behavior can model that of a function body in most programming languages.
+In addition, control flow may also not reach the end of a block or region, for
+example if a function call does not return. Such an operation prevents control
+flow from reaching later operations, but does not remove the structural
+requirement that the block must end with a terminator unless the enclosing
+operation opts out with `NoTerminator`.
Example:
@@ -565,6 +577,90 @@ func.func @accelerator_compute(i64, i1) -> i64 { // An SSACFG region
}
```
+#### Region-Exit Terminators
+
+A terminator that implements `RegionExitTerminatorOpInterface` terminates the
+current region and identifies the operation that may receive the region exit. A
+normal region exit can target the immediately containing operation. For
+example, `scf.yield` exits only its own immediately enclosing region and
+returns control to the parent operation.
+
+Some terminators instead request a breaking-control-flow event addressed to a
+specific ancestor. The terminator transfers control back only to its immediate
+parent operation; reaching the addressed ancestor requires each intermediate
+parent operation to collaborate by propagating the request outward. The way a
+terminator designates the ancestor is dialect-defined: it may be a builtin
+`token` operand, an integer count of region levels to exit, a named symbolic
+relationship to a parent operation, or another dialect-specific mechanism.
+
+Dialects must verify that every registered intermediate operation between a
+breaking terminator and its receiver defines `PropagateControlFlowBreak`, and
+that the receiver implements `HasBreakingControlFlowOpInterface`. These target
+and propagation constraints are checked by the dialect's terminator verifier;
+implementing `RegionExitTerminatorOpInterface` alone does not enforce them.
+Unregistered operations are conservatively treated as propagating exits.
+For example, a loop nested inside another loop can implement the receiver
+interface and carry the propagation trait simultaneously: it
+handles breaks whose token identifies itself, and propagates breaks whose token
+identifies an outer loop.
+
+Region-exit terminators may carry values to the receiver's continuation or to
+one of its region entries. For example, a break supplies a loop's result
+values, while a continue supplies loop-carried arguments at the body entry. The exact mapping between
+terminator operands and the receiving op's results is dialect-defined (the
+receiving op may also ignore operands entirely).
+
+Examples:
+
+```mlir
+// scf.yield is the standard immediate region-exit terminator.
+// It exits only its own immediately enclosing region.
+scf.if %cond {
+ scf.yield // returns control to scf.if
+}
+```
+
+```mlir
+// Trait legend:
+// [H] = HasBreakingControlFlowOpInterface (receives/catches the break)
+// [P] = PropagateControlFlowBreak (passes the break request upward)
+// [H][P] = both: handles breaks whose token identifies it, and propagates
+// breaks whose token identifies an outer loop
+scf.loop token(%outer) { // [H][P]
+ scf.loop token(%inner) { // [H][P]
+ scf.if %cond1 { // [P]
+ // Requests a break of the inner loop.
+ scf.break [%inner]
+ }
+ scf.if %cond2 { // [P]
+ // Requests a break of the outer loop. The scf.if and inner scf.loop
+ // both propagate the request upward.
+ scf.break [%outer]
+ }
+ scf.if %cond3 { // [P]
+ // Requests re-entry into the inner loop.
+ scf.continue [%inner]
+ }
+ }
+}
+return
+```
+
+```mlir
+// A loop that yields a result value on early exit.
+// scf.break carries operands that become the loop's results. (scf.continue
+// would instead carry operands that become the next iteration's iter_args,
+// but this loop has none.)
+%result = scf.loop token(%loop) -> f32 { // [H]
+ scf.if %found { // [P]
+ // %value becomes the loop result.
+ scf.break [%loop] %value : f32
+ }
+ // Re-enter the loop for the next iteration (no iter_args here).
+ scf.continue [%loop]
+}
+```
+
#### Operations with Multiple Regions
An operation containing multiple regions also completely determines the
diff --git a/mlir/include/mlir/Dialect/SCF/IR/SCF.h b/mlir/include/mlir/Dialect/SCF/IR/SCF.h
index 455c9c34cb81a3..1eefbed1f0a7bb 100644
--- a/mlir/include/mlir/Dialect/SCF/IR/SCF.h
+++ b/mlir/include/mlir/Dialect/SCF/IR/SCF.h
@@ -18,6 +18,7 @@
#include "mlir/Dialect/SCF/IR/SCFDialect.h"
#include "mlir/IR/Builders.h"
#include "mlir/IR/BuiltinTypes.h"
+#include "mlir/IR/RegionExitInterfaces.h"
#include "mlir/IR/RegionKindInterface.h"
#include "mlir/Interfaces/ControlFlowInterfaces.h"
#include "mlir/Interfaces/DestinationStyleOpInterface.h"
@@ -31,6 +32,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
@@ -113,6 +119,44 @@ SmallVector<Value> replaceAndCastForOpIterArg(RewriterBase &rewriter,
OpOperand &operand,
Value replacement,
const ValueTypeCastFnTy &castFn);
+namespace op_impl {
+/// An implicit terminator type for `if` operations. Empty regions still receive
+/// an `scf.yield`, but explicitly terminated regions may use any terminator
+/// implementing RegionExitTerminatorOpInterface so dialect-defined exits can
+/// propagate through `scf.if`.
+struct IfOpImplicitTerminatorType {
+ static bool classof(Operation *op) {
+ return isa<RegionExitTerminatorOpInterface>(op);
+ }
+
+ template <typename... Args>
+ static void build(Args &&...args) {
+ YieldOp::build(std::forward<Args>(args)...);
+ }
+ static constexpr StringLiteral getOperationName() {
+ return YieldOp::getOperationName();
+ }
+};
+struct LoopOpImplicitTerminatorType {
+ static bool classof(Operation *op) { return isa<ContinueOp, BreakOp>(op); }
+ static constexpr StringLiteral getOperationName() {
+ return ContinueOp::getOperationName();
+ }
+
+ /// Build the implicit `scf.continue` terminator. The control token consumed
+ /// by the terminator is the loop body's entry block argument #0; the
+ /// insertion block is guaranteed to be that body when the implicit
+ /// terminator is materialized. This keeps `scf.continue`'s public builders
+ /// free of any hidden block-argument assumption.
+ static void build(OpBuilder &builder, OperationState &state) {
+ Block *block = builder.getInsertionBlock();
+ assert(block && block->getNumArguments() > 0 &&
+ isa<TokenType>(block->getArgument(0).getType()) &&
+ "expected insertion block with a token-typed loop control argument");
+ state.addOperands(block->getArgument(0));
+ }
+};
+} // namespace op_impl
/// Helper function to compute the difference between two values. This is used
/// by the loop implementations to compute the trip count.
diff --git a/mlir/include/mlir/Dialect/SCF/IR/SCFOps.td b/mlir/include/mlir/Dialect/SCF/IR/SCFOps.td
index 6704e656a9ec8c..b4005bfab4f387 100644
--- a/mlir/include/mlir/Dialect/SCF/IR/SCFOps.td
+++ b/mlir/include/mlir/Dialect/SCF/IR/SCFOps.td
@@ -15,6 +15,8 @@
include "mlir/Interfaces/ControlFlowInterfaces.td"
include "mlir/Interfaces/LoopLikeInterface.td"
+include "mlir/IR/OpAsmInterface.td"
+include "mlir/IR/RegionExitInterfaces.td"
include "mlir/IR/RegionKindInterface.td"
include "mlir/Dialect/SCF/IR/DeviceMappingInterface.td"
include "mlir/Interfaces/DestinationStyleOpInterface.td"
@@ -149,6 +151,183 @@ 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,
+ HasNestedTerminators<["ContinueOp", "BreakOp"]>
+ ]> {
+ let summary = "Loop until a break operation";
+ let description = [{
+ The `loop` operation represents a loop that executes until a `break` is reached.
+
+ The loop consists of (1) a set of loop-carried values which are initialized by
+ `initArgs` and updated by each iteration of the loop, and
+ (2) a region which represents the loop body.
+
+ The loop will execute the body of the loop until a `break` is executed.
+
+ Each control path of the loop must be terminated by:
+
+ - a `continue` that yields the next iteration's value for each loop carried variable.
+ - a `break` that terminates the loop and yields the final loop results.
+ - a nested terminator requesting to propagate the control-flow to a parent.
+
+ The control token may only be used as the target operand of `scf.break` or
+ `scf.continue`. Both terminators can be nested under dialect-defined
+ operations that propagate breaking control flow.
+
+ Example:
+ ```mlir
+ %result = scf.loop token(%loop) iter_args(%arg = %init) : i32 -> i32 {
+ scf.if %done {
+ scf.break [%loop] %arg : i32
+ }
+ %next = arith.addi %arg, %one : i32
+ scf.continue [%loop] %next : i32
+ }
+ ```
+
+ This operation does not currently implement `LoopLikeOpInterface`, so
+ transformations that require that interface, such as loop invariant code
+ motion, do not apply to it.
+ }];
+
+ let arguments = (ins Variadic<AnyType>:$initArgs);
+ let results = (outs Variadic<AnyType>:$resultValues);
+ let regions = (region SizedRegion<1>:$region);
+
+ let extraClassDeclaration = [{
+ /// Return the iteration values of the loop region.
+ Block::BlockArgListType getRegionIterArgs() {
+ return getRegion().getArguments().drop_front();
+ }
+
+ /// Return the `index`-th region iteration value.
+ BlockArgument getRegionIterArg(unsigned index) {
+ return getRegionIterArgs()[index];
+ }
+
+ /// Return the loop control token.
+ BlockArgument getControlToken() {
+ return getRegion().getArgument(0);
+ }
+
+ /// Returns the number of region arguments for loop-carried values.
+ unsigned getNumRegionIterArgs() {
+ return getRegion().getNumArguments() - 1;
+ }
+
+ /// Returns the loop block body
+ Block *getBody() { return &getRegion().front(); }
+ }];
+
+ let hasCustomAssemblyFormat = 1;
+ let hasVerifier = 1;
+ let hasRegionVerifier = 1;
+ let hasCanonicalizer = 1;
+}
+
+//===----------------------------------------------------------------------===//
+// BreakOp
+//===----------------------------------------------------------------------===//
+
+def BreakOp : SCF_Op<"break", [
+ Terminator,
+ DeclareOpInterfaceMethods<RegionExitTerminatorOpInterface>,
+ DeclareOpInterfaceMethods<RegionBranchTerminatorOpInterface,
+ ["getMutableSuccessorOperands"]>
+ ]> {
+ let summary = "Break from loop";
+ let description = [{
+ The `break` operation is a terminator that exits one or more nested
+ regions and terminates the `scf.loop` that defined its control token.
+
+ The `break` may yield any number of operands; their types must match the
+ result types of the target `scf.loop`.
+
+ Example: break out of the immediately enclosing loop:
+ ```mlir
+ scf.loop token(%loop) -> i32 {
+ scf.break [%loop] %result : i32
+ }
+ ```
+
+ Example: break out of a loop through an enclosing `scf.if`:
+ ```mlir
+ scf.loop token(%loop) {
+ scf.if %cond {
+ scf.break [%loop]
+ }
+ scf.continue [%loop]
+ }
+ ```
+ }];
+
+ let arguments = (ins Token:$targetToken, Variadic<AnyType>:$args);
+ let assemblyFormat = [{
+ ` ` `[` $targetToken `]` ($args^ `:` type($args))? attr-dict
+ }];
+ let hasVerifier = 1;
+}
+
+//===----------------------------------------------------------------------===//
+// ContinueOp
+//===----------------------------------------------------------------------===//
+
+def ContinueOp : SCF_Op<"continue", [
+ Terminator,
+ DeclareOpInterfaceMethods<RegionExitTerminatorOpInterface>,
+ DeclareOpInterfaceMethods<RegionBranchTerminatorOpInterface,
+ ["getMutableSuccessorOperands"]>
+ ]> {
+ let summary = "Continue to next loop iteration";
+ let description = [{
+ The `continue` operation is a terminator that re-enters a `scf.loop`
+ for its next iteration. The target loop is the one that defined the control
+ token operand.
+
+ The operands of `continue` become the loop-carried values (iter_args) for
+ the next iteration; their types must match the loop's iter_arg types.
+
+ Example: continue the immediately enclosing loop:
+ ```mlir
+ scf.loop token(%loop) iter_args(%i = %init) : i32 {
+ %next = arith.addi %i, %one : i32
+ scf.continue [%loop] %next : i32
+ }
+ ```
+
+ Example: continue an outer loop from inside a nested `scf.if`:
+ ```mlir
+ scf.loop token(%outer) iter_args(%counter = %init) : i64 {
+ scf.loop token(%inner) iter_args(%inner_arg = %counter) : i64 {
+ scf.if %restart_outer {
+ scf.continue [%outer] %inner_arg : i64
+ }
+ scf.continue [%inner] %inner_arg : i64
+ }
+ scf.continue [%outer] %counter : i64
+ }
+ ```
+ }];
+
+ let arguments = (ins Token:$targetToken, Variadic<AnyType>:$args);
+ let assemblyFormat = [{
+ ` ` `[` $targetToken `]` ($args^ `:` type($args))? attr-dict
+ }];
+ let hasVerifier = 1;
+}
//===----------------------------------------------------------------------===//
// ForOp
@@ -717,8 +896,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
@@ -754,8 +933,14 @@ def IfOp : SCF_Op<"if", [DeclareOpInterfaceMethods<RegionBranchOpInterface, [
block. In case the `scf.if` produces results, the "else" region must also
have exactly 1 block.
- The blocks are always terminated with `scf.yield`. If `scf.if` defines no
- values, the `scf.yield` can be left out, and will be inserted implicitly.
+ A block normally terminates with `scf.yield`. It may instead terminate with
+ a `RegionExitTerminatorOpInterface` operation that exits an enclosing
+ receiver through this `scf.if`, such as `scf.break` or `scf.continue`.
+ Such a branch does not contribute values to the `scf.if` results. If both
+ branches exit enclosing receivers, the continuation is unreachable: explicit
+ result types are permitted, while the callback builder infers no results.
+ If `scf.if` defines no values, `scf.yield` can be left out and is inserted
+ implicitly.
Otherwise, it must be explicit.
Example:
@@ -801,9 +986,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;
@@ -937,6 +1130,7 @@ def ParallelOp : SCF_Op<"parallel",
def ReduceOp : SCF_Op<"reduce", [
Terminator, HasParent<"ParallelOp">, RecursiveMemoryEffects,
+ DeclareOpInterfaceMethods<RegionExitTerminatorOpInterface>,
DeclareOpInterfaceMethods<PromotableRegionOpInterface>,
DeclareOpInterfaceMethods<RegionBranchTerminatorOpInterface,
["getMutableSuccessorOperands"]>]> {
@@ -1238,6 +1432,7 @@ def IndexSwitchOp : SCF_Op<"index_switch", [RecursiveMemoryEffects,
//===----------------------------------------------------------------------===//
def YieldOp : SCF_Op<"yield", [Pure, ReturnLike, Terminator,
+ DeclareOpInterfaceMethods<RegionExitTerminatorOpInterface>,
ParentOneOf<["ExecuteRegionOp", "ForOp", "IfOp", "IndexSwitchOp",
"WhileOp"]>]> {
let summary = "loop yield and termination operation";
diff --git a/mlir/include/mlir/IR/CMakeLists.txt b/mlir/include/mlir/IR/CMakeLists.txt
index 15518901b901ab..f821423ff809c5 100644
--- a/mlir/include/mlir/IR/CMakeLists.txt
+++ b/mlir/include/mlir/IR/CMakeLists.txt
@@ -3,6 +3,7 @@ set(LLVM_TARGET_DEFINITIONS SymbolInterfaces.td)
mlir_tablegen(SymbolInterfacesAttrInterface.h.inc -gen-attr-interface-decls)
mlir_tablegen(SymbolInterfacesAttrInterface.cpp.inc -gen-attr-interface-defs)
add_mlir_interface(RegionKindInterface)
+add_mlir_interface(RegionExitInterfaces)
add_mlir_type_interface(QuantStorageTypeInterface)
diff --git a/mlir/include/mlir/IR/Dominance.h b/mlir/include/mlir/IR/Dominance.h
index 70924a2e9ae59d..01a9eccbaac86c 100644
--- a/mlir/include/mlir/IR/Dominance.h
+++ b/mlir/include/mlir/IR/Dominance.h
@@ -123,6 +123,11 @@ class DominanceInfoBase {
Block::iterator bIt,
bool enclosingOk = true) const;
+ /// Return true if `op` contains a nested breaking terminator that escapes
+ /// through it to an ancestor. This cache is invalidated together with
+ /// dominance info.
+ bool hasEscapingBreakingControlFlowCached(Operation *op) const;
+
/// A mapping of regions to their base dominator tree and a cached
/// "hasSSADominance" bit. This map does not contain dominator trees for
/// single block CFG regions, but we do want to cache the "hasSSADominance"
@@ -131,6 +136,10 @@ class DominanceInfoBase {
///
mutable DenseMap<Region *, llvm::PointerIntPair<DomTree *, 1, bool>>
dominanceInfos;
+
+ /// Cached escaping breaking-terminator query results. This is stable for the
+ /// lifetime of the analyzed IR and cleared on dominance invalidation.
+ mutable DenseMap<Operation *, bool> breakingControlFlowOpsCache;
};
extern template class DominanceInfoBase</*IsPostDom=*/true>;
diff --git a/mlir/include/mlir/IR/OpBase.td b/mlir/include/mlir/IR/OpBase.td
index a130d5847a1c72..655ace1967ad83 100644
--- a/mlir/include/mlir/IR/OpBase.td
+++ b/mlir/include/mlir/IR/OpBase.td
@@ -134,6 +134,13 @@ class SingleBlockImplicitTerminatorImpl<string op>
class SingleBlockImplicitTerminator<string op>
: TraitList<[SingleBlock, SingleBlockImplicitTerminatorImpl<op>]>;
+// This operation has nested regions with the supplied list of terminator
+// operations.
+class HasNestedTerminators<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 32f757ebceaff0..5f0a3aa7425ebf 100644
--- a/mlir/include/mlir/IR/OpDefinition.h
+++ b/mlir/include/mlir/IR/OpDefinition.h
@@ -1364,6 +1364,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/RegionExitInterfaces.h b/mlir/include/mlir/IR/RegionExitInterfaces.h
new file mode 100644
index 00000000000000..da78e648e94c06
--- /dev/null
+++ b/mlir/include/mlir/IR/RegionExitInterfaces.h
@@ -0,0 +1,124 @@
+//===- RegionExitInterfaces.h - Region exit interfaces ----------*- C++ -*-===//
+//
+// 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
+//
+//===----------------------------------------------------------------------===//
+
+#ifndef MLIR_IR_REGIONEXITINTERFACES_H
+#define MLIR_IR_REGIONEXITINTERFACES_H
+
+#include "mlir/IR/OpDefinition.h"
+
+namespace mlir {
+namespace OpTrait {
+/// Indicates that this operation is transparent to breaking control flow:
+/// a terminator (e.g. scf.break / scf.continue) can propagate through
+/// this op on its way to the addressed HasBreakingControlFlowOpInterface
+/// ancestor. The op does NOT consume the break; it simply passes it upward.
+/// All intermediate ops that receive and forward the break request must carry
+/// this trait.
+template <typename ConcreteType>
+class PropagateControlFlowBreak
+ : public TraitBase<ConcreteType, PropagateControlFlowBreak> {
+public:
+ static LogicalResult verifyTrait(Operation *op) {
+ // Verify the operation has regions and can handle breaking control flow
+ if (op->getNumRegions() == 0)
+ return op->emitOpError(
+ "operation with PropagateControlFlowBreak trait must have regions");
+ return success();
+ }
+};
+} // namespace OpTrait
+
+/// Summary of breaking terminator operations nested under an op.
+struct NestedBreakingControlFlowInfo {
+ /// Breaking terminators that may target the queried op directly. These
+ /// transfer control to the queried op rather than only to an intermediate
+ /// PropagateControlFlowBreak op.
+ SmallVector<Operation *> predecessors;
+
+ /// True if any direct predecessor targets the queried op from below the
+ /// queried op's immediate region.
+ bool hasNestedPredecessors = false;
+
+ /// True if any breaking terminator under the queried op targets one of the
+ /// queried op's ancestors, escaping through the queried op.
+ bool hasEscapingBreakingControlFlow = false;
+};
+
+/// Collect nested breaking-control-flow information for `op` with a single walk
+/// of its nested region tree.
+NestedBreakingControlFlowInfo getNestedBreakingControlFlowInfo(Operation *op);
+
+/// Return true if `op` (which implements HasBreakingControlFlowOpInterface)
+/// contains at least one breaking terminator that directly targets it from a
+/// nested region. Such a terminator is a "nested predecessor" of `op` because
+/// control flow may re-enter or exit `op` through a request propagated from a
+/// deeply nested site rather than only through an immediately enclosing
+/// terminator.
+bool hasNestedPredecessors(Operation *op);
+
+/// Return true if `op` contains any breaking terminator that would "break
+/// through" `op` towards an outer HasBreakingControlFlowOpInterface ancestor.
+/// This is used to detect whether an op's post-dominance is disrupted by an
+/// early-exit request that propagates through it.
+bool hasEscapingBreakingControlFlow(Operation *op);
+} // namespace mlir
+
+#include "mlir/IR/RegionExitInterfaces.h.inc"
+
+namespace mlir {
+namespace detail {
+/// Implementation helper for visitNestedBreakingControlFlowOps. Walks the
+/// regions of `op` and invokes `callback` for every breaking terminator that
+/// either targets `op` or propagates further upward through `op`.
+/// The `nestedLevel` argument passed to the callback is the 1-based depth of
+/// the terminator relative to `op`'s outermost region.
+void visitNestedBreakingControlFlowOpsImpl(
+ Operation *op,
+ function_ref<WalkResult(RegionExitTerminatorOpInterface, int nestedLevel)>
+ callback);
+} // namespace detail
+
+/// Walk all breaking terminators that are relevant to breaking control
+/// flow inside `op` (see visitNestedBreakingControlFlowOpsImpl). The callback
+/// receives the terminator op and its 1-based nesting level. Callbacks
+/// returning WalkResult support early termination via WalkResult::interrupt();
+/// void-returning callbacks always continue.
+template <typename CallbackT>
+void visitNestedBreakingControlFlowOps(Operation *op, CallbackT &&callback) {
+ using RetT =
+ decltype(callback(std::declval<Operation *>(), std::declval<int>()));
+ if constexpr (std::is_same_v<RetT, WalkResult>) {
+ detail::visitNestedBreakingControlFlowOpsImpl(op, callback);
+ } else {
+ detail::visitNestedBreakingControlFlowOpsImpl(
+ op, [&](RegionExitTerminatorOpInterface visitedOp, int nestedLevel) {
+ callback(visitedOp, nestedLevel);
+ return WalkResult::advance();
+ });
+ }
+}
+
+/// Return true if the given region may contain breaking control flow — either
+/// because its parent op propagates breaks (PropagateControlFlowBreak) or
+/// because it is the body of a HasBreakingControlFlowOpInterface op. Used to
+/// decide whether post-dominance analysis must account for early-exit paths.
+inline bool mightHaveBreakingControlFlow(Region *region) {
+ Operation *parentOp = region->getParentOp();
+ return (!parentOp->isRegistered() ||
+ parentOp->hasTrait<OpTrait::PropagateControlFlowBreak>() ||
+ isa<HasBreakingControlFlowOpInterface>(parentOp));
+}
+
+/// Return the HasBreakingControlFlowOpInterface operations potentially
+/// addressed by the given terminator. Returns an empty vector for non-breaking
+/// terminators or malformed target designators.
+SmallVector<HasBreakingControlFlowOpInterface>
+findPotentialBreakTargets(Operation *terminator);
+} // namespace mlir
+
+#endif // MLIR_IR_REGIONEXITINTERFACES_H
diff --git a/mlir/include/mlir/IR/RegionExitInterfaces.td b/mlir/include/mlir/IR/RegionExitInterfaces.td
new file mode 100644
index 00000000000000..fb93b5592f4d4f
--- /dev/null
+++ b/mlir/include/mlir/IR/RegionExitInterfaces.td
@@ -0,0 +1,95 @@
+//===- RegionExitInterfaces.td - Region exit interfaces ----*- tablegen -*-===//
+//
+// 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
+//
+//===----------------------------------------------------------------------===//
+
+#ifndef MLIR_IR_REGIONEXITINTERFACES
+#define MLIR_IR_REGIONEXITINTERFACES
+
+include "mlir/IR/OpBase.td"
+
+// Indicates that this op may propagate a breaking control-flow event from a
+// nested region upward to the addressed HasBreakingControlFlowOpInterface
+// operation. The op does NOT consume the break itself; it is merely transparent
+// to it. All ops that sit between a breaking terminator and the
+// HasBreakingControlFlowOpInterface ancestor that will ultimately receive the
+// break must carry this trait.
+def PropagateControlFlowBreak : NativeOpTrait<"PropagateControlFlowBreak">;
+
+def HasBreakingControlFlowOpInterface : OpInterface<"HasBreakingControlFlowOpInterface"> {
+ let description = [{
+ Interface for operations that act as the target of a breaking control-flow
+ event (e.g. `scf.break` or `scf.continue`). Every intermediate op must
+ carry the `PropagateControlFlowBreak` trait. The terminator may identify the
+ receiver with any dialect-defined mechanism, such as a builtin token value,
+ an attribute, or a symbolic relationship to an enclosing operation.
+ }];
+ let cppNamespace = "::mlir";
+
+ let methods = [
+ StaticInterfaceMethod<
+ /*desc=*/[{
+ Return true if this operation accepts the given terminator operation
+ as a breaking-control-flow predecessor. This can be implemented through
+ the `HasNestedTerminators<[...]>` trait, otherwise an explicit implementation
+ must be provided (e.g. `return true;` to accept all, or a type check to
+ restrict).
+ }],
+ /*retTy=*/"bool",
+ /*methodName=*/"acceptsTerminator",
+ /*args=*/(ins "Operation *":$op)
+ >
+ ];
+ let verifyWithRegions = 1;
+ let verify = [{
+ for (Operation *terminator :
+ ::mlir::getNestedBreakingControlFlowInfo($_op).predecessors)
+ if (!ConcreteOp::acceptsTerminator(terminator))
+ return $_op->emitOpError("does not accept terminator: ")
+ << OpWithFlags(terminator, OpPrintingFlags().skipRegions());
+ return success();
+ }];
+}
+
+// OpInterface for terminators that leave their current region and identify
+// potential receiving operations using a dialect-defined mechanism. A normal
+// region return may simply target the immediately containing operation; a
+// breaking-control-flow region exit may target an ancestor operation that
+// implements HasBreakingControlFlowOpInterface.
+def RegionExitTerminatorOpInterface : OpInterface<"RegionExitTerminatorOpInterface"> {
+ let description = [{
+ Interface for terminators that leave their current region. Implemented by
+ ops such as `scf.yield`, `scf.break`, and `scf.continue`.
+ }];
+ let cppNamespace = "::mlir";
+
+ let methods = [
+ InterfaceMethod<
+ /*desc=*/[{
+ Return the operations that may receive this region-terminating control
+ flow event. Receivers for breaking control flow must implement
+ `HasBreakingControlFlowOpInterface`; normal region exits may return the
+ immediate parent operation. This set must be a sound over-approximation:
+ it must contain every possible target in every defined execution.
+ Omitting a defined target makes control-flow analyses unsound. Dialects
+ may exclude choices whose execution has undefined behavior.
+ }],
+ /*retTy=*/"::llvm::SmallVector<::mlir::Operation *>",
+ /*methodName=*/"getPotentialTargets",
+ /*args=*/(ins),
+ /*methodBody=*/[{}]
+ >
+ ];
+
+ let verify = [{
+ static_assert(ConcreteOp::template hasTrait<OpTrait::IsTerminator>(),
+ "expected operation to be a terminator");
+ return success();
+ }];
+}
+
+
+#endif // MLIR_IR_REGIONEXITINTERFACES
diff --git a/mlir/include/mlir/Interfaces/ControlFlowInterfaces.h b/mlir/include/mlir/Interfaces/ControlFlowInterfaces.h
index 48b80e2059e119..6073af5ac0866a 100644
--- a/mlir/include/mlir/Interfaces/ControlFlowInterfaces.h
+++ b/mlir/include/mlir/Interfaces/ControlFlowInterfaces.h
@@ -193,9 +193,10 @@ using RegionBranchInverseSuccessorMapping =
DenseMap<Value, SmallVector<OpOperand *>>;
/// This class represents a successor of a region. A region successor can either
-/// target another region or target an ancestor operation (at the moment,
-/// limited to the immediate parent operation). In the latter case, the control
-/// flow branches after/out of the target operation.
+/// target a region or an ancestor operation. In the latter case, control flow
+/// branches after/out of the target operation. Propagated exits may target an
+/// operation other than the immediate parent, or a region owned by such an
+/// ancestor. Query getRegionBranchSuccessorOwner to identify that owner.
class RegionSuccessor {
public:
/// Initialize a successor that branches to a region of the parent operation.
@@ -446,6 +447,11 @@ inline llvm::raw_ostream &operator<<(llvm::raw_ostream &os,
OpPrintingFlags().skipRegions())
<< ">";
}
+
+/// Return the RegionBranchOpInterface operation that owns `successor` and
+/// defines its successor inputs.
+RegionBranchOpInterface
+getRegionBranchSuccessorOwner(RegionSuccessor successor);
} // namespace mlir
#endif // MLIR_INTERFACES_CONTROLFLOWINTERFACES_H
diff --git a/mlir/include/mlir/Interfaces/ControlFlowInterfaces.td b/mlir/include/mlir/Interfaces/ControlFlowInterfaces.td
index da79de52807d16..3be6ec1d332a39 100644
--- a/mlir/include/mlir/Interfaces/ControlFlowInterfaces.td
+++ b/mlir/include/mlir/Interfaces/ControlFlowInterfaces.td
@@ -235,7 +235,15 @@ def RegionBranchOpInterface : OpInterface<"RegionBranchOpInterface"> {
When `point = RegionBranchPoint::parent()`, this method returns the
region successors when entering the operation. Otherwise, this method
returns the successor regions when branching from the region indicated
- by `point`.
+ by `point`. A propagated region exit may target an ancestor operation
+ or one of its regions instead of a region/result owned by this op.
+ Query `getRegionBranchSuccessorOwner(successor).getSuccessorInputs(...)`
+ for those inputs. Clients that mutate inputs must check ownership.
+
+ For a `HasBreakingControlFlowOpInterface` receiver, `point` may also be
+ a nested terminator that targets this op or an ancestor through
+ propagating operations. Ordinary propagators enumerate only immediate
+ terminators; the receiver accounts for the nested incoming edges.
Example: In the above example, this method returns the region of the
`scf.for` and `parent` for either region branch point. An implementation
@@ -269,7 +277,9 @@ def RegionBranchOpInterface : OpInterface<"RegionBranchOpInterface"> {
Return all successor inputs for the given region successor. If the
given region successor is a region, then the returned values are block
arguments. Otherwise, if the given region successor is an operation,
- the returned values are op results.
+ the returned values are op results. This method is queried on the
+ owner of the successor, which may differ from the op that enumerated
+ the edge when control flow propagates to an ancestor.
}],
"::mlir::ValueRange", "getSuccessorInputs",
(ins "::mlir::RegionSuccessor":$successor),
@@ -392,7 +402,10 @@ def RegionBranchOpInterface : OpInterface<"RegionBranchOpInterface"> {
/// successor operand could be forwarded to multiple successor inputs.
/// Operands that are not forwarded are not added to the map. Unless a
/// specific region branch point is specified, this function takes into
- /// account all possible region branch points.
+ /// account all possible region branch points. Inputs may belong to an
+ /// ancestor operation or its regions. Such entries describe only the
+ /// edges visible from this op, not the owner's complete predecessor set;
+ /// callers must not use them to rewrite another owner's inputs.
void getSuccessorOperandInputMapping(
::mlir::RegionBranchSuccessorMapping &mapping,
std::optional<::mlir::RegionBranchPoint> src = std::nullopt);
@@ -403,7 +416,9 @@ def RegionBranchOpInterface : OpInterface<"RegionBranchOpInterface"> {
::mlir::RegionBranchInverseSuccessorMapping &mapping);
/// Return all possible region branch points: the region branch op itself
- /// and all region branch terminators.
+ /// and its immediate region branch terminators. Receivers implementing
+ /// HasBreakingControlFlowOpInterface additionally enumerate nested breaking
+ /// terminators targeting them or an ancestor through propagating ops.
::llvm::SmallVector<::mlir::RegionBranchPoint> getAllRegionBranchPoints();
}];
}
@@ -442,7 +457,8 @@ def RegionBranchTerminatorOpInterface :
>,
InterfaceMethod<[{
Returns all potential region successors that are branched to after this
- terminator based on the given constant operands.
+ terminator based on the given constant operands. A propagated exit may
+ target an ancestor operation or a region owned by that ancestor.
This method also receives the constant operands of this op (one entry
per operand, "null" if the operand has no/unknown constant value). The
diff --git a/mlir/lib/Analysis/DataFlow/DeadCodeAnalysis.cpp b/mlir/lib/Analysis/DataFlow/DeadCodeAnalysis.cpp
index 37aaf1f634a643..cf703785c7e082 100644
--- a/mlir/lib/Analysis/DataFlow/DeadCodeAnalysis.cpp
+++ b/mlir/lib/Analysis/DataFlow/DeadCodeAnalysis.cpp
@@ -509,11 +509,15 @@ void DeadCodeAnalysis::visitRegionTerminator(Operation *op,
if (!operands)
return;
- SmallVector<RegionSuccessor> successors;
auto terminator = dyn_cast<RegionBranchTerminatorOpInterface>(op);
if (!terminator)
return;
+
+ SmallVector<RegionSuccessor> successors;
+ // Use operand-aware resolution to prune dead successors via constant folding
+ // (e.g. a scf.while condition known to be false).
terminator.getSuccessorRegions(*operands, successors);
+
visitRegionBranchEdges(branch, op, successors);
}
@@ -540,8 +544,10 @@ void DeadCodeAnalysis::visitRegionBranchEdges(
auto *predecessors = getOrCreate<PredecessorState>(point);
propagateIfChanged(
predecessors,
- predecessors->join(predecessorOp,
- regionBranchOp.getSuccessorInputs(successor)));
+ predecessors->join(
+ predecessorOp,
+ getRegionBranchSuccessorOwner(successor).getSuccessorInputs(
+ successor)));
LDBG() << "Added region branch as predecessor for successor: " << *point;
}
}
diff --git a/mlir/lib/Analysis/DataFlow/DenseAnalysis.cpp b/mlir/lib/Analysis/DataFlow/DenseAnalysis.cpp
index a2a9018b851ea0..dc30aae167ac79 100644
--- a/mlir/lib/Analysis/DataFlow/DenseAnalysis.cpp
+++ b/mlir/lib/Analysis/DataFlow/DenseAnalysis.cpp
@@ -636,6 +636,9 @@ void AbstractDenseBackwardDataFlowAnalysis::visitRegionBranchOperation(
branch.getSuccessorRegions(branchPoint, successors);
LDBG() << " Processing " << successors.size() << " successor regions";
for (const RegionSuccessor &successor : successors) {
+ RegionBranchOpInterface successorOwner =
+ getRegionBranchSuccessorOwner(successor);
+ assert(successorOwner && "expected RegionBranchOpInterface owner");
const AbstractDenseLattice *after;
if (successor.isOperation()) {
LDBG() << " Successor is operation";
@@ -643,7 +646,7 @@ void AbstractDenseBackwardDataFlowAnalysis::visitRegionBranchOperation(
getProgramPointAfter(successor.getSuccessorOp()));
} else if (successor.getSuccessor()->empty()) {
LDBG() << " Successor is empty region";
- after = getLatticeFor(point, getProgramPointAfter(branch));
+ after = getLatticeFor(point, getProgramPointAfter(successorOwner));
} else {
Region *successorRegion = successor.getSuccessor();
assert(!successorRegion->empty() && "unexpected empty successor region");
@@ -662,7 +665,7 @@ void AbstractDenseBackwardDataFlowAnalysis::visitRegionBranchOperation(
}
LDBG() << " After state: " << *after;
- visitRegionBranchControlFlowTransfer(branch, branchPoint, successor, *after,
- before);
+ visitRegionBranchControlFlowTransfer(successorOwner, branchPoint, successor,
+ *after, before);
}
}
diff --git a/mlir/lib/Analysis/DataFlow/SparseAnalysis.cpp b/mlir/lib/Analysis/DataFlow/SparseAnalysis.cpp
index bc61c19832d297..90453f0e5fd78f 100644
--- a/mlir/lib/Analysis/DataFlow/SparseAnalysis.cpp
+++ b/mlir/lib/Analysis/DataFlow/SparseAnalysis.cpp
@@ -665,14 +665,23 @@ void AbstractSparseBackwardDataFlowAnalysis::
// non-contiguous in the presence of multiple successors.
BitVector unaccounted(terminator->getNumOperands(), true);
- RegionBranchSuccessorMapping mapping;
- branch.getSuccessorOperandInputMapping(mapping,
- RegionBranchPoint(terminator));
- for (const auto &[operand, inputs] : mapping) {
- for (Value input : inputs) {
- meet(getLatticeElement(operand->get()),
+ SmallVector<RegionSuccessor> successors;
+ branch.getSuccessorRegions(RegionBranchPoint(terminator), successors);
+
+ for (RegionSuccessor successor : successors) {
+ RegionBranchOpInterface successorOwner =
+ getRegionBranchSuccessorOwner(successor);
+ assert(successorOwner && "expected RegionBranchOpInterface owner");
+ OperandRange operands =
+ branch.getSuccessorOperands(RegionBranchPoint(terminator), successor);
+ ValueRange inputs = successorOwner.getSuccessorInputs(successor);
+ assert(operands.size() == inputs.size() &&
+ "expected the same number of operands and inputs");
+ MutableArrayRef<OpOperand> opOperands(operands.getBase(), operands.size());
+ for (const auto &[operand, input] : llvm::zip_equal(opOperands, inputs)) {
+ meet(getLatticeElement(operand.get()),
*getLatticeElementFor(getProgramPointAfter(terminator), input));
- unaccounted.reset(operand->getOperandNumber());
+ unaccounted.reset(operand.getOperandNumber());
}
}
diff --git a/mlir/lib/Conversion/SCFToControlFlow/SCFToControlFlow.cpp b/mlir/lib/Conversion/SCFToControlFlow/SCFToControlFlow.cpp
index 4d3cb36f0367ca..4c915c9f8643e6 100644
--- a/mlir/lib/Conversion/SCFToControlFlow/SCFToControlFlow.cpp
+++ b/mlir/lib/Conversion/SCFToControlFlow/SCFToControlFlow.cpp
@@ -12,6 +12,7 @@
//===----------------------------------------------------------------------===//
#include "mlir/Conversion/SCFToControlFlow/SCFToControlFlow.h"
+#include "mlir/IR/RegionExitInterfaces.h"
#include "mlir/Dialect/Arith/IR/Arith.h"
#include "mlir/Dialect/ControlFlow/IR/ControlFlowOps.h"
@@ -311,6 +312,18 @@ 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 void copyLLVMDialectAttrs(Operation *from, Operation *to) {
@@ -322,6 +335,66 @@ static void copyLLVMDialectAttrs(Operation *from, Operation *to) {
to->setDiscardableAttrs(llvmAttrs);
}
+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.
@@ -411,58 +484,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
@@ -740,10 +762,87 @@ LogicalResult ForallLowering::matchAndRewrite(ForallOp forallOp,
return scf::forallToParallelLoop(rewriter, forallOp);
}
+LogicalResult
+LoopOpLowering::matchAndRewrite(LoopOp loopOp, OpAdaptor adaptor,
+ ConversionPatternRewriter &rewriter) const {
+ {
+ OpBuilder::InsertionGuard guard(rewriter);
+ if (failed(rewriter.legalize(&loopOp.getRegion())))
+ return rewriter.notifyMatchFailure(loopOp,
+ "failed to convert nested region");
+ }
+
+ // Handle degenerate case before modifying any IR.
+ Region &bodyRegion = loopOp.getRegion();
+ if (bodyRegion.empty() || bodyRegion.front().empty()) {
+ rewriter.eraseOp(loopOp);
+ return success();
+ }
+
+ NestedBreakingControlFlowInfo nestedBreakingControlFlow =
+ getNestedBreakingControlFlowInfo(loopOp);
+ if (nestedBreakingControlFlow.hasNestedPredecessors)
+ return rewriter.notifyMatchFailure(loopOp,
+ "loop op with nested predecessors");
+
+ // Check token users before changing the block signature. Erasures may be
+ // deferred by dialect conversion, so use_empty() is not sufficient here.
+ BlockArgument controlToken = loopOp.getControlToken();
+ if (llvm::any_of(controlToken.getUses(), [&](OpOperand &use) {
+ return use.getOperandNumber() != 0 ||
+ !llvm::is_contained(nestedBreakingControlFlow.predecessors,
+ use.getOwner());
+ }))
+ return rewriter.notifyMatchFailure(loopOp,
+ "unsupported control token user");
+
+ Location loc = loopOp.getLoc();
+ Block *initBlock = loopOp->getBlock();
+ Block *remainingBlock = rewriter.splitBlock(initBlock, loopOp->getIterator());
+ Block *continueBlock =
+ rewriter.createBlock(remainingBlock, loopOp.getResultTypes(),
+ SmallVector<Location>(loopOp.getNumResults(), loc));
+ rewriter.mergeBlocks(remainingBlock, continueBlock);
+
+ // Drop the token using a signature conversion so rollback restores both the
+ // block and its arguments. Preserve the loop-carried block arguments.
+ Block *loopBody = &bodyRegion.front();
+ TypeConverter::SignatureConversion signature(loopBody->getNumArguments());
+ for (BlockArgument arg : loopBody->getArguments().drop_front())
+ signature.addInputs(arg.getArgNumber(), arg.getType());
+ loopBody = rewriter.applySignatureConversion(loopBody, signature);
+ if (!loopBody)
+ return failure();
+ rewriter.inlineRegionBefore(bodyRegion, continueBlock);
+
+ for (Operation *predecessor : nestedBreakingControlFlow.predecessors) {
+ // Erasure may be deferred. Keep the new CFG branch last in the block so
+ // enclosing loop conversions do not revisit an erased SCF terminator.
+ rewriter.setInsertionPointAfter(predecessor);
+ if (auto breakOp = dyn_cast<scf::BreakOp>(predecessor))
+ cf::BranchOp::create(rewriter, breakOp.getLoc(), continueBlock,
+ breakOp.getArgs());
+ else {
+ auto contOp = cast<scf::ContinueOp>(predecessor);
+ cf::BranchOp::create(rewriter, contOp.getLoc(), loopBody,
+ contOp.getArgs());
+ }
+ rewriter.modifyOpInPlace(predecessor,
+ [&] { predecessor->setOperands({}); });
+ rewriter.eraseOp(predecessor);
+ }
+
+ rewriter.setInsertionPointToEnd(initBlock);
+ cf::BranchOp::create(rewriter, loc, loopBody, adaptor.getOperands());
+ rewriter.replaceOp(loopOp, continueBlock->getArguments());
+ return success();
+}
+
void mlir::populateSCFToControlFlowConversionPatterns(
RewritePatternSet &patterns) {
- patterns.add<ForallLowering, ForLowering, IfLowering, ParallelLowering,
- WhileLowering, ExecuteRegionLowering, IndexSwitchLowering>(
+ patterns.add<IfLowering>(patterns.getContext(), /*benefit=*/2);
+ patterns.add<ForallLowering, ForLowering, ParallelLowering, WhileLowering,
+ ExecuteRegionLowering, IndexSwitchLowering, LoopOpLowering>(
patterns.getContext());
patterns.add<DoWhileLowering>(patterns.getContext(), /*benefit=*/2);
}
@@ -755,7 +854,8 @@ void SCFToControlFlowPass::runOnOperation() {
// Configure conversion to lower out SCF operations.
ConversionTarget target(getContext());
target.addIllegalOp<scf::ForallOp, scf::ForOp, scf::IfOp, scf::IndexSwitchOp,
- scf::ParallelOp, scf::WhileOp, scf::ExecuteRegionOp>();
+ scf::LoopOp, scf::ParallelOp, scf::WhileOp,
+ scf::ExecuteRegionOp>();
target.markUnknownOpDynamicallyLegal([](Operation *) { return true; });
ConversionConfig config;
config.allowPatternRollback = allowPatternRollback;
diff --git a/mlir/lib/Dialect/SCF/IR/SCF.cpp b/mlir/lib/Dialect/SCF/IR/SCF.cpp
index 8bbad457ef716c..2b8b98996fc9b2 100644
--- a/mlir/lib/Dialect/SCF/IR/SCF.cpp
+++ b/mlir/lib/Dialect/SCF/IR/SCF.cpp
@@ -23,6 +23,7 @@
#include "mlir/IR/Operation.h"
#include "mlir/IR/OperationSupport.h"
#include "mlir/IR/PatternMatch.h"
+#include "mlir/IR/RegionExitInterfaces.h"
#include "mlir/Interfaces/FunctionInterfaces.h"
#include "mlir/Interfaces/ParallelCombiningOpInterface.h"
#include "mlir/Interfaces/ValueBoundsOpInterface.h"
@@ -288,6 +289,10 @@ ValueRange ExecuteRegionOp::getSuccessorInputs(RegionSuccessor successor) {
: ValueRange();
}
+SmallVector<Operation *> YieldOp::getPotentialTargets() {
+ return {getOperation()->getParentOp()};
+}
+
//===----------------------------------------------------------------------===//
// ConditionOp
//===----------------------------------------------------------------------===//
@@ -317,6 +322,304 @@ void ConditionOp::getSuccessorRegions(
regions.push_back(RegionSuccessor(whileOp.getOperation()));
}
+//===----------------------------------------------------------------------===//
+// Control Flow Op Utilities
+//===----------------------------------------------------------------------===//
+
+static LoopOp getLoopTargetFromToken(Value target) {
+ auto targetArg = cast<BlockArgument>(target);
+ auto loopOp = cast<LoopOp>(targetArg.getOwner()->getParentOp());
+ assert(targetArg == loopOp.getControlToken() &&
+ "expected target token to be the loop control token");
+ return loopOp;
+}
+
+static LogicalResult verifyLoopTerminatorTarget(Operation *terminator,
+ Value target) {
+ auto targetArg = dyn_cast<BlockArgument>(target);
+ if (!targetArg)
+ return terminator->emitOpError()
+ << "target token must be an entry block argument of an scf.loop";
+
+ Block *targetBlock = targetArg.getOwner();
+ auto loopOp = dyn_cast_or_null<LoopOp>(targetBlock->getParentOp());
+ if (!loopOp || targetArg != loopOp.getControlToken())
+ return terminator->emitOpError()
+ << "target token must be the control token of an scf.loop";
+
+ if (!loopOp->isProperAncestor(terminator))
+ return terminator->emitOpError(
+ "target token must be defined by an enclosing 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();
+ }
+
+ return success();
+}
+
+static bool terminatorPropagatesThrough(Operation *terminator, Operation *op) {
+ if (isa<YieldOp>(terminator))
+ return false;
+ for (HasBreakingControlFlowOpInterface target :
+ findPotentialBreakTargets(terminator)) {
+ // Callback builders infer types before the containing op is attached.
+ if (!op || (target.getOperation() != op &&
+ target.getOperation()->isProperAncestor(op)))
+ return true;
+ }
+ return false;
+}
+
+static bool appendPropagatedTerminatorSuccessors(
+ RegionBranchTerminatorOpInterface terminator, Operation *op,
+ SmallVectorImpl<RegionSuccessor> ®ions) {
+ if (isa<YieldOp>(terminator.getOperation()))
+ return false;
+ size_t oldSize = regions.size();
+ for (HasBreakingControlFlowOpInterface target :
+ findPotentialBreakTargets(terminator)) {
+ Operation *targetOp = target.getOperation();
+ if (targetOp == op || !targetOp->isProperAncestor(op))
+ continue;
+ auto targetBranch = dyn_cast<RegionBranchOpInterface>(targetOp);
+ assert(targetBranch && "expected breaking target to model region branch");
+ targetBranch.getSuccessorRegions(RegionBranchPoint(terminator), regions);
+ }
+ return regions.size() != oldSize;
+}
+
+LogicalResult BreakOp::verify() {
+ return verifyLoopTerminatorTarget(getOperation(), getTargetToken());
+}
+
+SmallVector<Operation *> BreakOp::getPotentialTargets() {
+ LoopOp loopOp = getLoopTargetFromToken(getTargetToken());
+ return {loopOp.getOperation()};
+}
+
+MutableOperandRange
+BreakOp::getMutableSuccessorOperands(RegionSuccessor point) {
+ return MutableOperandRange(getOperation(), /*start=*/1,
+ /*length=*/getOperation()->getNumOperands() - 1);
+}
+
+LogicalResult ContinueOp::verify() {
+ return verifyLoopTerminatorTarget(getOperation(), getTargetToken());
+}
+
+SmallVector<Operation *> ContinueOp::getPotentialTargets() {
+ LoopOp loopOp = getLoopTargetFromToken(getTargetToken());
+ return {loopOp.getOperation()};
+}
+
+MutableOperandRange
+ContinueOp::getMutableSuccessorOperands(RegionSuccessor point) {
+ return MutableOperandRange(getOperation(), /*start=*/1,
+ /*length=*/getOperation()->getNumOperands() - 1);
+}
+
+//===----------------------------------------------------------------------===//
+// LoopOp
+//===----------------------------------------------------------------------===//
+
+LogicalResult LoopOp::verify() {
+ if (getBody()->getNumArguments() == 0)
+ return emitOpError("expected a leading control token argument");
+ return success();
+}
+
+LogicalResult LoopOp::verifyRegions() {
+ if (!isa<TokenType>(getControlToken().getType()))
+ return emitOpError("first region argument must be a control token");
+ for (OpOperand &use : getControlToken().getUses()) {
+ if (!isa<BreakOp, ContinueOp>(use.getOwner()) ||
+ use.getOperandNumber() != 0)
+ return emitOpError("control token may only be used as the target of "
+ "scf.break or scf.continue");
+ }
+ return success();
+}
+
+void LoopOp::print(OpAsmPrinter &p) {
+ p << " token(" << getControlToken() << ") ";
+ bool hasIters = !getInitArgs().empty();
+ bool hasReturn = !getResultTypes().empty();
+
+ if (hasIters) {
+ p << "iter_args(";
+ llvm::interleaveComma(
+ llvm::zip(getRegionIterArgs(), getInitArgs()), p,
+ [&](auto it) { p << std::get<0>(it) << " = " << std::get<1>(it); });
+ p << ") : ";
+ p << getInitArgs().getTypes();
+ p << " ";
+ }
+ if (hasReturn) {
+ p << "-> ";
+ p << getResultTypes();
+ p << " ";
+ }
+
+ Operation *terminator = getRegion().front().getTerminator();
+ // Elide the terminator only when it is the trivial implicit terminator:
+ // a `scf.continue` that targets this loop's control token and carries no
+ // iter values. A single-operand `scf.continue` targeting an *outer* loop
+ // must be printed explicitly, otherwise the implicit terminator rebuilt on
+ // parse would silently retarget it to this loop.
+ auto continueOp = dyn_cast<ContinueOp>(terminator);
+ bool printBlockTerminators = !continueOp || !continueOp.getArgs().empty() ||
+ continueOp.getTargetToken() != getControlToken();
+ p.printRegion(getRegion(), /*printEntryBlockArgs=*/false,
+ printBlockTerminators);
+ p.printOptionalAttrDict((*this)->getAttrs());
+}
+
+ParseResult LoopOp::parse(OpAsmParser &parser, OperationState &result) {
+ SmallVector<OpAsmParser::Argument, 4> regionArgs;
+ SmallVector<OpAsmParser::Argument, 4> iterRegionArgs;
+ SmallVector<OpAsmParser::UnresolvedOperand, 4> iterOperands;
+ SmallVector<Type, 4> iterTypes;
+
+ OpAsmParser::Argument controlToken;
+ if (parser.parseKeyword("token") || parser.parseLParen() ||
+ parser.parseArgument(controlToken) || parser.parseRParen())
+ return failure();
+ controlToken.type = parser.getBuilder().getType<TokenType>();
+ regionArgs.push_back(controlToken);
+
+ if (failed(parser.parseOptionalKeyword("iter_args"))) {
+ // no iter_args, but can still have a return type
+ if (succeeded(parser.parseOptionalArrow()))
+ if (parser.parseTypeList(result.types))
+ return failure();
+ } else {
+ // iter_args are present and must have colon followed by types
+ if (parser.parseAssignmentList(iterRegionArgs, iterOperands) ||
+ parser.parseColon() || parser.parseTypeList(iterTypes))
+ return failure();
+ if (iterRegionArgs.size() != iterTypes.size())
+ return parser.emitError(parser.getCurrentLocation(),
+ "found different number of iter_args and types");
+ // check for optional result type(s)
+ if (succeeded(parser.parseOptionalArrow()))
+ if (parser.parseTypeList(result.types))
+ return failure();
+ // Set region argument types for loop body
+ for (auto [regionArg, type] : llvm::zip_equal(iterRegionArgs, iterTypes)) {
+ regionArg.type = type;
+ }
+ llvm::append_range(regionArgs, iterRegionArgs);
+ }
+
+ // Parse region and attr dict.
+ Region *body = result.addRegion();
+ if (parser.parseRegion(*body, regionArgs))
+ return failure();
+ LoopOp::ensureTerminator(*body, parser.getBuilder(),
+ parser.getEncodedSourceLoc(parser.getNameLoc()));
+ if (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 &&
+ appendPropagatedTerminatorSuccessors(terminator, getOperation(), regions))
+ return;
+
+ if (isa<ContinueOp>(terminator)) {
+ regions.push_back(RegionSuccessor(&getRegion()));
+ return;
+ }
+ assert(isa<BreakOp>(terminator) && "expected continue or break terminator");
+
+ regions.push_back(RegionSuccessor(getOperation()));
+}
+
+OperandRange LoopOp::getEntrySuccessorOperands(RegionSuccessor successor) {
+ return getInitArgs();
+}
+
+ValueRange LoopOp::getSuccessorInputs(RegionSuccessor successor) {
+ return successor.isOperation() ? ValueRange(getResults())
+ : ValueRange(getRegionIterArgs());
+}
+
+namespace {
+/// Rewriting pattern that erases loops that have a single iteration.
+struct SimplifyTrivialLoops : public OpRewritePattern<LoopOp> {
+ using OpRewritePattern<LoopOp>::OpRewritePattern;
+
+ LogicalResult matchAndRewrite(LoopOp op,
+ PatternRewriter &rewriter) const override {
+ // Terminator must be a break.
+ auto breakOp = dyn_cast<BreakOp>(op.getBody()->getTerminator());
+ if (!breakOp)
+ return rewriter.notifyMatchFailure(op, "loop terminator isn't a break");
+ SmallVector<HasBreakingControlFlowOpInterface> targets =
+ findPotentialBreakTargets(breakOp);
+ if (targets.size() != 1 ||
+ targets.front().getOperation() != op.getOperation())
+ return rewriter.notifyMatchFailure(
+ op, "loop terminator targets another loop");
+
+ // If it has nested predecessors, it can't be trivially simplified.
+ if (hasNestedPredecessors(op))
+ return rewriter.notifyMatchFailure(op, "has nested predecessors");
+
+ // Great: it is a single iteration loop, we can simplify it.
+ Block *body = op.getBody();
+ SmallVector<Value> replacements;
+ for (Value value : breakOp.getArgs()) {
+ if (auto blockArg = dyn_cast<BlockArgument>(value);
+ blockArg && blockArg.getOwner() == body) {
+ if (blockArg.getArgNumber() == 0)
+ return rewriter.notifyMatchFailure(
+ op, "loop terminator cannot yield the control token");
+ replacements.push_back(op.getInitArgs()[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.getInitArgs());
+ rewriter.replaceOp(op, replacements);
+
+ return success();
+ }
+};
+} // namespace
+
+void LoopOp::getCanonicalizationPatterns(RewritePatternSet &results,
+ MLIRContext *context) {
+ results.add<SimplifyTrivialLoops>(context);
+}
+
//===----------------------------------------------------------------------===//
// ForOp
//===----------------------------------------------------------------------===//
@@ -1972,13 +2275,26 @@ IfOp::inferReturnTypes(MLIRContext *ctx, std::optional<Location> loc,
Region *r = &adaptor.getThenRegion();
if (r->empty())
return failure();
- Block &b = r->front();
- if (b.empty())
- return failure();
- auto yieldOp = llvm::dyn_cast<YieldOp>(b.back());
- if (!yieldOp)
+ Block *b = &r->front();
+ if (b->empty())
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())) {
+ // Neither branch reaches this op's continuation. Preserve explicit types
+ // when verifying attached IR; an unattached callback builder infers none.
+ if (Operation *parent = terminator->getParentOp())
+ llvm::append_range(inferredReturnTypes, parent->getResultTypes());
+ return success();
+ }
+ }
+ TypeRange types = terminator->getOperandTypes();
llvm::append_range(inferredReturnTypes, types);
return success();
}
@@ -2103,7 +2419,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()) {
@@ -2133,6 +2451,14 @@ 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 (appendPropagatedTerminatorSuccessors(terminator, getOperation(),
+ regions))
+ return;
+ }
regions.push_back(RegionSuccessor(getOperation()));
return;
}
@@ -2217,9 +2543,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)) {
@@ -2263,10 +2594,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();
@@ -2439,18 +2772,17 @@ 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())) {
+ for (auto [trueResult, falseResult, opResult] : llvm::zip(
+ thenYield.getResults(), elseYield.getResults(), op.getResults())) {
if (trueResult == falseResult) {
if (!opResult.use_empty()) {
opResult.replaceAllUsesWith(trueResult);
@@ -2459,13 +2791,13 @@ struct ReplaceIfYieldWithConditionOrValue : public OpRewritePattern<IfOp> {
continue;
}
- BoolAttr trueYield, falseYield;
- if (!matchPattern(trueResult, m_Constant(&trueYield)) ||
- !matchPattern(falseResult, m_Constant(&falseYield)))
+ BoolAttr thenConst, elseConst;
+ if (!matchPattern(trueResult, m_Constant(&thenConst)) ||
+ !matchPattern(falseResult, m_Constant(&elseConst)))
continue;
- bool trueVal = trueYield.getValue();
- bool falseVal = falseYield.getValue();
+ bool trueVal = thenConst.getValue();
+ bool falseVal = elseConst.getValue();
if (!trueVal && falseVal) {
if (!opResult.use_empty()) {
Dialect *constDialect = trueResult.getDefiningOp()->getDialect();
@@ -2557,14 +2889,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(
@@ -2592,16 +2935,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(),
@@ -2614,18 +2957,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);
}
}
@@ -2653,7 +2995,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(),
@@ -2689,21 +3032,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> thenYield(op.thenTerminator()->getOperands());
SmallVector<Value> elseYield;
if (op.elseBlock())
- llvm::append_range(elseYield, op.elseYield().getOperands());
+ llvm::append_range(elseYield, op.elseTerminator()->getOperands());
// A list of indices for which we should upgrade the value yielded
// in the else to a select.
@@ -2719,13 +3073,14 @@ struct CombineNestedIfs : public OpRewritePattern<IfOp> {
for (const auto &tup : llvm::enumerate(thenYield)) {
if (tup.value().getDefiningOp() == nestedIf) {
auto nestedIdx = llvm::cast<OpResult>(tup.value()).getResultNumber();
- if (nestedIf.elseYield().getOperand(nestedIdx) !=
+ if (nestedIf.elseTerminator()->getOperand(nestedIdx) !=
elseYield[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);
+ thenYield[tup.index()] =
+ nestedIf.thenTerminator()->getOperand(nestedIdx);
continue;
}
@@ -2761,7 +3116,7 @@ struct CombineNestedIfs : public OpRewritePattern<IfOp> {
rewriter.mergeBlocks(nestedIf.thenBlock(), newIfBlock);
rewriter.setInsertionPointToEnd(newIf.thenBlock());
- rewriter.replaceOpWithNewOp<YieldOp>(newIf.thenYield(), thenYield);
+ rewriter.replaceOpWithNewOp<YieldOp>(newIf.thenTerminator(), thenYield);
if (!elseYield.empty()) {
rewriter.createBlock(&newIf.getElseRegion());
rewriter.setInsertionPointToEnd(newIf.elseBlock());
@@ -2772,13 +3127,69 @@ struct CombineNestedIfs : public OpRewritePattern<IfOp> {
}
};
+/// 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;
+ }
+}
+
+/// 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();
+ 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 = IfOp::create(rewriter, op.getLoc(), TypeRange{},
+ op.getCondition(), /*addThenBlock=*/false,
+ /*addElseBlock=*/false);
+ newIfOp->setAttrs(op->getAttrs());
+ rewriter.inlineRegionBefore(op.getThenRegion(), newIfOp.getThenRegion(),
+ newIfOp.getThenRegion().begin());
+ rewriter.eraseOp(op);
+
+ return success();
+ }
+};
} // namespace
void IfOp::getCanonicalizationPatterns(RewritePatternSet &results,
MLIRContext *context) {
results.add<CombineIfs, CombineNestedIfs, ConditionPropagation,
ConvertTrivialIfToSelect, RemoveEmptyElseBranch,
- ReplaceIfYieldWithConditionOrValue>(context);
+ ReplaceIfYieldWithConditionOrValue,
+ SimplifyIfWithBreakingControlFlowInBothBranches>(context);
populateRegionBranchOpInterfaceCanonicalizationPatterns(
results, IfOp::getOperationName());
populateRegionBranchOpInterfaceInliningPattern(results,
@@ -2786,14 +3197,13 @@ void IfOp::getCanonicalizationPatterns(RewritePatternSet &results,
}
Block *IfOp::thenBlock() { return &getThenRegion().back(); }
-YieldOp IfOp::thenYield() { return cast<YieldOp>(&thenBlock()->back()); }
+
Block *IfOp::elseBlock() {
Region &r = getElseRegion();
if (r.empty())
return nullptr;
return &r.back();
}
-YieldOp IfOp::elseYield() { return cast<YieldOp>(&elseBlock()->back()); }
//===----------------------------------------------------------------------===//
// ParallelOp
@@ -3249,6 +3659,10 @@ ReduceOp::getMutableSuccessorOperands(RegionSuccessor point) {
return MutableOperandRange(getOperation(), /*start=*/0, /*length=*/0);
}
+SmallVector<Operation *> ReduceOp::getPotentialTargets() {
+ return {getOperation()->getParentOp()};
+}
+
//===----------------------------------------------------------------------===//
// ReduceReturnOp
//===----------------------------------------------------------------------===//
@@ -3519,9 +3933,9 @@ struct WhileMoveIfDown : public OpRewritePattern<scf::WhileOp> {
continue;
size_t ifOpIdx = it.getIndex();
rewriter.replaceAllUsesWith(op.getAfterArguments()[idx],
- ifOp.thenYield()->getOperand(ifOpIdx));
+ ifOp.thenTerminator()->getOperand(ifOpIdx));
unsigned argIdx = idx;
- Value elseValue = ifOp.elseYield()->getOperand(ifOpIdx);
+ Value elseValue = ifOp.elseTerminator()->getOperand(ifOpIdx);
rewriter.modifyOpInPlace(conditionOp, [&] {
conditionOp.getArgsMutable()[argIdx].assign(elseValue);
});
@@ -3568,7 +3982,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 171ab3f19e09af..6f0aae4fa41df7 100644
--- a/mlir/lib/Dialect/SCF/IR/ValueBoundsOpInterfaceImpl.cpp
+++ b/mlir/lib/Dialect/SCF/IR/ValueBoundsOpInterfaceImpl.cpp
@@ -192,8 +192,22 @@ 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];
+ auto thenYield = dyn_cast<YieldOp>(ifOp.thenTerminator());
+ auto elseYield = dyn_cast<YieldOp>(ifOp.elseTerminator());
+ // A branch that exits an enclosing loop contributes no value to the if.
+ if (!thenYield || !elseYield) {
+ auto yield = thenYield ? thenYield : elseYield;
+ if (!yield)
+ return;
+ Value yielded = yield.getOperand(resultNum);
+ if (dim)
+ cstr.bound(value)[*dim] == cstr.getExpr(yielded, dim);
+ else
+ cstr.bound(value) == yielded;
+ return;
+ }
+ Value thenValue = thenYield.getOperand(resultNum);
+ Value elseValue = elseYield.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 35357b150e60f5..bd1c5fcbe3e744 100644
--- a/mlir/lib/Dialect/SCF/Transforms/BufferizableOpInterfaceImpl.cpp
+++ b/mlir/lib/Dialect/SCF/Transforms/BufferizableOpInterfaceImpl.cpp
@@ -227,10 +227,12 @@ 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);
- return {{thenOperand, BufferRelation::Equivalent, /*isDefinite=*/false},
- {elseOperand, BufferRelation::Equivalent, /*isDefinite=*/false}};
+ AliasingOpOperandList aliases;
+ for (Operation *terminator : {ifOp.thenTerminator(), ifOp.elseTerminator()})
+ if (auto yield = dyn_cast<scf::YieldOp>(terminator))
+ aliases.addAlias({&yield->getOpOperand(resultNum),
+ BufferRelation::Equivalent, /*isDefinite=*/false});
+ return aliases;
}
LogicalResult bufferize(Operation *op, RewriterBase &rewriter,
@@ -273,35 +275,30 @@ struct IfOpInterface
const BufferizationState &state,
SmallVector<Value> &invocationStack) const {
auto ifOp = cast<scf::IfOp>(op);
- auto thenYieldOp = cast<scf::YieldOp>(ifOp.thenBlock()->getTerminator());
- auto elseYieldOp = cast<scf::YieldOp>(ifOp.elseBlock()->getTerminator());
- assert(value.getDefiningOp() == op && "invalid valid");
-
- // Determine buffer types of the true/false branches.
+ assert(value.getDefiningOp() == op && "invalid value");
auto opResult = cast<OpResult>(value);
- auto thenValue = thenYieldOp.getOperand(opResult.getResultNumber());
- auto elseValue = elseYieldOp.getOperand(opResult.getResultNumber());
- BufferLikeType thenBufferType, elseBufferType;
- if (isa<BufferLikeType>(thenValue.getType())) {
- // True branch was already bufferized.
- thenBufferType = cast<BufferLikeType>(thenValue.getType());
- } else {
- auto maybeBufferType = bufferization::getBufferType(
- thenValue, options, state, invocationStack);
- if (failed(maybeBufferType))
- return failure();
- thenBufferType = *maybeBufferType;
- }
- if (isa<BufferLikeType>(elseValue.getType())) {
- // False branch was already bufferized.
- elseBufferType = cast<BufferLikeType>(elseValue.getType());
- } else {
- auto maybeBufferType = bufferization::getBufferType(
- elseValue, options, state, invocationStack);
- if (failed(maybeBufferType))
+ SmallVector<BufferLikeType> bufferTypes;
+ for (Operation *terminator :
+ {ifOp.thenTerminator(), ifOp.elseTerminator()}) {
+ auto yield = dyn_cast<scf::YieldOp>(terminator);
+ if (!yield)
+ continue;
+ Value yielded = yield.getOperand(opResult.getResultNumber());
+ if (auto type = dyn_cast<BufferLikeType>(yielded.getType())) {
+ bufferTypes.push_back(type);
+ continue;
+ }
+ auto type = bufferization::getBufferType(yielded, options, state,
+ invocationStack);
+ if (failed(type))
return failure();
- elseBufferType = *maybeBufferType;
+ bufferTypes.push_back(*type);
}
+ if (bufferTypes.empty())
+ return failure();
+ if (bufferTypes.size() == 1)
+ return bufferTypes.front();
+ auto thenBufferType = bufferTypes[0], elseBufferType = bufferTypes[1];
// Best case: Both branches have the exact same buffer type.
if (thenBufferType == elseBufferType)
diff --git a/mlir/lib/Dialect/XeGPU/Utils/XeGPUUtils.cpp b/mlir/lib/Dialect/XeGPU/Utils/XeGPUUtils.cpp
index a12588945ded36..9d8b62ce5a8ccc 100644
--- a/mlir/lib/Dialect/XeGPU/Utils/XeGPUUtils.cpp
+++ b/mlir/lib/Dialect/XeGPU/Utils/XeGPUUtils.cpp
@@ -983,11 +983,17 @@ xegpu::precomputeLoopBlockArgTypes(Operation *topLevelOp,
}
if (auto ifOp = dyn_cast<scf::IfOp>(op)) {
// Each result and its then/else yield operands share one position and
- // must convert identically; derive all from the result's layout.
- scf::YieldOp thenYield = ifOp.thenYield();
- scf::YieldOp elseYield = ifOp.elseBlock() ? ifOp.elseYield() : nullptr;
+ // must convert identically; derive all from the result's layout. A
+ // branch that exits early does not yield a value to this `scf.if`, so
+ // only record operands from branches that actually end with scf.yield.
+ scf::YieldOp thenYield = dyn_cast<scf::YieldOp>(ifOp.thenTerminator());
+ scf::YieldOp elseYield;
+ if (ifOp.elseBlock())
+ elseYield = dyn_cast<scf::YieldOp>(ifOp.elseTerminator());
for (auto [idx, res] : llvm::enumerate(ifOp.getResults())) {
- SmallVector<Value> dests{res, thenYield.getOperand(idx)};
+ SmallVector<Value> dests{res};
+ if (thenYield)
+ dests.push_back(thenYield.getOperand(idx));
if (elseYield)
dests.push_back(elseYield.getOperand(idx));
recordTypes(res, dests);
diff --git a/mlir/lib/IR/CMakeLists.txt b/mlir/lib/IR/CMakeLists.txt
index 01c80dafcd21a7..882421448cc405 100644
--- a/mlir/lib/IR/CMakeLists.txt
+++ b/mlir/lib/IR/CMakeLists.txt
@@ -33,6 +33,7 @@ add_mlir_library(MLIRIR
PatternMatch.cpp
QuantStorageTypeInterface.cpp
Region.cpp
+ RegionExitInterfaces.cpp
RegionKindInterface.cpp
Remarks.cpp
SymbolTable.cpp
@@ -69,6 +70,7 @@ add_mlir_library(MLIRIR
MLIRDialectFoldInterfaceIncGen
MLIROpAsmInterfaceIncGen
MLIRQuantStorageTypeInterfaceIncGen
+ MLIRRegionExitInterfacesIncGen
MLIRRegionKindInterfaceIncGen
MLIRSideEffectInterfacesIncGen
MLIRSymbolInterfacesIncGen
diff --git a/mlir/lib/IR/Dominance.cpp b/mlir/lib/IR/Dominance.cpp
index 79fb41f2e6b30a..ee785a30bcd938 100644
--- a/mlir/lib/IR/Dominance.cpp
+++ b/mlir/lib/IR/Dominance.cpp
@@ -13,9 +13,14 @@
#include "mlir/IR/Dominance.h"
#include "mlir/IR/Operation.h"
+#include "mlir/IR/RegionExitInterfaces.h"
#include "mlir/IR/RegionKindInterface.h"
+#include "llvm/ADT/SmallPtrSet.h"
+#include "llvm/Support/DebugLog.h"
#include "llvm/Support/GenericDomTreeConstruction.h"
+#define DEBUG_TYPE "dominance"
+
using namespace mlir;
using namespace mlir::detail;
@@ -38,6 +43,7 @@ void DominanceInfoBase<IsPostDom>::invalidate() {
for (auto entry : dominanceInfos)
delete entry.second.getPointer();
dominanceInfos.clear();
+ breakingControlFlowOpsCache.clear();
}
template <bool IsPostDom>
@@ -53,6 +59,7 @@ void DominanceInfoBase<IsPostDom>::invalidate(Region *region, bool recursive) {
region->walk([&](Region *r) { invalidate(r); });
else
invalidate(region);
+ breakingControlFlowOpsCache.clear();
}
/// Return the dom tree and "hasSSADominance" bit for the given region. The
@@ -255,6 +262,15 @@ static bool isBeforeInBlock(Block *block, Block::iterator a,
return a->isBeforeInBlock(&*b);
}
+template <bool IsPostDom>
+bool DominanceInfoBase<IsPostDom>::hasEscapingBreakingControlFlowCached(
+ Operation *op) const {
+ auto [it, inserted] = breakingControlFlowOpsCache.try_emplace(op, false);
+ if (inserted)
+ it->second = hasEscapingBreakingControlFlow(op);
+ return it->second;
+}
+
template <bool IsPostDom>
bool DominanceInfoBase<IsPostDom>::properlyDominatesImpl(
Block *aBlock, Block::iterator aIt, Block *bBlock, Block::iterator bIt,
@@ -288,6 +304,23 @@ bool DominanceInfoBase<IsPostDom>::properlyDominatesImpl(
return true;
}
+ auto hasEscapingBreakingControlFlowInRange =
+ [&](Block *block, Block::iterator begin, Block::iterator end) {
+ assert(begin == block->end() || begin->getBlock() == block);
+ assert(end == block->end() || end->getBlock() == block);
+ for (auto it = begin; it != end; ++it) {
+ Operation *op = &*it;
+ if ((op->hasTrait<OpTrait::PropagateControlFlowBreak>() ||
+ !op->isRegistered()) &&
+ hasEscapingBreakingControlFlowCached(op)) {
+ LDBG() << "Breaking control flow: "
+ << OpWithFlags(op, OpPrintingFlags().skipRegions());
+ return true;
+ }
+ }
+ return false;
+ };
+
// Ok, they are in the same region now.
if (aBlock == bBlock) {
// Dominance changes based on the region type. In a region with SSA
@@ -295,15 +328,63 @@ bool DominanceInfoBase<IsPostDom>::properlyDominatesImpl(
// regions kinds, uses and defs can come in any order inside a block.
if (!hasSSADominance(aBlock))
return true;
+
if constexpr (IsPostDom) {
- return isBeforeInBlock(aBlock, bIt, aIt);
+ if (!isBeforeInBlock(aBlock, bIt, aIt))
+ return false;
+ if (aIt != aBlock->end() && bIt != bBlock->end() &&
+ mightHaveBreakingControlFlow(aBlock->getParent()) &&
+ hasEscapingBreakingControlFlowInRange(aBlock, bIt, aIt))
+ return false;
+ return true;
} else {
return isBeforeInBlock(aBlock, aIt, bIt);
}
}
// If the blocks are different, use DomTree to resolve the query.
- return getDomTree(aRegion).properlyDominates(aBlock, bBlock);
+ bool properlyDominates =
+ getDomTree(aRegion).properlyDominates(aBlock, bBlock);
+ if (!IsPostDom || !properlyDominates ||
+ !mightHaveBreakingControlFlow(aRegion))
+ return properlyDominates;
+
+ // A nested breaking terminator can bypass the containing block's ordinary
+ // CFG terminator, so the block post-dominator tree is not enough for
+ // operation post-dominance. Look for escaping breaks on paths from bBlock to
+ // aBlock. This requires a traversal per query, only for multi-block regions
+ // whose parent permits breaking control flow.
+ if (bIt != bBlock->end() &&
+ hasEscapingBreakingControlFlowInRange(bBlock, bIt, bBlock->end()))
+ return false;
+
+ SmallPtrSet<Block *, 16> visited;
+ SmallVector<Block *> worklist;
+ auto enqueueSuccessor = [&](Block *successor) {
+ if (!successor || successor == aBlock || successor->getParent() != aRegion)
+ return;
+ if (visited.insert(successor).second)
+ worklist.push_back(successor);
+ };
+ for (Block *successor : bBlock->getSuccessors())
+ enqueueSuccessor(successor);
+
+ while (!worklist.empty()) {
+ Block *block = worklist.pop_back_val();
+ if (hasEscapingBreakingControlFlowInRange(block, block->begin(),
+ block->end()))
+ return false;
+ for (Block *successor : block->getSuccessors())
+ enqueueSuccessor(successor);
+ }
+
+ if (aIt != aBlock->end()) {
+ Block::iterator endIt = std::next(aIt);
+ if (hasEscapingBreakingControlFlowInRange(aBlock, aBlock->begin(), endIt))
+ return false;
+ }
+
+ return true;
}
/// Return true if the specified block is reachable from the entry block of
diff --git a/mlir/lib/IR/RegionExitInterfaces.cpp b/mlir/lib/IR/RegionExitInterfaces.cpp
new file mode 100644
index 00000000000000..90f35fb532bc5b
--- /dev/null
+++ b/mlir/lib/IR/RegionExitInterfaces.cpp
@@ -0,0 +1,171 @@
+//===- RegionExitInterfaces.cpp - Region exit interfaces ------------------===//
+//
+// 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/IR/RegionExitInterfaces.h"
+#include "mlir/Support/WalkResult.h"
+
+using namespace mlir;
+
+#include "mlir/IR/RegionExitInterfaces.cpp.inc"
+
+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();
+ 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;
+ }
+
+ // 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 block terminators nested within `rootOp` through propagators,
+/// receivers, and unregistered operations. 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 walkNestedBlockTerminators(
+ 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();
+
+ // A registered op without either capability is a propagation boundary.
+ if (op->isRegistered() &&
+ !op->hasTrait<OpTrait::PropagateControlFlowBreak>() &&
+ !isa<HasBreakingControlFlowOpInterface>(op))
+ continue;
+
+ // Recursively visit the nested regions.
+ for (Region &nestedRegion : op->getRegions()) {
+ if (nestedRegion.empty())
+ continue;
+ worklist.push_back({&nestedRegion, currentNestedLevel + 1});
+ }
+ }
+}
+
+NestedBreakingControlFlowInfo
+mlir::getNestedBreakingControlFlowInfo(Operation *op) {
+ NestedBreakingControlFlowInfo info;
+ walkNestedBlockTerminators(op, [&](Operation *visitedOp, int nestedLevel) {
+ bool directlyTargetsOp = false;
+ bool escapesThroughOp = false;
+ for (HasBreakingControlFlowOpInterface target :
+ findPotentialBreakTargets(visitedOp)) {
+ if (!target)
+ continue;
+ Operation *targetOp = target.getOperation();
+ if (targetOp == op)
+ directlyTargetsOp = true;
+ else if (targetOp->isProperAncestor(op))
+ escapesThroughOp = true;
+ }
+
+ if (directlyTargetsOp) {
+ info.predecessors.push_back(visitedOp);
+ if (nestedLevel > 1)
+ info.hasNestedPredecessors = true;
+ }
+ if (escapesThroughOp)
+ info.hasEscapingBreakingControlFlow = true;
+
+ return WalkResult::advance();
+ });
+ return info;
+}
+
+bool mlir::hasNestedPredecessors(Operation *op) {
+ return getNestedBreakingControlFlowInfo(op).hasNestedPredecessors;
+}
+
+bool mlir::hasEscapingBreakingControlFlow(Operation *op) {
+ return getNestedBreakingControlFlowInfo(op).hasEscapingBreakingControlFlow;
+}
+
+void mlir::detail::visitNestedBreakingControlFlowOpsImpl(
+ Operation *op,
+ function_ref<WalkResult(RegionExitTerminatorOpInterface, int nestedLevel)>
+ callback) {
+ walkNestedBlockTerminators(op, [&](Operation *visitedOp, int nestedLevel) {
+ for (HasBreakingControlFlowOpInterface target :
+ findPotentialBreakTargets(visitedOp)) {
+ if (target && (target.getOperation() == op ||
+ target.getOperation()->isProperAncestor(op)))
+ return callback(cast<RegionExitTerminatorOpInterface>(visitedOp),
+ nestedLevel);
+ }
+ return WalkResult::advance();
+ });
+}
+
+SmallVector<HasBreakingControlFlowOpInterface>
+mlir::findPotentialBreakTargets(Operation *terminator) {
+ auto breakingTerminator =
+ dyn_cast<RegionExitTerminatorOpInterface>(terminator);
+ if (!breakingTerminator)
+ return {};
+ SmallVector<HasBreakingControlFlowOpInterface> targets;
+ for (Operation *target : breakingTerminator.getPotentialTargets()) {
+ if (!target)
+ continue;
+ if (auto breakTarget = dyn_cast<HasBreakingControlFlowOpInterface>(target))
+ targets.push_back(breakTarget);
+ }
+ return targets;
+}
diff --git a/mlir/lib/IR/Verifier.cpp b/mlir/lib/IR/Verifier.cpp
index 11771e78d5f20e..4209a36b1c5609 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 3f4f5d7691df4d..e57d8910da6967 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/RegionExitInterfaces.h"
#include "mlir/Interfaces/ControlFlowInterfaces.h"
#include "llvm/ADT/EquivalenceClasses.h"
#include "llvm/ADT/PointerIntPair.h"
@@ -169,11 +170,17 @@ LogicalResult detail::verifyRegionBranchOpInterface(Operation *op) {
SmallVector<RegionSuccessor> successors;
regionInterface.getSuccessorRegions(branchPoint, successors);
for (const RegionSuccessor &successor : successors) {
+ RegionBranchOpInterface successorOwner =
+ getRegionBranchSuccessorOwner(successor);
+ if (!successorOwner)
+ return regionInterface->emitOpError()
+ << "has region successor not owned by a "
+ "RegionBranchOpInterface";
// Helper function that print the region branch point and the region
// successor.
auto emitRegionEdgeError = [&]() {
InFlightDiagnostic diag =
- regionInterface->emitOpError("along control flow edge from ");
+ successorOwner->emitOpError("along control flow edge from ");
if (branchPoint.isParent()) {
diag << "parent";
diag.attachNote(op->getLoc()) << "region branch point";
@@ -196,7 +203,7 @@ LogicalResult detail::verifyRegionBranchOpInterface(Operation *op) {
// Verify number of successor operands and successor inputs.
OperandRange succOperands =
regionInterface.getSuccessorOperands(branchPoint, successor);
- ValueRange succInputs = regionInterface.getSuccessorInputs(successor);
+ ValueRange succInputs = successorOwner.getSuccessorInputs(successor);
if (succOperands.size() != succInputs.size()) {
return emitRegionEdgeError()
<< ": region branch point has " << succOperands.size()
@@ -211,7 +218,7 @@ LogicalResult detail::verifyRegionBranchOpInterface(Operation *op) {
llvm::enumerate(llvm::zip(succOperandTypes, succInputTypes))) {
Type succOperandType = std::get<0>(typesIdx.value());
Type succInputType = std::get<1>(typesIdx.value());
- if (!regionInterface.areTypesCompatible(succOperandType, succInputType))
+ if (!successorOwner.areTypesCompatible(succOperandType, succInputType))
return emitRegionEdgeError()
<< ": successor operand type #" << typesIdx.index() << " "
<< succOperandType << " should match successor input type #"
@@ -486,6 +493,14 @@ RegionBranchOpInterface::getNonSuccessorInputs(RegionSuccessor successor) {
return results;
}
+RegionBranchOpInterface
+mlir::getRegionBranchSuccessorOwner(RegionSuccessor successor) {
+ if (Operation *successorOp = successor.getSuccessorOp())
+ return dyn_cast<RegionBranchOpInterface>(successorOp);
+ return dyn_cast<RegionBranchOpInterface>(
+ successor.getSuccessor()->getParentOp());
+}
+
static MutableArrayRef<OpOperand> operandsToOpOperands(OperandRange &operands) {
return MutableArrayRef<OpOperand>(operands.getBase(), operands.size());
}
@@ -497,11 +512,14 @@ getSuccessorOperandInputMapping(RegionBranchOpInterface branchOp,
SmallVector<RegionSuccessor> successors;
branchOp.getSuccessorRegions(src, successors);
for (RegionSuccessor dst : successors) {
+ RegionBranchOpInterface successorOwner = getRegionBranchSuccessorOwner(dst);
+ assert(successorOwner && "expected RegionBranchOpInterface owner");
OperandRange operands = branchOp.getSuccessorOperands(src, dst);
- assert(operands.size() == branchOp.getSuccessorInputs(dst).size() &&
+ ValueRange inputs = successorOwner.getSuccessorInputs(dst);
+ assert(operands.size() == inputs.size() &&
"expected the same number of operands and inputs");
- for (const auto &[operand, input] : llvm::zip_equal(
- operandsToOpOperands(operands), branchOp.getSuccessorInputs(dst)))
+ for (const auto &[operand, input] :
+ llvm::zip_equal(operandsToOpOperands(operands), inputs))
mapping[&operand].push_back(input);
}
}
@@ -548,6 +566,26 @@ RegionBranchOpInterface::getAllRegionBranchPoints() {
branchPoints.push_back(RegionBranchPoint(terminator));
}
}
+
+ Operation *op = getOperation();
+ if (!isa<HasBreakingControlFlowOpInterface>(op))
+ return branchPoints;
+
+ // Only receivers need nested branch points. Intermediate propagators do not
+ // receive values from nested exits: the addressed receiver verifies those
+ // edges and enumerates their predecessors.
+ visitNestedBreakingControlFlowOps(
+ op, [&](Operation *nestedTerminator, int nestedLevel) {
+ // Immediate region terminators are already covered above. Deeper
+ // terminators may add a propagated control-flow edge through this op.
+ if (nestedLevel <= 1)
+ return;
+ auto terminator =
+ dyn_cast<RegionBranchTerminatorOpInterface>(nestedTerminator);
+ if (!terminator)
+ return;
+ branchPoints.push_back(RegionBranchPoint(terminator));
+ });
return branchPoints;
}
@@ -652,6 +690,12 @@ static bool isDefinedBefore(Operation *regionBranchOp, Value a, Value b) {
return true;
}
+static bool isSuccessorInputOwnedBy(Operation *regionBranchOp, Value value) {
+ if (Operation *definingOp = value.getDefiningOp())
+ return definingOp == regionBranchOp;
+ return value.getParentRegion()->getParentOp() == regionBranchOp;
+}
+
/// Compute all non-successor-input values that a successor input could have
/// based on the given successor input to successor operand mapping.
///
@@ -708,8 +752,8 @@ static LogicalResult computeReachableValuesFromSuccessorInput(
if (visited.insert(operand->get()).second)
worklist.push_back(operand->get());
}
- // Note: The result does not contain any successor inputs. (Therefore,
- // `value` is also guaranteed to be excluded.)
+ // The result does not contain any keys in inputToOperands. In particular,
+ // `value` is guaranteed to be excluded.
return success();
}
@@ -750,11 +794,19 @@ struct MakeRegionBranchOpSuccessorInputsDead : public RewritePattern {
auto regionBranchOp = cast<RegionBranchOpInterface>(op);
RegionBranchInverseSuccessorMapping inputToOperands;
regionBranchOp.getSuccessorInputOperandMapping(inputToOperands);
+ // An ancestor input's predecessor set is incomplete here. Treat that input
+ // as an incoming value instead of following its partial mapping, including
+ // when it is reached transitively from one of this op's own inputs.
+ inputToOperands.remove_if([&](const auto &entry) {
+ return !isSuccessorInputOwnedBy(op, entry.first);
+ });
// Try to replace the uses of each successor input one-by-one.
bool changed = false;
const bool isIsolated = op->hasTrait<OpTrait::IsIsolatedFromAbove>();
for (Value value : inputToOperands.keys()) {
+ if (!isSuccessorInputOwnedBy(regionBranchOp, value))
+ continue;
// Nothing to do for successor inputs that are already dead.
if (value.use_empty())
continue;
@@ -911,8 +963,9 @@ struct RemoveDeadRegionBranchOpSuccessorInputs : public RewritePattern {
bool allDead = true;
for (auto memberIt = tiedSuccessorInputs.member_begin(**it);
memberIt != tiedSuccessorInputs.member_end(); ++memberIt) {
- // Iterate over all values in the set and check their liveness.
- if (!memberIt->use_empty()) {
+ // A tied set can only be removed if all inputs belong to this op.
+ // An ancestor-owned input may have other predecessors we cannot edit.
+ if (!isSuccessorInputOwnedBy(op, *memberIt) || !memberIt->use_empty()) {
allDead = false;
break;
}
@@ -1082,6 +1135,10 @@ struct RemoveDuplicateSuccessorInputUses : public RewritePattern {
// Total complexity: O(n * k * max(log k, log n)). For each input, sorting
// the signature costs O(k log k) and the std::map lookup costs O(k log n).
for (Value input : inputs) {
+ // Propagated exits expose inputs owned by an ancestor. Their complete
+ // predecessor set is only known to that owner.
+ if (!isSuccessorInputOwnedBy(op, input))
+ continue;
// Gather the predecessor value for each predecessor (region branch
// point) and sort them to form this input's signature.
Signature sig;
@@ -1171,6 +1228,9 @@ computeSingleAcyclicRegionBranchPath(RegionBranchOpInterface op) {
// through the region branch op.
return {};
}
+ // Never inline or rewrite regions/results owned by another operation.
+ if (getRegionBranchSuccessorOwner(successors.front()) != op)
+ return {};
path.push_back(successors.front());
if (successors.front().isOperation()) {
// Found path that ends after an operation.
diff --git a/mlir/lib/Transforms/Utils/CMakeLists.txt b/mlir/lib/Transforms/Utils/CMakeLists.txt
index c9f2fd75d25cfd..70b28df0d762b5 100644
--- a/mlir/lib/Transforms/Utils/CMakeLists.txt
+++ b/mlir/lib/Transforms/Utils/CMakeLists.txt
@@ -17,6 +17,7 @@ add_mlir_library(MLIRTransformUtils
DEPENDS
MLIRTransformsDialectInterfaceIncGen
+ MLIRRegionExitInterfacesIncGen
LINK_LIBS PUBLIC
MLIRAnalysisBase
diff --git a/mlir/lib/Transforms/Utils/InliningUtils.cpp b/mlir/lib/Transforms/Utils/InliningUtils.cpp
index ce4d6ad71b6d3a..a23c8bbaa88769 100644
--- a/mlir/lib/Transforms/Utils/InliningUtils.cpp
+++ b/mlir/lib/Transforms/Utils/InliningUtils.cpp
@@ -11,6 +11,7 @@
//===----------------------------------------------------------------------===//
#include "mlir/Transforms/InliningUtils.h"
+#include "mlir/IR/RegionExitInterfaces.h"
#include "mlir/IR/Builders.h"
#include "mlir/IR/BuiltinOps.h"
@@ -286,6 +287,13 @@ 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 (mightHaveBreakingControlFlow(src) &&
+ hasEscapingBreakingControlFlow(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/early-exit-backward.mlir b/mlir/test/Analysis/DataFlow/early-exit-backward.mlir
new file mode 100644
index 00000000000000..1b8bdaae6d46b5
--- /dev/null
+++ b/mlir/test/Analysis/DataFlow/early-exit-backward.mlir
@@ -0,0 +1,46 @@
+// RUN: mlir-opt %s -test-liveness-analysis -split-input-file 2>&1 | FileCheck %s --check-prefix=LIVE
+// RUN: mlir-opt %s -test-next-access -split-input-file | FileCheck %s --check-prefix=ACCESS
+
+// LIVE-LABEL: test_tag: dead:
+// LIVE: operand #0: live
+// LIVE: operand #1: not live
+// LIVE: result #0: not live
+// LIVE-LABEL: test_tag: live:
+// LIVE: operand #0: live
+// LIVE: operand #1: live
+// LIVE: result #0: live
+func.func @liveness(%cond : i1, %x : i32, %y : i32) -> i32 {
+ %0 = scf.loop token(%loop) iter_args(%a = %x) : i32 -> i32 {
+ %dead = arith.addi %a, %y {tag = "dead"} : i32
+ %live = arith.addi %a, %a {tag = "live"} : i32
+ scf.if %cond {
+ scf.break [%loop] %live : i32
+ }
+ scf.continue [%loop] %a : i32
+ }
+ return %0 : i32
+}
+
+// -----
+
+// ACCESS-LABEL: func.func @next_access
+// The first entry is for the scalar store value, whose next memory access is
+// unknown. The memref's next access is either the inner load or the loop store.
+// ACCESS: name = "before", next_access = ["unknown", ["inloop", "inif"]]
+// ACCESS: name = "inif", next_access = ["unknown"]
+// ACCESS: name = "inloop", next_access = ["unknown", "unknown"]
+// ACCESS: name = "after", next_access = ["unknown"]
+func.func @next_access(%cond : i1, %m : memref<i32>) {
+ %c = arith.constant 0 : i32
+ memref.store %c, %m[] {name = "before"} : memref<i32>
+ scf.loop token(%loop) {
+ scf.if %cond {
+ memref.load %m[] {name = "inif"} : memref<i32>
+ scf.break [%loop]
+ }
+ memref.store %c, %m[] {name = "inloop"} : memref<i32>
+ scf.continue [%loop]
+ }
+ memref.load %m[] {name = "after"} : memref<i32>
+ return
+}
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 00000000000000..5171610a091a6a
--- /dev/null
+++ b/mlir/test/Analysis/DataFlow/test-dead-code-analysis-early-exit.mlir
@@ -0,0 +1,96 @@
+// RUN: mlir-opt -test-dead-code-analysis --split-input-file 2>&1 %s | FileCheck %s
+
+// DeadCodeAnalysis follows explicit ancestor-owned region successors through
+// propagating ops and records the break predecessor at its target loop's exit.
+
+// -----
+
+// A loop whose only exit is scf.break propagated through scf.if.
+// The loop's exit point (op_preds) must show the break as a predecessor.
+
+// CHECK-LABEL: loop:
+// CHECK: region #0
+// CHECK: ^bb0 = live
+// CHECK: region_preds: (all) predecessors:
+// CHECK: %0 = scf.loop
+// CHECK: scf.continue
+// CHECK: op_preds: (all) predecessors:
+// CHECK: scf.break
+func.func @test_break_through_if(%cond: i1) -> i32 {
+ %result = scf.loop token(%loop) -> i32 {
+ scf.if %cond {
+ %c42 = arith.constant 42 : i32
+ scf.break [%loop] %c42 : i32
+ }
+ scf.continue [%loop]
+ } {tag = "loop"}
+ return %result : i32
+}
+
+// -----
+
+// For comparison: a loop with a direct scf.break.
+
+// CHECK-LABEL: loop:
+// CHECK: region #0
+// CHECK: ^bb0 = live
+// CHECK: region_preds: (all) predecessors:
+// CHECK: %0 = scf.loop
+// CHECK: op_preds: (all) predecessors:
+// CHECK: scf.break
+func.func @test_direct_break(%cond: i1) -> i32 {
+ %result = scf.loop token(%loop) -> i32 {
+ scf.if %cond {
+ }
+ %c42 = arith.constant 42 : i32
+ scf.break [%loop] %c42 : i32
+ } {tag = "loop"}
+ return %result : i32
+}
+
+// -----
+
+// A loop whose only exit is scf.break propagated through two nested scf.if
+// ops. The loop's op_preds must show the break. The intermediate if ops do not
+// appear in op_preds for the break — the break bypasses them without yielding
+// values to them, which is correct.
+
+// CHECK-LABEL: inner_if:
+// CHECK: region #0
+// CHECK: ^bb0 = live
+// CHECK: region_preds: (all) predecessors:
+// CHECK: scf.if {{.*}} {tag = "inner_if"}
+// CHECK: region #1
+// CHECK: op_preds: (all) predecessors:
+// CHECK: scf.if {{.*}} {tag = "inner_if"}
+
+// CHECK-LABEL: outer_if:
+// CHECK: region #0
+// CHECK: ^bb0 = live
+// CHECK: region_preds: (all) predecessors:
+// CHECK: scf.if {{.*}} {tag = "outer_if"}
+// CHECK: region #1
+// CHECK: op_preds: (all) predecessors:
+// CHECK: scf.if {{.*}} {tag = "outer_if"}
+
+// CHECK-LABEL: loop:
+// CHECK: region #0
+// CHECK: ^bb0 = live
+// CHECK: region_preds: (all) predecessors:
+// CHECK: %0 = scf.loop
+// CHECK: scf.continue
+// CHECK: op_preds: (all) predecessors:
+// CHECK: scf.break
+func.func @test_break_through_nested_ifs(%cond1: i1, %cond2: i1) -> i32 {
+ %result = scf.loop token(%loop) -> i32 {
+ scf.if %cond1 {
+ scf.if %cond2 {
+ %c99 = arith.constant 99 : i32
+ scf.break [%loop] %c99 : i32
+ } {tag = "inner_if"}
+ scf.continue [%loop]
+ } {tag = "outer_if"}
+ scf.continue [%loop]
+ } {tag = "loop"}
+ return %result : i32
+}
diff --git a/mlir/test/Analysis/test-dominance.mlir b/mlir/test/Analysis/test-dominance.mlir
index a926a8271200a3..5627c9994b0e82 100644
--- a/mlir/test/Analysis/test-dominance.mlir
+++ b/mlir/test/Analysis/test-dominance.mlir
@@ -680,3 +680,63 @@ func.func @func_loop_nested_region(
// CHECK: ^{{.*}}
// CHECK: }
// CHECK: }
+
+
+// -----
+
+// CHECK-LABEL: Testing : func_loop_early_exit
+func.func @func_loop_early_exit(%cond : i1, %arg0 : index) -> index {
+ %0 = scf.loop token(%outer) -> index {
+ scf.loop token(%inner) {
+ scf.if %cond {
+ scf.break [%outer] %arg0 : index
+ }
+ "test.foo"() : () -> ()
+ scf.break [%inner] {test.print_dominance = true}
+ }
+ }
+ return %0 : index
+}
+
+// CHECK: postdominates(scf.break {{.*}} {test.print_dominance = true} {{.*}} scf.break
+// CHECK-SAME: = 0
+// CHECK: postdominates(scf.break {{.*}} {test.print_dominance = true} {{.*}} scf.if
+// CHECK-SAME: = 0
+// CHECK: postdominates(scf.break {{.*}} {test.print_dominance = true} {{.*}} "test.foo"
+// CHECK-SAME: = 1
+// CHECK: postdominates(scf.break {{.*}} {test.print_dominance = true} {{.*}} scf.break
+// CHECK-SAME: = 1
+
+// -----
+
+// CHECK-LABEL: Testing : func_loop_early_exit_cross_block
+func.func @func_loop_early_exit_cross_block(%cond : i1, %arg0 : index) -> index {
+ %0 = scf.loop token(%outer) -> index {
+ scf.loop token(%inner) {
+ test.propagate_control_flow_break {
+ ^bb0:
+ cf.cond_br %cond, ^bb1, ^bb2
+ ^bb1:
+ scf.if %cond {
+ scf.break [%outer] %arg0 : index
+ }
+ cf.br ^bb2
+ ^bb2:
+ "test.foo"() {test.print_dominance = true} : () -> ()
+ "test.return"() : () -> ()
+ }
+ scf.break [%inner]
+ }
+ }
+ return %0 : index
+}
+
+// CHECK: --- PostDominanceInfo ---
+// CHECK: postdominates("test.foo"() {test.print_dominance = true} {{.*}} cf.cond_br
+// CHECK-SAME: = 0
+// CHECK: postdominates("test.foo"() {test.print_dominance = true} {{.*}} scf.if
+// CHECK-SAME: = 0
+// CHECK: postdominates("test.foo"() {test.print_dominance = true} {{.*}} cf.br
+// CHECK-SAME: = 1
+// CHECK: postdominates("test.foo"() {test.print_dominance = true} {{.*}} "test.foo"
+// CHECK-SAME: = 1
diff --git a/mlir/test/Conversion/SCFToControlFlow/convert-early-exit-to-cfg.mlir b/mlir/test/Conversion/SCFToControlFlow/convert-early-exit-to-cfg.mlir
new file mode 100644
index 00000000000000..425bb7cb87c04c
--- /dev/null
+++ b/mlir/test/Conversion/SCFToControlFlow/convert-early-exit-to-cfg.mlir
@@ -0,0 +1,196 @@
+// RUN: mlir-opt -convert-scf-to-cf -split-input-file %s | FileCheck %s
+// RUN: mlir-opt -convert-scf-to-cf="allow-pattern-rollback=false" -split-input-file %s | FileCheck %s
+
+
+func.func @loop_break(%cond : i1) {
+ // CHECK: test.op1
+ "test.op1"() : () -> ()
+ // CHECK-NEXT: cf.br [[LOOP1_ENTRY:.*]]
+ // CHECK-NEXT: [[LOOP1_ENTRY]]
+ scf.loop token(%loop) {
+ // CHECK-NEXT: test.op2
+ "test.op2"() : () -> ()
+ // CHECK-NEXT: cf.cond_br %arg0, [[IF_ENTRY:.*]], [[IF_CONTINUE:.*]]
+ // CHECK-NEXT: [[IF_ENTRY]]
+ scf.if %cond {
+ "test.op3"() : () -> ()
+ scf.break [%loop] loc("break1")
+ }
+ "test.op3"() : () -> ()
+ } loc("loop1")
+ "test.op4"() : () -> ()
+ return
+}
+
+// -----
+
+// Bug regression test: IfLowering was using thenTerminator instead
+// of elseTerminator, causing the else block's scf.yield to not be replaced with
+// a cf.branch when the then-branch has a scf.break.
+// CHECK-LABEL: func @if_break_then_yield_else
+func.func @if_break_then_yield_else(%cond : i1) {
+ // CHECK: cf.br ^[[LOOP:.*]]
+ scf.loop token(%loop) {
+ // CHECK: ^[[LOOP]]:
+ // CHECK: cf.cond_br %arg0, ^[[THEN:.*]], ^[[ELSE:.*]]
+ scf.if %cond {
+ // CHECK: ^[[THEN]]:
+ // CHECK-NEXT: cf.br ^[[EXIT:.*]]
+ scf.break [%loop]
+ } else {
+ // CHECK: ^[[ELSE]]:
+ // CHECK-NEXT: cf.br ^[[AFTER_IF:.*]]
+ }
+ // CHECK: ^[[AFTER_IF]]:
+ // CHECK-NEXT: cf.br ^[[LOOP]]
+ scf.continue [%loop]
+ }
+ // CHECK: ^[[EXIT]]:
+ // CHECK-NEXT: return
+ return
+}
+
+// -----
+
+// CHECK-LABEL: func @nested_loops_and_ifs(
+// CHECK-SAME: %[[COND1:.*]]: i1,
+// CHECK-SAME: %[[COND2:.*]]: i1
+func.func @nested_loops_and_ifs(%cond1 : i1, %cond2 : i1) {
+ // CHECK: test.op1
+ "test.op1"() : () -> ()
+ // CHECK-NEXT: cf.br ^[[OUTER_LOOP_ENTRY:.*]]
+ scf.loop token(%outer) {
+ // CHECK-NEXT: ^[[OUTER_LOOP_ENTRY]]:
+ // CHECK-NEXT: cf.cond_br %[[COND1]], ^[[IF1_THEN_BLOCK:.*]], ^[[IF1_EXIT:.*]]
+ scf.if %cond1 {
+ // CHECK-NEXT: ^[[IF1_THEN_BLOCK]]:
+ // CHECK-NEXT: test.op2
+ "test.op2"() : () -> ()
+ // CHECK-NEXT: cf.br ^[[INNER_LOOP_ENTRY:.*]]
+ scf.loop token(%inner) {
+ // CHECK-NEXT: ^[[INNER_LOOP_ENTRY]]:
+ // CHECK-NEXT: test.op3
+ "test.op3"() : () -> ()
+ // CHECK-NEXT: cf.cond_br %[[COND1]], ^[[IF2_THEN_BLOCK:.*]], ^[[IF2_EXIT:.*]]
+ scf.if %cond1 {
+ // CHECK-NEXT: ^[[IF2_THEN_BLOCK]]:
+ // CHECK-NEXT: test.op4
+ "test.op4"() : () -> ()
+ // CHECK-NEXT: cf.br ^[[INNER_LOOP_ENTRY]]
+ scf.continue [%inner] loc("continue1")
+ }
+ // CHECK-NEXT: ^[[IF2_EXIT]]:
+ // CHECK-NEXT: test.op5
+ "test.op5"() : () -> ()
+ // CHECK-NEXT: cf.cond_br %[[COND2]], ^[[IF3_THEN_BLOCK:.*]], ^[[IF3_EXIT:.*]]
+ scf.if %cond2 {
+ // CHECK-NEXT: ^[[IF3_THEN_BLOCK]]:
+ // CHECK-NEXT: test.op6
+ "test.op6"() : () -> ()
+ // CHECK-NEXT: cf.br ^[[FUNC_EXIT:.*]]
+ scf.break [%outer] loc("break2")
+ }
+ // CHECK-NEXT: ^[[IF3_EXIT]]:
+ // CHECK-NEXT: test.op7
+ "test.op7"() : () -> ()
+ // CHECK-NEXT: cf.br ^[[INNER_LOOP_ENTRY]]
+ scf.continue [%inner] loc("continue2")
+ } loc("loop3")
+ // CHECK-NEXT: ^[[AFTER_INNER_LOOP:.*]]:
+ // CHECK-NEXT: test.op8
+ "test.op8"() : () -> ()
+ // CHECK-NEXT: cf.br ^[[IF1_EXIT]]
+ } loc("if1")
+ // CHECK-NEXT: ^[[IF1_EXIT]]:
+ // CHECK-NEXT: cf.br ^[[OUTER_LOOP_ENTRY]]
+ scf.continue [%outer] loc("continue3")
+ } loc("loop2")
+ // CHECK-NEXT: ^[[FUNC_EXIT]]:
+ // CHECK-NEXT: test.op9
+ "test.op9"() : () -> ()
+ // CHECK-NEXT: return
+ return
+}
+
+// -----
+
+// CHECK-LABEL: func @loop_with_iter_args
+// CHECK-SAME: %[[INIT:.*]]: i32, %[[COND:.*]]: i1
+func.func @loop_with_iter_args(%init: i32, %cond: i1) -> i32 {
+ // CHECK: cf.br ^[[LOOP:.*]](%[[INIT]] : i32)
+ %result = scf.loop token(%loop) iter_args(%arg = %init) : i32 -> i32 {
+ // CHECK: ^[[LOOP]](%[[ARG:.*]]: i32):
+ // CHECK: cf.cond_br %[[COND]], ^[[THEN:.*]], ^[[ELSE:.*]]
+ scf.if %cond {
+ // CHECK: ^[[THEN]]:
+ // CHECK-NEXT: cf.br ^[[EXIT:.*]](%[[ARG]] : i32)
+ scf.break [%loop] %arg : i32
+ }
+ // CHECK: ^[[ELSE]]:
+ // CHECK-NEXT: cf.br ^[[LOOP]](%[[ARG]] : i32)
+ scf.continue [%loop] %arg : i32
+ }
+ // CHECK: ^[[EXIT]](%[[RES:.*]]: i32):
+ // CHECK-NEXT: return %[[RES]]
+ return %result : i32
+}
+
+// -----
+
+// CHECK-LABEL: func.func @if_results_with_break(
+// CHECK-SAME: %[[COND:.*]]: i1, %[[X:.*]]: i32, %[[Y:.*]]: i32
+func.func @if_results_with_break(%cond: i1, %x: i32, %y: i32) -> i32 {
+ %r = scf.loop token(%token) -> i32 {
+ // CHECK: cf.cond_br %[[COND]], ^[[THEN:.*]], ^[[ELSE:.*]]
+ // CHECK: ^[[THEN]]:
+ // CHECK-NEXT: cf.br ^[[EXIT:.*]](%[[X]] : i32)
+ // CHECK: ^[[ELSE]]:
+ // CHECK-NEXT: cf.br ^[[JOIN:.*]](%[[Y]] : i32)
+ // CHECK: ^[[JOIN]](%[[VALUE:.*]]: i32):
+ %v = scf.if %cond -> i32 {
+ scf.break [%token] %x : i32
+ } else {
+ scf.yield %y : i32
+ }
+ // CHECK: cf.br ^[[EXIT]](%[[VALUE]] : i32)
+ scf.break [%token] %v : i32
+ }
+ // CHECK: ^[[EXIT]](%[[RESULT:.*]]: i32):
+ // CHECK-NEXT: return %[[RESULT]] : i32
+ return %r : i32
+}
+
+// -----
+
+// Multiple result types, an inner loop in the else branch, and payload-bearing
+// exits to the outer loop must preserve all branch operands.
+// CHECK-LABEL: func.func @multi_results_else_loop(
+// CHECK-SAME: %[[COND:.*]]: i1, %[[X:.*]]: i32, %[[Y:.*]]: f32
+func.func @multi_results_else_loop(%cond: i1, %x: i32, %y: f32) -> (i32, f32) {
+ // CHECK: cf.br ^[[BODY:.*]](%[[X]], %[[Y]] : i32, f32)
+ // CHECK: ^[[BODY]](%[[A:.*]]: i32, %[[B:.*]]: f32):
+ %r:2 = scf.loop token(%outer) iter_args(%a = %x, %b = %y) : i32, f32 -> i32, f32 {
+ // CHECK: cf.cond_br %[[COND]], ^[[THEN:.*]], ^[[ELSE:.*]]
+ scf.if %cond {
+ // CHECK: ^[[THEN]]:
+ // CHECK-NEXT: cf.br ^[[BODY]](%[[A]], %[[B]] : i32, f32)
+ scf.continue [%outer] %a, %b : i32, f32
+ } else {
+ // CHECK: ^[[ELSE]]:
+ // CHECK: cf.cond_br %[[COND]], ^[[BREAK_OUTER:.*]], ^{{.*}}
+ scf.loop token(%inner) {
+ scf.if %cond {
+ // CHECK: ^[[BREAK_OUTER]]:
+ // CHECK-NEXT: cf.br ^[[EXIT:.*]](%[[A]], %[[B]] : i32, f32)
+ scf.break [%outer] %a, %b : i32, f32
+ }
+ scf.break [%inner]
+ }
+ }
+ // CHECK: cf.br ^[[EXIT]](%[[A]], %[[B]] : i32, f32)
+ scf.break [%outer] %a, %b : i32, f32
+ }
+ // CHECK: ^[[EXIT]](%[[R0:.*]]: i32, %[[R1:.*]]: f32):
+ // CHECK-NEXT: return %[[R0]], %[[R1]] : i32, f32
+ return %r#0, %r#1 : i32, f32
+}
diff --git a/mlir/test/Dialect/SCF/early-exit-bufferize.mlir b/mlir/test/Dialect/SCF/early-exit-bufferize.mlir
new file mode 100644
index 00000000000000..c46f10bec9b97c
--- /dev/null
+++ b/mlir/test/Dialect/SCF/early-exit-bufferize.mlir
@@ -0,0 +1,43 @@
+// RUN: mlir-opt %s -one-shot-bufferize="allow-unknown-ops bufferize-function-boundaries" -split-input-file | FileCheck %s
+
+// A break carries a token before its payload. Only yields contribute aliases
+// and buffer types for an if result.
+// CHECK-LABEL: func.func @break_then(
+// CHECK-SAME: %[[COND:.*]]: i1, %[[INPUT:.*]]: memref<5xf32
+// CHECK: scf.if %[[COND]] -> (memref<5xf32
+// CHECK: scf.break
+// CHECK: } else {
+// CHECK: scf.yield %[[INPUT]] : memref<5xf32
+func.func @break_then(%cond: i1, %input: tensor<5xf32>) {
+ scf.loop token(%token) {
+ %r = scf.if %cond -> tensor<5xf32> {
+ scf.break [%token]
+ } else {
+ scf.yield %input : tensor<5xf32>
+ }
+ "test.use"(%r) : (tensor<5xf32>) -> ()
+ scf.break [%token]
+ }
+ return
+}
+
+// -----
+
+// CHECK-LABEL: func.func @break_else(
+// CHECK-SAME: %[[COND:.*]]: i1, %[[INPUT:.*]]: memref<5xf32
+// CHECK: scf.if %[[COND]] -> (memref<5xf32
+// CHECK: scf.yield %[[INPUT]] : memref<5xf32
+// CHECK: } else {
+// CHECK: scf.break
+func.func @break_else(%cond: i1, %input: tensor<5xf32>) {
+ scf.loop token(%token) {
+ %r = scf.if %cond -> tensor<5xf32> {
+ scf.yield %input : tensor<5xf32>
+ } else {
+ scf.break [%token]
+ }
+ "test.use"(%r) : (tensor<5xf32>) -> ()
+ scf.break [%token]
+ }
+ return
+}
diff --git a/mlir/test/Dialect/SCF/early-exit-invalid.mlir b/mlir/test/Dialect/SCF/early-exit-invalid.mlir
new file mode 100644
index 00000000000000..3fd581f6990be8
--- /dev/null
+++ b/mlir/test/Dialect/SCF/early-exit-invalid.mlir
@@ -0,0 +1,221 @@
+// RUN: mlir-opt %s --split-input-file --verify-diagnostics
+
+func.func @loop_result_mismatch(%value : f32) {
+ // expected-error @+1 {{'scf.loop' op along control flow edge from Operation scf.break to Operation scf.loop: successor operand type #0 'f32' should match successor input type #0 'i32'}}
+ %result = scf.loop token(%loop) -> i32 {
+ scf.break [%loop] %value : f32 // expected-note {{region branch point}}
+ }
+ return
+}
+
+// -----
+
+// A PropagateControlFlowBreak op that is immediately nested in a
+// RegionBranchOpInterface can contain a region terminator that targets an
+// ancestor of that RegionBranchOpInterface. The intermediate RegionBranchOp
+// must expose this as a possible propagated edge so the addressed receiver can
+// verify the break operands against its results.
+func.func @outer_loop_result_mismatch_through_if_in_inner_loop(%cond : i1,
+ %value : f32) {
+ // expected-error @+1 {{'scf.loop' op along control flow edge from Operation scf.break to Operation scf.loop: successor operand type #0 'f32' should match successor input type #0 'i32'}}
+ %result = scf.loop token(%outer) -> i32 {
+ scf.loop token(%inner) {
+ scf.if %cond {
+ scf.break [%outer] %value : f32 // expected-note {{region branch point}}
+ }
+ scf.continue [%inner]
+ }
+ scf.continue [%outer]
+ }
+ return
+}
+
+// -----
+
+func.func @loop_result_number_mismatch(%value : f32) {
+ // expected-error @+1 {{'scf.loop' op along control flow edge from Operation scf.break to Operation scf.loop: region branch point has 1 operands, but region successor needs 2 inputs}}
+ %result:2 = scf.loop token(%loop) -> f32, f32 {
+ scf.break [%loop] %value : f32 // expected-note {{region branch point}}
+ }
+ return
+}
+
+// -----
+
+func.func @loop_continue_mismatch(%init : i32, %value : f32) {
+ // expected-error @+1 {{'scf.loop' op along control flow edge from Operation scf.continue to Region #0: successor operand type #0 'f32' should match successor input type #0 'i32'}}
+ scf.loop token(%loop) iter_args(%next = %init) : i32 {
+ scf.continue [%loop] %value : f32 // expected-note {{region branch point}}
+ }
+ return
+}
+
+
+// -----
+
+func.func @loop_iterargs_type_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_count_mismatch(%init : i32, %value : f32) {
+ // expected-error @+2 {{'scf.loop' op along control flow edge from parent to Region #0: region branch point has 1 operands, but region successor needs 2 inputs}}
+ // expected-note @+1 {{region branch point}}
+ "scf.loop"(%init) ({
+ ^body(%token : token, %next : i32, %next2 : f32):
+ scf.continue [%token] %init : i32
+ }) : (i32) -> ()
+ return
+}
+
+// -----
+
+// scf.for lacks PropagateControlFlowBreak, so it cannot be an intermediate
+// parent for a break targeting an enclosing loop token.
+func.func @break_through_for_missing_trait(%lb: index, %ub: index, %step: index, %cond: i1) {
+ scf.loop token(%loop) {
+ scf.for %i = %lb to %ub step %step {
+ scf.if %cond {
+ // expected-error @+1 {{target token crosses an op that does not have the PropagateControlFlowBreak trait}}
+ scf.break [%loop]
+ }
+ }
+ }
+ return
+}
+
+// -----
+
+// scf.while lacks PropagateControlFlowBreak, so break through it is rejected.
+// The scf.while verifier catches this first as the after region must terminate
+// with scf.yield.
+func.func @break_through_while(%cond: i1) {
+ %init = arith.constant true
+ scf.loop token(%loop) {
+ // expected-error @+1 {{'scf.while' op expects the 'after' region to terminate with 'scf.yield'}}
+ scf.while(%arg = %init) : (i1) -> i1 {
+ scf.condition(%arg) %arg : i1
+ } do {
+ ^bb0(%arg2: i1):
+ // expected-note @+1 {{terminator here}}
+ scf.break [%loop]
+ }
+ }
+ return
+}
+
+// -----
+
+// Verify that a break without a token target is rejected.
+func.func @break_without_token() {
+ scf.loop token(%loop) {
+ // expected-error @+1 {{'scf.break' op expected 1 or more operands}}
+ "scf.break"() : () -> ()
+ }
+ return
+}
+
+// -----
+
+// Verify that a loop whose first region argument is not a control token is
+// rejected. The break targets the outer loop so the inner loop does not need a
+// self-targeting terminator, isolating the control-token check.
+func.func @loop_non_token_control(%cond: i1) {
+ "scf.loop"() ({
+ ^bb0(%outer: token):
+ // expected-error @+1 {{'scf.loop' op first region argument must be a control token}}
+ "scf.loop"() ({
+ ^bb1(%bad: i32):
+ "scf.break"(%outer) : (token) -> ()
+ }) : () -> ()
+ "scf.continue"(%outer) : (token) -> ()
+ }) : () -> ()
+ return
+}
+
+// -----
+
+func.func @non_terminator_token_use() {
+ // expected-error @+1 {{control token may only be used as the target of scf.break or scf.continue}}
+ scf.loop token(%token) {
+ "test.use"(%token) : (token) -> ()
+ scf.break [%token]
+ }
+ return
+}
+
+// -----
+
+func.func @no_control_token_argument() {
+ scf.loop token(%outer) {
+ // expected-error @+1 {{expected a leading control token argument}}
+ "scf.loop"() ({
+ scf.break [%outer]
+ }) : () -> ()
+ scf.break [%outer]
+ }
+ return
+}
+
+// -----
+
+func.func @token_unreg() {
+ %t = "test.make_token"() : () -> token
+ scf.loop token(%loop) {
+ // expected-error @+1 {{target token must be an entry block argument of an scf.loop}}
+ scf.break [%t]
+ }
+ return
+}
+
+// -----
+
+func.func @token_testregion3(%c : i1) {
+ scf.loop token(%loop) {
+ "test.token.region"() ({
+ ^bb0(%t: token):
+ scf.if %c {
+ // expected-error @+1 {{target token must be the control token of an scf.loop}}
+ scf.break [%t]
+ }
+ "test.return"() : () -> ()
+ }) : () -> ()
+ }
+ return
+}
+
+// -----
+
+func.func @iterargs_parse(%x : i32, %y : i32) {
+ // expected-error @+1 {{found different number of iter_args and types}}
+ scf.loop token(%loop) iter_args(%a = %x, %b = %y) : i32 {
+ scf.continue [%loop] %a, %b : i32, i32
+ }
+ return
+}
+
+// -----
+
+func.func @implicit_iterargs(%init : i32) {
+ // expected-error @+2 {{region branch point has 0 operands, but region successor needs 1 inputs}}
+ // expected-note @+1 {{region branch point}}
+ scf.loop token(%loop) iter_args(%a = %init) : i32 {
+ }
+ return
+}
+
+// -----
+
+func.func @count_mismatch_continue(%x : i32) {
+ scf.loop token(%loop) { // expected-error {{region branch point has 1 operands, but region successor needs 0 inputs}}
+ scf.continue [%loop] %x : i32 // expected-note {{region branch point}}
+ }
+ return
+}
diff --git a/mlir/test/Dialect/SCF/early-exit.mlir b/mlir/test/Dialect/SCF/early-exit.mlir
new file mode 100644
index 00000000000000..07002d7eb07462
--- /dev/null
+++ b/mlir/test/Dialect/SCF/early-exit.mlir
@@ -0,0 +1,125 @@
+// RUN: mlir-opt --print-region-branch-op-interface %s --split-input-file | FileCheck %s
+// RUN: mlir-opt %s --mlir-print-debuginfo --mlir-print-op-generic --split-input-file | mlir-opt --print-region-branch-op-interface --split-input-file | FileCheck %s
+
+
+func.func @loop_break(%cond : i1) {
+ // CHECK: Found RegionBranchOpInterface operation: scf.loop {{.*}} {...} loc("loop1")
+ // CHECK: - Successor is region #0
+ // CHECK: - Found 2 predecessor(s)
+ // CHECK: - Predecessor is scf.break {{.*}} loc("break1")
+ // CHECK: - Predecessor is scf.continue
+ scf.loop token(%loop) {
+ scf.if %cond {
+ scf.break [%loop] loc("break1")
+ }
+ } loc("loop1")
+ return
+}
+
+// -----
+
+func.func @loop_continue(%cond1 : i1, %cond2 : i1) {
+ // CHECK: Found RegionBranchOpInterface operation: scf.loop {{.*}} {...} loc("loop2")
+ // CHECK: - Successor is region #0
+ // CHECK: - Found 2 predecessor(s)
+ // CHECK: - Predecessor is scf.break {{.*}} loc("break2")
+ // CHECK: - Predecessor is scf.continue
+ scf.loop token(%outer) {
+ // CHECK: Found RegionBranchOpInterface operation: scf.loop {{.*}} {...} loc("loop3")
+ // CHECK: - Successor is region #0
+ // CHECK: - Found 2 predecessor(s)
+ // CHECK: - Predecessor is scf.continue {{.*}} loc("continue1")
+ // CHECK: - Predecessor is scf.continue
+ scf.loop token(%inner) {
+ scf.if %cond1 {
+ scf.continue [%inner] loc("continue1")
+ }
+ scf.if %cond2 {
+ scf.break [%outer] loc("break2")
+ }
+ } loc("loop3")
+ } loc("loop2")
+ return
+}
+
+// -----
+
+// CHECK-LABEL: func @loop_with_results(
+// CHECK-SAME: %[[VALUE:.*]]: f32
+// CHECK: %[[RESULT:.*]] = scf.loop token(%[[TOKEN:.*]]) -> f32 {
+// CHECK: scf.break [%[[TOKEN]]] %[[VALUE]] : f32
+// CHECK: return %[[RESULT]] : f32
+func.func @loop_with_results(%value : f32) -> f32 {
+ %result = scf.loop token(%loop) -> f32 {
+ scf.break [%loop] %value : f32
+ }
+ return %result : f32
+}
+
+// -----
+
+// CHECK-LABEL: func @loop_continue_iterargs(
+// CHECK-SAME: %[[INIT:.*]]: i32
+// CHECK: scf.loop token(%[[TOKEN:.*]]) iter_args(%[[ARG:.*]] = %[[INIT]]) : i32 {
+// CHECK: scf.continue [%[[TOKEN]]] %[[ARG]] : i32
+func.func @loop_continue_iterargs(%init : i32) {
+ scf.loop token(%loop) iter_args(%next = %init) : i32 {
+ scf.continue [%loop] %next : i32
+ }
+ return
+}
+
+// -----
+
+// A single-operand continue that targets an *outer* loop must be printed
+// explicitly. If the printer elided it as the trivial implicit terminator, the
+// parser would rebuild it as a continue of the inner loop, silently changing
+// the target.
+// CHECK-LABEL: func @continue_outer_from_inner
+func.func @continue_outer_from_inner() {
+ // CHECK: scf.loop token(%[[OUTER:.*]]) {
+ scf.loop token(%outer) {
+ // CHECK: scf.loop token(%[[INNER:.*]]) {
+ scf.loop token(%inner) {
+ // CHECK: scf.continue [%[[OUTER]]]
+ scf.continue [%outer]
+ }
+ }
+ return
+}
+
+// -----
+
+// CHECK-LABEL: func @iter_args_and_results(
+// CHECK-SAME: %[[INIT:.*]]: i32, %[[COND:.*]]: i1
+// CHECK: %[[RESULT:.*]] = scf.loop token(%[[TOKEN:.*]]) iter_args(%[[ARG:.*]] = %[[INIT]]) : i32 -> i32 {
+// CHECK: scf.if %[[COND]] {
+// CHECK: scf.break [%[[TOKEN]]] %[[ARG]] : i32
+// CHECK: scf.continue [%[[TOKEN]]] %[[ARG]] : i32
+// CHECK: return %[[RESULT]] : i32
+func.func @iter_args_and_results(%init: i32, %cond: i1) -> i32 {
+ %r = scf.loop token(%token) iter_args(%arg = %init) : i32 -> i32 {
+ scf.if %cond {
+ scf.break [%token] %arg : i32
+ }
+ scf.continue [%token] %arg : i32
+ }
+ return %r : i32
+}
+
+// -----
+
+// Dialect-defined propagators may host SCF loop exits.
+// CHECK-LABEL: func @break_through_dialect_op(
+// CHECK: scf.loop token(%[[TOKEN:.*]]) {
+// CHECK: test.propagate_control_flow_break
+// CHECK: scf.break [%[[TOKEN]]]
+func.func @break_through_dialect_op() {
+ scf.loop token(%token) {
+ test.propagate_control_flow_break {
+ scf.break [%token]
+ }
+ scf.break [%token]
+ }
+ return
+}
diff --git a/mlir/test/Dialect/SCF/if-callback-builder.mlir b/mlir/test/Dialect/SCF/if-callback-builder.mlir
new file mode 100644
index 00000000000000..0b16fb1b3209eb
--- /dev/null
+++ b/mlir/test/Dialect/SCF/if-callback-builder.mlir
@@ -0,0 +1,51 @@
+// RUN: mlir-opt %s -test-scf-if-callback-builder -split-input-file | FileCheck %s
+
+// CHECK-LABEL: @callback_break_then
+// CHECK: scf.if %{{.*}} -> (i32) {
+// CHECK: scf.break
+// CHECK: } else {
+// CHECK: scf.yield %{{.*}} : i32
+func.func @callback_break_then(%cond: i1, %x: i32) -> i32 {
+ %r = scf.loop token(%token) -> i32 {
+ %v = scf.if %cond -> i32 {
+ scf.break [%token] %x : i32
+ } else {
+ scf.yield %x : i32
+ }
+ scf.break [%token] %v : i32
+ }
+ return %r : i32
+}
+
+// -----
+
+// CHECK-LABEL: @callback_continue_else
+// CHECK: scf.if %{{.*}} -> (i32) {
+// CHECK: scf.yield %{{.*}} : i32
+// CHECK: } else {
+// CHECK: scf.continue
+func.func @callback_continue_else(%cond: i1, %x: i32) -> i32 {
+ %r = scf.loop token(%token) iter_args(%arg = %x) : i32 -> i32 {
+ %v = scf.if %cond -> i32 {
+ scf.yield %arg : i32
+ } else {
+ scf.continue [%token] %arg : i32
+ }
+ scf.break [%token] %v : i32
+ }
+ return %r : i32
+}
+
+// -----
+
+// CHECK-LABEL: @callback_break_without_else
+// CHECK: scf.if %{{.*}} {
+// CHECK: scf.break
+func.func @callback_break_without_else(%cond: i1) {
+ scf.loop token(%token) {
+ scf.if %cond {
+ scf.break [%token]
+ }
+ }
+ return
+}
diff --git a/mlir/test/Dialect/SCF/loop_canonicalize.mlir b/mlir/test/Dialect/SCF/loop_canonicalize.mlir
new file mode 100644
index 00000000000000..bbe471a9bb24de
--- /dev/null
+++ b/mlir/test/Dialect/SCF/loop_canonicalize.mlir
@@ -0,0 +1,496 @@
+// RUN: mlir-opt %s -pass-pipeline='builtin.module(func.func(canonicalize{test-convergence}))' -split-input-file | FileCheck %s
+
+// CHECK-LABEL: func @fold_single_iteration_loop1
+func.func @fold_single_iteration_loop1(%arg0 : index) -> index {
+ // CHECK-NOT: loop
+ %0 = scf.loop token(%loop) -> index {
+ scf.break [%loop] %arg0 : index
+ }
+ return %0 : index
+}
+
+// -----
+
+// CHECK-LABEL: func @fold_single_iteration_loop_with_propagating_control_flow
+// CHECK-SAME: %[[COND:.*]]: i1, %[[VALUE:.*]]: index
+// CHECK: %[[RESULT:.*]] = scf.loop token(%[[TOKEN:.*]]) -> index {
+// CHECK-NOT: scf.loop
+// CHECK: scf.if %[[COND]] {
+// CHECK: scf.break [%[[TOKEN]]] %[[VALUE]] : index
+// CHECK: return %[[RESULT]] : index
+func.func @fold_single_iteration_loop_with_propagating_control_flow(%cond : i1, %arg0 : index) -> index {
+ %0 = scf.loop token(%outer) -> index {
+ scf.loop token(%inner) {
+ scf.if %cond {
+ scf.break [%outer] %arg0 : index
+ }
+ scf.break [%inner]
+ }
+ }
+ return %0 : index
+}
+
+// -----
+
+// CHECK-LABEL: func @loop_not_combine_ifs
+func.func @loop_not_combine_ifs(%arg0 : i1, %arg2: i64) {
+ // Verify that we don't combine ifs when terminator mismatches
+ scf.loop token(%loop) {
+ // CHECK: scf.if
+ %res = scf.if %arg0 -> i32 {
+ %v = "test.firstCodeTrue"() : () -> i32
+ scf.yield %v : i32
+ } else {
+ %v2 = "test.firstCodeFalse"() : () -> i32
+ scf.break [%loop]
+ }
+ // CHECK: scf.if
+ %res2 = scf.if %arg0 -> i32 {
+ %v = "test.secondCodeTrue"() : () -> i32
+ scf.yield %v : i32
+ } else {
+ %v2 = "test.secondCodeFalse"() : () -> i32
+ scf.continue [%loop]
+ }
+ }
+ return
+}
+
+// -----
+
+// CHECK-LABEL: func @loop_not_combine_break_ifs
+func.func @loop_not_combine_break_ifs(%arg0 : i1, %arg2: i64) {
+ // CombineIfs is yield-only, so breaking terminators remain separate.
+ scf.loop token(%loop) {
+ // CHECK: scf.if
+ %res = scf.if %arg0 -> i32 {
+ %v = "test.firstCodeTrue"() : () -> i32
+ scf.yield %v : i32
+ } else {
+ %v2 = "test.firstCodeFalse"() : () -> i32
+ scf.break [%loop]
+ }
+ // CHECK: scf.if
+ %res2 = scf.if %arg0 -> i32 {
+ %v = "test.secondCodeTrue"() : () -> i32
+ scf.yield %v : i32
+ } else {
+ %v2 = "test.secondCodeFalse"() : () -> i32
+ scf.break [%loop]
+ }
+ }
+ return
+}
+
+// -----
+
+// CHECK-LABEL: @do_not_merge_nested_if_with_breaking_control_flow1
+func.func @do_not_merge_nested_if_with_breaking_control_flow1(%arg0: i1, %arg1: i1) {
+// The outer if then terminator isn't a yield, blocking the merge.
+// CHECK: scf.loop
+// CHECK: scf.if
+// CHECK: scf.if
+ scf.loop token(%loop) {
+ scf.if %arg0 {
+ scf.if %arg1 {
+ "test.op"() : () -> ()
+ scf.yield
+ }
+ scf.break [%loop]
+ }
+ }
+ return
+}
+
+// -----
+
+// CHECK-LABEL: @do_not_merge_nested_if_with_breaking_control_flow2
+func.func @do_not_merge_nested_if_with_breaking_control_flow2(%arg0: i1, %arg1: i1) {
+// The outer if else block terminator isn't a yield, blocking the merge.
+// CHECK: scf.loop
+// CHECK: scf.if
+// CHECK: scf.if
+// CHECK: else
+// CHECK-NEXT: scf.break
+ scf.loop token(%loop) {
+ scf.if %arg0 {
+ scf.if %arg1 {
+ "test.op"() : () -> ()
+ scf.yield
+ }
+ scf.yield
+ } else {
+ scf.break [%loop]
+ }
+ }
+ return
+}
+
+// -----
+
+// CHECK-LABEL: @do_not_merge_nested_if_with_breaking_control_flow3
+func.func @do_not_merge_nested_if_with_breaking_control_flow3(%arg0: i1, %arg1: i1) {
+// The nested if then block terminator isn't a yield, blocking the merge.
+// CHECK: scf.loop
+// CHECK: scf.if
+// CHECK: scf.if
+// CHECK: test.op
+// CHECK-NEXT: scf.break
+ scf.loop token(%loop) {
+ scf.if %arg0 {
+ scf.if %arg1 {
+ "test.op"() : () -> ()
+ scf.break [%loop]
+ }
+ }
+ }
+ return
+}
+
+// -----
+
+// CHECK-LABEL: @do_not_merge_nested_if_with_breaking_control_flow4
+func.func @do_not_merge_nested_if_with_breaking_control_flow4(%arg0: i1, %arg1: i1) {
+// The nested if else block terminator isn't a yield, blocking the merge.
+// CHECK: scf.loop
+// CHECK: scf.if
+// CHECK: scf.if
+// CHECK: else
+// CHECK-NEXT: scf.break
+ scf.loop token(%loop) {
+ scf.if %arg0 {
+ scf.if %arg1 {
+ "test.op"() : () -> ()
+ } else {
+ scf.break [%loop]
+ }
+ }
+ }
+ return
+}
+
+// -----
+
+// CHECK-LABEL: func @do_not_convert_if_to_select1
+func.func @do_not_convert_if_to_select1(%cond: i1, %arg0 : index, %arg1 : index) -> index {
+ %loop_res = scf.loop token(%loop) -> index {
+ // Inner then terminator is not a yield, blocking transform to select.
+ // CHECK: scf.if
+ %0 = scf.if %cond -> index {
+ scf.break [%loop] %arg0 : index
+ } else {
+ scf.yield %arg1 : index
+ }
+ scf.break [%loop] %0 : index
+ }
+ return %loop_res : index
+}
+
+// -----
+
+// CHECK-LABEL: func @do_not_convert_if_to_select2
+func.func @do_not_convert_if_to_select2(%cond: i1, %arg0 : index, %arg1 : index) -> index {
+ %loop_res = scf.loop token(%loop) -> index {
+ // Inner then terminator is not a yield, blocking transform to select.
+ // CHECK: scf.if
+ %0 = scf.if %cond -> index {
+ scf.yield %arg0 : index
+ } else {
+ scf.break [%loop] %arg1 : index
+ }
+ scf.break [%loop] %0 : index
+ }
+ return %loop_res : index
+}
+
+
+// -----
+
+// Verify that removing the unused results of an if with nested breaking control flow
+// operation works.
+// CHECK-LABEL: func @remove_unused_if_results1
+func.func @remove_unused_if_results1(%cond : i1, %arg0 : index) -> index {
+ // CHECK: scf.loop
+ %0 = scf.loop token(%loop) -> index {
+ // CHECK: %[[FOO:.*]]:3 = "test.foo"
+ %foo:3 = "test.foo" () : () -> (i32, i64, index)
+ // CHECK-NOT: %[[RES:.*]] = scf.if
+ // CHECK: scf.if
+ %res:3 = scf.if %cond -> (i32, i64, index) {
+ // CHECK: scf.yield
+ scf.yield %foo#0, %foo#1, %foo#2 : i32, i64, index
+ } else {
+ // CHECK: scf.break {{.*}} %[[FOO]]#2 : index
+ scf.break [%loop] %foo#2 : index
+ }
+ // CHECK: "test.op"(%[[FOO]]#1)
+ "test.op"(%res#1) : (i64) -> ()
+ }
+ return %0 : index
+}
+
+// -----
+
+// Verify that removing the unused results of an if with nested breaking control flow
+// operation works.
+// CHECK-LABEL: func @simplify_if_with_breaking_controlflow_in_both_branches
+func.func @simplify_if_with_breaking_controlflow_in_both_branches(%cond : i1, %cond2 : i1, %arg0 : index) -> index {
+ // CHECK: scf.loop
+ %0 = scf.loop token(%loop) -> index {
+ // CHECK: %[[FOO:.*]] = "test.foo"
+ %foo = "test.foo" () : () -> (index)
+ // CHECK: scf.if
+ scf.if %cond {
+ // CHECK: scf.break {{.*}} %[[FOO]] : index
+ scf.break [%loop] %foo : index
+ // CHECK-NOT: else
+ } else {
+ // CHECK: %[[BAR:.*]] = "test.bar"
+ %bar = "test.bar" () : () -> (index)
+ // CHECK: scf.if
+ scf.if %cond2 {
+ // verify that this is correctly updated when inlining the parent region.
+ // CHECK: scf.break {{.*}} %[[BAR]] : index
+ scf.break [%loop] %bar : index
+ }
+ scf.continue [%loop]
+ }
+ "test.op"() : () -> ()
+ }
+ return %0 : index
+}
+
+// -----
+
+// Verify simplification when both branches have continue.
+// CHECK-LABEL: func @simplify_if_with_continue_in_both_branches
+func.func @simplify_if_with_continue_in_both_branches(%cond : i1, %init : i32) {
+ // CHECK: scf.loop
+ scf.loop token(%loop) iter_args(%arg = %init) : i32 {
+ // CHECK: scf.if
+ scf.if %cond {
+ // CHECK: scf.continue {{.*}} %{{.*}} : i32
+ scf.continue [%loop] %arg : i32
+ } else {
+ scf.continue [%loop] %arg : i32
+ }
+ // The dead code after the if should be removed.
+ // CHECK-NOT: test.op
+ "test.op"() : () -> ()
+ scf.continue [%loop] %arg : i32
+ }
+ return
+}
+
+// -----
+
+// Verify simplification when both branches have break.
+// CHECK-LABEL: func @simplify_if_with_break_in_both_branches
+func.func @simplify_if_with_break_in_both_branches(%cond : i1) -> index {
+ // CHECK: scf.loop
+ %0 = scf.loop token(%loop) -> index {
+ %c0 = arith.constant 0 : index
+ %c1 = arith.constant 1 : index
+ // CHECK: scf.if
+ scf.if %cond {
+ // CHECK: scf.break
+ scf.break [%loop] %c0 : index
+ } else {
+ scf.break [%loop] %c1 : index
+ }
+ // Dead code after the if.
+ // CHECK-NOT: test.op
+ "test.op"() : () -> ()
+ scf.continue [%loop]
+ }
+ return %0 : index
+}
+
+// -----
+
+// A propagating if sees only one of the loop's two break predecessors. It must
+// not merge loop results just because that one predecessor repeats an operand.
+// CHECK-LABEL: func.func @preserve_distinct_loop_results(
+// CHECK-SAME: %[[COND:.*]]: i1, %[[X:.*]]: i32, %[[Y:.*]]: i32
+func.func @preserve_distinct_loop_results(%cond: i1, %x: i32, %y: i32) -> (i32, i32) {
+ // CHECK: %[[RESULT:.*]]:2 = scf.loop token(%[[TOKEN:.*]]) -> i32, i32 {
+ %r:2 = scf.loop token(%t) -> i32, i32 {
+ // CHECK: scf.if %[[COND]] {
+ scf.if %cond {
+ // CHECK: scf.break [%[[TOKEN]]] %[[X]], %[[X]] : i32, i32
+ scf.break [%t] %x, %x : i32, i32
+ }
+ // CHECK: scf.break [%[[TOKEN]]] %[[X]], %[[Y]] : i32, i32
+ scf.break [%t] %x, %y : i32, i32
+ }
+ // CHECK: return %[[RESULT]]#0, %[[RESULT]]#1 : i32, i32
+ return %r#0, %r#1 : i32, i32
+}
+
+// -----
+
+// CHECK-LABEL: func.func @dead_loop_result(
+// CHECK: scf.loop token(%[[TOKEN:.*]]) -> i32
+// CHECK: scf.break [%[[TOKEN]]] %{{.*}} : i32
+func.func @dead_loop_result(%c: i1, %x: i32) {
+ %r = scf.loop token(%t) -> i32 {
+ %v = scf.if %c -> i32 {
+ scf.yield %x : i32
+ } else {
+ scf.break [%t] %x : i32
+ }
+ "test.use"(%v) : (i32) -> ()
+ scf.continue [%t]
+ }
+ return
+}
+
+// -----
+
+// CHECK-LABEL: func.func @dead_iter(
+// CHECK: scf.loop token(%[[TOKEN:.*]]) iter_args(%{{.*}} = %{{.*}}) : i32 {
+// CHECK: scf.continue [%[[TOKEN]]] %{{.*}} : i32
+func.func @dead_iter(%c: i1, %x: i32) {
+ scf.loop token(%t) iter_args(%a = %x) : i32 {
+ scf.if %c {
+ scf.continue [%t] %x : i32
+ }
+ scf.break [%t]
+ }
+ return
+}
+
+// -----
+
+// CHECK-LABEL: func.func @const_if_continue(
+// CHECK: scf.loop token(%[[TOKEN:[^ )]+]])
+// CHECK: scf.if %{{.*}} {
+// CHECK: scf.continue [%[[TOKEN]]] %{{.*}} : i32
+func.func @const_if_continue(%init: i32) -> i32 {
+ %r = scf.loop token(%t) iter_args(%a = %init) : i32 -> i32 {
+ %true = arith.constant true
+ %one = arith.constant 1 : i32
+ %b = arith.addi %a, %one : i32
+ scf.if %true {
+ scf.continue [%t] %b : i32
+ }
+ "test.side_effect"() : () -> ()
+ scf.break [%t] %a : i32
+ }
+ return %r : i32
+}
+
+// -----
+
+// CHECK-LABEL: func.func @const_if_break(
+// CHECK: scf.loop token(%[[TOKEN:[^ )]+]])
+// CHECK: scf.if %{{.*}} {
+// CHECK: scf.break [%[[TOKEN]]] %{{.*}} : i32
+func.func @const_if_break(%init: i32) -> i32 {
+ %r = scf.loop token(%t) iter_args(%a = %init) : i32 -> i32 {
+ %true = arith.constant true
+ %one = arith.constant 1 : i32
+ %b = arith.addi %a, %one : i32
+ scf.if %true {
+ scf.break [%t] %b : i32
+ }
+ "test.side_effect"() : () -> ()
+ scf.continue [%t] %a : i32
+ }
+ return %r : i32
+}
+
+// -----
+
+// CHECK-LABEL: func.func @const_if_continue0(
+// CHECK: scf.loop token(%[[TOKEN:[^ )]+]])
+// CHECK: scf.if %{{.*}} {
+// CHECK: scf.continue [%[[TOKEN]]]
+func.func @const_if_continue0() {
+ scf.loop token(%t) {
+ %true = arith.constant true
+ scf.if %true {
+ "test.a"() : () -> ()
+ scf.continue [%t]
+ }
+ "test.b"() : () -> ()
+ scf.break [%t]
+ }
+ return
+}
+
+// -----
+
+// CHECK-LABEL: func.func @const_if_break0(
+// CHECK: scf.loop token(%[[TOKEN:[^ )]+]])
+// CHECK: scf.if %{{.*}} {
+// CHECK: scf.break [%[[TOKEN]]]
+func.func @const_if_break0() {
+ scf.loop token(%t) {
+ %true = arith.constant true
+ scf.if %true {
+ "test.a"() : () -> ()
+ scf.break [%t]
+ }
+ "test.b"() : () -> ()
+ scf.continue [%t]
+ }
+ return
+}
+
+// -----
+
+// CHECK-LABEL: func.func @fold_iter_arg_to_init(
+// CHECK-SAME: %[[INIT:.*]]: i32
+// CHECK-NOT: scf.loop
+// CHECK: return %[[INIT]] : i32
+func.func @fold_iter_arg_to_init(%init: i32) -> i32 {
+ %result = scf.loop token(%token) iter_args(%arg = %init) : i32 -> i32 {
+ scf.break [%token] %arg : i32
+ }
+ return %result : i32
+}
+
+// -----
+
+// CHECK-LABEL: func.func @both_break_with_results(
+// CHECK: scf.loop
+// CHECK: scf.if
+// CHECK: scf.break
+// CHECK-NOT: test.dead
+// CHECK: return
+func.func @both_break_with_results(%cond: i1, %x: index) -> index {
+ %result = scf.loop token(%token) -> index {
+ %r = scf.if %cond -> index {
+ scf.break [%token] %x : index
+ } else {
+ scf.break [%token] %x : index
+ }
+ "test.dead"(%r) : (index) -> ()
+ scf.break [%token] %r : index
+ }
+ return %result : index
+}
+
+// -----
+
+// Following an ancestor input's partial predecessor mapping is also unsound
+// when the starting input belongs to this if. Its result is the current iter
+// arg, not the next-iteration value forwarded by the continue in the else.
+// CHECK-LABEL: func.func @transitive_successor_input_ownership(
+// CHECK-SAME: %[[COND:.*]]: i1, %[[INIT:.*]]: i32, %[[NEXT:.*]]: i32
+func.func @transitive_successor_input_ownership(%cond: i1, %init: i32, %next: i32) {
+ // CHECK: scf.loop token(%[[TOKEN:.*]]) iter_args(%[[ARG:.*]] = %[[INIT]]) : i32 {
+ scf.loop token(%token) iter_args(%arg = %init) : i32 {
+ %r = scf.if %cond -> i32 {
+ scf.yield %arg : i32
+ } else {
+ // CHECK: scf.continue [%[[TOKEN]]] %[[NEXT]] : i32
+ scf.continue [%token] %next : i32
+ }
+ // CHECK: "test.use"(%[[ARG]]) : (i32) -> ()
+ "test.use"(%r) : (i32) -> ()
+ scf.break [%token]
+ }
+ return
+}
diff --git a/mlir/test/Dialect/SCF/value-bounds-op-interface-impl.mlir b/mlir/test/Dialect/SCF/value-bounds-op-interface-impl.mlir
index ab3e0b3989cb16..64f12dc99d7bd4 100644
--- a/mlir/test/Dialect/SCF/value-bounds-op-interface-impl.mlir
+++ b/mlir/test/Dialect/SCF/value-bounds-op-interface-impl.mlir
@@ -442,3 +442,61 @@ func.func @scf_for_unsigned_induction_var() {
}
return
}
+
+// -----
+
+// CHECK-LABEL: func @if_break_then(
+// CHECK-SAME: %[[COND:.*]]: i1, %[[SIZE:.*]]: index
+// CHECK: "test.some_use"(%[[SIZE]])
+func.func @if_break_then(%cond: i1, %size: index) {
+ scf.loop token(%token) {
+ %r = scf.if %cond -> index {
+ scf.break [%token]
+ } else {
+ scf.yield %size : index
+ }
+ %bound = "test.reify_bound"(%r) {type = "LB"} : (index) -> index
+ "test.some_use"(%bound) : (index) -> ()
+ scf.break [%token]
+ }
+ return
+}
+
+// -----
+
+// CHECK-LABEL: func @if_break_else(
+// CHECK-SAME: %[[COND:.*]]: i1, %[[SIZE:.*]]: index
+// CHECK: "test.some_use"(%[[SIZE]])
+func.func @if_break_else(%cond: i1, %size: index) {
+ scf.loop token(%token) {
+ %r = scf.if %cond -> index {
+ scf.yield %size : index
+ } else {
+ scf.break [%token]
+ }
+ %bound = "test.reify_bound"(%r) {type = "LB"} : (index) -> index
+ "test.some_use"(%bound) : (index) -> ()
+ scf.break [%token]
+ }
+ return
+}
+
+// -----
+
+// CHECK-LABEL: func @if_break_shaped_result(
+// CHECK-SAME: %[[COND:.*]]: i1, %[[INPUT:.*]]: tensor<?xf32>
+// CHECK: %[[DIM:.*]] = tensor.dim %[[INPUT]], %{{.*}} : tensor<?xf32>
+// CHECK: "test.some_use"(%[[DIM]])
+func.func @if_break_shaped_result(%cond: i1, %input: tensor<?xf32>) {
+ scf.loop token(%token) {
+ %r = scf.if %cond -> tensor<?xf32> {
+ scf.break [%token]
+ } else {
+ scf.yield %input : tensor<?xf32>
+ }
+ %bound = "test.reify_bound"(%r) {type = "EQ", dim = 0} : (tensor<?xf32>) -> index
+ "test.some_use"(%bound) : (index) -> ()
+ scf.break [%token]
+ }
+ return
+}
diff --git a/mlir/test/IR/early-exit-test-dialect-invalid.mlir b/mlir/test/IR/early-exit-test-dialect-invalid.mlir
new file mode 100644
index 00000000000000..bbfbb867e7abd9
--- /dev/null
+++ b/mlir/test/IR/early-exit-test-dialect-invalid.mlir
@@ -0,0 +1,142 @@
+// RUN: mlir-opt %s --split-input-file --verify-diagnostics
+
+func.func @non_index_depth(%depth: i32) {
+ test.breakable_loop {
+ // expected-error @+1 {{operand #0 must be index}}
+ "test.dynamic_break"(%depth) : (i32) -> ()
+ }
+ return
+}
+
+// -----
+
+func.func @no_reachable_loop(%depth: index) {
+ // expected-error @+1 {{must be nested inside a reachable test.breakable_loop}}
+ test.dynamic_break %depth
+}
+
+// -----
+
+func.func @zero_depth() {
+ %c0 = arith.constant 0 : index
+ test.breakable_loop {
+ // expected-error @+1 {{depth must be positive}}
+ test.dynamic_continue %c0
+ }
+ return
+}
+
+// -----
+
+func.func @too_deep() {
+ %c2 = arith.constant 2 : index
+ test.breakable_loop {
+ // expected-error @+1 {{constant depth exceeds the number of reachable test.breakable_loop operations}}
+ test.dynamic_break %c2
+ }
+ return
+}
+
+// -----
+
+func.func @blocked_parent() {
+ %c1 = arith.constant 1 : index
+ test.breakable_loop {
+ "test.any_cond"() ({
+ // expected-error @+1 {{depth target crosses an op that does not have the PropagateControlFlowBreak trait}}
+ test.dynamic_break %c1
+ }) : () -> ()
+ test.dynamic_continue %c1
+ }
+ return
+}
+
+// -----
+
+func.func @dynamic_depth_blocked_outer(%depth: index) {
+ %c1 = arith.constant 1 : index
+ test.breakable_loop {
+ test.single_no_terminator_custom_asm_op {
+ test.breakable_loop {
+ // expected-error @+1 {{dynamic depth may target across an op that does not have the PropagateControlFlowBreak trait}}
+ test.dynamic_break %depth
+ }
+ }
+ test.dynamic_continue %c1
+ }
+ return
+}
+
+// -----
+
+// Constant depths select one concrete loop, so payload mismatches are still
+// diagnosed against that selected target.
+func.func @constant_break_payload_mismatch(%value: f32) -> i32 {
+ %c1 = arith.constant 1 : index
+ %result = test.breakable_loop -> i32 {
+ // expected-error @+1 {{payload operand #0 has type 'f32', but target results}}
+ test.dynamic_break %c1 %value : f32
+ }
+ return %result : i32
+}
+
+// -----
+
+// Dynamic depths filter out incompatible targets. This is still invalid when
+// filtering leaves no target that can accept the terminator payload.
+func.func @dynamic_break_no_compatible_target(%depth: index, %value: f32) -> i32 {
+ %result = test.breakable_loop -> i32 {
+ // expected-error @+1 {{dynamic depth has no compatible test.breakable_loop target}}
+ test.dynamic_break %depth %value : f32
+ }
+ return %result : i32
+}
+
+// -----
+
+func.func @unsupported_loop_terminator() {
+ // expected-error @+1 {{expected test.dynamic_break or test.dynamic_continue}}
+ test.breakable_loop {
+ "test.return"() : () -> ()
+ }
+ return
+}
+
+// -----
+
+func.func @payload_count_mismatch(%x: i32) -> (i32, i32) {
+ %one = arith.constant 1 : index
+ %r:2 = test.breakable_loop -> i32, i32 {
+ // expected-error @+1 {{has 1 payload operands, but target results expects 2}}
+ test.dynamic_break %one %x : i32
+ }
+ return %r#0, %r#1 : i32, i32
+}
+
+// -----
+
+func.func @negative_depth() {
+ %negative = arith.constant -1 : index
+ test.breakable_loop {
+ // expected-error @+1 {{depth must be positive}}
+ test.dynamic_break %negative
+ }
+ return
+}
+
+// -----
+
+func.func @constant_depth_crosses_blocker() {
+ %two = arith.constant 2 : index
+ test.breakable_loop {
+ test.single_no_terminator_custom_asm_op {
+ test.breakable_loop {
+ // expected-error @+1 {{depth target crosses an op that does not have the PropagateControlFlowBreak trait}}
+ test.dynamic_break %two
+ }
+ }
+ %one = arith.constant 1 : index
+ test.dynamic_continue %one
+ }
+ return
+}
diff --git a/mlir/test/IR/early-exit-test-dialect.mlir b/mlir/test/IR/early-exit-test-dialect.mlir
new file mode 100644
index 00000000000000..3ea7f0e8154b88
--- /dev/null
+++ b/mlir/test/IR/early-exit-test-dialect.mlir
@@ -0,0 +1,124 @@
+// RUN: mlir-opt %s --split-input-file | FileCheck %s --check-prefix=PRINT
+// RUN: mlir-opt --print-region-branch-op-interface %s --split-input-file | FileCheck %s --check-prefix=BRANCH
+
+// PRINT-LABEL: func.func @immediate_continue(
+func.func @immediate_continue(%depth: index) {
+ // PRINT: test.breakable_loop
+ test.breakable_loop {
+ // PRINT: test.dynamic_continue %arg0
+ test.dynamic_continue %depth
+ }
+ return
+}
+
+// -----
+
+// PRINT-LABEL: func.func @break_with_result(
+func.func @break_with_result(%depth: index, %value: i32) -> i32 {
+ // PRINT: test.breakable_loop -> i32
+ %result = test.breakable_loop -> i32 {
+ // PRINT: test.dynamic_break %arg0 %arg1 : i32
+ test.dynamic_break %depth %value : i32
+ }
+ return %result : i32
+}
+
+// -----
+
+// PRINT-LABEL: func.func @continue_iter_args(
+func.func @continue_iter_args(%depth: index, %init: i32) {
+ // PRINT: test.breakable_loop iter_args(%{{.*}} = %arg1) : i32
+ test.breakable_loop iter_args(%iter = %init) : i32 {
+ // PRINT: test.dynamic_continue %arg0 %{{.*}} : i32
+ test.dynamic_continue %depth %iter : i32
+ }
+ return
+}
+
+// -----
+
+// PRINT-LABEL: func.func @constant_depth_selects_outer(
+func.func @constant_depth_selects_outer(%outer_init: i32, %inner_init: f32) {
+ %c1 = arith.constant 1 : index
+ %c2 = arith.constant 2 : index
+ test.breakable_loop iter_args(%outer = %outer_init) : i32 {
+ test.breakable_loop iter_args(%inner = %inner_init) : f32 {
+ // PRINT: test.dynamic_continue %c2 %{{.*}} : i32
+ test.dynamic_continue %c2 %outer : i32
+ }
+ test.dynamic_continue %c1 %outer : i32
+ }
+ return
+}
+
+// -----
+
+// PRINT-LABEL: func.func @through_scf_if(
+func.func @through_scf_if(%cond: i1, %depth: index) {
+ %c1 = arith.constant 1 : index
+ // BRANCH-LABEL: Found RegionBranchOpInterface operation: test.breakable_loop {{.*}} loc("loop")
+ // BRANCH: - Successor is region #0
+ // BRANCH: - Found 2 predecessor(s)
+ // BRANCH: - Predecessor is test.dynamic_break {{.*}} loc("if_break")
+ // BRANCH: - Predecessor is test.dynamic_continue {{.*}} loc("after_if")
+ test.breakable_loop {
+ // PRINT: scf.if
+ scf.if %cond {
+ // PRINT: test.dynamic_break %arg1
+ test.dynamic_break %depth loc("if_break")
+ } loc("if")
+ test.dynamic_continue %c1 loc("after_if")
+ } loc("loop")
+ return
+}
+
+// -----
+
+// A dynamic depth can target either the immediately enclosing loop or an outer
+// compatible loop. The inner loop sees the dynamic break as a direct
+// predecessor, and the outer loop sees both the propagated dynamic break and
+// its own explicit continue.
+func.func @nested_dynamic(%depth: index) {
+ %c1 = arith.constant 1 : index
+ // BRANCH-LABEL: Found RegionBranchOpInterface operation: test.breakable_loop {{.*}} loc("outer")
+ // BRANCH: - Successor is region #0
+ // BRANCH: - Found 2 predecessor(s)
+ // BRANCH: - Predecessor is test.dynamic_break {{.*}} loc("dyn_break")
+ // BRANCH: - Predecessor is test.dynamic_continue {{.*}} loc("outer_continue")
+ test.breakable_loop {
+ // BRANCH-LABEL: Found RegionBranchOpInterface operation: test.breakable_loop {{.*}} loc("inner")
+ // BRANCH: - Successor is region #0
+ // BRANCH: - Found 1 predecessor(s)
+ // BRANCH: - Predecessor is test.dynamic_break {{.*}} loc("dyn_break")
+ test.breakable_loop {
+ test.dynamic_break %depth loc("dyn_break")
+ } loc("inner")
+ test.dynamic_continue %c1 loc("outer_continue")
+ } loc("outer")
+ return
+}
+
+// -----
+
+// The dynamic continue payload is f32. That makes the inner f32 loop a
+// potential target, but filters out the enclosing i32 loop. The outer loop
+// should only see its explicit continue predecessor.
+func.func @dynamic_continue_filters_incompatible_outer(%depth: index, %i: i32,
+ %f: f32) {
+ %c1 = arith.constant 1 : index
+ // BRANCH-LABEL: Found RegionBranchOpInterface operation: test.breakable_loop {{.*}} loc("filtered_outer")
+ // BRANCH: - Successor is region #0
+ // BRANCH: - Found 1 predecessor(s)
+ // BRANCH: - Predecessor is test.dynamic_continue {{.*}} loc("filtered_outer_continue")
+ test.breakable_loop iter_args(%outer = %i) : i32 {
+ // BRANCH-LABEL: Found RegionBranchOpInterface operation: test.breakable_loop {{.*}} loc("filtered_inner")
+ // BRANCH: - Successor is region #0
+ // BRANCH: - Found 1 predecessor(s)
+ // BRANCH: - Predecessor is test.dynamic_continue {{.*}} loc("filtered_inner_continue")
+ test.breakable_loop iter_args(%inner = %f) : f32 {
+ test.dynamic_continue %depth %inner : f32 loc("filtered_inner_continue")
+ } loc("filtered_inner")
+ test.dynamic_continue %c1 %outer : i32 loc("filtered_outer_continue")
+ } loc("filtered_outer")
+ return
+}
diff --git a/mlir/test/Integration/Dialect/SCF/early_exit.mlir b/mlir/test/Integration/Dialect/SCF/early_exit.mlir
new file mode 100644
index 00000000000000..5ba6dfa3ba4464
--- /dev/null
+++ b/mlir/test/Integration/Dialect/SCF/early_exit.mlir
@@ -0,0 +1,84 @@
+// 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 nested loop breaks and continues.
+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: Iteration 3, break out of inner loop
+// CHECK-NEXT: continue outer loop
+// CHECK-NEXT: Outer Loop Begin with counter: 3
+// CHECK-NEXT: Inner Loop Begin, counter: 4
+// CHECK-NEXT: continue inner loop
+// CHECK-NEXT: Inner Loop Begin, counter: 5
+// CHECK-NEXT: Last iteration, break out of outer loop
+// CHECK-NEXT: Outer loop finished with result: 4
+
+
+ %result = scf.loop token(%outer) iter_args(%counter_out = %counter_init) : i64 -> i64 {
+ // Outer loop iteration
+ vector.print str "Outer Loop Begin with counter: "
+ vector.print %counter_out : i64
+
+ %next_counter = scf.loop token(%inner) iter_args(%counter = %counter_out) : i64 -> 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, 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 %cond3 {
+ vector.print str "Iteration 3, break out of inner loop\n"
+ scf.break [%inner] %counter_update : i64
+ }
+ vector.print str "continue inner loop\n"
+ scf.continue [%inner] %counter_update : i64
+ }
+ vector.print str "continue outer loop\n"
+ scf.continue [%outer] %next_counter : 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 00000000000000..2c98becffa1a35
--- /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 to the inner loop) carries the constant 5.
+/// The early-exit path (break to the outer loop) 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 outer -5; } // early exit from both loops
+/// break 5; // normal exit from inner loop
+/// };
+/// // here a == 5 always
+/// if (a < 0) "dead"(); // dead code
+/// break a;
+/// }
+
+// CHECK-LABEL: func @early_exit_dead_code(
+// CHECK-DAG: %[[FALSE:.*]] = arith.constant false
+// CHECK-DAG: %[[C5:.*]] = arith.constant 5 : i32
+// CHECK: scf.loop
+// CHECK: scf.loop
+// CHECK: scf.if %[[FALSE]]
+// CHECK: "test.dead_op"
+// CHECK: scf.break {{.*}} %[[C5]] : i32
+func.func @early_exit_dead_code(%cond: i1) -> i32 {
+ %c0 = arith.constant 0 : i32
+ %c5 = arith.constant 5 : i32
+ %cm5 = arith.constant -5 : i32
+ %outer = scf.loop token(%outer_token) -> i32 {
+ %inner = scf.loop token(%inner_token) -> i32 {
+ scf.if %cond {
+ scf.break [%outer_token] %cm5 : i32
+ }
+ scf.break [%inner_token] %c5 : i32
+ }
+ %is_neg = arith.cmpi slt, %inner, %c0 : i32
+ scf.if %is_neg {
+ "test.dead_op"() : () -> ()
+ }
+ scf.break [%outer_token] %inner : i32
+ }
+ return %outer : i32
+}
+
+// -----
+
+/// For contrast: the same logic emulated with scf.while and a boolean flag
+/// instead of early exit. Without the break to the outer loop 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/Transforms/test-convert-breakable-loop-to-scf.mlir b/mlir/test/Transforms/test-convert-breakable-loop-to-scf.mlir
new file mode 100644
index 00000000000000..f33530e6d91d19
--- /dev/null
+++ b/mlir/test/Transforms/test-convert-breakable-loop-to-scf.mlir
@@ -0,0 +1,139 @@
+// RUN: mlir-opt %s --test-convert-breakable-loop-to-scf --split-input-file | FileCheck %s
+
+// CHECK-LABEL: func.func @immediate_continue(
+// CHECK-SAME: %[[INIT:.*]]: i32
+func.func @immediate_continue(%init: i32) -> i32 {
+ %c1 = arith.constant 1 : index
+ // CHECK: %[[RESULT:.*]] = scf.loop token(%[[TOKEN:.*]]) iter_args(%[[ITER:.*]] = %[[INIT]]) : i32 -> i32
+ %result = test.breakable_loop iter_args(%iter = %init) : i32 -> i32 {
+ // CHECK: scf.continue [%[[TOKEN]]] %[[ITER]] : i32
+ test.dynamic_continue %c1 %iter : i32
+ }
+ // CHECK: return %[[RESULT]] : i32
+ return %result : i32
+}
+
+// -----
+
+// CHECK-LABEL: func.func @constant_outer_break(
+func.func @constant_outer_break(%init: i32) -> i32 {
+ %c1 = arith.constant 1 : index
+ %c2 = arith.constant 2 : index
+ // CHECK: %[[RESULT:.*]] = scf.loop token(%[[OUTER_TOKEN:.*]]) iter_args(%[[OUTER_ARG:.*]] =
+ %result = test.breakable_loop iter_args(%outer = %init) : i32 -> i32 {
+ // CHECK: scf.loop token(%[[INNER_TOKEN:.*]]) {
+ test.breakable_loop {
+ // CHECK: scf.break [%[[OUTER_TOKEN]]] %[[OUTER_ARG]] : i32
+ test.dynamic_break %c2 %outer : i32
+ }
+ test.dynamic_break %c1 %outer : i32
+ }
+ // CHECK: return %[[RESULT]] : i32
+ return %result : i32
+}
+
+// -----
+
+// Dynamic break over two compatible loops lowers to a dispatch ladder with one
+// arm for depth 1 (inner) and one for depth 2 (outer), plus a deterministic
+// fallback for UB depth values.
+// CHECK-LABEL: func.func @dynamic_break_dispatch(
+// CHECK-SAME: %[[CHOOSE:.*]]: i1, %[[INIT:.*]]: i32
+func.func @dynamic_break_dispatch(%choose_outer: i1, %init: i32) -> i32 {
+ %c1 = arith.constant 1 : index
+ %c2 = arith.constant 2 : index
+ // CHECK: %[[DEPTH:.*]] = arith.select %[[CHOOSE]],
+ %depth = arith.select %choose_outer, %c2, %c1 : index
+ // CHECK: %[[RESULT:.*]] = scf.loop token(%[[OUTER_TOKEN:.*]]) iter_args(%[[OUTER_ARG:.*]] =
+ %result = test.breakable_loop iter_args(%outer = %init) : i32 -> i32 {
+ // CHECK: %[[INNER_RESULT:.*]] = scf.loop token(%[[INNER_TOKEN:.*]]) iter_args(%[[INNER_ARG:.*]] =
+ %inner_result = test.breakable_loop iter_args(%inner = %outer) : i32 -> i32 {
+ // CHECK: %[[DEPTH_ONE:.*]] = arith.constant 1 : index
+ // CHECK: %[[IS_ONE:.*]] = arith.cmpi eq, %[[DEPTH]], %[[DEPTH_ONE]] : index
+ // CHECK: scf.if %[[IS_ONE]] {
+ // CHECK: scf.break [%[[INNER_TOKEN]]] %[[INNER_ARG]] : i32
+ // CHECK: %[[DEPTH_TWO:.*]] = arith.constant 2 : index
+ // CHECK: %[[IS_TWO:.*]] = arith.cmpi eq, %[[DEPTH]], %[[DEPTH_TWO]] : index
+ // CHECK: scf.if %[[IS_TWO]] {
+ // CHECK: scf.break [%[[OUTER_TOKEN]]] %[[INNER_ARG]] : i32
+ // CHECK: scf.break [%[[INNER_TOKEN]]] %[[INNER_ARG]] : i32
+ test.dynamic_break %depth %inner : i32
+ }
+ test.dynamic_break %c1 %inner_result : i32
+ }
+ // CHECK: return %[[RESULT]] : i32
+ return %result : i32
+}
+
+// -----
+
+// The same dispatch shape is required for dynamic continue, except each arm
+// targets the next iteration of the selected loop.
+// CHECK-LABEL: func.func @dynamic_continue_dispatch(
+// CHECK-SAME: %[[CHOOSE:.*]]: i1, %[[INIT:.*]]: i32
+func.func @dynamic_continue_dispatch(%choose_outer: i1, %init: i32) -> i32 {
+ %c1 = arith.constant 1 : index
+ %c2 = arith.constant 2 : index
+ // CHECK: %[[DEPTH:.*]] = arith.select %[[CHOOSE]],
+ %depth = arith.select %choose_outer, %c2, %c1 : index
+ // CHECK: %[[RESULT:.*]] = scf.loop token(%[[OUTER_TOKEN:.*]]) iter_args(%[[OUTER_ARG:.*]] =
+ %result = test.breakable_loop iter_args(%outer = %init) : i32 -> i32 {
+ // CHECK: scf.loop token(%[[INNER_TOKEN:.*]]) iter_args(%[[INNER_ARG:.*]] =
+ test.breakable_loop iter_args(%inner = %outer) : i32 {
+ // CHECK: %[[DEPTH_ONE:.*]] = arith.constant 1 : index
+ // CHECK: %[[IS_ONE:.*]] = arith.cmpi eq, %[[DEPTH]], %[[DEPTH_ONE]] : index
+ // CHECK: scf.if %[[IS_ONE]] {
+ // CHECK: scf.continue [%[[INNER_TOKEN]]] %[[INNER_ARG]] : i32
+ // CHECK: %[[DEPTH_TWO:.*]] = arith.constant 2 : index
+ // CHECK: %[[IS_TWO:.*]] = arith.cmpi eq, %[[DEPTH]], %[[DEPTH_TWO]] : index
+ // CHECK: scf.if %[[IS_TWO]] {
+ // CHECK: scf.continue [%[[OUTER_TOKEN]]] %[[INNER_ARG]] : i32
+ // CHECK: scf.continue [%[[INNER_TOKEN]]] %[[INNER_ARG]] : i32
+ test.dynamic_continue %depth %inner : i32
+ }
+ test.dynamic_break %c1 %outer : i32
+ }
+ // CHECK: return %[[RESULT]] : i32
+ return %result : i32
+}
+
+// -----
+
+// The innermost loop carries f32, so it is not compatible with the i32 break
+// payload and must be filtered out. The remaining compatible targets are still
+// addressed by their original depths: 2 for the middle loop and 3 for the outer
+// loop.
+// CHECK-LABEL: func.func @dynamic_break_dispatch_filtered_depths(
+// CHECK-SAME: %[[CHOOSE:.*]]: i1, %[[INIT:.*]]: i32, %[[F_INIT:.*]]: f32
+func.func @dynamic_break_dispatch_filtered_depths(%choose_outer: i1,
+ %init: i32, %inner_init: f32)
+ -> i32 {
+ %c1 = arith.constant 1 : index
+ %c2 = arith.constant 2 : index
+ %c3 = arith.constant 3 : index
+ // CHECK: %[[DEPTH:.*]] = arith.select %[[CHOOSE]],
+ %depth = arith.select %choose_outer, %c3, %c2 : index
+ // CHECK: %[[RESULT:.*]] = scf.loop token(%[[OUTER_TOKEN:.*]]) iter_args(%[[OUTER_ARG:.*]] =
+ %result = test.breakable_loop iter_args(%outer = %init) : i32 -> i32 {
+ // CHECK: %[[MIDDLE_RESULT:.*]] = scf.loop token(%[[MIDDLE_TOKEN:.*]]) iter_args(%[[MIDDLE_ARG:.*]] =
+ %middle_result = test.breakable_loop iter_args(%middle = %outer) : i32 -> i32 {
+ // CHECK: scf.loop token(%[[INNER_TOKEN:.*]]) iter_args(%[[INNER_ARG:.*]] =
+ test.breakable_loop iter_args(%inner = %inner_init) : f32 {
+ // CHECK: %[[DEPTH_TWO:.*]] = arith.constant 2 : index
+ // CHECK: %[[IS_TWO:.*]] = arith.cmpi eq, %[[DEPTH]], %[[DEPTH_TWO]] : index
+ // CHECK: scf.if %[[IS_TWO]] {
+ // CHECK: scf.break [%[[MIDDLE_TOKEN]]] %[[MIDDLE_ARG]] : i32
+ // CHECK: %[[DEPTH_THREE:.*]] = arith.constant 3 : index
+ // CHECK: %[[IS_THREE:.*]] = arith.cmpi eq, %[[DEPTH]], %[[DEPTH_THREE]] : index
+ // CHECK: scf.if %[[IS_THREE]] {
+ // CHECK: scf.break [%[[OUTER_TOKEN]]] %[[MIDDLE_ARG]] : i32
+ // CHECK: scf.break [%[[MIDDLE_TOKEN]]] %[[MIDDLE_ARG]] : i32
+ test.dynamic_break %depth %middle : i32
+ }
+ test.dynamic_break %c1 %middle : i32
+ }
+ test.dynamic_break %c1 %middle_result : i32
+ }
+ // CHECK: return %[[RESULT]] : i32
+ return %result : i32
+}
diff --git a/mlir/test/lib/Dialect/SCF/TestSCFUtils.cpp b/mlir/test/lib/Dialect/SCF/TestSCFUtils.cpp
index 0778947a01c629..8b6e667933366e 100644
--- a/mlir/test/lib/Dialect/SCF/TestSCFUtils.cpp
+++ b/mlir/test/lib/Dialect/SCF/TestSCFUtils.cpp
@@ -17,6 +17,7 @@
#include "mlir/Dialect/SCF/Transforms/Patterns.h"
#include "mlir/Dialect/SCF/Utils/Utils.h"
#include "mlir/IR/Builders.h"
+#include "mlir/IR/IRMapping.h"
#include "mlir/IR/PatternMatch.h"
#include "mlir/Pass/Pass.h"
#include "mlir/Transforms/GreedyPatternRewriteDriver.h"
@@ -86,6 +87,39 @@ struct TestSCFForUtilsPass
}
};
+// Exercise result type inference while the callback-built regions are detached.
+struct TestSCFIfCallbackBuilderPass
+ : public PassWrapper<TestSCFIfCallbackBuilderPass,
+ OperationPass<func::FuncOp>> {
+ MLIR_DEFINE_EXPLICIT_INTERNAL_INLINE_TYPE_ID(TestSCFIfCallbackBuilderPass)
+
+ StringRef getArgument() const final { return "test-scf-if-callback-builder"; }
+ StringRef getDescription() const final {
+ return "Rebuild scf.if operations using the callback builder";
+ }
+
+ void runOnOperation() override {
+ getOperation().walk([&](scf::IfOp op) {
+ IRRewriter rewriter(op);
+ IRMapping mapping;
+ auto thenBuilder = [&](OpBuilder &b, Location) {
+ for (Operation &nested : *op.thenBlock())
+ b.clone(nested, mapping);
+ };
+ auto elseBuilder = [&](OpBuilder &b, Location) {
+ for (Operation &nested : *op.elseBlock())
+ b.clone(nested, mapping);
+ };
+ function_ref<void(OpBuilder &, Location)> elseCallback = nullptr;
+ if (!op.getElseRegion().empty())
+ elseCallback = elseBuilder;
+ auto replacement = scf::IfOp::create(
+ rewriter, op.getLoc(), op.getCondition(), thenBuilder, elseCallback);
+ rewriter.replaceOp(op, replacement.getResults());
+ });
+ }
+};
+
struct TestSCFIfUtilsPass
: public PassWrapper<TestSCFIfUtilsPass, OperationPass<ModuleOp>> {
MLIR_DEFINE_EXPLICIT_INTERNAL_INLINE_TYPE_ID(TestSCFIfUtilsPass)
@@ -337,6 +371,7 @@ namespace test {
void registerTestSCFUtilsPass() {
PassRegistration<TestSCFForUtilsPass>();
PassRegistration<TestSCFIfUtilsPass>();
+ PassRegistration<TestSCFIfCallbackBuilderPass>();
PassRegistration<TestSCFPipeliningPass>();
PassRegistration<TestSplitForOpAtPointPass>();
}
diff --git a/mlir/test/lib/Dialect/Test/TestDialect.h b/mlir/test/lib/Dialect/Test/TestDialect.h
index bcf3b55d33cb91..f0aa1d6316e35c 100644
--- a/mlir/test/lib/Dialect/Test/TestDialect.h
+++ b/mlir/test/lib/Dialect/Test/TestDialect.h
@@ -58,6 +58,9 @@ class RewritePatternSet;
#include "TestOpsDialect.h.inc"
namespace test {
+/// Return the payload forwarded by a dynamic break/continue terminator.
+mlir::OperandRange
+getBreakableLoopTerminatorPayload(mlir::Operation *terminator);
//===----------------------------------------------------------------------===//
// TestDialect version utilities
diff --git a/mlir/test/lib/Dialect/Test/TestOpDefs.cpp b/mlir/test/lib/Dialect/Test/TestOpDefs.cpp
index ef1fc725b53cda..699c178fb195bc 100644
--- a/mlir/test/lib/Dialect/Test/TestOpDefs.cpp
+++ b/mlir/test/lib/Dialect/Test/TestOpDefs.cpp
@@ -10,6 +10,8 @@
#include "TestOps.h"
#include "mlir/Dialect/Bufferization/IR/Bufferization.h"
#include "mlir/Dialect/Tensor/IR/Tensor.h"
+#include "mlir/IR/Matchers.h"
+#include "mlir/IR/RegionExitInterfaces.h"
#include "mlir/IR/Verifier.h"
#include "mlir/Interfaces/FunctionImplementation.h"
#include "mlir/Interfaces/MemorySlotInterfaces.h"
@@ -848,6 +850,351 @@ void AnyCondOp::getRegionInvocationBounds(
invocationBounds.emplace_back(1, 1);
}
+//===----------------------------------------------------------------------===//
+// TestBreakableLoopOp / TestDynamicBreakOp / TestDynamicContinueOp
+//===----------------------------------------------------------------------===//
+
+static std::optional<int64_t> getStaticDepth(Value depth) {
+ APInt value;
+ if (!matchPattern(depth, m_ConstantInt(&value)))
+ return std::nullopt;
+ return value.getSExtValue();
+}
+
+static SmallVector<TestBreakableLoopOp>
+getReachableBreakableLoops(Operation *terminator) {
+ SmallVector<TestBreakableLoopOp> loops;
+ for (Operation *currentOp = terminator->getParentOp(); currentOp;
+ currentOp = currentOp->getParentOp()) {
+ if (auto loopOp = dyn_cast<TestBreakableLoopOp>(currentOp))
+ loops.push_back(loopOp);
+ if (!currentOp->mightHaveTrait<OpTrait::PropagateControlFlowBreak>())
+ break;
+ }
+ return loops;
+}
+
+static Operation *getPropagationBlockingParent(Operation *terminator) {
+ for (Operation *currentOp = terminator->getParentOp(); currentOp;
+ currentOp = currentOp->getParentOp()) {
+ if (currentOp->mightHaveTrait<OpTrait::PropagateControlFlowBreak>())
+ continue;
+ for (Operation *parentOp = currentOp->getParentOp(); parentOp;
+ parentOp = parentOp->getParentOp())
+ if (isa<TestBreakableLoopOp>(parentOp))
+ return currentOp;
+ return nullptr;
+ }
+ return nullptr;
+}
+
+static SmallVector<TestBreakableLoopOp>
+getResolvedBreakableLoops(Operation *terminator, Value depth) {
+ SmallVector<TestBreakableLoopOp> loops =
+ getReachableBreakableLoops(terminator);
+ std::optional<int64_t> staticDepth = getStaticDepth(depth);
+ if (!staticDepth)
+ return loops;
+ if (*staticDepth <= 0 || static_cast<uint64_t>(*staticDepth) > loops.size())
+ return {};
+ return {loops[*staticDepth - 1]};
+}
+
+OperandRange test::getBreakableLoopTerminatorPayload(Operation *terminator) {
+ if (auto breakOp = dyn_cast<TestDynamicBreakOp>(terminator))
+ return breakOp.getArgs();
+ return cast<TestDynamicContinueOp>(terminator).getArgs();
+}
+
+static bool isContinueTerminator(Operation *terminator) {
+ if (isa<TestDynamicContinueOp>(terminator))
+ return true;
+ assert(isa<TestDynamicBreakOp>(terminator) &&
+ "expected test.dynamic_break or test.dynamic_continue");
+ return false;
+}
+
+static ValueRange getTargetExpectedValues(TestBreakableLoopOp target,
+ bool isContinue) {
+ return isContinue ? ValueRange(target.getRegionIterArgs())
+ : ValueRange(target.getResults());
+}
+
+static bool isTerminatorPayloadCompatible(OperandRange args,
+ ValueRange expectedValues) {
+ if (args.size() != expectedValues.size())
+ return false;
+ return llvm::all_of(llvm::zip(args, expectedValues), [](auto values) {
+ return std::get<0>(values).getType() == std::get<1>(values).getType();
+ });
+}
+
+static bool isPotentialTerminatorTarget(Operation *terminator,
+ OperandRange args,
+ TestBreakableLoopOp target,
+ bool isContinue) {
+ return target.acceptsTerminator(terminator) &&
+ isTerminatorPayloadCompatible(
+ args, getTargetExpectedValues(target, isContinue));
+}
+
+static SmallVector<Operation *>
+getPotentialBreakableLoopTargets(Operation *terminator, Value depth) {
+ SmallVector<Operation *> targets;
+ OperandRange args = getBreakableLoopTerminatorPayload(terminator);
+ bool isContinue = isContinueTerminator(terminator);
+ for (TestBreakableLoopOp loopOp :
+ getResolvedBreakableLoops(terminator, depth)) {
+ if (!isPotentialTerminatorTarget(terminator, args, loopOp, isContinue))
+ continue;
+ targets.push_back(loopOp.getOperation());
+ }
+ return targets;
+}
+
+static void
+appendTestBreakableLoopSuccessor(TestBreakableLoopOp loopOp, bool isContinue,
+ SmallVectorImpl<RegionSuccessor> ®ions) {
+ regions.push_back(isContinue ? RegionSuccessor(&loopOp.getBody())
+ : RegionSuccessor(loopOp.getOperation()));
+}
+
+static LogicalResult verifyTerminatorPayload(Operation *terminator,
+ OperandRange args,
+ ValueRange expectedValues,
+ TestBreakableLoopOp target,
+ StringRef targetKind) {
+ if (args.size() != expectedValues.size())
+ return terminator->emitOpError()
+ << "has " << args.size() << " payload operands, but target "
+ << targetKind << " expects " << expectedValues.size();
+
+ for (auto [index, argAndExpected] :
+ llvm::enumerate(llvm::zip(args, expectedValues))) {
+ Value arg = std::get<0>(argAndExpected);
+ Value expected = std::get<1>(argAndExpected);
+ if (arg.getType() == expected.getType())
+ continue;
+ return terminator->emitOpError()
+ << "payload operand #" << index << " has type " << arg.getType()
+ << ", but target " << targetKind << " of "
+ << OpWithFlags(target, OpPrintingFlags().skipRegions())
+ << " expects " << expected.getType();
+ }
+ return success();
+}
+
+static LogicalResult verifyDynamicTerminator(Operation *terminator, Value depth,
+ OperandRange args,
+ bool isContinue) {
+ SmallVector<TestBreakableLoopOp> reachableLoops =
+ getReachableBreakableLoops(terminator);
+ Operation *blockingOp = getPropagationBlockingParent(terminator);
+ if (reachableLoops.empty()) {
+ if (blockingOp)
+ return terminator->emitOpError()
+ << "depth target crosses an op that does not have the "
+ "PropagateControlFlowBreak trait: "
+ << OpWithFlags(blockingOp, OpPrintingFlags().skipRegions());
+ return terminator->emitOpError()
+ << "must be nested inside a reachable test.breakable_loop";
+ }
+
+ SmallVector<TestBreakableLoopOp> targets = reachableLoops;
+ std::optional<int64_t> staticDepth = getStaticDepth(depth);
+ if (staticDepth) {
+ if (*staticDepth <= 0)
+ return terminator->emitOpError() << "depth must be positive";
+ if (static_cast<uint64_t>(*staticDepth) > reachableLoops.size()) {
+ if (blockingOp)
+ return terminator->emitOpError()
+ << "depth target crosses an op that does not have the "
+ "PropagateControlFlowBreak trait: "
+ << OpWithFlags(blockingOp, OpPrintingFlags().skipRegions());
+ return terminator->emitOpError()
+ << "constant depth exceeds the number of reachable "
+ "test.breakable_loop operations";
+ }
+ targets.clear();
+ targets.push_back(reachableLoops[*staticDepth - 1]);
+ } else if (blockingOp) {
+ return terminator->emitOpError()
+ << "dynamic depth may target across an op that does not have the "
+ "PropagateControlFlowBreak trait: "
+ << OpWithFlags(blockingOp, OpPrintingFlags().skipRegions());
+ }
+
+ if (!staticDepth) {
+ for (TestBreakableLoopOp target : targets)
+ if (isPotentialTerminatorTarget(terminator, args, target, isContinue))
+ return success();
+ return terminator->emitOpError()
+ << "dynamic depth has no compatible test.breakable_loop target";
+ }
+
+ for (TestBreakableLoopOp target : targets) {
+ if (!target.acceptsTerminator(terminator))
+ return target.emitOpError("does not accept terminator: ")
+ << OpWithFlags(terminator, OpPrintingFlags().skipRegions());
+ ValueRange expectedValues = getTargetExpectedValues(target, isContinue);
+ if (failed(verifyTerminatorPayload(terminator, args, expectedValues, target,
+ isContinue ? "loop-carried values"
+ : "results")))
+ return failure();
+ }
+ return success();
+}
+
+void TestBreakableLoopOp::print(OpAsmPrinter &p) {
+ p << " ";
+ if (!getInitArgs().empty()) {
+ p << "iter_args(";
+ llvm::interleaveComma(
+ llvm::zip(getRegionIterArgs(), getInitArgs()), p,
+ [&](auto it) { p << std::get<0>(it) << " = " << std::get<1>(it); });
+ p << ") : " << getInitArgs().getTypes() << " ";
+ }
+ if (!getResultTypes().empty())
+ p << "-> " << getResultTypes() << " ";
+ p.printRegion(getBody(), /*printEntryBlockArgs=*/false,
+ /*printBlockTerminators=*/true);
+ p.printOptionalAttrDict((*this)->getAttrs());
+}
+
+ParseResult TestBreakableLoopOp::parse(OpAsmParser &parser,
+ OperationState &result) {
+ SmallVector<OpAsmParser::Argument, 4> regionArgs;
+ SmallVector<OpAsmParser::Argument, 4> iterRegionArgs;
+ SmallVector<OpAsmParser::UnresolvedOperand, 4> iterOperands;
+ SmallVector<Type, 4> iterTypes;
+
+ if (failed(parser.parseOptionalKeyword("iter_args"))) {
+ if (succeeded(parser.parseOptionalArrow()))
+ if (parser.parseTypeList(result.types))
+ return failure();
+ } else {
+ if (parser.parseAssignmentList(iterRegionArgs, iterOperands) ||
+ parser.parseColon() || parser.parseTypeList(iterTypes))
+ return failure();
+ if (iterRegionArgs.size() != iterTypes.size())
+ return parser.emitError(parser.getCurrentLocation(),
+ "found different number of iter_args and types");
+ if (succeeded(parser.parseOptionalArrow()))
+ if (parser.parseTypeList(result.types))
+ return failure();
+ for (auto [regionArg, type] : llvm::zip_equal(iterRegionArgs, iterTypes))
+ regionArg.type = type;
+ llvm::append_range(regionArgs, iterRegionArgs);
+ }
+
+ Region *body = result.addRegion();
+ if (parser.parseRegion(*body, regionArgs) ||
+ parser.parseOptionalAttrDict(result.attributes))
+ return failure();
+
+ return parser.resolveOperands(iterOperands, iterTypes, parser.getNameLoc(),
+ result.operands);
+}
+
+LogicalResult TestBreakableLoopOp::verifyRegions() {
+ if (getBody().empty())
+ return emitOpError("region cannot be empty");
+ if (getBody().front().getNumArguments() != getNumOperands())
+ return emitOpError("expected the region to have one argument per "
+ "loop-carried value (")
+ << getNumOperands() << " expected, but got "
+ << getBody().front().getNumArguments() << ")";
+ for (auto [index, argAndOperand] :
+ llvm::enumerate(llvm::zip(getRegionIterArgs(), getOperands()))) {
+ Type argType = std::get<0>(argAndOperand).getType();
+ Type operandType = std::get<1>(argAndOperand).getType();
+ if (argType != operandType)
+ return emitOpError() << "types mismatch between " << index
+ << "th iter operand (" << operandType
+ << ") and defined region argument (" << argType
+ << ")";
+ }
+ for (Block &block : getBody())
+ if (!acceptsTerminator(block.getTerminator()))
+ return emitOpError(
+ "expected test.dynamic_break or test.dynamic_continue");
+ return success();
+}
+
+void TestBreakableLoopOp::getSuccessorRegions(
+ RegionBranchPoint point, SmallVectorImpl<RegionSuccessor> ®ions) {
+ if (point.isParent()) {
+ regions.push_back(RegionSuccessor(&getBody()));
+ return;
+ }
+
+ Operation *terminator = point.getTerminatorPredecessorOrNull();
+ Value depth;
+ bool isContinue = false;
+ if (auto continueOp = dyn_cast<TestDynamicContinueOp>(terminator)) {
+ depth = continueOp.getDepth();
+ isContinue = true;
+ } else if (auto breakOp = dyn_cast<TestDynamicBreakOp>(terminator)) {
+ depth = breakOp.getDepth();
+ } else {
+ // The region verifier diagnoses unsupported terminators.
+ return;
+ }
+
+ bool isImmediateBranch = terminator->getParentOp() == getOperation();
+ for (Operation *target :
+ getPotentialBreakableLoopTargets(terminator, depth)) {
+ if (target == getOperation()) {
+ appendTestBreakableLoopSuccessor(*this, isContinue, regions);
+ continue;
+ }
+ if (!isImmediateBranch || !target->isProperAncestor(getOperation()))
+ continue;
+ appendTestBreakableLoopSuccessor(cast<TestBreakableLoopOp>(target),
+ isContinue, regions);
+ }
+}
+
+OperandRange
+TestBreakableLoopOp::getEntrySuccessorOperands(RegionSuccessor successor) {
+ return getInitArgs();
+}
+
+ValueRange TestBreakableLoopOp::getSuccessorInputs(RegionSuccessor successor) {
+ return successor.isOperation() ? ValueRange(getResults())
+ : ValueRange(getRegionIterArgs());
+}
+
+LogicalResult TestDynamicBreakOp::verify() {
+ return verifyDynamicTerminator(getOperation(), getDepth(), getArgs(),
+ /*isContinue=*/false);
+}
+
+SmallVector<Operation *> TestDynamicBreakOp::getPotentialTargets() {
+ return getPotentialBreakableLoopTargets(getOperation(), getDepth());
+}
+
+MutableOperandRange
+TestDynamicBreakOp::getMutableSuccessorOperands(RegionSuccessor point) {
+ return MutableOperandRange(getOperation(), /*start=*/1,
+ /*length=*/getOperation()->getNumOperands() - 1);
+}
+
+LogicalResult TestDynamicContinueOp::verify() {
+ return verifyDynamicTerminator(getOperation(), getDepth(), getArgs(),
+ /*isContinue=*/true);
+}
+
+SmallVector<Operation *> TestDynamicContinueOp::getPotentialTargets() {
+ return getPotentialBreakableLoopTargets(getOperation(), getDepth());
+}
+
+MutableOperandRange
+TestDynamicContinueOp::getMutableSuccessorOperands(RegionSuccessor point) {
+ return MutableOperandRange(getOperation(), /*start=*/1,
+ /*length=*/getOperation()->getNumOperands() - 1);
+}
+
//===----------------------------------------------------------------------===//
// SingleBlockImplicitTerminatorOp
//===----------------------------------------------------------------------===//
diff --git a/mlir/test/lib/Dialect/Test/TestOps.h b/mlir/test/lib/Dialect/Test/TestOps.h
index 38313017219906..c72c6152df1573 100644
--- a/mlir/test/lib/Dialect/Test/TestOps.h
+++ b/mlir/test/lib/Dialect/Test/TestOps.h
@@ -29,6 +29,7 @@
#include "mlir/IR/ExtensibleDialect.h"
#include "mlir/IR/OpDefinition.h"
#include "mlir/IR/OpImplementation.h"
+#include "mlir/IR/RegionExitInterfaces.h"
#include "mlir/IR/RegionKindInterface.h"
#include "mlir/IR/SymbolTable.h"
#include "mlir/Interfaces/CallInterfaces.h"
diff --git a/mlir/test/lib/Dialect/Test/TestOps.td b/mlir/test/lib/Dialect/Test/TestOps.td
index fc096a9e61d710..085d07004a4346 100644
--- a/mlir/test/lib/Dialect/Test/TestOps.td
+++ b/mlir/test/lib/Dialect/Test/TestOps.td
@@ -19,6 +19,7 @@ include "mlir/Interfaces/FunctionInterfaces.td"
include "mlir/IR/OpBase.td"
include "mlir/IR/OpAsmInterface.td"
include "mlir/IR/PatternBase.td"
+include "mlir/IR/RegionExitInterfaces.td"
include "mlir/IR/RegionKindInterface.td"
include "mlir/IR/SymbolInterfaces.td"
include "mlir/Interfaces/CallInterfaces.td"
@@ -2464,6 +2465,12 @@ def TestMergeBlocksOp : TEST_Op<"merge_blocks"> {
let results = (outs Variadic<AnyType>:$result);
}
+def TestPropagateControlFlowBreakOp
+ : TEST_Op<"propagate_control_flow_break", [PropagateControlFlowBreak]> {
+ let regions = (region AnyRegion:$body);
+ let assemblyFormat = "attr-dict-with-keyword $body";
+}
+
def TestRemappedValueRegionOp : TEST_Op<"remapped_value_region",
[SingleBlock]> {
let summary = "remapped_value_region operation";
@@ -2993,6 +3000,73 @@ def LoopWithExtraResultOp : TEST_Op<"loop_with_extra_result",
}];
}
+def TestBreakableLoopOp : TEST_Op<"breakable_loop", [
+ AutomaticAllocationScope,
+ RecursiveMemoryEffects,
+ PropagateControlFlowBreak,
+ SingleBlock,
+ DeclareOpInterfaceMethods<RegionBranchOpInterface,
+ ["getEntrySuccessorOperands", "getSuccessorInputs"]>,
+ HasBreakingControlFlowOpInterface,
+ HasNestedTerminators<["TestDynamicContinueOp", "TestDynamicBreakOp"]>
+ ]> {
+ let summary = "test loop with dynamically addressed break/continue";
+ let description = [{
+ Test dialect loop operation that models breaking control flow with
+ dynamically selected targets. Its terminators use an SSA index value to
+ count enclosing `test.breakable_loop` operations. A depth of one targets the
+ immediately enclosing test loop, a depth of two targets the next enclosing
+ test loop, and so on. A constant depth must identify a reachable target with
+ compatible payload types. A dynamic depth may evaluate to any reachable depth
+ with compatible payload types; evaluating to zero, to an incompatible target,
+ or to a depth greater than the number of reachable loops is undefined
+ behavior.
+ }];
+
+ let arguments = (ins Variadic<AnyType>:$initArgs);
+ let results = (outs Variadic<AnyType>:$results);
+ let regions = (region SizedRegion<1>:$body);
+
+ let extraClassDeclaration = [{
+ ::mlir::Block::BlockArgListType getRegionIterArgs() {
+ return getBody().getArguments();
+ }
+ }];
+
+ let hasCustomAssemblyFormat = 1;
+ let hasRegionVerifier = 1;
+}
+
+def TestDynamicBreakOp : TEST_Op<"dynamic_break", [
+ Terminator, RegionExitTerminatorOpInterface,
+ DeclareOpInterfaceMethods<RegionBranchTerminatorOpInterface,
+ ["getMutableSuccessorOperands"]>
+ ]> {
+ let summary = "dynamically targeted break from test.breakable_loop";
+ let arguments = (ins Index:$depth, Variadic<AnyType>:$args);
+ let assemblyFormat = "$depth ($args^ `:` type($args))? attr-dict";
+ let extraClassDeclaration = [{
+ ::llvm::SmallVector<::mlir::Operation *>
+ getPotentialTargets();
+ }];
+ let hasVerifier = 1;
+}
+
+def TestDynamicContinueOp : TEST_Op<"dynamic_continue", [
+ Terminator, RegionExitTerminatorOpInterface,
+ DeclareOpInterfaceMethods<RegionBranchTerminatorOpInterface,
+ ["getMutableSuccessorOperands"]>
+ ]> {
+ let summary = "dynamically targeted continue of test.breakable_loop";
+ let arguments = (ins Index:$depth, Variadic<AnyType>:$args);
+ let assemblyFormat = "$depth ($args^ `:` type($args))? attr-dict";
+ let extraClassDeclaration = [{
+ ::llvm::SmallVector<::mlir::Operation *>
+ getPotentialTargets();
+ }];
+ let hasVerifier = 1;
+}
+
def TestNoTerminatorOp : TEST_Op<"switch_with_no_break", [
NoTerminator,
DeclareOpInterfaceMethods<RegionBranchOpInterface>
diff --git a/mlir/test/lib/Dialect/Test/TestPatterns.cpp b/mlir/test/lib/Dialect/Test/TestPatterns.cpp
index 0417d6662814e0..51ce2151c973e6 100644
--- a/mlir/test/lib/Dialect/Test/TestPatterns.cpp
+++ b/mlir/test/lib/Dialect/Test/TestPatterns.cpp
@@ -16,17 +16,21 @@
#include "mlir/Dialect/ControlFlow/Transforms/StructuralTypeConversions.h"
#include "mlir/Dialect/Func/IR/FuncOps.h"
#include "mlir/Dialect/Func/Transforms/FuncConversions.h"
+#include "mlir/Dialect/SCF/IR/SCF.h"
#include "mlir/Dialect/SCF/Transforms/Patterns.h"
#include "mlir/Dialect/Tensor/IR/Tensor.h"
#include "mlir/IR/BuiltinAttributes.h"
+#include "mlir/IR/BuiltinTypes.h"
#include "mlir/IR/Matchers.h"
#include "mlir/IR/PatternMatch.h"
+#include "mlir/IR/RegionExitInterfaces.h"
#include "mlir/IR/Visitors.h"
#include "mlir/Pass/Pass.h"
#include "mlir/Transforms/DialectConversion.h"
#include "mlir/Transforms/FoldUtils.h"
#include "mlir/Transforms/GreedyPatternRewriteDriver.h"
#include "mlir/Transforms/WalkPatternRewriteDriver.h"
+#include "llvm/ADT/DenseMap.h"
#include "llvm/ADT/ScopeExit.h"
#include <cstdint>
@@ -2656,6 +2660,208 @@ struct TestFoldTypeConvertingOp
signalPassFailure();
}
};
+
+class TestConvertBreakableLoopToSCFPass
+ : public PassWrapper<TestConvertBreakableLoopToSCFPass, OperationPass<>> {
+public:
+ MLIR_DEFINE_EXPLICIT_INTERNAL_INLINE_TYPE_ID(
+ TestConvertBreakableLoopToSCFPass)
+
+ StringRef getArgument() const final {
+ return "test-convert-breakable-loop-to-scf";
+ }
+ StringRef getDescription() const final {
+ return "Convert test.breakable_loop operations to scf.loop operations";
+ }
+
+ void runOnOperation() override;
+
+ struct TargetInfo {
+ Operation *target;
+ int64_t depth;
+ };
+ using TargetMap = llvm::DenseMap<Operation *, SmallVector<TargetInfo>>;
+ using TokenMap = llvm::DenseMap<Operation *, Value>;
+
+private:
+ LogicalResult rewriteLoop(TestBreakableLoopOp loopOp,
+ const TargetMap &terminatorTargets,
+ TokenMap &loopTokens);
+ LogicalResult rewriteTerminator(Operation *terminator,
+ const TargetMap &terminatorTargets,
+ const TokenMap &loopTokens,
+ PatternRewriter &rewriter);
+};
+
+static void appendTerminatorTargets(
+ Operation *terminator,
+ SmallVectorImpl<TestConvertBreakableLoopToSCFPass::TargetInfo> &targets,
+ ArrayRef<Operation *> potentialTargets) {
+ for (Operation *targetOp : potentialTargets) {
+ int64_t depth = 0;
+ for (Operation *currentOp = terminator->getParentOp(); currentOp;
+ currentOp = currentOp->getParentOp()) {
+ if (isa<TestBreakableLoopOp>(currentOp)) {
+ ++depth;
+ if (currentOp == targetOp) {
+ targets.push_back({targetOp, depth});
+ break;
+ }
+ }
+ if (!currentOp->mightHaveTrait<OpTrait::PropagateControlFlowBreak>())
+ break;
+ }
+ }
+}
+
+static bool isBreakableLoopTerminator(Operation *op) {
+ return isa<TestDynamicBreakOp, TestDynamicContinueOp>(op);
+}
+
+static Value getTerminatorDepth(Operation *terminator) {
+ if (auto breakOp = dyn_cast<TestDynamicBreakOp>(terminator))
+ return breakOp.getDepth();
+ return cast<TestDynamicContinueOp>(terminator).getDepth();
+}
+
+static void createSCFTerminator(PatternRewriter &rewriter, Location loc,
+ Operation *terminator, Value targetToken,
+ ValueRange args) {
+ if (isa<TestDynamicBreakOp>(terminator)) {
+ scf::BreakOp::create(rewriter, loc, targetToken, args);
+ return;
+ }
+ scf::ContinueOp::create(rewriter, loc, targetToken, args);
+}
+
+static std::optional<Value>
+lookupToken(Operation *terminator, Operation *target,
+ const TestConvertBreakableLoopToSCFPass::TokenMap &loopTokens) {
+ auto it = loopTokens.find(target);
+ if (it != loopTokens.end())
+ return it->second;
+ terminator->emitOpError()
+ << "cannot convert because target loop has not been converted";
+ return std::nullopt;
+}
+
+LogicalResult TestConvertBreakableLoopToSCFPass::rewriteTerminator(
+ Operation *terminator, const TargetMap &terminatorTargets,
+ const TokenMap &loopTokens, PatternRewriter &rewriter) {
+ auto targetIt = terminatorTargets.find(terminator);
+ if (targetIt == terminatorTargets.end() || targetIt->second.empty())
+ return terminator->emitOpError()
+ << "cannot convert terminator without a target loop";
+
+ ArrayRef<TargetInfo> targets = targetIt->second;
+ Location loc = terminator->getLoc();
+ SmallVector<Value> args(getBreakableLoopTerminatorPayload(terminator));
+
+ if (targets.size() == 1) {
+ std::optional<Value> token =
+ lookupToken(terminator, targets.front().target, loopTokens);
+ if (!token)
+ return failure();
+ rewriter.setInsertionPoint(terminator);
+ createSCFTerminator(rewriter, loc, terminator, *token, args);
+ rewriter.eraseOp(terminator);
+ return success();
+ }
+
+ rewriter.setInsertionPoint(terminator);
+ Value depth = getTerminatorDepth(terminator);
+ for (TargetInfo target : targets) {
+ std::optional<Value> token =
+ lookupToken(terminator, target.target, loopTokens);
+ if (!token)
+ return failure();
+
+ Value depthValue =
+ arith::ConstantIndexOp::create(rewriter, loc, target.depth);
+ Value isTarget = arith::CmpIOp::create(
+ rewriter, loc, arith::CmpIPredicate::eq, depth, depthValue);
+ auto ifOp = scf::IfOp::create(rewriter, loc, TypeRange{}, isTarget,
+ /*addThenBlock=*/true,
+ /*addElseBlock=*/false);
+ {
+ OpBuilder::InsertionGuard guard(rewriter);
+ rewriter.setInsertionPointToStart(ifOp.thenBlock());
+ createSCFTerminator(rewriter, loc, terminator, *token, args);
+ }
+ rewriter.setInsertionPointAfter(ifOp);
+ }
+
+ // Dynamic depths are verified to be able to target any listed loop. If the
+ // value falls outside that set, selects an incompatible target, or otherwise
+ // has undefined behavior, the fallback may choose any valid target. Use the
+ // first compatible target to keep the generated IR deterministic.
+ std::optional<Value> fallbackToken =
+ lookupToken(terminator, targets.front().target, loopTokens);
+ if (!fallbackToken)
+ return failure();
+ createSCFTerminator(rewriter, loc, terminator, *fallbackToken, args);
+ rewriter.eraseOp(terminator);
+ return success();
+}
+
+LogicalResult TestConvertBreakableLoopToSCFPass::rewriteLoop(
+ TestBreakableLoopOp loopOp, const TargetMap &terminatorTargets,
+ TokenMap &loopTokens) {
+ PatternRewriter rewriter(loopOp.getContext());
+ rewriter.setInsertionPoint(loopOp);
+
+ auto scfLoop = scf::LoopOp::create(
+ rewriter, loopOp.getLoc(), loopOp.getResultTypes(), loopOp.getInitArgs());
+ scfLoop->setAttrs(loopOp->getAttrs());
+ scfLoop.getRegion().takeBody(loopOp.getBody());
+ scfLoop.getBody()->insertArgument(
+ /*index=*/0u, TokenType::get(loopOp.getContext()), loopOp.getLoc());
+ loopTokens[loopOp.getOperation()] = scfLoop.getControlToken();
+
+ SmallVector<Operation *> terminators;
+ scfLoop.getRegion().walk([&](Operation *op) {
+ if (!isBreakableLoopTerminator(op))
+ return;
+ if (op->getParentOfType<TestBreakableLoopOp>())
+ return;
+ terminators.push_back(op);
+ });
+
+ for (Operation *terminator : terminators)
+ if (failed(rewriteTerminator(terminator, terminatorTargets, loopTokens,
+ rewriter)))
+ return failure();
+
+ rewriter.replaceOp(loopOp, scfLoop.getResults());
+ return success();
+}
+
+void TestConvertBreakableLoopToSCFPass::runOnOperation() {
+ SmallVector<TestBreakableLoopOp> loops;
+ TargetMap terminatorTargets;
+ getOperation()->walk<WalkOrder::PreOrder>([&](Operation *op) {
+ if (auto loopOp = dyn_cast<TestBreakableLoopOp>(op)) {
+ loops.push_back(loopOp);
+ return;
+ }
+ if (auto breakOp = dyn_cast<TestDynamicBreakOp>(op)) {
+ appendTerminatorTargets(op, terminatorTargets[op],
+ breakOp.getPotentialTargets());
+ return;
+ }
+ if (auto continueOp = dyn_cast<TestDynamicContinueOp>(op))
+ appendTerminatorTargets(op, terminatorTargets[op],
+ continueOp.getPotentialTargets());
+ });
+
+ TokenMap loopTokens;
+ for (TestBreakableLoopOp loopOp : loops) {
+ if (failed(rewriteLoop(loopOp, terminatorTargets, loopTokens))) {
+ signalPassFailure();
+ return;
+ }
+ }
+}
} // namespace
//===----------------------------------------------------------------------===//
@@ -2688,6 +2894,8 @@ void registerPatternsTestPass() {
PassRegistration<TestSelectiveReplacementPatternDriver>();
PassRegistration<TestFoldTypeConvertingOp>();
+
+ PassRegistration<TestConvertBreakableLoopToSCFPass>();
}
} // namespace test
} // namespace mlir
diff --git a/mlir/test/lib/IR/TestDominance.cpp b/mlir/test/lib/IR/TestDominance.cpp
index 60d37476b5bf91..9c5076f6ac8094 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);
@@ -73,6 +82,27 @@ class DominanceTest {
template <typename DominanceT>
void printDominance(DominanceT &dominanceInfo,
bool 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 6a21ed10eec6f3..3aa5097b7ed206 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 00000000000000..8e003942e41c0d
--- /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 00000000000000..4aa9798efb4734
--- /dev/null
+++ b/mlir/test/lib/Interfaces/RegionBranchOpInterface/TestRegionBranchOpInterface.cpp
@@ -0,0 +1,78 @@
+//===- TestRegionBranchOpInterface.cpp - RegionBranch test pass -----------===//
+//
+// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
+// See https://llvm.org/LICENSE.txt for license information.
+// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+//
+//===----------------------------------------------------------------------===//
+
+#include "mlir/Dialect/Func/IR/FuncOps.h"
+#include "mlir/IR/BuiltinOps.h"
+#include "mlir/IR/RegionExitInterfaces.h"
+#include "mlir/Interfaces/ControlFlowInterfaces.h"
+#include "mlir/Pass/Pass.h"
+#include "llvm/Support/raw_ostream.h"
+
+using namespace mlir;
+
+namespace {
+/// Test pass that prints RegionBranchOpInterface successor information and
+/// nested breaking-control-flow predecessors.
+struct PrintRegionBranchOpInterfacePass
+ : public PassWrapper<PrintRegionBranchOpInterfacePass, OperationPass<>> {
+ MLIR_DEFINE_EXPLICIT_INTERNAL_INLINE_TYPE_ID(PrintRegionBranchOpInterfacePass)
+
+ StringRef getArgument() const final {
+ return "print-region-branch-op-interface";
+ }
+ StringRef getDescription() const final {
+ return "Print control-flow edges represented by "
+ "mlir::RegionBranchOpInterface";
+ }
+
+ void runOnOperation() override {
+ Operation *op = getOperation();
+ op->walk<WalkOrder::PreOrder>([&](RegionBranchOpInterface branchOp) {
+ llvm::outs() << "Found RegionBranchOpInterface operation: "
+ << OpWithFlags(
+ branchOp,
+ OpPrintingFlags().skipRegions().enableDebugInfo())
+ << "\n";
+ SmallVector<RegionSuccessor> regions;
+ branchOp.getSuccessorRegions(RegionBranchPoint::parent(), regions);
+ for (auto &successor : regions) {
+ if (successor.isOperation()) {
+ llvm::outs() << " - Successor is operation "
+ << successor.getSuccessorOp()->getName() << "\n";
+ } else {
+ llvm::outs() << " - Successor is region #"
+ << successor.getSuccessor()->getRegionNumber() << "\n";
+ }
+ }
+ if (auto breakingControlFlowOp =
+ dyn_cast<HasBreakingControlFlowOpInterface>(
+ branchOp.getOperation())) {
+ llvm::outs() << " - Collecting all nested predecessors\n";
+ auto predecessors =
+ getNestedBreakingControlFlowInfo(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-runner/test-convert-breakable-loop-to-scf.mlir b/mlir/test/mlir-runner/test-convert-breakable-loop-to-scf.mlir
new file mode 100644
index 00000000000000..d2a9ccfc846502
--- /dev/null
+++ b/mlir/test/mlir-runner/test-convert-breakable-loop-to-scf.mlir
@@ -0,0 +1,145 @@
+// RUN: mlir-opt %s -pass-pipeline="builtin.module(func.func(test-convert-breakable-loop-to-scf,convert-scf-to-cf,canonicalize,convert-arith-to-llvm),convert-func-to-llvm,convert-cf-to-llvm,reconcile-unrealized-casts)" -o %t.mlir
+// RUN: mlir-runner %t.mlir -e immediate_continue -entry-point-result=i32 | FileCheck %s --check-prefix=IMMEDIATE
+// RUN: mlir-runner %t.mlir -e break_with_result -entry-point-result=i32 | FileCheck %s --check-prefix=BREAK
+// RUN: mlir-runner %t.mlir -e constant_outer_break -entry-point-result=i32 | FileCheck %s --check-prefix=OUTER-BREAK
+// RUN: mlir-runner %t.mlir -e constant_outer_continue -entry-point-result=i32 | FileCheck %s --check-prefix=OUTER-CONTINUE
+// RUN: mlir-runner %t.mlir -e dynamic_break_depth_1 -entry-point-result=i32 | FileCheck %s --check-prefix=DYNAMIC-ONE
+// RUN: mlir-runner %t.mlir -e dynamic_break_depth_2 -entry-point-result=i32 | FileCheck %s --check-prefix=DYNAMIC-TWO
+// RUN: mlir-runner %t.mlir -e dynamic_filtered_depth_2 -entry-point-result=i32 | FileCheck %s --check-prefix=FILTERED-TWO
+// RUN: mlir-runner %t.mlir -e dynamic_filtered_depth_3 -entry-point-result=i32 | FileCheck %s --check-prefix=FILTERED-THREE
+// XFAIL: system-aix
+
+func.func @immediate_continue() -> i32 {
+ %depth = arith.constant 1 : index
+ %c0 = arith.constant 0 : i32
+ %c1 = arith.constant 1 : i32
+ %c4 = arith.constant 4 : i32
+ %result = test.breakable_loop iter_args(%i = %c0) : i32 -> i32 {
+ %done = arith.cmpi eq, %i, %c4 : i32
+ scf.if %done {
+ test.dynamic_break %depth %i : i32
+ }
+ %next = arith.addi %i, %c1 : i32
+ test.dynamic_continue %depth %next : i32
+ }
+ return %result : i32
+}
+// IMMEDIATE: 4
+
+func.func @break_with_result() -> i32 {
+ %depth = arith.constant 1 : index
+ %value = arith.constant 17 : i32
+ %result = test.breakable_loop -> i32 {
+ test.dynamic_break %depth %value : i32
+ }
+ return %result : i32
+}
+// BREAK: 17
+
+func.func @constant_outer_break() -> i32 {
+ %inner_depth = arith.constant 1 : index
+ %outer_depth = arith.constant 2 : index
+ %outer_value = arith.constant 42 : i32
+ %fallback = arith.constant 13 : i32
+ %result = test.breakable_loop -> i32 {
+ test.breakable_loop {
+ test.dynamic_break %outer_depth %outer_value : i32
+ }
+ test.dynamic_break %inner_depth %fallback : i32
+ }
+ return %result : i32
+}
+// OUTER-BREAK: 42
+
+func.func @constant_outer_continue() -> i32 {
+ %inner_depth = arith.constant 1 : index
+ %outer_depth = arith.constant 2 : index
+ %c0 = arith.constant 0 : i32
+ %c1 = arith.constant 1 : i32
+ %c3 = arith.constant 3 : i32
+ %result = test.breakable_loop iter_args(%i = %c0) : i32 -> i32 {
+ %done = arith.cmpi eq, %i, %c3 : i32
+ scf.if %done {
+ test.dynamic_break %inner_depth %i : i32
+ }
+ test.breakable_loop {
+ %next = arith.addi %i, %c1 : i32
+ test.dynamic_continue %outer_depth %next : i32
+ }
+ test.dynamic_break %inner_depth %i : i32
+ }
+ return %result : i32
+}
+// OUTER-CONTINUE: 3
+
+// Distinguish the two compatible dynamic break targets at runtime. Depth 1
+// breaks the inner loop, then the outer loop adds one; depth 2 skips that outer
+// post-inner add by breaking directly to the outer loop result.
+func.func @dynamic_break(%choose_outer: i1) -> i32 {
+ %inner_depth = arith.constant 1 : index
+ %outer_depth = arith.constant 2 : index
+ %depth = arith.select %choose_outer, %outer_depth, %inner_depth : index
+ %value = arith.constant 7 : i32
+ %one = arith.constant 1 : i32
+ %result = test.breakable_loop -> i32 {
+ %inner_result = test.breakable_loop -> i32 {
+ test.dynamic_break %depth %value : i32
+ }
+ %after_inner = arith.addi %inner_result, %one : i32
+ test.dynamic_break %inner_depth %after_inner : i32
+ }
+ return %result : i32
+}
+
+func.func @dynamic_break_depth_1() -> i32 {
+ %choose_outer = arith.constant false
+ %result = func.call @dynamic_break(%choose_outer) : (i1) -> i32
+ return %result : i32
+}
+// DYNAMIC-ONE: 8
+
+func.func @dynamic_break_depth_2() -> i32 {
+ %choose_outer = arith.constant true
+ %result = func.call @dynamic_break(%choose_outer) : (i1) -> i32
+ return %result : i32
+}
+// DYNAMIC-TWO: 7
+
+// The innermost f32 loop is incompatible with the i32 break payload and is
+// filtered out. Depth 2 therefore targets the middle loop and observes the
+// outer post-middle add, while depth 3 targets the outer loop directly.
+func.func @dynamic_break_filtered(%choose_outer: i1) -> i32 {
+ %inner_depth = arith.constant 1 : index
+ %middle_depth = arith.constant 2 : index
+ %outer_depth = arith.constant 3 : index
+ %depth = arith.select %choose_outer, %outer_depth, %middle_depth : index
+ %value = arith.constant 5 : i32
+ %one = arith.constant 1 : i32
+ %f0 = arith.constant 0.000000e+00 : f32
+ %result = test.breakable_loop -> i32 {
+ %middle_result = test.breakable_loop -> i32 {
+ test.breakable_loop iter_args(%inner = %f0) : f32 {
+ test.dynamic_break %depth %value : i32
+ }
+ %after_inner = arith.addi %value, %one : i32
+ test.dynamic_break %inner_depth %after_inner : i32
+ }
+ %after_middle = arith.addi %middle_result, %one : i32
+ test.dynamic_break %inner_depth %after_middle : i32
+ }
+ return %result : i32
+}
+
+func.func @dynamic_filtered_depth_2() -> i32 {
+ %choose_outer = arith.constant false
+ %result = func.call @dynamic_break_filtered(%choose_outer) : (i1) -> i32
+ return %result : i32
+}
+// FILTERED-TWO: 6
+
+func.func @dynamic_filtered_depth_3() -> i32 {
+ %choose_outer = arith.constant true
+ %result = func.call @dynamic_break_filtered(%choose_outer) : (i1) -> i32
+ return %result : i32
+}
+// FILTERED-THREE: 5
diff --git a/mlir/tools/mlir-opt/CMakeLists.txt b/mlir/tools/mlir-opt/CMakeLists.txt
index b34f7fdbe484f4..51304cc7d6862a 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 bf8efb86372eb6..67e45e02c2a116 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 registerShapeFunctionTestPasses();
void registerSideEffectTestPasses();
@@ -190,6 +191,7 @@ static void registerTestPasses() {
registerPassManagerTestPass();
registerPrintSpirvAvailabilityPass();
registerRegionTestPasses();
+ registerRegionBranchOpInterfaceTestPasses();
registerShapeFunctionTestPasses();
registerSideEffectTestPasses();
registerSliceAnalysisTestPass();
diff --git a/utils/bazel/llvm-project-overlay/mlir/BUILD.bazel b/utils/bazel/llvm-project-overlay/mlir/BUILD.bazel
index 41f6f6998f21b9..5f33ae88fddc78 100644
--- a/utils/bazel/llvm-project-overlay/mlir/BUILD.bazel
+++ b/utils/bazel/llvm-project-overlay/mlir/BUILD.bazel
@@ -109,6 +109,7 @@ exports_files(glob(["include/**/*.td"]))
deps = [":OpBaseTdFiles"],
)
for name in [
+ "RegionExitInterfaces",
"RegionKindInterface",
"SymbolInterfaces",
]
@@ -473,6 +474,7 @@ cc_library(
":InferIntRangeInterfaceIncGen",
":OpAsmInterfaceIncGen",
":QuantStorageTypeInterfaceIncGen",
+ ":RegionExitInterfacesIncGen",
":RegionKindInterfaceIncGen",
":SideEffectInterfacesIncGen",
":Support",
@@ -1438,6 +1440,7 @@ td_library(
"include/mlir/IR/OpBase.td",
"include/mlir/IR/PatternBase.td",
"include/mlir/IR/Properties.td",
+ "include/mlir/IR/RegionExitInterfaces.td",
"include/mlir/IR/RegionKindInterface.td",
"include/mlir/IR/SymbolInterfaces.td",
"include/mlir/IR/TensorEncoding.td",
More information about the Mlir-commits
mailing list