[Mlir-commits] [mlir] [MLIR][SCCP] Fix in-place folds leaking into IR during simulation (PR #213933)

[Berke Ates]

https://github.com/Berke-Ates updated https://github.com/llvm/llvm-project/pull/213933

>From 1110408464244b8e4380fda9a9eb6c88aed9312c Mon Sep 17 00:00:00 2001
From: Berke-Ates <berke at ates.ch>
Date: Tue, 4 Aug 2026 15:10:37 +0200
Subject: [PATCH] [MLIR][SCCP] Fix in-place folds leaking into IR during
 simulation

Co-Authored-By: Claude Fable 5 <noreply at anthropic.com>
---
 .../DataFlow/ConstantPropagationAnalysis.cpp  | 24 ++++++++--------
 mlir/test/Transforms/sccp.mlir                | 28 +++++++++++++++++++
 2 files changed, 41 insertions(+), 11 deletions(-)

diff --git a/mlir/lib/Analysis/DataFlow/ConstantPropagationAnalysis.cpp b/mlir/lib/Analysis/DataFlow/ConstantPropagationAnalysis.cpp
index dbf68ac575dbe..e6e7491b041ee 100644
--- a/mlir/lib/Analysis/DataFlow/ConstantPropagationAnalysis.cpp
+++ b/mlir/lib/Analysis/DataFlow/ConstantPropagationAnalysis.cpp
@@ -72,21 +72,23 @@ LogicalResult SparseConstantPropagation::visitOperation(
   SmallVector<Value, 8> originalOperands(op->getOperands());
   DictionaryAttr originalAttrs = op->getAttrDictionary();
 
-  // Simulate the result of folding this operation to a constant. If folding
-  // fails or was an in-place fold, mark the results as overdefined.
+  // Simulate the result of folding this operation to a constant.
   SmallVector<OpFoldResult, 8> foldResults;
   foldResults.reserve(op->getNumResults());
-  if (failed(op->fold(constantOperands, foldResults))) {
-    setAllToEntryStates(results);
-    return success();
-  }
+  LogicalResult folded = op->fold(constantOperands, foldResults);
 
-  // If the folding was in-place, mark the results as overdefined and reset
-  // the operation. We don't allow in-place folds as the desire here is for
-  // simulated execution, and not general folding.
-  if (foldResults.empty()) {
+  // `fold` can mutate the operation in place and still return an out-of-place
+  // result, so the mutation must be reverted regardless of the outcome.
+  // Only write the operands back if the fold changed them, as `setOperands`
+  // relinks use-lists even for identical values.
+  if (!llvm::equal(op->getOperands(), originalOperands))
     op->setOperands(originalOperands);
-    op->setAttrs(originalAttrs);
+  op->setAttrs(originalAttrs);
+
+  // If folding failed or was in-place, mark the results as overdefined. We
+  // don't allow in-place folds here: the goal is simulated execution, not
+  // general folding.
+  if (failed(folded) || foldResults.empty()) {
     setAllToEntryStates(results);
     return success();
   }
diff --git a/mlir/test/Transforms/sccp.mlir b/mlir/test/Transforms/sccp.mlir
index 251a74dc20647..80273c86090a7 100644
--- a/mlir/test/Transforms/sccp.mlir
+++ b/mlir/test/Transforms/sccp.mlir
@@ -335,3 +335,31 @@ func.func @fold_to_non_operand_value(%x: i64, %cond: i1) -> i64 {
   %cast2 = builtin.unrealized_conversion_cast %cast1 : index to i64
   return %cast2 : i64
 }
+
+// -----
+
+// Regression test: SCCP must revert in-place folds. The `vector.extract`
+// folder rewrites a constant dynamic position into a static one in place, but
+// the constants SCCP feeds into `fold` are speculative: on the first visit of
+// ^bb1 %iv is Constant 0, before the back edge widens it to overdefined.
+// Without the revert, the dynamic extract would permanently read element 0.
+
+// CHECK-LABEL: func @no_inplace_extract_fold_of_speculative_constant
+func.func @no_inplace_extract_fold_of_speculative_constant(%a: f32, %b: f32) -> f32 {
+  %c0_i32 = arith.constant 0 : i32
+  %c1_i32 = arith.constant 1 : i32
+  %c2_i32 = arith.constant 2 : i32
+  %v = vector.from_elements %a, %b : vector<2xf32>
+  cf.br ^bb1(%c0_i32 : i32)
+// CHECK: ^bb1(%[[IV:.*]]: i32):
+^bb1(%iv: i32):
+  // CHECK: %[[IDX:.*]] = arith.index_cast %[[IV]] : i32 to index
+  %idx = arith.index_cast %iv : i32 to index
+  // CHECK: vector.extract %{{.*}}[%[[IDX]]] : f32 from vector<2xf32>
+  %e = vector.extract %v[%idx] : f32 from vector<2xf32>
+  %next = arith.addi %iv, %c1_i32 : i32
+  %cond = arith.cmpi ne, %next, %c2_i32 : i32
+  cf.cond_br %cond, ^bb1(%next : i32), ^bb2
+^bb2:
+  return %e : f32
+}