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

[Tobias Gysi]

================
@@ -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>
----------------
gysit wrote:

I would either add the insert to the test above or drop this second test since I would assume that it doesn't check anything in addition to the test above?

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