[Mlir-commits] [mlir] [mlir][linalg] Fix rank-reducing insert_slice vectorization writing to wrong dims (PR #226490)
Dhairyashil R G
llvmlistbot at llvm.org
Fri Sep 25 06:33:50 PDT 2026
https://github.com/dhairyashilRG created https://github.com/llvm/llvm-project/pull/226490
Vectorizing a rank-reducing `tensor.insert_slice` took the vector shape from the non-dropped result dims, but wrote the vector back with an implied minor-identity permutation map, i.e. onto the trailing result dims. When the dropped unit dims are not the trailing ones, those are the wrong dims.
Inserting a `tensor<5xi32>` into `tensor<5x3xi32>` at `[0, 2] [5, 1] [1, 1]` fills column 2, one element per row. Vectorized, the write went along dim 1, so only the first element landed and the other four were discarded by the out-of-bounds check. With user-provided vector sizes the write was masked with `dim(1) - index(0)` and stored past the end of the row instead.
Fixed by passing a permutation map built from the non-dropped result dims, and by making the masking path of `createReadOrMaskedRead/createWriteOrMaskedWrite` honour such a map: mask dims follow the transfer's mask layout (`inferTransferOpMaskType`) and each size is `dim - index` for the source dim it covers.
Note for out-of-tree callers of `createWriteOrMaskedWrite`: when the destination has more dims than the vector and the write indices are non-zero, a mask size is now `dim - index` for the same dim, where before it paired `destSizes[diff + idx]` with `writeIndices[idx]`. No in-tree caller was affected.
Three of the four rank-reducing cases in insert-slice.mlir were checking the old write; their expectations now pin the permutation map. The fourth (a dropped leading dim) is unchanged, because minor identity was already correct there. A new integration test runs the vectorized and un-vectorized forms and requires them to agree.
Assisted-by: Claude
>From 44393b16eb08686bafb865261f757d987ca55b48 Mon Sep 17 00:00:00 2001
From: Dhairyashil R G <dhairyashil25 at gmail.com>
Date: Fri, 25 Sep 2026 18:05:00 +0530
Subject: [PATCH] [mlir][linalg] Fix rank-reducing insert_slice vectorization
writing to wrong dims
Vectorizing a rank-reducing `tensor.insert_slice` took the vector shape from the
non-dropped result dims, but wrote the vector back with an implied minor-identity
permutation map, i.e. onto the *trailing* result dims. When the dropped unit dims
are not the trailing ones, those are the wrong dims.
Inserting a `tensor<5xi32>` into `tensor<5x3xi32>` at `[0, 2] [5, 1] [1, 1]`
fills column 2, one element per row. Vectorized, the write went along dim 1, so
only the first element landed and the other four were discarded by the
out-of-bounds check. With user-provided vector sizes the write was masked with
`dim(1) - index(0)` and stored past the end of the row instead.
Fixed by passing a permutation map built from the non-dropped result dims, and by
making the masking path of `createReadOrMaskedRead`/`createWriteOrMaskedWrite`
honour such a map: mask dims follow the transfer's mask layout
(`inferTransferOpMaskType`) and each size is `dim - index` for the source dim it
covers.
Note for out-of-tree callers of `createWriteOrMaskedWrite`: when the destination
has more dims than the vector and the write indices are non-zero, a mask size is
now `dim - index` for the same dim, where before it paired `destSizes[diff + idx]`
with `writeIndices[idx]`. No in-tree caller was affected.
Three of the four rank-reducing cases in insert-slice.mlir were checking the old
write; their expectations now pin the permutation map. The fourth (a dropped
leading dim) is unchanged, because minor identity was already correct there. A
new integration test runs the vectorized and un-vectorized forms and requires
them to agree.
Assisted-by: Claude
---
.../Linalg/Transforms/Vectorization.cpp | 17 ++++-
mlir/lib/Dialect/Vector/Utils/VectorUtils.cpp | 76 ++++++++++++++-----
.../Linalg/vectorization/insert-slice.mlir | 50 +++++++++++-
.../vectorize-insert-slice-rank-reducing.mlir | 53 +++++++++++++
4 files changed, 173 insertions(+), 23 deletions(-)
create mode 100644 mlir/test/Integration/Dialect/Linalg/CPU/vectorize-insert-slice-rank-reducing.mlir
diff --git a/mlir/lib/Dialect/Linalg/Transforms/Vectorization.cpp b/mlir/lib/Dialect/Linalg/Transforms/Vectorization.cpp
index 5c649473b87f1a..8f14fa39693aa4 100644
--- a/mlir/lib/Dialect/Linalg/Transforms/Vectorization.cpp
+++ b/mlir/lib/Dialect/Linalg/Transforms/Vectorization.cpp
@@ -3067,9 +3067,20 @@ vectorizeAsInsertSliceOp(RewriterBase &rewriter, tensor::InsertSliceOp sliceOp,
// Create write
auto writeIndices =
getValueOrCreateConstantIndexOp(rewriter, loc, sliceOp.getMixedOffsets());
- Operation *write =
- vector::createWriteOrMaskedWrite(rewriter, loc, read, sliceOp.getDest(),
- writeIndices, inputVectorSizes.empty());
+ // For a rank-reducing slice the vector does not cover the trailing result
+ // dims, so state which result dim each vector dim writes to.
+ AffineMap writeMap;
+ if (sourceType.getRank() != resultType.getRank()) {
+ SmallVector<AffineExpr> exprs;
+ for (int64_t resultDim : resultDimsForSourceDims)
+ exprs.push_back(rewriter.getAffineDimExpr(resultDim));
+ writeMap = AffineMap::get(resultType.getRank(), /*symbolCount=*/0, exprs,
+ rewriter.getContext());
+ }
+
+ Operation *write = vector::createWriteOrMaskedWrite(
+ rewriter, loc, read, sliceOp.getDest(), writeIndices,
+ inputVectorSizes.empty(), writeMap);
// 4. Finalize
newResults.push_back(write->getResult(0));
diff --git a/mlir/lib/Dialect/Vector/Utils/VectorUtils.cpp b/mlir/lib/Dialect/Vector/Utils/VectorUtils.cpp
index c38213850c6ea5..a93c17ecbd0889 100644
--- a/mlir/lib/Dialect/Vector/Utils/VectorUtils.cpp
+++ b/mlir/lib/Dialect/Vector/Utils/VectorUtils.cpp
@@ -441,6 +441,44 @@ static SmallVector<bool> computeInBoundsFromPermutationMap(
return inBounds;
}
+/// Computes the mask type and sizes for a transfer of `vecType` at `indices`.
+///
+/// Mask dimensions follow the transfer's mask layout (see
+/// `inferTransferOpMaskType`): one per source dimension read by the permutation
+/// map, in increasing dimension order, broadcasts dropped. Each size is
+/// `dim - index`. Returns a null type for a pure broadcast: nothing to mask.
+static VectorType
+computeTransferMask(OpBuilder &builder, Location loc,
+ ArrayRef<OpFoldResult> sourceSizes, ArrayRef<Value> indices,
+ AffineMap permutationMap, VectorType vecType,
+ SmallVectorImpl<OpFoldResult> &maskSizes) {
+ SmallVector<int64_t> sourceDims;
+ if (permutationMap) {
+ for (AffineExpr expr : permutationMap.getResults())
+ if (auto dimExpr = dyn_cast<AffineDimExpr>(expr))
+ sourceDims.push_back(dimExpr.getPosition());
+ llvm::sort(sourceDims);
+ if (sourceDims.empty())
+ return VectorType();
+ } else {
+ // No map implies a minor identity: the trailing source dimensions.
+ int64_t rankDiff = sourceSizes.size() - vecType.getRank();
+ for (int64_t i = 0, e = vecType.getRank(); i < e; ++i)
+ sourceDims.push_back(rankDiff + i);
+ }
+
+ for (int64_t sourceDim : sourceDims) {
+ Value dim =
+ getValueOrCreateConstantIndexOp(builder, loc, sourceSizes[sourceDim]);
+ maskSizes.push_back(
+ builder.createOrFold<arith::SubIOp>(loc, dim, indices[sourceDim]));
+ }
+
+ if (permutationMap)
+ return vector::inferTransferOpMaskType(vecType, permutationMap);
+ return vecType.cloneWith(/*shape=*/{}, builder.getI1Type());
+}
+
Value vector::createReadOrMaskedRead(OpBuilder &builder, Location loc,
Value source,
const VectorType &vecToReadTy,
@@ -514,13 +552,19 @@ Value vector::createReadOrMaskedRead(OpBuilder &builder, Location loc,
? memref::getMixedSizes(builder, loc, source)
: tensor::getMixedSizes(builder, loc, source);
- if (isMaskTriviallyFoldable(mixedSourceDims, indices, sourceShape,
- vecToReadShape))
+ // isMaskTriviallyFoldable pairs indices with dimensions assuming a minor
+ // identity, so it does not apply when a permutation map is present.
+ if (!permutationMap && isMaskTriviallyFoldable(mixedSourceDims, indices,
+ sourceShape, vecToReadShape))
return transferReadOp;
- auto maskType = vecToReadTy.cloneWith(/*shape=*/{}, builder.getI1Type());
- Value mask =
- vector::CreateMaskOp::create(builder, loc, maskType, mixedSourceDims);
+ SmallVector<OpFoldResult> maskSizes;
+ VectorType maskType =
+ computeTransferMask(builder, loc, mixedSourceDims, indices,
+ permutationMap, vecToReadTy, maskSizes);
+ if (!maskType)
+ return transferReadOp;
+ Value mask = vector::CreateMaskOp::create(builder, loc, maskType, maskSizes);
return mlir::vector::maskOperation(builder, transferReadOp, mask)
->getResult(0);
}
@@ -586,8 +630,8 @@ Operation *vector::createWriteOrMaskedWrite(OpBuilder &builder, Location loc,
return write;
// Compute the mask and mask the write Op.
- auto writeMaskType = VectorType::get(vecToStoreShape, builder.getI1Type(),
- vecToStoreType.getScalableDims());
+ VectorType writeMaskType = VectorType::get(
+ vecToStoreShape, builder.getI1Type(), vecToStoreType.getScalableDims());
SmallVector<OpFoldResult> destSizes =
isa<MemRefType>(dest.getType())
@@ -600,18 +644,16 @@ Operation *vector::createWriteOrMaskedWrite(OpBuilder &builder, Location loc,
maskSizes = SmallVector<OpFoldResult>(destSizes.end() - vecToStoreRank,
destSizes.end());
} else {
- size_t diff = destShape.size() - vecToStoreRank;
- for (int64_t idx = 0; idx < vecToStoreRank; idx++) {
- auto value =
- getValueOrCreateConstantIndexOp(builder, loc, destSizes[diff + idx]);
- auto neg =
- builder.createOrFold<arith::SubIOp>(loc, value, writeIndices[idx]);
- maskSizes.push_back(OpFoldResult(neg));
- }
+ writeMaskType =
+ computeTransferMask(builder, loc, destSizes, writeIndices,
+ permutationMap, vecToStoreType, maskSizes);
+ if (!writeMaskType)
+ return write;
}
- if (isMaskTriviallyFoldable(maskSizes, writeIndices, destShape,
- vecToStoreShape))
+ // See the note in createReadOrMaskedRead.
+ if (!permutationMap && isMaskTriviallyFoldable(maskSizes, writeIndices,
+ destShape, vecToStoreShape))
return write;
Value maskForWrite =
diff --git a/mlir/test/Dialect/Linalg/vectorization/insert-slice.mlir b/mlir/test/Dialect/Linalg/vectorization/insert-slice.mlir
index 34e251c527e316..bc8807a5649322 100644
--- a/mlir/test/Dialect/Linalg/vectorization/insert-slice.mlir
+++ b/mlir/test/Dialect/Linalg/vectorization/insert-slice.mlir
@@ -97,13 +97,16 @@ func.func private @insert_slice_dynamic_src_dim_non_leading_unit_dim_dropped(
return %res : tensor<8x1x4xi32>
}
+// CHECK-DAG: #[[$MAP_0_2:.*]] = affine_map<(d0, d1, d2) -> (d0, d2)>
+
// 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>
+// Dim 1 is dropped, so the source dims map to result dims 0 and 2.
+// CHECK: %[[RES:.*]] = vector.transfer_write %[[READ]], %[[INIT]][%{{.*}}, %{{.*}}, %{{.*}}] {{.*}}permutation_map = #[[$MAP_0_2]]{{.*}} : vector<8x4xi32>, tensor<8x1x4xi32>
// CHECK: return %[[RES]] : tensor<8x1x4xi32>
module attributes {transform.with_named_sequence} {
@@ -131,6 +134,8 @@ func.func private @insert_slice_dynamic_src_dim_leading_unit_dim_dropped(
// 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>
+// Dim 0 is dropped, so the source dims are already the trailing result dims.
+// CHECK-NOT: permutation_map
// CHECK: %[[RES:.*]] = vector.transfer_write %[[READ]], %[[INIT]][%{{.*}}, %{{.*}}, %{{.*}}] {{.*}} : vector<8x4xi32>, tensor<1x8x4xi32>
// CHECK: return %[[RES]] : tensor<1x8x4xi32>
@@ -153,13 +158,16 @@ func.func private @insert_slice_dynamic_src_dim_trailing_unit_dim_dropped(
return %res : tensor<8x4x1xi32>
}
+// CHECK-DAG: #[[$MAP_0_1:.*]] = affine_map<(d0, d1, d2) -> (d0, d1)>
+
// 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>
+// Dim 2 is dropped, so the source dims map to result dims 0 and 1.
+// CHECK: %[[RES:.*]] = vector.transfer_write %[[READ]], %[[INIT]][%{{.*}}, %{{.*}}, %{{.*}}] {{.*}}permutation_map = #[[$MAP_0_1]]{{.*}} : vector<8x4xi32>, tensor<8x4x1xi32>
// CHECK: return %[[RES]] : tensor<8x4x1xi32>
module attributes {transform.with_named_sequence} {
@@ -181,13 +189,16 @@ func.func private @insert_slice_dynamic_src_dim_leading_and_trailing_unit_dims_d
return %res : tensor<1x8x4x1xi32>
}
+// CHECK-DAG: #[[$MAP_1_2:.*]] = affine_map<(d0, d1, d2, d3) -> (d1, d2)>
+
// 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>
+// Dims 0 and 3 are dropped, so the source dims map to result dims 1 and 2.
+// CHECK: %[[RES:.*]] = vector.transfer_write %[[READ]], %[[INIT]][%{{.*}}, %{{.*}}, %{{.*}}, %{{.*}}] {{.*}}permutation_map = #[[$MAP_1_2]]{{.*}} : vector<8x4xi32>, tensor<1x8x4x1xi32>
// CHECK: return %[[RES]] : tensor<1x8x4x1xi32>
module attributes {transform.with_named_sequence} {
@@ -334,3 +345,36 @@ 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
+ }
+ }
diff --git a/mlir/test/Integration/Dialect/Linalg/CPU/vectorize-insert-slice-rank-reducing.mlir b/mlir/test/Integration/Dialect/Linalg/CPU/vectorize-insert-slice-rank-reducing.mlir
new file mode 100644
index 00000000000000..da96d327324154
--- /dev/null
+++ b/mlir/test/Integration/Dialect/Linalg/CPU/vectorize-insert-slice-rank-reducing.mlir
@@ -0,0 +1,53 @@
+// DEFINE: %{entry_point} = main
+// DEFINE: %{lower} = -one-shot-bufferize=bufferize-function-boundaries \
+// DEFINE: -buffer-deallocation-pipeline -convert-linalg-to-loops \
+// DEFINE: -convert-vector-to-scf -lower-affine
+// DEFINE: %{run} = mlir-runner -e %{entry_point} -entry-point-result=void \
+// DEFINE: -shared-libs=%mlir_runner_utils,%mlir_c_runner_utils
+
+/// End-to-end test for vectorizing a rank-reducing tensor.insert_slice. Both
+/// runs use the same lowering and must print the same thing.
+
+/// Reference: the schedule is erased, so the insert_slice is not vectorized.
+// RUN: mlir-opt %s -test-transform-dialect-erase-schedule %{lower} \
+// RUN: | mlir-opt -test-lower-to-llvm | %{run} | FileCheck %s
+
+/// Vectorized.
+// RUN: mlir-opt %s -transform-interpreter -test-transform-dialect-erase-schedule %{lower} \
+// RUN: | mlir-opt -test-lower-to-llvm | %{run} | FileCheck %s
+
+func.func @main() {
+ %pad = arith.constant 0 : i32
+ %src = arith.constant dense<[1, 2, 3, 4, 5]> : tensor<5xi32>
+ %empty = tensor.empty() : tensor<5x3xi32>
+ %init = linalg.fill ins(%pad : i32) outs(%empty : tensor<5x3xi32>) -> tensor<5x3xi32>
+
+ /// Dim 1 of the slice is dropped, so the source's only dim corresponds to
+ /// dim 0 of the result. Writing along dim 1 instead keeps just the first
+ /// element.
+ %res = tensor.insert_slice %src into %init[0, 2] [5, 1] [1, 1]
+ : tensor<5xi32> into tensor<5x3xi32>
+
+ %res_cast = tensor.cast %res : tensor<5x3xi32> to tensor<*xi32>
+
+ // CHECK: Unranked Memref base@ = 0x{{.*}} rank = 2 offset = 0 sizes = [5, 3] strides = [3, 1] data =
+ // CHECK-NEXT: [0, 0, 1]
+ // CHECK-NEXT: [0, 0, 2]
+ // CHECK-NEXT: [0, 0, 3]
+ // CHECK-NEXT: [0, 0, 4]
+ // CHECK-NEXT: [0, 0, 5]
+ call @printMemrefI32(%res_cast) : (tensor<*xi32>) -> ()
+
+ return
+}
+
+module @transforms attributes { transform.with_named_sequence } {
+ transform.named_sequence @__transform_main(%module: !transform.any_op {transform.readonly}) {
+ %slice = transform.structured.match ops{["tensor.insert_slice"]} in %module
+ : (!transform.any_op) -> !transform.any_op
+ transform.structured.vectorize %slice : !transform.any_op
+ transform.yield
+ }
+}
+
+func.func private @printMemrefI32(%ptr : tensor<*xi32>)
More information about the Mlir-commits
mailing list