[Mlir-commits] [mlir] [MLIR] Fix control-flow sinking through nested regions (PR #217168)
llvmlistbot at llvm.org
llvmlistbot at llvm.org
Thu Aug 20 02:37:56 PDT 2026
https://github.com/MarkVeerasingam updated https://github.com/llvm/llvm-project/pull/217168
>From 034303fc1a07cb61f319100a2cc5146bab46b129 Mon Sep 17 00:00:00 2001
From: MarkVeerasingam <markveer70 at gmail.com>
Date: Wed, 19 Aug 2026 00:36:03 +0100
Subject: [PATCH 1/7] Fix control-flow sinking through nested regions
---
.../Transforms/Utils/ControlFlowSinkUtils.cpp | 23 ++++++++++--------
mlir/test/Dialect/SCF/control-flow-sink.mlir | 24 +++++++++++++++++++
2 files changed, 37 insertions(+), 10 deletions(-)
diff --git a/mlir/lib/Transforms/Utils/ControlFlowSinkUtils.cpp b/mlir/lib/Transforms/Utils/ControlFlowSinkUtils.cpp
index b546d3adef08f..3d30beeaef303 100644
--- a/mlir/lib/Transforms/Utils/ControlFlowSinkUtils.cpp
+++ b/mlir/lib/Transforms/Utils/ControlFlowSinkUtils.cpp
@@ -80,8 +80,11 @@ bool Sinker::allUsersDominatedBy(Operation *op, Region *region) {
"expected op to be defined outside the region");
return llvm::all_of(op->getUsers(), [&](Operation *user) {
// The user is dominated by the region if its containing block is dominated
- // by the region's entry block.
- return domInfo.dominates(®ion->front(), user->getBlock());
+ // by the region's entry block. Additionally, allow users that are in
+ // descendant regions of the region (e.g., nested loops) since those
+ // should still permit sinking into the outer region.
+ return domInfo.dominates(®ion->front(), user->getBlock()) ||
+ region->isAncestor(user->getParentRegion());
});
}
@@ -91,12 +94,13 @@ void Sinker::tryToSinkPredecessors(Operation *user, Region *region,
<< OpWithFlags(user, OpPrintingFlags().skipRegions());
for (Value value : user->getOperands()) {
Operation *op = value.getDefiningOp();
- // Ignore block arguments and ops that are already inside the region.
- if (!op || op->getParentRegion() == region)
+ // Ignore block arguments and ops already contained in the target region,
+ // including ops in nested regions. Only consider defs outside the target
+ // region.
+ if (!op || region->isAncestor(op->getParentRegion()))
continue;
LDBG() << "Try to sink:\n"
<< OpWithFlags(op, OpPrintingFlags().skipRegions());
-
// If the op's users are all in the region and it can be moved, then do so.
if (allUsersDominatedBy(op, region) && shouldMoveIntoRegion(op, region)) {
moveIntoRegion(op, region);
@@ -108,12 +112,11 @@ void Sinker::tryToSinkPredecessors(Operation *user, Region *region,
}
void Sinker::sinkRegion(Region *region) {
- // Initialize the work queue with all the ops in the region.
+ // Initialize the work queue with all the ops in the region, including
+ // nested regions. Seed from all operations so that uses in non-entry and
+ // nested blocks trigger sinking.
std::vector<Operation *> stack;
- for (Block &block : *region)
- // Seed from all block so that uses in non-entry blocks and trigger sinking
- for (Operation &op : block)
- stack.push_back(&op);
+ region->walk([&](Operation *op) { stack.push_back(op); });
// Process all the ops depth-first. This ensures that nodes of subgraphs are
// sunk in the correct order.
diff --git a/mlir/test/Dialect/SCF/control-flow-sink.mlir b/mlir/test/Dialect/SCF/control-flow-sink.mlir
index 6a3e2b651ee8a..ea3a3a678969d 100644
--- a/mlir/test/Dialect/SCF/control-flow-sink.mlir
+++ b/mlir/test/Dialect/SCF/control-flow-sink.mlir
@@ -80,3 +80,27 @@ func.func @test_scf_execute_region_multiblock_sink(%arg0: i32, %arg1: i32) {
}
return
}
+
+// -----
+
+func.func private @sink_i32(i32)
+
+// CHECK-LABEL: @test_scf_if_sink_through_loop
+// CHECK-SAME: (%[[ARG0:.*]]: i1, %[[ARG1:.*]]: index, %[[ARG2:.*]]: i32, %[[ARG3:.*]]: i32)
+// CHECK: scf.if %[[ARG0]]
+// CHECK: %[[V0:.*]] = arith.muli %[[ARG2]], %[[ARG3]]
+// CHECK: scf.for
+// CHECK: call @sink_i32(%[[V0]])
+
+func.func @test_scf_if_sink_through_loop(
+ %arg0: i1, %arg1: index, %arg2: i32, %arg3: i32) {
+ %c0 = arith.constant 0 : index
+ %c1 = arith.constant 1 : index
+ %0 = arith.muli %arg2, %arg3 : i32
+ scf.if %arg0 {
+ scf.for %arg4 = %c0 to %arg1 step %c1 {
+ func.call @sink_i32(%0) : (i32) -> ()
+ }
+ }
+ return
+}
>From 7480e202a69f70beede23c006e00ab076832ad91 Mon Sep 17 00:00:00 2001
From: MarkVeerasingam <markveer70 at gmail.com>
Date: Wed, 19 Aug 2026 14:23:49 +0100
Subject: [PATCH 2/7] Address control-flow sink review feedback
---
mlir/lib/Transforms/Utils/ControlFlowSinkUtils.cpp | 7 ++-----
1 file changed, 2 insertions(+), 5 deletions(-)
diff --git a/mlir/lib/Transforms/Utils/ControlFlowSinkUtils.cpp b/mlir/lib/Transforms/Utils/ControlFlowSinkUtils.cpp
index 3d30beeaef303..9bdc25b10ba95 100644
--- a/mlir/lib/Transforms/Utils/ControlFlowSinkUtils.cpp
+++ b/mlir/lib/Transforms/Utils/ControlFlowSinkUtils.cpp
@@ -80,11 +80,8 @@ bool Sinker::allUsersDominatedBy(Operation *op, Region *region) {
"expected op to be defined outside the region");
return llvm::all_of(op->getUsers(), [&](Operation *user) {
// The user is dominated by the region if its containing block is dominated
- // by the region's entry block. Additionally, allow users that are in
- // descendant regions of the region (e.g., nested loops) since those
- // should still permit sinking into the outer region.
- return domInfo.dominates(®ion->front(), user->getBlock()) ||
- region->isAncestor(user->getParentRegion());
+ // by the region's entry block
+ return domInfo.dominates(®ion->front(), user->getBlock());
});
}
>From 440bc0fbe3815abd5caf9d5363b39b84ccc5e16e Mon Sep 17 00:00:00 2001
From: MarkVeerasingam <markveer70 at gmail.com>
Date: Wed, 19 Aug 2026 14:30:20 +0100
Subject: [PATCH 3/7] Add control-flow sink regression test
---
mlir/test/Dialect/SCF/control-flow-sink.mlir | 21 ++++++++++++++++++++
1 file changed, 21 insertions(+)
diff --git a/mlir/test/Dialect/SCF/control-flow-sink.mlir b/mlir/test/Dialect/SCF/control-flow-sink.mlir
index ea3a3a678969d..083591f16d615 100644
--- a/mlir/test/Dialect/SCF/control-flow-sink.mlir
+++ b/mlir/test/Dialect/SCF/control-flow-sink.mlir
@@ -104,3 +104,24 @@ func.func @test_scf_if_sink_through_loop(
}
return
}
+
+// CHECK-LABEL: @test_scf_if_sink_through_loop_with_external_use
+// CHECK: %[[V0:.*]] = arith.muli %{{.*}}, %{{.*}}
+// CHECK: scf.if
+// CHECK: scf.for
+// CHECK: call @sink_i32(%[[V0]])
+// CHECK: call @sink_i32(%[[V0]])
+
+func.func @test_scf_if_sink_through_loop_with_external_use(
+ %arg0: i1, %arg1: index, %arg2: i32, %arg3: i32) {
+ %c0 = arith.constant 0 : index
+ %c1 = arith.constant 1 : index
+ %0 = arith.muli %arg2, %arg3 : i32
+ scf.if %arg0 {
+ scf.for %arg4 = %c0 to %arg1 step %c1 {
+ func.call @sink_i32(%0) : (i32) -> ()
+ }
+ }
+ func.call @sink_i32(%0) : (i32) -> ()
+ return
+}
>From 0ed47276e1118bd0cc98a2df5795ae173cb8aa4d Mon Sep 17 00:00:00 2001
From: MarkVeerasingam <markveer70 at gmail.com>
Date: Wed, 19 Aug 2026 15:36:07 +0100
Subject: [PATCH 4/7] Clean up control-flow sink comments
---
mlir/lib/Transforms/Utils/ControlFlowSinkUtils.cpp | 2 +-
1 file changed, 1 insertion(+), 1 deletion(-)
diff --git a/mlir/lib/Transforms/Utils/ControlFlowSinkUtils.cpp b/mlir/lib/Transforms/Utils/ControlFlowSinkUtils.cpp
index 9bdc25b10ba95..c045d414caadc 100644
--- a/mlir/lib/Transforms/Utils/ControlFlowSinkUtils.cpp
+++ b/mlir/lib/Transforms/Utils/ControlFlowSinkUtils.cpp
@@ -80,7 +80,7 @@ bool Sinker::allUsersDominatedBy(Operation *op, Region *region) {
"expected op to be defined outside the region");
return llvm::all_of(op->getUsers(), [&](Operation *user) {
// The user is dominated by the region if its containing block is dominated
- // by the region's entry block
+ // by the region's entry block.
return domInfo.dominates(®ion->front(), user->getBlock());
});
}
>From 9e9dd05ca138fbbea6418530f85a31ae6b9692a6 Mon Sep 17 00:00:00 2001
From: MarkVeerasingam <122292447+MarkVeerasingam at users.noreply.github.com>
Date: Wed, 19 Aug 2026 15:36:51 +0100
Subject: [PATCH 5/7] Apply suggestion from @FedericoBruzzone
Co-authored-by: Federico Bruzzone <federico.bruzzone.i at gmail.com>
---
mlir/lib/Transforms/Utils/ControlFlowSinkUtils.cpp | 2 +-
1 file changed, 1 insertion(+), 1 deletion(-)
diff --git a/mlir/lib/Transforms/Utils/ControlFlowSinkUtils.cpp b/mlir/lib/Transforms/Utils/ControlFlowSinkUtils.cpp
index c045d414caadc..b72c5584e8e16 100644
--- a/mlir/lib/Transforms/Utils/ControlFlowSinkUtils.cpp
+++ b/mlir/lib/Transforms/Utils/ControlFlowSinkUtils.cpp
@@ -92,7 +92,7 @@ void Sinker::tryToSinkPredecessors(Operation *user, Region *region,
for (Value value : user->getOperands()) {
Operation *op = value.getDefiningOp();
// Ignore block arguments and ops already contained in the target region,
- // including ops in nested regions. Only consider defs outside the target
+ // including ops in nested regions.
// region.
if (!op || region->isAncestor(op->getParentRegion()))
continue;
>From 2327167798bd3f754776a8367c736ea6d8729443 Mon Sep 17 00:00:00 2001
From: MarkVeerasingam <122292447+MarkVeerasingam at users.noreply.github.com>
Date: Wed, 19 Aug 2026 15:37:44 +0100
Subject: [PATCH 6/7] Update mlir/lib/Transforms/Utils/ControlFlowSinkUtils.cpp
Co-authored-by: Federico Bruzzone <federico.bruzzone.i at gmail.com>
---
mlir/lib/Transforms/Utils/ControlFlowSinkUtils.cpp | 3 +--
1 file changed, 1 insertion(+), 2 deletions(-)
diff --git a/mlir/lib/Transforms/Utils/ControlFlowSinkUtils.cpp b/mlir/lib/Transforms/Utils/ControlFlowSinkUtils.cpp
index b72c5584e8e16..beaaa1fc358b8 100644
--- a/mlir/lib/Transforms/Utils/ControlFlowSinkUtils.cpp
+++ b/mlir/lib/Transforms/Utils/ControlFlowSinkUtils.cpp
@@ -110,8 +110,7 @@ void Sinker::tryToSinkPredecessors(Operation *user, Region *region,
void Sinker::sinkRegion(Region *region) {
// Initialize the work queue with all the ops in the region, including
- // nested regions. Seed from all operations so that uses in non-entry and
- // nested blocks trigger sinking.
+ // nested regions.
std::vector<Operation *> stack;
region->walk([&](Operation *op) { stack.push_back(op); });
>From ca914b024e36e6325ef8989a528385bb2e2c2403 Mon Sep 17 00:00:00 2001
From: MarkVeerasingam <markveer70 at gmail.com>
Date: Thu, 20 Aug 2026 10:37:37 +0100
Subject: [PATCH 7/7] fixed redundant comment
---
mlir/lib/Transforms/Utils/ControlFlowSinkUtils.cpp | 1 -
1 file changed, 1 deletion(-)
diff --git a/mlir/lib/Transforms/Utils/ControlFlowSinkUtils.cpp b/mlir/lib/Transforms/Utils/ControlFlowSinkUtils.cpp
index beaaa1fc358b8..af633f2ce4a95 100644
--- a/mlir/lib/Transforms/Utils/ControlFlowSinkUtils.cpp
+++ b/mlir/lib/Transforms/Utils/ControlFlowSinkUtils.cpp
@@ -93,7 +93,6 @@ void Sinker::tryToSinkPredecessors(Operation *user, Region *region,
Operation *op = value.getDefiningOp();
// Ignore block arguments and ops already contained in the target region,
// including ops in nested regions.
- // region.
if (!op || region->isAncestor(op->getParentRegion()))
continue;
LDBG() << "Try to sink:\n"
More information about the Mlir-commits
mailing list