[Mlir-commits] [mlir] [mlir][scf] Fix WhileMoveIfDown with duplicated scf.condition operands (PR #219458)

[llvmlistbot at llvm.org]

llvmorg-github-actions[bot] wrote:


<!--LLVM PR SUMMARY COMMENT-->

@llvm/pr-subscribers-mlir-scf

Author: Samarth Narang (snarang181)

<details>
<summary>Changes</summary>

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

---
Full diff: https://github.com/llvm/llvm-project/pull/219458.diff


2 Files Affected:

- (modified) mlir/lib/Dialect/SCF/IR/SCF.cpp (+22-8) 
- (modified) mlir/test/Dialect/SCF/canonicalize.mlir (+41) 


``````````diff
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 {

``````````

</details>


https://github.com/llvm/llvm-project/pull/219458