[Mlir-commits] [mlir] [mlir][linalg] Check the strides and the dropped dims before vectorizing insert_slice (PR #221293)

[Hung-Kuan Tseng]

https://github.com/Tim096 updated https://github.com/llvm/llvm-project/pull/221293

>From 7ce815586e1a5ec5f33eb6c3d3a5aef1529c833a Mon Sep 17 00:00:00 2001
From: Hung-Kuan Tseng <tseng.tim096 at gmail.com>
Date: Thu, 8 Oct 2026 18:17:24 +0800
Subject: [PATCH] [mlir][linalg] Reject non-unit strides when vectorizing
 insert_slice

`vectorizeAsInsertSliceOp` writes the source back with a
`vector.transfer_write` at the slice offsets, which cannot express a
stride. An insert with stride 3 was written to rows 0 and 1 instead of
rows 0 and 3. Bail out in the precondition unless all strides are 1.

Part of #227821.
---
 .../Linalg/Transforms/Vectorization.cpp       | 10 ++++-
 .../Linalg/vectorization/unsupported.mlir     | 37 +++++++++++++++++++
 2 files changed, 46 insertions(+), 1 deletion(-)

diff --git a/mlir/lib/Dialect/Linalg/Transforms/Vectorization.cpp b/mlir/lib/Dialect/Linalg/Transforms/Vectorization.cpp
index 0b6bc6b68db33..8ffe8da456357 100644
--- a/mlir/lib/Dialect/Linalg/Transforms/Vectorization.cpp
+++ b/mlir/lib/Dialect/Linalg/Transforms/Vectorization.cpp
@@ -66,7 +66,8 @@ vectorizeConvolution(RewriterBase &rewriter, LinalgOp convOp,
 ///   * inferred from the static dims in the input and output tensors.
 /// Bails out if:
 ///   * vector sizes are not user-provided, and
-///   * at least one dim is dynamic (in both the input and output tensors).
+///   * at least one dim is dynamic (in both the input and output tensors),
+/// or if the strides are not 1.
 ///
 /// Before:
 ///     !t_in_type = tensor<1x2x3xf32>
@@ -1992,6 +1993,13 @@ vectorizeInsertSliceOpPrecondition(tensor::InsertSliceOp sliceOp,
   if (!VectorType::isValidElementType(sourceType.getElementType()))
     return failure();
 
+  // The write is a vector.transfer_write at the slice offsets, which cannot
+  // express a stride.
+  if (!sliceOp.hasUnitStride()) {
+    LDBG() << "Only unit strides are supported";
+    return failure();
+  }
+
   // Get the pad value.
   // TransferReadOp (which is used to vectorize InsertSliceOp), requires a
   // scalar padding value. Note that:
diff --git a/mlir/test/Dialect/Linalg/vectorization/unsupported.mlir b/mlir/test/Dialect/Linalg/vectorization/unsupported.mlir
index b5ee8da25adba..66ff11e02e18f 100644
--- a/mlir/test/Dialect/Linalg/vectorization/unsupported.mlir
+++ b/mlir/test/Dialect/Linalg/vectorization/unsupported.mlir
@@ -411,3 +411,40 @@ module attributes {transform.with_named_sequence} {
     transform.yield
   }
 }
+
+// -----
+
+// A stride of 3 along the first dim cannot be expressed by the
+// vector.transfer_write that the insert is vectorized to.
+
+func.func @insert_slice_non_unit_stride(%src: tensor<2x4xf32>, %dest: tensor<8x4xf32>) -> tensor<8x4xf32> {
+  // expected-error @+1 {{Attempted to vectorize, but failed}}
+  %res = tensor.insert_slice %src into %dest[0, 0] [2, 4] [3, 1] : tensor<2x4xf32> into tensor<8x4xf32>
+  return %res : tensor<8x4xf32>
+}
+
+module attributes {transform.with_named_sequence} {
+  transform.named_sequence @__transform_main(%arg1: !transform.any_op {transform.readonly}) {
+    %0 = transform.structured.match ops{["tensor.insert_slice"]} in %arg1 : (!transform.any_op) -> !transform.any_op
+    transform.structured.vectorize %0 : !transform.any_op
+    transform.yield
+  }
+}
+
+// -----
+
+// Same with a dynamic stride, which may not be 1.
+
+func.func @insert_slice_dynamic_stride(%src: tensor<2x4xf32>, %dest: tensor<8x4xf32>, %stride: index) -> tensor<8x4xf32> {
+  // expected-error @+1 {{Attempted to vectorize, but failed}}
+  %res = tensor.insert_slice %src into %dest[0, 0] [2, 4] [%stride, 1] : tensor<2x4xf32> into tensor<8x4xf32>
+  return %res : tensor<8x4xf32>
+}
+
+module attributes {transform.with_named_sequence} {
+  transform.named_sequence @__transform_main(%arg1: !transform.any_op {transform.readonly}) {
+    %0 = transform.structured.match ops{["tensor.insert_slice"]} in %arg1 : (!transform.any_op) -> !transform.any_op
+    transform.structured.vectorize %0 : !transform.any_op
+    transform.yield
+  }
+}