[Mlir-commits] [mlir] [mlir][scf] Fix WhileMoveIfDown with duplicated scf.condition operands (PR #219458)
Samarth Narang
llvmlistbot at llvm.org
Sat Sep 5 02:56:19 PDT 2026
https://github.com/snarang181 updated https://github.com/llvm/llvm-project/pull/219458
>From fb836d80a43743fa9a9f62fe1bb88d998036a641 Mon Sep 17 00:00:00 2001
From: Samarth Narang <sn2910 at nyu.edu>
Date: Fri, 28 Aug 2026 12:49:41 +0000
Subject: [PATCH 1/2] [mlir][scf] Fix WhileMoveIfDown with duplicated
scf.condition operands
WhileMoveIfDown replaces an scf.if in the before region of an scf.while by
forwarding the then value to the after-region arguments and the else value to
the while results.
The rewrite classified conditionOp.getArgs() one position at a time while
already mutating them: replaceAllUsesWith(ifResult, elseValue) rewrites every
use at once, including condition operands not yet visited. When the same
scf.if result is forwarded to several condition operands - which is legal, as
neither the scf.condition verifier nor the region-branch contract requires
distinct operands - the first position rewrote all copies to the else value,
so later positions were no longer recognized as scf.if results and their
after-region arguments kept the else value even though the after region only
executes when the condition is true. The result was verifier-clean IR that
silently computed wrong values.
Classify all condition operands before mutating: record the operand-to-result
mapping first, then replace the after-region arguments (distinct block
arguments, so they cannot interfere), and only then collapse each scf.if
result to its else value once.
Fixes https://github.com/llvm/llvm-project/issues/219456
---
mlir/lib/Dialect/SCF/IR/SCF.cpp | 30 +++++++++++++-----
mlir/test/Dialect/SCF/canonicalize.mlir | 41 +++++++++++++++++++++++++
2 files changed, 63 insertions(+), 8 deletions(-)
diff --git a/mlir/lib/Dialect/SCF/IR/SCF.cpp b/mlir/lib/Dialect/SCF/IR/SCF.cpp
index 3388fa490f996..78e621e1a2960 100644
--- a/mlir/lib/Dialect/SCF/IR/SCF.cpp
+++ b/mlir/lib/Dialect/SCF/IR/SCF.cpp
@@ -3474,18 +3474,32 @@ struct WhileMoveIfDown : public OpRewritePattern<scf::WhileOp> {
// Replace uses of ifOp results in the conditionOp with the yielded values
// from the ifOp branches.
+ //
+ // The same ifOp result may be forwarded to several condition operands, so
+ // classify every operand before mutating anything: replacing an ifOp result
+ // rewrites *all* of its uses at once, including condition operands that
+ // have not been visited yet, which would hide them from this scan.
+ SmallVector<std::pair<size_t, size_t>> conditionToIfResult;
for (auto [idx, arg] : llvm::enumerate(conditionOp.getArgs())) {
auto it = llvm::find(ifOp->getResults(), arg);
- if (it != ifOp->getResults().end()) {
- size_t ifOpIdx = it.getIndex();
- Value thenValue = ifOp.thenYield()->getOperand(ifOpIdx);
- Value elseValue = ifOp.elseYield()->getOperand(ifOpIdx);
-
- rewriter.replaceAllUsesWith(ifOp->getResults()[ifOpIdx], elseValue);
- rewriter.replaceAllUsesWith(op.getAfterArguments()[idx], thenValue);
- }
+ if (it != ifOp->getResults().end())
+ conditionToIfResult.emplace_back(idx, it.getIndex());
}
+ // The after-region arguments are distinct block arguments, so these
+ // replacements cannot interfere with one another.
+ for (auto [idx, ifOpIdx] : conditionToIfResult)
+ rewriter.replaceAllUsesWith(op.getAfterArguments()[idx],
+ ifOp.thenYield()->getOperand(ifOpIdx));
+
+ // Any remaining use of an ifOp result is on the false path, i.e. a result
+ // of the while op; collapse each one to its else value exactly once.
+ llvm::SmallDenseSet<size_t> collapsedIfResults;
+ for (auto [idx, ifOpIdx] : conditionToIfResult)
+ if (collapsedIfResults.insert(ifOpIdx).second)
+ rewriter.replaceAllUsesWith(ifOp->getResults()[ifOpIdx],
+ ifOp.elseYield()->getOperand(ifOpIdx));
+
// Collect additional used values from before region.
SetVector<Value> additionalUsedValuesSet;
visitUsedValuesDefinedAbove(ifOp.getThenRegion(), [&](OpOperand *operand) {
diff --git a/mlir/test/Dialect/SCF/canonicalize.mlir b/mlir/test/Dialect/SCF/canonicalize.mlir
index e89ad45867a57..4766b1f1a7f45 100644
--- a/mlir/test/Dialect/SCF/canonicalize.mlir
+++ b/mlir/test/Dialect/SCF/canonicalize.mlir
@@ -1110,6 +1110,47 @@ func.func @while_move_if_down() -> i32 {
// -----
+// The same scf.if result may be forwarded to several scf.condition operands.
+// Every corresponding after-region argument must receive the *then* value,
+// while the while op's results receive the *else* value.
+
+// CHECK-LABEL: @while_move_if_down_duplicate_forward
+func.func @while_move_if_down_duplicate_forward() -> i32 {
+ %0:2 = scf.while () : () -> (i32, i32) {
+ %else_value = "test.get_some_value0" () : () -> (i32)
+ %condition = "test.condition"() : () -> i1
+ %res = scf.if %condition -> (i32) {
+ %then_value = "test.get_some_value1" () : () -> (i32)
+ scf.yield %then_value : i32
+ } else {
+ scf.yield %else_value : i32
+ }
+ scf.condition(%condition) %res, %res : i32, i32
+ } do {
+ ^bb0(%first: i32, %second: i32):
+ "test.use0" (%first) : (i32) -> ()
+ "test.use1" (%second) : (i32) -> ()
+ scf.yield
+ }
+ return %0#1 : i32
+}
+// CHECK: %[[WHILE_RES:.*]] = scf.while : () -> i32 {
+// CHECK: %[[else_value:.*]] = "test.get_some_value0"() : () -> i32
+// CHECK: %[[condition:.*]] = "test.condition"() : () -> i1
+// The while result is the else value.
+// CHECK: scf.condition(%[[condition]]) %[[else_value]] : i32
+// CHECK: } do {
+// CHECK: ^bb0(%{{.*}}: i32):
+// CHECK: %[[then_value:.*]] = "test.get_some_value1"() : () -> i32
+// Both forwarded positions must observe the then value, not the else value.
+// CHECK: "test.use0"(%[[then_value]]) : (i32) -> ()
+// CHECK: "test.use1"(%[[then_value]]) : (i32) -> ()
+// CHECK: scf.yield
+// CHECK: }
+// CHECK: return %[[WHILE_RES]] : i32
+
+// -----
+
// CHECK-LABEL: @while_cond_true
func.func @while_cond_true() -> i1 {
%0 = scf.while () : () -> i1 {
>From cf4be2515e755f43ac3d1f503a6757153e7d4411 Mon Sep 17 00:00:00 2001
From: Samarth Narang <sn2910 at nyu.edu>
Date: Mon, 31 Aug 2026 20:32:25 +0000
Subject: [PATCH 2/2] Address review: drop the redundant set and document why
the replacement is safe
The tracking set is unnecessary: the first replaceAllUsesWith leaves the ifOp
result without uses, so repeating it for a result forwarded to several
condition operands is a no-op. Replace the third loop with an unconditional
call and note that ifOp is conditionOp->getPrevNode() and, per the assertion
above, conditionOp is its only user, which is what makes replacing all uses
safe here.
---
mlir/lib/Dialect/SCF/IR/SCF.cpp | 12 +++++++-----
1 file changed, 7 insertions(+), 5 deletions(-)
diff --git a/mlir/lib/Dialect/SCF/IR/SCF.cpp b/mlir/lib/Dialect/SCF/IR/SCF.cpp
index 78e621e1a2960..2d2702c471b2e 100644
--- a/mlir/lib/Dialect/SCF/IR/SCF.cpp
+++ b/mlir/lib/Dialect/SCF/IR/SCF.cpp
@@ -3493,12 +3493,14 @@ struct WhileMoveIfDown : public OpRewritePattern<scf::WhileOp> {
ifOp.thenYield()->getOperand(ifOpIdx));
// Any remaining use of an ifOp result is on the false path, i.e. a result
- // of the while op; collapse each one to its else value exactly once.
- llvm::SmallDenseSet<size_t> collapsedIfResults;
+ // of the while op, so it becomes the else value. `ifOp` is
+ // `conditionOp->getPrevNode()` and the assertion above establishes that
+ // `conditionOp` is its only user, so replacing all uses is safe here.
+ // Repeating a replacement for a result forwarded more than once is a no-op
+ // because the first one leaves it without uses.
for (auto [idx, ifOpIdx] : conditionToIfResult)
- if (collapsedIfResults.insert(ifOpIdx).second)
- rewriter.replaceAllUsesWith(ifOp->getResults()[ifOpIdx],
- ifOp.elseYield()->getOperand(ifOpIdx));
+ rewriter.replaceAllUsesWith(ifOp->getResults()[ifOpIdx],
+ ifOp.elseYield()->getOperand(ifOpIdx));
// Collect additional used values from before region.
SetVector<Value> additionalUsedValuesSet;
More information about the Mlir-commits
mailing list