[Mlir-commits] [mlir] [mlir][linalg] Fix rank-reducing insert_slice vectorization writing to wrong dims (PR #226490)

[Andrzej Warzyński]

================
@@ -334,3 +345,133 @@ func.func private @insert_slice_non_zero_offset_for_dyn_dim(%source: tensor<?x3x
     transform.yield
   }
  }
+
+// -----
+
+/// A 1-D source inserted down a column: the dropped dim is the trailing one,
+/// so the source's only dim corresponds to result dim 0.
+
+func.func private @insert_slice_1d_source_into_column(
+    %source: tensor<5xi32>) -> tensor<5x3xi32> {
+  %pad = arith.constant 0 : i32
+  %empty = tensor.empty() : tensor<5x3xi32>
+  %init = linalg.fill ins(%pad : i32) outs(%empty : tensor<5x3xi32>) -> tensor<5x3xi32>
+  %res = tensor.insert_slice %source into %init[0, 2] [5, 1] [1, 1] : tensor<5xi32> into tensor<5x3xi32>
+  return %res : tensor<5x3xi32>
+}
+
+// CHECK-DAG: #[[$MAP_D0:.*]] = affine_map<(d0, d1) -> (d0)>
+
+// CHECK-LABEL:   func.func private @insert_slice_1d_source_into_column(
+// CHECK-SAME:      %[[SRC:.*]]: tensor<5xi32>) -> tensor<5x3xi32> {
+// CHECK-DAG:       %[[PAD:.*]] = arith.constant 0 : i32
+// CHECK:           %[[INIT:.*]] = linalg.fill ins(%[[PAD]] : i32) outs({{.*}} : tensor<5x3xi32>) -> tensor<5x3xi32>
+// CHECK:           %[[READ:.*]] = vector.transfer_read %[[SRC]][%{{.*}}], %[[PAD]] {{.*}} : tensor<5xi32>, vector<5xi32>
+// CHECK:           %[[C_2:.*]] = arith.constant 2 : index
+// CHECK:           %[[RES:.*]] = vector.transfer_write %[[READ]], %[[INIT]][%{{.*}}, %[[C_2]]] {{.*}}permutation_map = #[[$MAP_D0]]{{.*}} : vector<5xi32>, tensor<5x3xi32>
+// CHECK:           return %[[RES]] : tensor<5x3xi32>
+
+ module attributes {transform.with_named_sequence} {
+  transform.named_sequence @__transform_main(%arg0: !transform.any_op {transform.readonly}) {
+    %0 = transform.structured.match ops{["tensor.insert_slice"]} in %arg0 : (!transform.any_op) -> !transform.any_op
+    transform.structured.vectorize %0 : !transform.any_op
+    transform.yield
+  }
+ }
+
+// -----
+
+/// Rank-reducing with user-provided vector sizes, the dropped dim not being the
+/// trailing one. The vector is written to result dims 0 and 2, so each mask
+/// size is one of those dims minus its offset.
+
+func.func private @insert_slice_masked_non_trailing_unit_dim_dropped(
----------------
banach-space wrote:

This is breaking with the naming scheme used in this file (we don't use "masked" to mark tests that lead to masks). See our testing guide for suggestions on good naming scheme, e.g.
* https://mlir.llvm.org/getting_started/TestingGuide/#test-formatting-best-practices

Also, what's unique about this test? Is there no other test that exercises similar behaviour

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