[Mlir-commits] [mlir] [mlir][vector]add extractInsertFoldConstantOp fold function and apply it to extractOp and insertOp. (PR #124399)

[Andrzej Warzyński]

================
@@ -1473,6 +1473,22 @@ func.func @extract_scalar_from_vec_0d_index(%arg0: vector<index>) -> index {
 //       CHECK:   %[[T3:.*]] = builtin.unrealized_conversion_cast %[[T2]] : i64 to index
 //       CHECK:   return %[[T3]] : index
 
+
+// -----
+
+func.func @extract_scalar_from_vec_all_constant_idx(%arg : vector<32x1xi32>) -> i32 {
+  %0 = arith.constant 0 : index
+  %1 = vector.extract %arg[%0, %0] : i32 from vector<32x1xi32>
+  return %1 : i32
+}
+// CHECK-LABEL: func @extract_scalar_from_vec_all_constant_idx
+// CHECK-SAME: %[[ARG:.*]]: vector<32x1xi32>) -> i32 {
+// CHECK: %[[CAST:.*]] = builtin.unrealized_conversion_cast %[[ARG]] : vector<32x1xi32> to !llvm.array<32 x vector<1xi32>>
+// CHECK: %[[VEC_0:.*]] = llvm.extractvalue %[[CAST]][0] : !llvm.array<32 x vector<1xi32>>
+// CHECK: %[[C0:.*]] = llvm.mlir.constant(0 : i64) : i64
+// CHECK: %[[RES:.*]] = llvm.extractelement %[[VEC_0]]{{\[}}%[[C0]] : i64] : vector<1xi32>
+// CHECK:  return %[[RES]] : i32
----------------
banach-space wrote:

Thanks for updating this test! I have 2 more requests:
1. Please use similar indention style for CHECK lines as the other tests in the file, e.g.

```
// CHECK-LABEL: @extract_scalar_from_vec_0d_index(
//  CHECK-SAME:   %[[A:.*]]: vector<index>)
//       CHECK:   %[[T0:.*]] = builtin.unrealized_conversion_cast %[[A]] : vector<index> to vector<1xi64>
//       CHECK:   %[[T1:.*]] = llvm.mlir.constant(0 : i64) : i64
//       CHECK:   %[[T2:.*]] = llvm.extractelement %[[T0]][%[[T1]] : i64] : vector<1xi64>
//       CHECK:   %[[T3:.*]] = builtin.unrealized_conversion_cast %[[T2]] : i64 to index
//       CHECK:   return %[[T3]] : index
```

2. This test seems very similar to e.g. `extract_scalar_from_vec_2d_f32_outer_dynamic_idx`:
```mlir
func.func @extract_scalar_from_vec_2d_f32_outer_dynamic_idx(%arg0: vector<1x16xf32>, %arg1: index) -> f32 {
  %0 = vector.extract %arg0[%arg1, 0]: f32 from vector<1x16xf32>
  return %0 : f32
}
```

Please, could you make that more explicit, e.g.:

```mlir
func.func @extract_scalar_from_vec_2d_f32_dynamic_idxs_compile_time_const(%arg0: vector<1x16xf32>) -> f32 {
  %c0 = arith.constant 0 : index
  %0 = vector.extract %arg0[%c0, %c0]: f32 from vector<1x16xf32>
  return %0 : f32
}
```

This would make it super clear this is a special case of `@extract_scalar_from_vec_2d_f32` where:
* both indices are dynamic, but also
* both indices are compile-time constant.

Additional comment would also be appreciated (e.g. "Compile-time constants are folded away, hence the lowering works.")

Same comment for `@insert_scalar_from_vec_all_constant_idx`.

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