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

[Federico Bruzzone]

https://github.com/FedericoBruzzone created https://github.com/llvm/llvm-project/pull/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, 
2. trailing, and
3. both (leading + trailing).

It complements the test introduced in #219192.


>From 7f8ea63f3428c9e0d205961abe31eceeb7eae979 Mon Sep 17 00:00:00 2001
From: Federico Bruzzone <federico.bruzzone.i at gmail.com>
Date: Tue, 8 Sep 2026 13:43:11 +0200
Subject: [PATCH] [mlir][linalg] Add `insert_slice` vectorization tests for
 leading/trailing dropped dims

Signed-off-by: Federico Bruzzone <federico.bruzzone.i at gmail.com>
---
 .../Linalg/vectorization/insert-slice.mlir    | 84 +++++++++++++++++++
 1 file changed, 84 insertions(+)

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> {