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

[llvmlistbot at llvm.org]

Author: Berke Ates
Date: 2026-08-05T08:55:30+02:00
New Revision: dce57cf279d65b0ce871a77f1b493e5141a99379

URL: https://github.com/llvm/llvm-project/commit/dce57cf279d65b0ce871a77f1b493e5141a99379
DIFF: https://github.com/llvm/llvm-project/commit/dce57cf279d65b0ce871a77f1b493e5141a99379.diff

LOG: [MLIR][SCCP] Fix in-place folds leaking into IR during simulation (#213933)

SparseConstantPropagation restored the operation only when `fold` did
not return any fold results. But a folder can mutate the op in place and
still return out-of-place results or fail, e.g. `vector.extract` folds
constant dynamic positions into static ones before attempting further
folds.
The constants fed to `fold` are speculative lattice values, so the
mutation bakes a possibly-wrong constant into the IR. SCCP would
permanently replace a loop-carried dynamic index with its first lattice
value for example.

Fix: Restore the original operands and attributes after every fold call,
regardless of its outcome.

Co-authored-by: Claude Fable 5 <noreply at anthropic.com>

Added: 
    

Modified: 
    mlir/lib/Analysis/DataFlow/ConstantPropagationAnalysis.cpp
    mlir/test/Transforms/sccp.mlir

Removed: 
    


################################################################################
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
+}