[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:

This test is no different from the first one to me.

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