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

[llvmlistbot at llvm.org]

llvmorg-github-actions[bot] wrote:


<!--LLVM PR SUMMARY COMMENT-->

@llvm/pr-subscribers-mlir

Author: Berke Ates (Berke-Ates)

<details>
<summary>Changes</summary>

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.

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


2 Files Affected:

- (modified) mlir/lib/Analysis/DataFlow/ConstantPropagationAnalysis.cpp (+13-11) 
- (modified) mlir/test/Transforms/sccp.mlir (+49) 


``````````diff
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..9f2217524287e 100644
--- a/mlir/test/Transforms/sccp.mlir
+++ b/mlir/test/Transforms/sccp.mlir
@@ -335,3 +335,52 @@ 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)
+^bb1(%iv: i32):
+  %idx = arith.index_cast %iv : i32 to index
+  // CHECK: vector.extract %{{.*}}[%{{.*}}] : 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
+}
+
+// -----
+
+// The same in-place position rewrite happens in the `vector.insert` folder.
+
+// CHECK-LABEL: func @no_inplace_insert_fold_of_speculative_constant
+func.func @no_inplace_insert_fold_of_speculative_constant() -> vector<2xf32> {
+  %c0_i32 = arith.constant 0 : i32
+  %c1_i32 = arith.constant 1 : i32
+  %c2_i32 = arith.constant 2 : i32
+  %f = arith.constant 3.000000e+00 : f32
+  %init = arith.constant dense<[1.000000e+00, 2.000000e+00]> : vector<2xf32>
+  cf.br ^bb1(%c0_i32, %init : i32, vector<2xf32>)
+^bb1(%iv: i32, %acc: vector<2xf32>):
+  %idx = arith.index_cast %iv : i32 to index
+  // CHECK: vector.insert %{{.*}}, %{{.*}}[%{{.*}}] : f32 into vector<2xf32>
+  %ins = vector.insert %f, %acc[%idx] : f32 into vector<2xf32>
+  %next = arith.addi %iv, %c1_i32 : i32
+  %cond = arith.cmpi ne, %next, %c2_i32 : i32
+  cf.cond_br %cond, ^bb1(%next, %ins : i32, vector<2xf32>), ^bb2
+^bb2:
+  return %ins : vector<2xf32>
+}

``````````

</details>


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