[Mlir-commits] [mlir] [mlir][linalg] Support non-trailing dropped dims when vectorizing ran… (PR #219192)

[Federico Bruzzone]

https://github.com/FedericoBruzzone created https://github.com/llvm/llvm-project/pull/219192

…

`vectorizeAsInsertSliceOp` inferred a dynamic source dim's vector size from
`resultType.getDimSize(rankDiff + i)`, which only covers the case where the
rank-reducing dims dropped from `resultType` are the leading ones and
`source` maps to the trailing result dims.

Generalize this to the general case, where a dropped unit dim can be
anywhere, by using `sliceOp.getDroppedDims()` to map each source dim to
its actual corresponding (non-dropped) result dim.

>From 07392e83d0e1925dfd9d45a2c3bb78fc3dd22122 Mon Sep 17 00:00:00 2001
From: Federico Bruzzone <federico.bruzzone.i at gmail.com>
Date: Thu, 27 Aug 2026 13:58:03 +0200
Subject: [PATCH] [mlir][linalg] Support non-trailing dropped dims when
 vectorizing rank-reducing `InsertSliceOp`

Signed-off-by: Federico Bruzzone <federico.bruzzone.i at gmail.com>
---
 .../Linalg/Transforms/Vectorization.cpp       | 20 +++++++++----
 .../Linalg/vectorization/insert-slice.mlir    | 28 +++++++++++++++++++
 2 files changed, 43 insertions(+), 5 deletions(-)

diff --git a/mlir/lib/Dialect/Linalg/Transforms/Vectorization.cpp b/mlir/lib/Dialect/Linalg/Transforms/Vectorization.cpp
index 21ca3108efcd6..18d3d12e01cd1 100644
--- a/mlir/lib/Dialect/Linalg/Transforms/Vectorization.cpp
+++ b/mlir/lib/Dialect/Linalg/Transforms/Vectorization.cpp
@@ -2979,20 +2979,30 @@ vectorizeAsInsertSliceOp(RewriterBase &rewriter, tensor::InsertSliceOp sliceOp,
   }
 
   // 2. Get the vector shape
+  // Map each source dim to its corresponding (non-dropped) result dim: for a
+  // rank-reducing slice, dropped dims need not be the trailing ones.
+  llvm::SmallBitVector droppedDims = sliceOp.getDroppedDims();
+  SmallVector<int64_t> resultDimsForSourceDims;
+  resultDimsForSourceDims.reserve(sourceType.getRank());
+  for (int64_t resultDim = 0, end = resultType.getRank(); resultDim < end;
+       ++resultDim)
+    if (!droppedDims[resultDim])
+      resultDimsForSourceDims.push_back(resultDim);
+  assert(resultDimsForSourceDims.size() ==
+             static_cast<size_t>(sourceType.getRank()) &&
+         "expected one non-dropped result dim per source dim");
+
   SmallVector<int64_t> vecShape;
-  size_t rankDiff = resultType.getRank() - sourceType.getRank();
   for (int64_t i = 0, end = sourceType.getRank(); i < end; ++i) {
     if (!inputVectorSizes.empty()) {
       vecShape.push_back(inputVectorSizes[i]);
     } else if (!sourceType.isDynamicDim(i)) {
       vecShape.push_back(sourceType.getDimSize(i));
-    } else if (!resultType.isDynamicDim(i)) {
+    } else if (!resultType.isDynamicDim(resultDimsForSourceDims[i])) {
       // Source shape is not statically known, but result shape is.
       // Vectorize with size of result shape. This may be larger than the
       // source size.
-      // FIXME: Using rankDiff implies that the source tensor is inserted at
-      // the end of the destination tensor. However, that's not required.
-      vecShape.push_back(resultType.getDimSize(rankDiff + i));
+      vecShape.push_back(resultType.getDimSize(resultDimsForSourceDims[i]));
     } else {
       // Neither source nor result dim of padOp is static. Cannot vectorize
       // the copy.
diff --git a/mlir/test/Dialect/Linalg/vectorization/insert-slice.mlir b/mlir/test/Dialect/Linalg/vectorization/insert-slice.mlir
index 91ceea3b394c3..c2e4d83b14886 100644
--- a/mlir/test/Dialect/Linalg/vectorization/insert-slice.mlir
+++ b/mlir/test/Dialect/Linalg/vectorization/insert-slice.mlir
@@ -88,6 +88,34 @@ func.func private @insert_slice_dynamic_src_dim(%source: tensor<?x3x?x1xi32>, %s
 
 // -----
 
+func.func private @insert_slice_dynamic_src_dim_non_leading_unit_dim_dropped(
+    %source: tensor<?x4xi32>, %size: index) -> tensor<8x1x4xi32> {
+  %pad = arith.constant 0 : i32
+  %empty = tensor.empty() : tensor<8x1x4xi32>
+  %init = linalg.fill ins(%pad : i32) outs(%empty : tensor<8x1x4xi32>) -> tensor<8x1x4xi32>
+  %res = tensor.insert_slice %source into %init[0, 0, 0] [%size, 1, 4] [1, 1, 1] : tensor<?x4xi32> into tensor<8x1x4xi32>
+  return %res : tensor<8x1x4xi32>
+}
+
+// CHECK-LABEL:   func.func private @insert_slice_dynamic_src_dim_non_leading_unit_dim_dropped(
+// CHECK-SAME:      %[[SRC:.*]]: tensor<?x4xi32>,
+// CHECK-SAME:      %[[SIZE:.*]]: index) -> tensor<8x1x4xi32> {
+// CHECK-DAG:       %[[PAD:.*]] = arith.constant 0 : i32
+// CHECK:           %[[INIT:.*]] = linalg.fill ins(%[[PAD]] : i32) outs({{.*}} : tensor<8x1x4xi32>) -> tensor<8x1x4xi32>
+// CHECK:           %[[READ:.*]] = vector.transfer_read %[[SRC]][%{{.*}}, %{{.*}}], %[[PAD]] {{.*}} : tensor<?x4xi32>, vector<8x4xi32>
+// CHECK:           %[[RES:.*]] = vector.transfer_write %[[READ]], %[[INIT]][%{{.*}}, %{{.*}}, %{{.*}}] {{.*}} : vector<8x4xi32>, tensor<8x1x4xi32>
+// CHECK:           return %[[RES]] : tensor<8x1x4xi32>
+
+ 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> {