[Mlir-commits] [mlir] 152917c - [mlir][linalg] Add `insert_slice` vectorization tests for leading/trailing dropped dims (#221966)

[llvmlistbot at llvm.org]

Author: Federico Bruzzone
Date: 2026-09-11T13:45:35+02:00
New Revision: 152917c9555d4da0171e7d992b38196e2fb9f5aa

URL: https://github.com/llvm/llvm-project/commit/152917c9555d4da0171e7d992b38196e2fb9f5aa
DIFF: https://github.com/llvm/llvm-project/commit/152917c9555d4da0171e7d992b38196e2fb9f5aa.diff

LOG: [mlir][linalg] Add `insert_slice` vectorization tests for leading/trailing dropped dims (#221966)

Follow-up to #219192, addressing review feedback after merge.

This PR _simply_ add tests for vectorizing rank-reducing
`tensor.insert_slice` where the dropped unit dim is:
1. leading, e.g., `dest = 1x8x4, source = ?x4`,
2. trailing, e.g., `dest = 8x4x1, source = ?x4`, and
3. both  (leading + trailing), e.g., `dest = 1x8x4x1, source = ?x4`.
 
It complements the test introduced in #219192.

Signed-off-by: Federico Bruzzone <federico.bruzzone.i at gmail.com>

Added: 
    

Modified: 
    mlir/test/Dialect/Linalg/vectorization/insert-slice.mlir

Removed: 
    


################################################################################
diff  --git a/mlir/test/Dialect/Linalg/vectorization/insert-slice.mlir b/mlir/test/Dialect/Linalg/vectorization/insert-slice.mlir
index c2e4d83b14886..34e251c527e31 100644
--- a/mlir/test/Dialect/Linalg/vectorization/insert-slice.mlir
+++ b/mlir/test/Dialect/Linalg/vectorization/insert-slice.mlir
@@ -116,6 +116,90 @@ func.func private @insert_slice_dynamic_src_dim_non_leading_unit_dim_dropped(
 
 // -----
 
+func.func private @insert_slice_dynamic_src_dim_leading_unit_dim_dropped(
+    %source: tensor<?x4xi32>, %size: index) -> tensor<1x8x4xi32> {
+  %pad = arith.constant 0 : i32
+  %empty = tensor.empty() : tensor<1x8x4xi32>
+  %init = linalg.fill ins(%pad : i32) outs(%empty : tensor<1x8x4xi32>) -> tensor<1x8x4xi32>
+  %res = tensor.insert_slice %source into %init[0, 0, 0] [1, %size, 4] [1, 1, 1] : tensor<?x4xi32> into tensor<1x8x4xi32>
+  return %res : tensor<1x8x4xi32>
+}
+
+// CHECK-LABEL:   func.func private @insert_slice_dynamic_src_dim_leading_unit_dim_dropped(
+// CHECK-SAME:      %[[SRC:.*]]: tensor<?x4xi32>,
+// CHECK-SAME:      %[[SIZE:.*]]: index) -> tensor<1x8x4xi32> {
+// CHECK-DAG:       %[[PAD:.*]] = arith.constant 0 : i32
+// CHECK:           %[[INIT:.*]] = linalg.fill ins(%[[PAD]] : i32) outs({{.*}} : tensor<1x8x4xi32>) -> tensor<1x8x4xi32>
+// CHECK:           %[[READ:.*]] = vector.transfer_read %[[SRC]][%{{.*}}, %{{.*}}], %[[PAD]] {{.*}} : tensor<?x4xi32>, vector<8x4xi32>
+// CHECK:           %[[RES:.*]] = vector.transfer_write %[[READ]], %[[INIT]][%{{.*}}, %{{.*}}, %{{.*}}] {{.*}} : vector<8x4xi32>, tensor<1x8x4xi32>
+// CHECK:           return %[[RES]] : tensor<1x8x4xi32>
+
+ 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
+  }
+ }
+
+// -----
+
+func.func private @insert_slice_dynamic_src_dim_trailing_unit_dim_dropped(
+    %source: tensor<?x4xi32>, %size: index) -> tensor<8x4x1xi32> {
+  %pad = arith.constant 0 : i32
+  %empty = tensor.empty() : tensor<8x4x1xi32>
+  %init = linalg.fill ins(%pad : i32) outs(%empty : tensor<8x4x1xi32>) -> tensor<8x4x1xi32>
+  %res = tensor.insert_slice %source into %init[0, 0, 0] [%size, 4, 1] [1, 1, 1] : tensor<?x4xi32> into tensor<8x4x1xi32>
+  return %res : tensor<8x4x1xi32>
+}
+
+// CHECK-LABEL:   func.func private @insert_slice_dynamic_src_dim_trailing_unit_dim_dropped(
+// CHECK-SAME:      %[[SRC:.*]]: tensor<?x4xi32>,
+// CHECK-SAME:      %[[SIZE:.*]]: index) -> tensor<8x4x1xi32> {
+// CHECK-DAG:       %[[PAD:.*]] = arith.constant 0 : i32
+// CHECK:           %[[INIT:.*]] = linalg.fill ins(%[[PAD]] : i32) outs({{.*}} : tensor<8x4x1xi32>) -> tensor<8x4x1xi32>
+// CHECK:           %[[READ:.*]] = vector.transfer_read %[[SRC]][%{{.*}}, %{{.*}}], %[[PAD]] {{.*}} : tensor<?x4xi32>, vector<8x4xi32>
+// CHECK:           %[[RES:.*]] = vector.transfer_write %[[READ]], %[[INIT]][%{{.*}}, %{{.*}}, %{{.*}}] : vector<8x4xi32>, tensor<8x4x1xi32>
+// CHECK:           return %[[RES]] : tensor<8x4x1xi32>
+
+ 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
+  }
+ }
+
+// -----
+
+func.func private @insert_slice_dynamic_src_dim_leading_and_trailing_unit_dims_dropped(
+    %source: tensor<?x4xi32>, %size: index) -> tensor<1x8x4x1xi32> {
+  %pad = arith.constant 0 : i32
+  %empty = tensor.empty() : tensor<1x8x4x1xi32>
+  %init = linalg.fill ins(%pad : i32) outs(%empty : tensor<1x8x4x1xi32>) -> tensor<1x8x4x1xi32>
+  %res = tensor.insert_slice %source into %init[0, 0, 0, 0] [1, %size, 4, 1] [1, 1, 1, 1] : tensor<?x4xi32> into tensor<1x8x4x1xi32>
+  return %res : tensor<1x8x4x1xi32>
+}
+
+// CHECK-LABEL:   func.func private @insert_slice_dynamic_src_dim_leading_and_trailing_unit_dims_dropped(
+// CHECK-SAME:      %[[SRC:.*]]: tensor<?x4xi32>,
+// CHECK-SAME:      %[[SIZE:.*]]: index) -> tensor<1x8x4x1xi32> {
+// CHECK-DAG:       %[[PAD:.*]] = arith.constant 0 : i32
+// CHECK:           %[[INIT:.*]] = linalg.fill ins(%[[PAD]] : i32) outs({{.*}} : tensor<1x8x4x1xi32>) -> tensor<1x8x4x1xi32>
+// CHECK:           %[[READ:.*]] = vector.transfer_read %[[SRC]][%{{.*}}, %{{.*}}], %[[PAD]] {{.*}} : tensor<?x4xi32>, vector<8x4xi32>
+// CHECK:           %[[RES:.*]] = vector.transfer_write %[[READ]], %[[INIT]][%{{.*}}, %{{.*}}, %{{.*}}, %{{.*}}] : vector<8x4xi32>, tensor<1x8x4x1xi32>
+// CHECK:           return %[[RES]] : tensor<1x8x4x1xi32>
+
+ 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
+  }
+ }
+
+// -----
+
 // One of the _destination_ dimensions is dynamic (but _source_ dimensions are static).
 
 func.func private @insert_slice_dynamic_dest_dim(%source: tensor<?x3x?x1xi32>, %size: index) -> tensor<?x3xi32> {