[Mlir-commits] [mlir] [MLIR] `remove-dead-values` and RegionBranchOp canononicalizations fixes around fixes around the `IsolatedFromAbove` ops (PR #224831)
Ivan Butygin
llvmlistbot at llvm.org
Sat Sep 19 04:48:53 PDT 2026
https://github.com/Hardcode84 created https://github.com/llvm/llvm-project/pull/224831
I have a downstream op which is both `RegionBranchOpInterface` and `IsolatedFromAbove`, trying use `remove-dead-values` on it and especially trying to use this op as pass root uncovered various issues.
* `LivenessAnalysis` were treating yielded values as dead
* `RegionBranchOpInterface` canonicalization were just asserting on `IsolatedFromAbove`
* When rooting on `func.func` it would leave callers invalid.
* Subsequent canonicalization could reach outside the scope.
>From cf6fedc7d0d00237121bb0d754235252dd4b5657 Mon Sep 17 00:00:00 2001
From: Ivan Butygin <ivan.butygin at gmail.com>
Date: Sat, 19 Sep 2026 11:33:08 +0200
Subject: [PATCH 1/5] [MLIR] Preserve the root of remove-dead-values
---
.../mlir/Analysis/DataFlow/LivenessAnalysis.h | 2 +
.../Analysis/DataFlow/LivenessAnalysis.cpp | 7 ++
mlir/lib/Transforms/RemoveDeadValues.cpp | 47 +++++----
.../remove-dead-values-function-root.mlir | 17 ++++
.../remove-dead-values-nested-root.mlir | 96 +++++++++++++++++++
mlir/test/lib/Dialect/Test/TestOpDefs.cpp | 31 ++++++
mlir/test/lib/Dialect/Test/TestOps.td | 16 ++++
7 files changed, 195 insertions(+), 21 deletions(-)
create mode 100644 mlir/test/Transforms/remove-dead-values-function-root.mlir
create mode 100644 mlir/test/Transforms/remove-dead-values-nested-root.mlir
diff --git a/mlir/include/mlir/Analysis/DataFlow/LivenessAnalysis.h b/mlir/include/mlir/Analysis/DataFlow/LivenessAnalysis.h
index 90947d06e7999f..f27188e2cc75af 100644
--- a/mlir/include/mlir/Analysis/DataFlow/LivenessAnalysis.h
+++ b/mlir/include/mlir/Analysis/DataFlow/LivenessAnalysis.h
@@ -79,6 +79,8 @@ class LivenessAnalysis : public SparseBackwardDataFlowAnalysis<Liveness> {
public:
using SparseBackwardDataFlowAnalysis::SparseBackwardDataFlowAnalysis;
+ LogicalResult initialize(Operation *top) override;
+
LogicalResult visitOperation(Operation *op, ArrayRef<Liveness *> operands,
ArrayRef<const Liveness *> results) override;
diff --git a/mlir/lib/Analysis/DataFlow/LivenessAnalysis.cpp b/mlir/lib/Analysis/DataFlow/LivenessAnalysis.cpp
index 07e3c98e0bb603..42cf874cc96e81 100644
--- a/mlir/lib/Analysis/DataFlow/LivenessAnalysis.cpp
+++ b/mlir/lib/Analysis/DataFlow/LivenessAnalysis.cpp
@@ -48,6 +48,13 @@ ChangeResult Liveness::meet(const AbstractSparseLattice &other) {
// LivenessAnalysis
//===----------------------------------------------------------------------===//
+LogicalResult LivenessAnalysis::initialize(Operation *top) {
+ // Users outside the analysis root are not visited.
+ for (Value result : top->getResults())
+ setToExitState(getLatticeElement(result));
+ return SparseBackwardDataFlowAnalysis<Liveness>::initialize(top);
+}
+
/// For every value, liveness analysis determines whether or not it is "live".
///
/// A value is considered "live" iff it:
diff --git a/mlir/lib/Transforms/RemoveDeadValues.cpp b/mlir/lib/Transforms/RemoveDeadValues.cpp
index a4c2aa59b4421f..fd856711d8edf4 100644
--- a/mlir/lib/Transforms/RemoveDeadValues.cpp
+++ b/mlir/lib/Transforms/RemoveDeadValues.cpp
@@ -786,14 +786,14 @@ struct RemoveDeadValues
void RemoveDeadValues::runOnOperation() {
auto &la = getAnalysis<RunLivenessAnalysis>();
- Operation *module = getOperation();
+ Operation *root = getOperation();
// Build a symbol user map once up front so that processFuncOp can look up the
// callers of each function in O(1). Otherwise, each call would walk the
- // entire module to find the callers, making the pass O(numFunctions *
+ // entire root to find the callers, making the pass O(numFunctions *
// numOperations).
SymbolTableCollection symbolTableCollection;
- SymbolUserMap symbolUserMap(symbolTableCollection, module);
+ SymbolUserMap symbolUserMap(symbolTableCollection, root);
// Tracks values eligible for erasure - complements liveness analysis to
// identify "droppable" values.
@@ -803,7 +803,9 @@ void RemoveDeadValues::runOnOperation() {
// end of this pass.
RDVFinalCleanupList finalCleanupList;
- module->walk([&](Operation *op) {
+ root->walk([&](Operation *op) {
+ if (op == root)
+ return;
if (auto funcOp = dyn_cast<FunctionOpInterface>(op)) {
processFuncOp(funcOp, symbolUserMap, la, deadVals, finalCleanupList);
} else if (auto regionBranchOp = dyn_cast<RegionBranchOpInterface>(op)) {
@@ -821,27 +823,30 @@ void RemoveDeadValues::runOnOperation() {
}
});
- MLIRContext *context = module->getContext();
+ MLIRContext *context = root->getContext();
cleanUpDeadVals(context, finalCleanupList);
if (!canonicalize)
return;
- // Canonicalize all region branch ops.
- SmallVector<Operation *> opsToCanonicalize;
- module->walk([&](RegionBranchOpInterface regionBranchOp) {
- opsToCanonicalize.push_back(regionBranchOp.getOperation());
- });
- // Collect all canonicalization patterns for region branch ops.
- RewritePatternSet owningPatterns(context);
- DenseSet<RegisteredOperationName> populatedPatterns;
- for (Operation *op : opsToCanonicalize)
- if (std::optional<RegisteredOperationName> info = op->getRegisteredInfo())
- if (populatedPatterns.insert(*info).second)
- info->getCanonicalizationPatterns(owningPatterns, context);
- if (failed(applyOpPatternsGreedily(opsToCanonicalize,
- std::move(owningPatterns)))) {
- module->emitError("greedy pattern rewrite failed to converge");
- signalPassFailure();
+ for (Region ®ion : root->getRegions()) {
+ SmallVector<Operation *> opsToCanonicalize;
+ region.walk([&](RegionBranchOpInterface regionBranchOp) {
+ opsToCanonicalize.push_back(regionBranchOp.getOperation());
+ });
+ RewritePatternSet owningPatterns(context);
+ DenseSet<RegisteredOperationName> populatedPatterns;
+ for (Operation *op : opsToCanonicalize)
+ if (std::optional<RegisteredOperationName> info = op->getRegisteredInfo())
+ if (populatedPatterns.insert(*info).second)
+ info->getCanonicalizationPatterns(owningPatterns, context);
+ // Keep greedy worklist expansion below the pass root.
+ GreedyRewriteConfig config;
+ config.setScope(®ion);
+ if (failed(applyOpPatternsGreedily(opsToCanonicalize,
+ std::move(owningPatterns), config))) {
+ root->emitError("greedy pattern rewrite failed to converge");
+ return signalPassFailure();
+ }
}
}
diff --git a/mlir/test/Transforms/remove-dead-values-function-root.mlir b/mlir/test/Transforms/remove-dead-values-function-root.mlir
new file mode 100644
index 00000000000000..2cec260eb9640b
--- /dev/null
+++ b/mlir/test/Transforms/remove-dead-values-function-root.mlir
@@ -0,0 +1,17 @@
+// RUN: mlir-opt %s --pass-pipeline='builtin.module(func.func(remove-dead-values))' | FileCheck %s
+
+// CHECK-LABEL: func.func private @callee
+// CHECK-SAME: %[[ARG:.*]]: i32, %{{.*}}: i32) -> i32
+// CHECK-NEXT: %[[SUM:.*]] = arith.addi %[[ARG]], %[[ARG]] : i32
+// CHECK-NEXT: return %[[SUM]] : i32
+func.func private @callee(%arg: i32, %unused: i32) -> i32 {
+ %sum = arith.addi %arg, %arg : i32
+ return %sum : i32
+}
+
+// CHECK-LABEL: func.func @caller
+// CHECK: call @callee(%{{.*}}, %{{.*}}) : (i32, i32) -> i32
+func.func @caller(%x: i32) -> i32 {
+ %r = call @callee(%x, %x) : (i32, i32) -> i32
+ return %r : i32
+}
diff --git a/mlir/test/Transforms/remove-dead-values-nested-root.mlir b/mlir/test/Transforms/remove-dead-values-nested-root.mlir
new file mode 100644
index 00000000000000..6a42910d1e2d22
--- /dev/null
+++ b/mlir/test/Transforms/remove-dead-values-nested-root.mlir
@@ -0,0 +1,96 @@
+// RUN: mlir-opt %s --pass-pipeline='builtin.module(func.func(test.isolated_region_branch(remove-dead-values)))' > %t.parallel
+// RUN: FileCheck %s --check-prefixes=CHECK,CANON < %t.parallel
+// RUN: mlir-opt %s --mlir-disable-threading --pass-pipeline='builtin.module(func.func(test.isolated_region_branch(remove-dead-values)))' > %t.serial
+// RUN: diff %t.parallel %t.serial
+// RUN: mlir-opt %s --pass-pipeline='builtin.module(func.func(test.isolated_region_branch(remove-dead-values{canonicalize=false})))' | FileCheck %s
+
+// CHECK-LABEL: func.func @live_result
+// CHECK-SAME: %[[X:.*]]: i32
+// CHECK: %[[ROOT:.*]] = "test.isolated_region_branch"(%[[X]])
+// CHECK: ^bb0(%[[ARG:.*]]: i32):
+// CHECK-NEXT: %[[SUM:.*]] = arith.addi %[[ARG]], %[[ARG]] : i32
+// CHECK-NEXT: "test.isolated_region_yield"(%[[SUM]])
+// CHECK: return %[[ROOT]] : i32
+func.func @live_result(%x: i32) -> i32 {
+ %r = "test.isolated_region_branch"(%x) ({
+ ^bb0(%arg: i32):
+ %dead = arith.muli %arg, %arg : i32
+ %sum = arith.addi %arg, %arg : i32
+ "test.isolated_region_yield"(%sum) : (i32) -> ()
+ }) : (i32) -> i32
+ return %r : i32
+}
+
+// Unused root results and inputs are outside this pass's rewrite scope.
+// CHECK-LABEL: func.func @unused_result
+// CHECK-SAME: %[[X:.*]]: i32
+// CHECK: "test.isolated_region_branch"(%[[X]], %[[X]])
+// CHECK: ^bb0(%[[A:.*]]: i32, %[[UNUSED:.*]]: i32):
+// CHECK-NEXT: %[[SUM:.*]] = arith.addi %[[A]], %[[A]] : i32
+// CHECK-NEXT: "test.isolated_region_yield"(%[[SUM]])
+// CHECK: "test.isolated_region_branch"(%[[X]], %[[X]])
+// CHECK: "test.isolated_region_yield"
+// CHECK: return
+func.func @unused_result(%x: i32) {
+ %r = "test.isolated_region_branch"(%x, %x) ({
+ ^bb0(%arg: i32, %unused: i32):
+ %sum = arith.addi %arg, %arg : i32
+ "test.isolated_region_yield"(%sum) : (i32) -> ()
+ }) : (i32, i32) -> i32
+ %other = "test.isolated_region_branch"(%x, %x) ({
+ ^bb0(%arg: i32, %unused: i32):
+ "test.isolated_region_yield"(%arg) : (i32) -> ()
+ }) : (i32, i32) -> i32
+ return
+}
+
+// CHECK-LABEL: func.func @multiple_regions
+// CHECK: "test.isolated_region_branch"()
+// CHECK-NEXT: "test.isolated_region_yield"()
+// CHECK-NEXT: }, {
+// CHECK-NEXT: "test.isolated_region_yield"()
+// CHECK: return
+func.func @multiple_regions() {
+ "test.isolated_region_branch"() ({
+ %dead = arith.constant 1 : i32
+ "test.isolated_region_yield"() : () -> ()
+ }, {
+ %dead = arith.constant 2 : i32
+ "test.isolated_region_yield"() : () -> ()
+ }) : () -> ()
+ return
+}
+
+// CHECK-LABEL: func.func @effects
+// CHECK: "test.isolated_region_branch"
+// CHECK: ^bb0(%[[MEM:.*]]: memref<i32>, %[[VAL:.*]]: i32):
+// CHECK-NEXT: memref.store %[[VAL]], %[[MEM]][] : memref<i32>
+// CHECK-NEXT: "test.isolated_region_yield"()
+func.func @effects(%mem: memref<i32>, %x: i32) {
+ "test.isolated_region_branch"(%mem, %x) ({
+ ^bb0(%arg: memref<i32>, %val: i32):
+ %dead = arith.muli %val, %val : i32
+ memref.store %val, %arg[] : memref<i32>
+ "test.isolated_region_yield"() : () -> ()
+ }) : (memref<i32>, i32) -> ()
+ return
+}
+
+// CANON-LABEL: func.func @dead_carry
+// CANON: "test.isolated_region_branch"
+// CANON: scf.for {{.*}} iter_args(%[[CARRY:.*]] = %{{.*}}) -> (i32)
+// CANON: scf.yield %{{.*}} : i32
+// CANON: "test.isolated_region_yield"
+func.func @dead_carry(%lb: index, %ub: index, %step: index, %x: i32) -> i32 {
+ %r = "test.isolated_region_branch"(%lb, %ub, %step, %x) ({
+ ^bb0(%lower: index, %upper: index, %stride: index, %init: i32):
+ %loop:2 = scf.for %iv = %lower to %upper step %stride
+ iter_args(%live = %init, %dead = %init) -> (i32, i32) {
+ %sum = arith.addi %live, %live : i32
+ %unused = arith.muli %dead, %dead : i32
+ scf.yield %sum, %unused : i32, i32
+ }
+ "test.isolated_region_yield"(%loop#0) : (i32) -> ()
+ }) : (index, index, index, i32) -> i32
+ return %r : i32
+}
diff --git a/mlir/test/lib/Dialect/Test/TestOpDefs.cpp b/mlir/test/lib/Dialect/Test/TestOpDefs.cpp
index 3effcfa8f040ed..39a0c8eb213c69 100644
--- a/mlir/test/lib/Dialect/Test/TestOpDefs.cpp
+++ b/mlir/test/lib/Dialect/Test/TestOpDefs.cpp
@@ -1364,6 +1364,37 @@ SegmentedRegionBranchTerminatorOp::getMutableSuccessorOperands(
: getSecondMutable();
}
+//===----------------------------------------------------------------------===//
+// IsolatedRegionBranchOp
+//===----------------------------------------------------------------------===//
+
+void IsolatedRegionBranchOp::getSuccessorRegions(
+ RegionBranchPoint point, SmallVectorImpl<RegionSuccessor> ®ions) {
+ if (!point.isParent()) {
+ regions.emplace_back(getOperation());
+ return;
+ }
+ for (Region ®ion : getBranches())
+ regions.emplace_back(®ion);
+}
+
+ValueRange
+IsolatedRegionBranchOp::getSuccessorInputs(RegionSuccessor successor) {
+ return successor.isOperation()
+ ? ValueRange(getOutputs())
+ : ValueRange(successor.getSuccessor()->getArguments());
+}
+
+OperandRange
+IsolatedRegionBranchOp::getEntrySuccessorOperands(RegionSuccessor successor) {
+ return getInputs();
+}
+
+MutableOperandRange
+IsolatedRegionYieldOp::getMutableSuccessorOperands(RegionSuccessor successor) {
+ return getValuesMutable();
+}
+
//===----------------------------------------------------------------------===//
// LoopBlockOp
//===----------------------------------------------------------------------===//
diff --git a/mlir/test/lib/Dialect/Test/TestOps.td b/mlir/test/lib/Dialect/Test/TestOps.td
index bbf6c4bbfe6489..d9420346c8d13a 100644
--- a/mlir/test/lib/Dialect/Test/TestOps.td
+++ b/mlir/test/lib/Dialect/Test/TestOps.td
@@ -2828,6 +2828,22 @@ def SegmentedRegionBranchTerminatorOp
Variadic<AnyType>:$second, Optional<I1>:$exitCondition);
}
+def IsolatedRegionBranchOp : TEST_Op<"isolated_region_branch", [
+ IsolatedFromAbove, RecursiveMemoryEffects,
+ DeclareOpInterfaceMethods<RegionBranchOpInterface,
+ ["getEntrySuccessorOperands", "getSuccessorInputs"]>]> {
+ let arguments = (ins Variadic<AnyType>:$inputs);
+ let results = (outs Variadic<AnyType>:$outputs);
+ let regions = (region VariadicRegion<SizedRegion<1>>:$branches);
+}
+
+def IsolatedRegionYieldOp : TEST_Op<"isolated_region_yield", [
+ Pure, ReturnLike, Terminator,
+ DeclareOpInterfaceMethods<RegionBranchTerminatorOpInterface,
+ ["getMutableSuccessorOperands"]>]> {
+ let arguments = (ins Variadic<AnyType>:$values);
+}
+
def RegionIfYieldOp : TEST_Op<"region_if_yield",
[NoMemoryEffect, ReturnLike, Terminator]> {
let arguments = (ins Variadic<AnyType>:$results);
>From b29fe124e938170255fe5c54a90e2b95a356519f Mon Sep 17 00:00:00 2001
From: Ivan Butygin <ivan.butygin at gmail.com>
Date: Sat, 19 Sep 2026 11:35:01 +0200
Subject: [PATCH 2/5] [MLIR] Canonicalize isolated region branch operations
---
mlir/lib/Interfaces/ControlFlowInterfaces.cpp | 13 +--
.../isolated-region-branch-canonicalize.mlir | 97 +++++++++++++++++++
.../remove-dead-values-nested-root.mlir | 13 +++
mlir/test/lib/Dialect/Test/TestOpDefs.cpp | 6 ++
mlir/test/lib/Dialect/Test/TestOps.td | 1 +
5 files changed, 121 insertions(+), 9 deletions(-)
create mode 100644 mlir/test/Transforms/isolated-region-branch-canonicalize.mlir
diff --git a/mlir/lib/Interfaces/ControlFlowInterfaces.cpp b/mlir/lib/Interfaces/ControlFlowInterfaces.cpp
index d505ff0772a9cf..0ea76fc75bd074 100644
--- a/mlir/lib/Interfaces/ControlFlowInterfaces.cpp
+++ b/mlir/lib/Interfaces/ControlFlowInterfaces.cpp
@@ -722,9 +722,6 @@ struct MakeRegionBranchOpSuccessorInputsDead : public RewritePattern {
LogicalResult matchAndRewrite(Operation *op,
PatternRewriter &rewriter) const override {
- assert(!op->hasTrait<OpTrait::IsIsolatedFromAbove>() &&
- "isolated-from-above ops are not supported");
-
// Compute the mapping of successor inputs to successor operands.
auto regionBranchOp = cast<RegionBranchOpInterface>(op);
RegionBranchInverseSuccessorMapping inputToOperands;
@@ -733,6 +730,10 @@ struct MakeRegionBranchOpSuccessorInputsDead : public RewritePattern {
// Try to replace the uses of each successor input one-by-one.
bool changed = false;
for (Value value : inputToOperands.keys()) {
+ // Isolated entry arguments cannot capture values from above the op.
+ if (op->hasTrait<OpTrait::IsIsolatedFromAbove>() &&
+ isa<BlockArgument>(value))
+ continue;
// Nothing to do for successor inputs that are already dead.
if (value.use_empty())
continue;
@@ -845,9 +846,6 @@ struct RemoveDeadRegionBranchOpSuccessorInputs : public RewritePattern {
LogicalResult matchAndRewrite(Operation *op,
PatternRewriter &rewriter) const override {
- assert(!op->hasTrait<OpTrait::IsIsolatedFromAbove>() &&
- "isolated-from-above ops are not supported");
-
// Compute tied values: values that must come as a set. If you remove one,
// you must remove all. If a successor op operand is forwarded to two
// successor inputs %a and %b, both %a and %b are in the same set.
@@ -995,9 +993,6 @@ struct RemoveDuplicateSuccessorInputUses : public RewritePattern {
LogicalResult matchAndRewrite(Operation *op,
PatternRewriter &rewriter) const override {
- assert(!op->hasTrait<OpTrait::IsIsolatedFromAbove>() &&
- "isolated-from-above ops are not supported");
-
// Collect all successor inputs and sort them. When dropping the uses of a
// successor input, we'd like to also drop the uses of the same tied
// successor inputs. Otherwise, a set of tied successor inputs may not
diff --git a/mlir/test/Transforms/isolated-region-branch-canonicalize.mlir b/mlir/test/Transforms/isolated-region-branch-canonicalize.mlir
new file mode 100644
index 00000000000000..e872a8cebe6e0d
--- /dev/null
+++ b/mlir/test/Transforms/isolated-region-branch-canonicalize.mlir
@@ -0,0 +1,97 @@
+// RUN: mlir-opt %s --canonicalize | FileCheck %s
+// RUN: mlir-opt %s --remove-dead-values --canonicalize | FileCheck %s
+
+func.func private @sink(i32)
+
+// CHECK-LABEL: func.func @dead_input
+// CHECK: "test.isolated_region_branch"(%{{.*}}) ({
+// CHECK: ^bb0(%[[ARG:.*]]: i32):
+// CHECK: call @sink(%[[ARG]])
+func.func @dead_input(%x: i32, %y: i32) -> i32 {
+ %r = "test.isolated_region_branch"(%x, %y) ({
+ ^bb0(%arg: i32, %dead: i32):
+ call @sink(%arg) : (i32) -> ()
+ %sum = arith.addi %arg, %arg : i32
+ "test.isolated_region_yield"(%sum) : (i32) -> ()
+ }) : (i32, i32) -> i32
+ return %r : i32
+}
+
+// CHECK-LABEL: func.func @duplicate_inputs
+// CHECK: "test.isolated_region_branch"(%{{.*}}) ({
+// CHECK: ^bb0(%[[ARG:.*]]: i32):
+// CHECK: call @sink(%[[ARG]])
+// CHECK: call @sink(%[[ARG]])
+// CHECK: arith.addi %[[ARG]], %[[ARG]]
+func.func @duplicate_inputs(%x: i32) -> i32 {
+ %r = "test.isolated_region_branch"(%x, %x) ({
+ ^bb0(%a: i32, %b: i32):
+ call @sink(%a) : (i32) -> ()
+ call @sink(%b) : (i32) -> ()
+ %sum = arith.addi %a, %b : i32
+ "test.isolated_region_yield"(%sum) : (i32) -> ()
+ }) : (i32, i32) -> i32
+ return %r : i32
+}
+
+// CHECK-LABEL: func.func @duplicate_results
+// CHECK: %[[R:.*]] = "test.isolated_region_branch"
+// CHECK: "test.isolated_region_yield"(%{{.*}}) : (i32) -> ()
+// CHECK: return %[[R]], %[[R]] : i32, i32
+func.func @duplicate_results(%x: i32) -> (i32, i32) {
+ %r:2 = "test.isolated_region_branch"(%x) ({
+ ^bb0(%arg: i32):
+ %sum = arith.addi %arg, %arg : i32
+ "test.isolated_region_yield"(%sum, %sum) : (i32, i32) -> ()
+ }) : (i32) -> (i32, i32)
+ return %r#0, %r#1 : i32, i32
+}
+
+// CHECK-LABEL: func.func @forwarded_result
+// CHECK-SAME: %[[X:.*]]: i32
+// CHECK: "test.isolated_region_branch"(%[[X]]) ({
+// CHECK: ^bb0(%[[ARG:.*]]: i32):
+// CHECK: call @sink(%[[ARG]])
+// CHECK: "test.isolated_region_yield"() : () -> ()
+// CHECK: return %[[X]] : i32
+func.func @forwarded_result(%x: i32) -> i32 {
+ %r = "test.isolated_region_branch"(%x) ({
+ ^bb0(%arg: i32):
+ call @sink(%arg) : (i32) -> ()
+ "test.isolated_region_yield"(%arg) : (i32) -> ()
+ }) : (i32) -> i32
+ return %r : i32
+}
+
+// Keep replacement arguments in their own regions. Remove a shared input only
+// when its arguments are dead in all regions.
+// CHECK-LABEL: func.func @multiple_regions
+// CHECK-SAME: %[[X:.*]]: i32,
+// CHECK: %[[R:.*]] = "test.isolated_region_branch"(%[[X]]) ({
+// CHECK: ^bb0(%[[A:.*]]: i32):
+// CHECK: call @sink(%[[A]])
+// CHECK: call @sink(%[[A]])
+// CHECK: %[[SUM:.*]] = arith.addi %[[A]], %[[A]]
+// CHECK: "test.isolated_region_yield"(%[[SUM]]) : (i32) -> ()
+// CHECK: ^bb0(%[[B:.*]]: i32):
+// CHECK: call @sink(%[[B]])
+// CHECK: call @sink(%[[B]])
+// CHECK: %[[PRODUCT:.*]] = arith.muli %[[B]], %[[B]]
+// CHECK: "test.isolated_region_yield"(%[[PRODUCT]]) : (i32) -> ()
+// CHECK: return %[[R]], %[[R]] : i32, i32
+func.func @multiple_regions(%x: i32, %y: i32) -> (i32, i32) {
+ %r:2 = "test.isolated_region_branch"(%x, %x, %y) ({
+ ^bb0(%a: i32, %b: i32, %dead: i32):
+ call @sink(%a) : (i32) -> ()
+ call @sink(%b) : (i32) -> ()
+ %sum = arith.addi %a, %b : i32
+ "test.isolated_region_yield"(%sum, %sum) : (i32, i32) -> ()
+ }, {
+ ^bb0(%a: i32, %b: i32, %dead: i32):
+ call @sink(%a) : (i32) -> ()
+ call @sink(%b) : (i32) -> ()
+ %product = arith.muli %a, %b : i32
+ "test.isolated_region_yield"(%product, %product) : (i32, i32) -> ()
+ }) : (i32, i32, i32) -> (i32, i32)
+ return %r#0, %r#1 : i32, i32
+}
diff --git a/mlir/test/Transforms/remove-dead-values-nested-root.mlir b/mlir/test/Transforms/remove-dead-values-nested-root.mlir
index 6a42910d1e2d22..7e373386019dca 100644
--- a/mlir/test/Transforms/remove-dead-values-nested-root.mlir
+++ b/mlir/test/Transforms/remove-dead-values-nested-root.mlir
@@ -94,3 +94,16 @@ func.func @dead_carry(%lb: index, %ub: index, %step: index, %x: i32) -> i32 {
}) : (index, index, index, i32) -> i32
return %r : i32
}
+
+// CHECK-LABEL: func.func @isolated_dead_input
+// CHECK-SAME: %[[X:.*]]: i32
+// CHECK: "test.isolated_region_branch"(%[[X]], %[[X]])
+// CHECK: ^bb0(%[[ARG:.*]]: i32, %{{.*}}: i32):
+// CHECK-NEXT: "test.isolated_region_yield"(%[[ARG]])
+func.func @isolated_dead_input(%x: i32) -> i32 {
+ %r = "test.isolated_region_branch"(%x, %x) ({
+ ^bb0(%arg: i32, %dead: i32):
+ "test.isolated_region_yield"(%arg) : (i32) -> ()
+ }) : (i32, i32) -> i32
+ return %r : i32
+}
diff --git a/mlir/test/lib/Dialect/Test/TestOpDefs.cpp b/mlir/test/lib/Dialect/Test/TestOpDefs.cpp
index 39a0c8eb213c69..9fe2168e117374 100644
--- a/mlir/test/lib/Dialect/Test/TestOpDefs.cpp
+++ b/mlir/test/lib/Dialect/Test/TestOpDefs.cpp
@@ -1390,6 +1390,12 @@ IsolatedRegionBranchOp::getEntrySuccessorOperands(RegionSuccessor successor) {
return getInputs();
}
+void IsolatedRegionBranchOp::getCanonicalizationPatterns(
+ RewritePatternSet &patterns, MLIRContext *context) {
+ populateRegionBranchOpInterfaceCanonicalizationPatterns(patterns,
+ getOperationName());
+}
+
MutableOperandRange
IsolatedRegionYieldOp::getMutableSuccessorOperands(RegionSuccessor successor) {
return getValuesMutable();
diff --git a/mlir/test/lib/Dialect/Test/TestOps.td b/mlir/test/lib/Dialect/Test/TestOps.td
index d9420346c8d13a..c3e06dc968a270 100644
--- a/mlir/test/lib/Dialect/Test/TestOps.td
+++ b/mlir/test/lib/Dialect/Test/TestOps.td
@@ -2835,6 +2835,7 @@ def IsolatedRegionBranchOp : TEST_Op<"isolated_region_branch", [
let arguments = (ins Variadic<AnyType>:$inputs);
let results = (outs Variadic<AnyType>:$outputs);
let regions = (region VariadicRegion<SizedRegion<1>>:$branches);
+ let hasCanonicalizer = 1;
}
def IsolatedRegionYieldOp : TEST_Op<"isolated_region_yield", [
>From 9cedb73e1f7d9d7b8aa0c94dcaa3e66f9b1c69bd Mon Sep 17 00:00:00 2001
From: Ivan Butygin <ivan.butygin at gmail.com>
Date: Sat, 19 Sep 2026 12:37:59 +0200
Subject: [PATCH 3/5] [MLIR] Explain pass-root safety and hoist the isolation
check
---
mlir/lib/Analysis/DataFlow/LivenessAnalysis.cpp | 4 +++-
mlir/lib/Interfaces/ControlFlowInterfaces.cpp | 8 +++++---
mlir/lib/Transforms/RemoveDeadValues.cpp | 8 +++++++-
3 files changed, 15 insertions(+), 5 deletions(-)
diff --git a/mlir/lib/Analysis/DataFlow/LivenessAnalysis.cpp b/mlir/lib/Analysis/DataFlow/LivenessAnalysis.cpp
index 42cf874cc96e81..706f89dcd7c5d1 100644
--- a/mlir/lib/Analysis/DataFlow/LivenessAnalysis.cpp
+++ b/mlir/lib/Analysis/DataFlow/LivenessAnalysis.cpp
@@ -49,7 +49,9 @@ ChangeResult Liveness::meet(const AbstractSparseLattice &other) {
//===----------------------------------------------------------------------===//
LogicalResult LivenessAnalysis::initialize(Operation *top) {
- // Users outside the analysis root are not visited.
+ // Users outside the analysis root are not visited. Mark all root results live
+ // so that this missing information does not make their producers dead. For a
+ // region branch root, this keeps the yielded values and their producers live.
for (Value result : top->getResults())
setToExitState(getLatticeElement(result));
return SparseBackwardDataFlowAnalysis<Liveness>::initialize(top);
diff --git a/mlir/lib/Interfaces/ControlFlowInterfaces.cpp b/mlir/lib/Interfaces/ControlFlowInterfaces.cpp
index 0ea76fc75bd074..2055d0edff80ec 100644
--- a/mlir/lib/Interfaces/ControlFlowInterfaces.cpp
+++ b/mlir/lib/Interfaces/ControlFlowInterfaces.cpp
@@ -729,10 +729,12 @@ struct MakeRegionBranchOpSuccessorInputsDead : public RewritePattern {
// 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()) {
- // Isolated entry arguments cannot capture values from above the op.
- if (op->hasTrait<OpTrait::IsIsolatedFromAbove>() &&
- isa<BlockArgument>(value))
+ // Replacing an isolated entry argument with an outside operand would
+ // create an illegal capture, even if that operand dominates the op.
+ // Result replacements remain subject to the dominance check below.
+ if (isIsolated && isa<BlockArgument>(value))
continue;
// Nothing to do for successor inputs that are already dead.
if (value.use_empty())
diff --git a/mlir/lib/Transforms/RemoveDeadValues.cpp b/mlir/lib/Transforms/RemoveDeadValues.cpp
index fd856711d8edf4..35e8d306132b38 100644
--- a/mlir/lib/Transforms/RemoveDeadValues.cpp
+++ b/mlir/lib/Transforms/RemoveDeadValues.cpp
@@ -804,6 +804,9 @@ void RemoveDeadValues::runOnOperation() {
RDVFinalCleanupList finalCleanupList;
root->walk([&](Operation *op) {
+ // Do not erase the pass root or change its operands and results. In
+ // particular, the symbol user map cannot see callers outside the root, so
+ // changing a root function's signature would leave those calls invalid.
if (op == root)
return;
if (auto funcOp = dyn_cast<FunctionOpInterface>(op)) {
@@ -829,6 +832,8 @@ void RemoveDeadValues::runOnOperation() {
if (!canonicalize)
return;
+ // Process each root region separately. For operations in different root
+ // regions, the inferred common scope can contain the root itself.
for (Region ®ion : root->getRegions()) {
SmallVector<Operation *> opsToCanonicalize;
region.walk([&](RegionBranchOpInterface regionBranchOp) {
@@ -840,7 +845,8 @@ void RemoveDeadValues::runOnOperation() {
if (std::optional<RegisteredOperationName> info = op->getRegisteredInfo())
if (populatedPatterns.insert(*info).second)
info->getCanonicalizationPatterns(owningPatterns, context);
- // Keep greedy worklist expansion below the pass root.
+ // The greedy driver can add ancestors to its worklist. Set an explicit
+ // scope so that it cannot rewrite the root or operations outside the root.
GreedyRewriteConfig config;
config.setScope(®ion);
if (failed(applyOpPatternsGreedily(opsToCanonicalize,
>From c4530cc6a0f921e606e5b5f4649b9ee6df5b7f68 Mon Sep 17 00:00:00 2001
From: Ivan Butygin <ivan.butygin at gmail.com>
Date: Sat, 19 Sep 2026 12:54:25 +0200
Subject: [PATCH 4/5] [MLIR] Reuse dead-value cleanup patterns across root
regions
---
mlir/lib/Transforms/RemoveDeadValues.cpp | 22 ++++++++++++-------
.../remove-dead-values-nested-root.mlir | 13 +++++++++++
2 files changed, 27 insertions(+), 8 deletions(-)
diff --git a/mlir/lib/Transforms/RemoveDeadValues.cpp b/mlir/lib/Transforms/RemoveDeadValues.cpp
index 35e8d306132b38..fe7eed7ce54c69 100644
--- a/mlir/lib/Transforms/RemoveDeadValues.cpp
+++ b/mlir/lib/Transforms/RemoveDeadValues.cpp
@@ -832,6 +832,19 @@ void RemoveDeadValues::runOnOperation() {
if (!canonicalize)
return;
+ // Collect and freeze the patterns once for all root regions.
+ RewritePatternSet owningPatterns(context);
+ DenseSet<RegisteredOperationName> populatedPatterns;
+ root->walk([&](RegionBranchOpInterface regionBranchOp) {
+ Operation *op = regionBranchOp.getOperation();
+ if (op == root)
+ return;
+ if (std::optional<RegisteredOperationName> info = op->getRegisteredInfo())
+ if (populatedPatterns.insert(*info).second)
+ info->getCanonicalizationPatterns(owningPatterns, context);
+ });
+ FrozenRewritePatternSet patterns(std::move(owningPatterns));
+
// Process each root region separately. For operations in different root
// regions, the inferred common scope can contain the root itself.
for (Region ®ion : root->getRegions()) {
@@ -839,18 +852,11 @@ void RemoveDeadValues::runOnOperation() {
region.walk([&](RegionBranchOpInterface regionBranchOp) {
opsToCanonicalize.push_back(regionBranchOp.getOperation());
});
- RewritePatternSet owningPatterns(context);
- DenseSet<RegisteredOperationName> populatedPatterns;
- for (Operation *op : opsToCanonicalize)
- if (std::optional<RegisteredOperationName> info = op->getRegisteredInfo())
- if (populatedPatterns.insert(*info).second)
- info->getCanonicalizationPatterns(owningPatterns, context);
// The greedy driver can add ancestors to its worklist. Set an explicit
// scope so that it cannot rewrite the root or operations outside the root.
GreedyRewriteConfig config;
config.setScope(®ion);
- if (failed(applyOpPatternsGreedily(opsToCanonicalize,
- std::move(owningPatterns), config))) {
+ if (failed(applyOpPatternsGreedily(opsToCanonicalize, patterns, config))) {
root->emitError("greedy pattern rewrite failed to converge");
return signalPassFailure();
}
diff --git a/mlir/test/Transforms/remove-dead-values-nested-root.mlir b/mlir/test/Transforms/remove-dead-values-nested-root.mlir
index 7e373386019dca..20b6b4cc99657c 100644
--- a/mlir/test/Transforms/remove-dead-values-nested-root.mlir
+++ b/mlir/test/Transforms/remove-dead-values-nested-root.mlir
@@ -76,13 +76,26 @@ func.func @effects(%mem: memref<i32>, %x: i32) {
return
}
+// Remove dead loop results in both root regions with the shared patterns.
// CANON-LABEL: func.func @dead_carry
// CANON: "test.isolated_region_branch"
// CANON: scf.for {{.*}} iter_args(%[[CARRY:.*]] = %{{.*}}) -> (i32)
// CANON: scf.yield %{{.*}} : i32
// CANON: "test.isolated_region_yield"
+// CANON: scf.for {{.*}} iter_args(%{{.*}} = %{{.*}}) -> (i32)
+// CANON: scf.yield %{{.*}} : i32
+// CANON: "test.isolated_region_yield"
func.func @dead_carry(%lb: index, %ub: index, %step: index, %x: i32) -> i32 {
%r = "test.isolated_region_branch"(%lb, %ub, %step, %x) ({
+ ^bb0(%lower: index, %upper: index, %stride: index, %init: i32):
+ %loop:2 = scf.for %iv = %lower to %upper step %stride
+ iter_args(%live = %init, %dead = %init) -> (i32, i32) {
+ %sum = arith.addi %live, %live : i32
+ %unused = arith.muli %dead, %dead : i32
+ scf.yield %sum, %unused : i32, i32
+ }
+ "test.isolated_region_yield"(%loop#0) : (i32) -> ()
+ }, {
^bb0(%lower: index, %upper: index, %stride: index, %init: i32):
%loop:2 = scf.for %iv = %lower to %upper step %stride
iter_args(%live = %init, %dead = %init) -> (i32, i32) {
>From 56683df2501be519e13552de2d7fb25d8886e5bb Mon Sep 17 00:00:00 2001
From: Ivan Butygin <ivan.butygin at gmail.com>
Date: Sat, 19 Sep 2026 13:10:16 +0200
Subject: [PATCH 5/5] [MLIR] Check replacement regions at isolation boundaries
---
mlir/lib/Interfaces/ControlFlowInterfaces.cpp | 22 ++++---
.../isolated-region-branch-canonicalize.mlir | 59 +++++++++++++++++++
2 files changed, 73 insertions(+), 8 deletions(-)
diff --git a/mlir/lib/Interfaces/ControlFlowInterfaces.cpp b/mlir/lib/Interfaces/ControlFlowInterfaces.cpp
index 2055d0edff80ec..dc9068aa7a8d87 100644
--- a/mlir/lib/Interfaces/ControlFlowInterfaces.cpp
+++ b/mlir/lib/Interfaces/ControlFlowInterfaces.cpp
@@ -731,11 +731,6 @@ struct MakeRegionBranchOpSuccessorInputsDead : public RewritePattern {
bool changed = false;
const bool isIsolated = op->hasTrait<OpTrait::IsIsolatedFromAbove>();
for (Value value : inputToOperands.keys()) {
- // Replacing an isolated entry argument with an outside operand would
- // create an illegal capture, even if that operand dominates the op.
- // Result replacements remain subject to the dominance check below.
- if (isIsolated && isa<BlockArgument>(value))
- continue;
// Nothing to do for successor inputs that are already dead.
if (value.use_empty())
continue;
@@ -747,8 +742,19 @@ struct MakeRegionBranchOpSuccessorInputsDead : public RewritePattern {
/*maxReachableValues=*/1)) ||
reachableValues.empty())
continue;
- assert(*reachableValues.begin() != value &&
+ Value replacement = *reachableValues.begin();
+ assert(replacement != value &&
"successor inputs are supposed to be excluded");
+ // A value inside an isolated region cannot be replaced with a value from
+ // another region, even if the replacement dominates its uses. Op results
+ // are in the parent region and do not cross this isolation boundary.
+ // Successor inputs are direct region arguments or results of this op.
+ // Valid input IR already prevents captures across nested isolation
+ // boundaries, so only this op's boundary needs an additional check.
+ Region *valueRegion = value.getParentRegion();
+ if (isIsolated && valueRegion->getParentOp() == op &&
+ replacement.getParentRegion() != valueRegion)
+ continue;
// Do not replace `value` with the found reachable value if doing so
// would violate dominance. Example:
// %r = scf.execute_region ... {
@@ -758,9 +764,9 @@ struct MakeRegionBranchOpSuccessorInputsDead : public RewritePattern {
// use(%r)
// In the above example, reachableValues(%r) = {%a}, but %a cannot be
// used as a replacement for %r due to dominance / scope.
- if (!isDefinedBefore(regionBranchOp, *reachableValues.begin(), value))
+ if (!isDefinedBefore(regionBranchOp, replacement, value))
continue;
- rewriter.replaceAllUsesWith(value, *reachableValues.begin());
+ rewriter.replaceAllUsesWith(value, replacement);
changed = true;
}
return success(changed);
diff --git a/mlir/test/Transforms/isolated-region-branch-canonicalize.mlir b/mlir/test/Transforms/isolated-region-branch-canonicalize.mlir
index e872a8cebe6e0d..50ed4495d848a2 100644
--- a/mlir/test/Transforms/isolated-region-branch-canonicalize.mlir
+++ b/mlir/test/Transforms/isolated-region-branch-canonicalize.mlir
@@ -63,6 +63,65 @@ func.func @forwarded_result(%x: i32) -> i32 {
return %r : i32
}
+// An op result can be replaced with an input from an ancestor region. Keep the
+// region argument in the isolated body so that its uses do not capture %x.
+// CHECK-LABEL: func.func @forwarded_ancestor_result
+// CHECK-SAME: %{{.*}}: i1, %[[X:.*]]: i32
+// CHECK: %[[SELECTED:.*]] = arith.select %{{.*}}, %[[X]], %{{.*}} : i32
+// CHECK: scf.if
+// CHECK: "test.isolated_region_branch"(%[[X]]) ({
+// CHECK: ^bb0(%[[ARG:.*]]: i32):
+// CHECK: call @sink(%[[ARG]])
+// CHECK: "test.isolated_region_yield"() : () -> ()
+// CHECK: }) : (i32) -> ()
+// CHECK: return %[[SELECTED]] : i32
+func.func @forwarded_ancestor_result(%cond: i1, %x: i32) -> i32 {
+ %r = scf.if %cond -> i32 {
+ %forwarded = "test.isolated_region_branch"(%x) ({
+ ^bb0(%arg: i32):
+ func.call @sink(%arg) : (i32) -> ()
+ "test.isolated_region_yield"(%arg) : (i32) -> ()
+ }) : (i32) -> i32
+ scf.yield %forwarded : i32
+ } else {
+ %zero = arith.constant 0 : i32
+ scf.yield %zero : i32
+ }
+ return %r : i32
+}
+
+// Keep both isolation boundaries when an inner isolated op is nested in an if.
+// CHECK-LABEL: func.func @nested_isolation
+// CHECK-SAME: %[[COND:.*]]: i1, %[[X:.*]]: i32
+// CHECK: "test.isolated_region_branch"(%[[COND]], %[[X]]) ({
+// CHECK: ^bb0(%[[C:.*]]: i1, %[[A:.*]]: i32):
+// CHECK: scf.if %[[C]]
+// CHECK: "test.isolated_region_branch"(%[[A]]) ({
+// CHECK: ^bb0(%[[B:.*]]: i32):
+// CHECK: call @sink(%[[B]])
+// CHECK: "test.isolated_region_yield"() : () -> ()
+// CHECK: }) : (i32) -> ()
+// CHECK: "test.isolated_region_yield"() : () -> ()
+// CHECK: }) : (i1, i32) -> ()
+// CHECK: return %[[X]] : i32
+func.func @nested_isolation(%cond: i1, %x: i32) -> i32 {
+ %r = "test.isolated_region_branch"(%cond, %x) ({
+ ^bb0(%c: i1, %a: i32):
+ %selected = scf.if %c -> i32 {
+ %inner = "test.isolated_region_branch"(%a) ({
+ ^bb0(%b: i32):
+ func.call @sink(%b) : (i32) -> ()
+ "test.isolated_region_yield"(%b) : (i32) -> ()
+ }) : (i32) -> i32
+ scf.yield %inner : i32
+ } else {
+ scf.yield %a : i32
+ }
+ "test.isolated_region_yield"(%selected) : (i32) -> ()
+ }) : (i1, i32) -> i32
+ return %r : i32
+}
+
// Keep replacement arguments in their own regions. Remove a shared input only
// when its arguments are dead in all regions.
// CHECK-LABEL: func.func @multiple_regions
More information about the Mlir-commits
mailing list