[Mlir-commits] [mlir] [mlir][vector] Add foldInsertUseChain folder function to insert op (PR #147045)

[Mehdi Amini]

================
@@ -3446,3 +3446,50 @@ func.func @fold_insert_constant_indices(%arg : vector<4x1xi32>) -> vector<4x1xi3
   %res = vector.insert %1, %arg[%0, %0] : i32 into vector<4x1xi32>
   return %res : vector<4x1xi32>
 }
+
+// -----
+
+// CHECK-LABEL: @fold_insert_use_chain(
+//  CHECK-SAME:   %[[DEST:.*]]: vector<4x4xf32>,
+//  CHECK-SAME:   %[[VAL:.*]]: f32,
+//  CHECK-SAME:   %[[POS:.*]]: index) -> vector<4x4xf32> {
+//  CHECK-NEXT:   %[[RES:.*]] = vector.insert %[[VAL]], %[[DEST]] {{\[}}%[[POS]], 0] : f32 into vector<4x4xf32>
+//  CHECK-NEXT:   return %[[RES]] : vector<4x4xf32>
+func.func @fold_insert_use_chain(%arg : vector<4x4xf32>, %value : f32, %pos: index) -> vector<4x4xf32> {
+  %v_0 = vector.insert %value, %arg[%pos, 0] : f32 into vector<4x4xf32>
+  %v_1 = vector.insert %value, %v_0[%pos, 0] : f32 into vector<4x4xf32>
+  %v_2 = vector.insert %value, %v_1[%pos, 0] : f32 into vector<4x4xf32>
+  return %v_2 : vector<4x4xf32>  
+}
+
+// -----
+
+// CHECK-LABEL: @no_fold_insert_use_chain(
+//  CHECK-SAME:   %[[DEST_0:.*]]: vector<4xf32>,
+//  CHECK-SAME:   %[[VAL:.*]]: f32) -> vector<4xf32> {
+//       CHECK:   %[[DEST_1:.*]] = vector.insert %[[VAL]], %[[DEST_0]] [0] : f32 into vector<4xf32>
+//       CHECK:   %[[RES:.*]] = vector.insert %[[VAL]], %[[DEST_1]] [1] : f32 into vector<4xf32>
+//       CHECK:   return %[[RES]] : vector<4xf32>
+func.func @no_fold_insert_use_chain(%v : vector<4xf32>, %value : f32) -> vector<4xf32> {
+  %v_0 = vector.insert %value, %v[0] : f32 into vector<4xf32>
+  %v_2 = vector.insert %value, %v_0[1] : f32 into vector<4xf32>
+  return %v_2 : vector<4xf32>  
+}
+
+// -----
+
+// CHECK-LABEL: @fold_insert_use_chain_add_float(
+//  CHECK-SAME:   %[[DEST:.*]]: vector<4xf32>,
+//  CHECK-SAME:   %[[VAL:.*]]: f32) -> vector<4xf32> {
+//       CHECK:   %{{.*}} = vector.insert %[[VAL]], %[[DEST]] [0] : f32 into vector<4xf32>
+//       CHECK:   %[[LHS:.*]] = arith.addf %{{.*}}, %{{.*}} : vector<4xf32>
+//       CHECK:   %[[RHS:.*]] = vector.insert %[[VAL]], %[[DEST]] [0] : f32 into vector<4xf32>
+//       CHECK:   %[[RES:.*]] = arith.addf %[[LHS]], %[[RHS]] : vector<4xf32>
+//       CHECK:   return %[[RES]] : vector<4xf32>
+func.func @fold_insert_use_chain_add_float(%v : vector<4xf32>, %value : f32) -> vector<4xf32> {
+  %v_0 = vector.insert %value, %v[0] : f32 into vector<4xf32>
+  %v_1 = arith.addf %v_0, %v_0 : vector<4xf32>
----------------
joker-eph wrote:

Oh I see now: when you had a canonicalizer you were checking if the first insert had multiple uses and didn't perform the transformation in this case. So you're trying to reproduce this here and showing that the transform applies even with multiple uses?
It does not seem necessary for the folder to me, but if you really want to keep it, document this as a comment.

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