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

[Dhairyashil R G]

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

Removed: the next test checks the same things (the mask sizes and the map on the write). The two remaining tests are renamed to follow the file's scheme, and each covers something no other test does:
- `@insert_slice_non_leading_unit_dim_dropped_dynamic_dest_dim`: the mask sizes and the "is a mask needed" check use the dims the map selects, not the trailing ones.
- `@insert_slice_leading_unit_dim_dropped_non_zero_offset`: each mask size is the dim minus its own index.

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