[Mlir-commits] [mlir] [mlir][scf] Fix WhileMoveIfDown with duplicated scf.condition operands (PR #219458)
Samarth Narang
llvmlistbot at llvm.org
Sat Sep 5 11:00:13 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/3] [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/3] 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;
>From 424fc7798bc3344a09537b1e2e95ef383f441d4b Mon Sep 17 00:00:00 2001
From: Samarth Narang <sn2910 at nyu.edu>
Date: Sat, 5 Sep 2026 17:59:54 +0000
Subject: [PATCH 3/3] Address review: assign the else value into the condition
operand directly
Rather than replacing all uses of the ifOp result and ordering the replacements
around the interference that causes, assign into the specific condition operand.
A result forwarded to several operands is then handled naturally, and the
two-phase classification is no longer needed.
---
mlir/lib/Dialect/SCF/IR/SCF.cpp | 36 +++++++++++++--------------------
1 file changed, 14 insertions(+), 22 deletions(-)
diff --git a/mlir/lib/Dialect/SCF/IR/SCF.cpp b/mlir/lib/Dialect/SCF/IR/SCF.cpp
index 2d2702c471b2e..d83fab0745040 100644
--- a/mlir/lib/Dialect/SCF/IR/SCF.cpp
+++ b/mlir/lib/Dialect/SCF/IR/SCF.cpp
@@ -3473,34 +3473,26 @@ struct WhileMoveIfDown : public OpRewritePattern<scf::WhileOp> {
Location loc = op.getLoc();
// Replace uses of ifOp results in the conditionOp with the yielded values
- // from the ifOp branches.
+ // from the ifOp branches: the after-region argument takes the `then` value,
+ // while the condition operand -- which becomes a while result once the
+ // condition is false -- takes the `else` value.
//
// 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;
+ // assign into the specific operand instead of replacing all uses of the
+ // ifOp result, which would also rewrite the operands not yet visited.
for (auto [idx, arg] : llvm::enumerate(conditionOp.getArgs())) {
auto it = llvm::find(ifOp->getResults(), arg);
- 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)
+ if (it == ifOp->getResults().end())
+ continue;
+ size_t ifOpIdx = it.getIndex();
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, 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)
- rewriter.replaceAllUsesWith(ifOp->getResults()[ifOpIdx],
- ifOp.elseYield()->getOperand(ifOpIdx));
+ unsigned argIdx = idx;
+ Value elseValue = ifOp.elseYield()->getOperand(ifOpIdx);
+ rewriter.modifyOpInPlace(conditionOp, [&] {
+ conditionOp.getArgsMutable()[argIdx].assign(elseValue);
+ });
+ }
// Collect additional used values from before region.
SetVector<Value> additionalUsedValuesSet;
More information about the Mlir-commits
mailing list