[Mlir-commits] [mlir] [mlir][linalg] Update createWriteOrMaskedWrite (PR #174810)
Andrzej Warzyński
llvmlistbot at llvm.org
Fri Jan 16 07:18:51 PST 2026
https://github.com/banach-space updated https://github.com/llvm/llvm-project/pull/174810
>From f5570b1f6d87c2843b45207c50e85253c61d84bd Mon Sep 17 00:00:00 2001
From: Andrzej Warzynski <andrzej.warzynski at arm.com>
Date: Wed, 7 Jan 2026 17:04:34 +0000
Subject: [PATCH 1/3] [mlir][linalg] Update createWriteOrMaskedWrite
`createWriteOrMaskedWrite` is used extensively in the Linalg vectorizer.
When a write uses non-zero indices, the helper currently computes mask
sizes as if the write started at 0 (`size = dim(d)`), which can produce
incorrect `vector.create_mask` operands for the generated
`vector.transfer_write`. Instead, the mask size should be computed as
`size = dim(d) - write_index(d)`.
EXAMPLE
-------
Let`s use this example to illustrate:
```mlir
%res = tensor.insert_slice
%src into %dest[0, %c2] [5, 1] [1, 1] : tensor<5x1xi32> into tensor<?x3xi32>
```
This op is vectorized as a pair of `vector.transfer_read` +
`vector.transfer_write` ops. When calculating the mask for the
vector.transfer_write operation, the write indices should be taken into
account (*):
```mlir
%dim = tensor.dim %dest, %c_0 : tensor<?x3xi32>
// Subtract dim-0 and idx-0
%mask_size_0 = arith.subi %dim, %c0 : index
%x3 = arith.constant 3 : index
// Subtract dim-1 and idx-1
%mask_size_1 = arith.subi %c3, %c2 : index
%mask = vector.create_mask %mask_size_0, %mask_size_1 : vector<8x1xi1>
```
Currently, the mask is incorrectly computed as:
```mlir
%dim = tensor.dim %dest, %c_0 : tensor<?x3xi32>
%mask = vector.create_mask %dim, %c3 : vector<8x1xi1>
```
This PR fixes that.
(*) Vectorized with `transform.structured.vectorize %0 vector_sizes [8, 1]`
---
.../Linalg/Transforms/Vectorization.cpp | 20 ++++++--
.../Linalg/vectorization/insert-slice.mlir | 49 ++++++++++++-------
2 files changed, 49 insertions(+), 20 deletions(-)
diff --git a/mlir/lib/Dialect/Linalg/Transforms/Vectorization.cpp b/mlir/lib/Dialect/Linalg/Transforms/Vectorization.cpp
index 2b76c24334c0a..a0ac74db4ab1a 100644
--- a/mlir/lib/Dialect/Linalg/Transforms/Vectorization.cpp
+++ b/mlir/lib/Dialect/Linalg/Transforms/Vectorization.cpp
@@ -1711,7 +1711,8 @@ createWriteOrMaskedWrite(OpBuilder &builder, Location loc, Value vecToStore,
}
// If missing, initialize the write indices to 0.
- assert((writeIndices.empty() ||
+ bool useDefaultWriteIdxs = writeIndices.empty();
+ assert((useDefaultWriteIdxs ||
writeIndices.size() == static_cast<size_t>(destRank)) &&
"Invalid number of write indices!");
if (writeIndices.empty()) {
@@ -1742,8 +1743,21 @@ createWriteOrMaskedWrite(OpBuilder &builder, Location loc, Value vecToStore,
isa<MemRefType>(dest.getType())
? memref::getMixedSizes(builder, loc, dest)
: tensor::getMixedSizes(builder, loc, dest);
- SmallVector<OpFoldResult> maskSizes(destSizes.end() - vecToStoreRank,
- destSizes.end());
+
+ // Compute sizes for write-mask
+ SmallVector<OpFoldResult> maskSizes;
+ if (useDefaultWriteIdxs) {
+ 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 = arith::SubIOp::create(builder, loc, value, writeIndices[idx]);
+ maskSizes.push_back(OpFoldResult(neg));
+ }
+ }
if (isMaskTriviallyFoldable(maskSizes, writeIndices, destShape,
vecToStoreShape))
diff --git a/mlir/test/Dialect/Linalg/vectorization/insert-slice.mlir b/mlir/test/Dialect/Linalg/vectorization/insert-slice.mlir
index 0563c21f220eb..8f96b240468ef 100644
--- a/mlir/test/Dialect/Linalg/vectorization/insert-slice.mlir
+++ b/mlir/test/Dialect/Linalg/vectorization/insert-slice.mlir
@@ -26,13 +26,17 @@ func.func private @insert_slice_static_sizes(%source: tensor<?x3x?x1xi32>) -> te
// CHECK-DAG: %[[C_1:.*]] = arith.constant 1 : index
// CHECK: %[[MASK_READ:.*]] = vector.create_mask %[[C_5]], %[[C_1]] : vector<8x1xi1>
// CHECK: %[[READ:.*]] = vector.mask %[[MASK_READ]] { vector.transfer_read %[[SRC_SLICE]][%[[C_0]], %[[C_0]]], %[[PAD]] {{.*}} : tensor<5x1xi32>, vector<8x1xi32> } : vector<8x1xi1> -> vector<8x1xi32>
-// CHECK: %[[C_0_1:.*]] = arith.constant 0 : index
+// CHECK: %[[C_0_2:.*]] = arith.constant 0 : index
// CHECK: %[[C_5_1:.*]] = arith.constant 5 : index
+// CHECK: %[[SUBI_0:.*]] = arith.subi %[[C_5_1]], %[[C_0_2]] : index
// CHECK: %[[C_3:.*]] = arith.constant 3 : index
-// CHECK: %[[MASK_WRITE:.*]] = vector.create_mask %[[C_5_1]], %[[C_3]] : vector<8x1xi1>
-// CHECK: %[[RES:.*]] = vector.mask %[[MASK_WRITE]] { vector.transfer_write %[[READ]], %[[INIT]][%[[C_0_1]], %[[C_2]]] {in_bounds = [true, true]} : vector<8x1xi32>, tensor<5x3xi32> } : vector<8x1xi1> -> tensor<5x3xi32>
+// CHECK: %[[SUBI_1:.*]] = arith.subi %[[C_3]], %[[C_2]] : index
+// CHECK: %[[MASK_WRITE:.*]] = vector.create_mask %[[SUBI_0]], %[[SUBI_1]] : vector<8x1xi1>
+// CHECK: %[[RES:.*]] = vector.mask %[[MASK_WRITE]] { vector.transfer_write %[[READ]], %[[INIT]][%[[C_0_2]], %[[C_2]]] {in_bounds = [true, true]} : vector<8x1xi32>, tensor<5x3xi32> } : vector<8x1xi1> -> 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
@@ -69,11 +73,13 @@ func.func private @insert_slice_dynamic_src_dim(%source: tensor<?x3x?x1xi32>, %s
// CHECK-DAG: %[[D0:.*]] = tensor.dim %[[SRC_SLICE]], %[[C_0_2]] : tensor<?x1xi32>
// CHECK: %[[MASK:.*]] = vector.create_mask %[[D0]], %[[C_1]] : vector<8x1xi1>
// CHECK: %[[READ:.*]] = vector.mask %[[MASK]] { vector.transfer_read %[[SRC_SLICE]][%[[C_0_1]], %[[C_0_1]]], %[[PAD]] {{.*}} : tensor<?x1xi32>, vector<8x1xi32> } : vector<8x1xi1> -> vector<8x1xi32>
-// CHECK: %[[C_0_1:.*]] = arith.constant 0 : index
-// CHECK: %[[C_5_1:.*]] = arith.constant 5 : index
-// CHECK: %[[C_3:.*]] = arith.constant 3 : index
-// CHECK: %[[MASK_WRITE:.*]] = vector.create_mask %[[C_5_1]], %[[C_3]] : vector<8x1xi1>
-// CHECK: %[[RES:.*]] = vector.mask %[[MASK_WRITE]] { vector.transfer_write %[[READ]], %[[INIT]][%[[C_0_1]], %[[C_2]]] {in_bounds = [true, true]} : vector<8x1xi32>, tensor<5x3xi32> } : vector<8x1xi1> -> tensor<5x3xi32>
+// CHECK: %[[C_0_2:.*]] = arith.constant 0 : index
+// CHECK: %[[C_5_2:.*]] = arith.constant 5 : index
+// CHECK: %[[SUBI_0:.*]] = arith.subi %[[C_5_2]], %[[C_0_2]] : index
+// CHECK: %[[C_3_1:.*]] = arith.constant 3 : index
+// CHECK: %[[SUBI_1:.*]] = arith.subi %[[C_3_1]], %[[C_2]] : index
+// CHECK: %[[MASK_WRITE:.*]] = vector.create_mask %[[SUBI_0]], %[[SUBI_1]] : vector<8x1xi1>
+// CHECK: %[[RES:.*]] = vector.mask %[[MASK_WRITE]] { vector.transfer_write %[[READ]], %[[INIT]][%[[C_0_2]], %[[C_2]]] {in_bounds = [true, true]} : vector<8x1xi32>, tensor<5x3xi32> } : vector<8x1xi1> -> tensor<5x3xi32>
// CHECK: return %[[RES]] : tensor<5x3xi32>
module attributes {transform.with_named_sequence} {
@@ -111,12 +117,16 @@ func.func private @insert_slice_dynamic_dest_dim(%source: tensor<?x3x?x1xi32>, %
// CHECK-DAG: %[[C_0:.*]] = arith.constant 0 : index
// CHECK-DAG: %[[C_0_1:.*]] = arith.constant 0 : index
// CHECK: %[[READ:.*]] = vector.mask %[[MASK]] { vector.transfer_read %[[SRC_SLICE]][%[[C_0_1]], %[[C_0_1]]], %[[PAD]] {{.*}} : tensor<5x1xi32>, vector<8x1xi32> } : vector<8x1xi1> -> vector<8x1xi32>
-// CHECK: %[[C_0_1:.*]] = arith.constant 0 : index
+
// CHECK: %[[C_0_2:.*]] = arith.constant 0 : index
-// CHECK: %[[DIM:.*]] = tensor.dim %[[INIT]], %[[C_0_2]] : tensor<?x3xi32>
+// CHECK: %[[C_0_3:.*]] = arith.constant 0 : index
+// CHECK: %[[DIM_0:.*]] = tensor.dim %[[INIT]], %[[C_0_3]] : tensor<?x3xi32>
+// CHECK: %[[SUBI_0:.*]] = arith.subi %[[DIM_0]], %[[C_0_2]] : index
// CHECK: %[[C_3:.*]] = arith.constant 3 : index
-// CHECK: %[[MASK_WRITE:.*]] = vector.create_mask %[[DIM]], %[[C_3]] : vector<8x1xi1>
-// CHECK: %[[RES:.*]] = vector.mask %[[MASK_WRITE]] { vector.transfer_write %[[READ]], %[[INIT]][%[[C_0_1]], %[[C_2]]] {in_bounds = [true, true]} : vector<8x1xi32>, tensor<?x3xi32> } : vector<8x1xi1> -> tensor<?x3xi32>
+// CHECK: %[[SUBI_1:.*]] = arith.subi %[[C_3]], %[[C_2]] : index
+// CHECK: %[[MASK_WRITE:.*]] = vector.create_mask %[[SUBI_0]], %[[SUBI_1]] : vector<8x1xi1>
+
+// CHECK: %[[RES:.*]] = vector.mask %[[MASK_WRITE]] { vector.transfer_write %[[READ]], %[[INIT]][%[[C_0_2]], %[[C_2]]] {in_bounds = [true, true]} : vector<8x1xi32>, tensor<?x3xi32> } : vector<8x1xi1> -> tensor<?x3xi32>
// CHECK: return %[[RES]] : tensor<?x3xi32>
module attributes {transform.with_named_sequence} {
@@ -155,12 +165,17 @@ func.func private @insert_slice_dynamic_source_and_dest_dim(%source: tensor<?x3x
// CHECK: %[[C1:.*]] = arith.constant 1 : index
// CHECK: %[[MASK:.*]] = vector.create_mask %[[D0]], %[[C1]] : vector<8x1xi1>
// CHECK: %[[READ:.*]] = vector.mask %[[MASK]] { vector.transfer_read %[[SRC_SLICE]][%[[C0_1]], %[[C0_1]]], %[[PAD]] {{.*}} : tensor<?x1xi32>, vector<8x1xi32> } : vector<8x1xi1> -> vector<8x1xi32>
-// CHECK: %[[C_0_1:.*]] = arith.constant 0 : index
-// CHECK: %[[C_0_2:.*]] = arith.constant 0 : index
-// CHECK: %[[DIM:.*]] = tensor.dim %[[INIT]], %[[C_0_2]] : tensor<?x3xi32>
+
+// CHECK: %[[C_0_3:.*]] = arith.constant 0 : index
+// CHECK: %[[C_0_4:.*]] = arith.constant 0 : index
+// CHECK: %[[DIM_1:.*]] = tensor.dim %[[INIT]], %[[C_0_4]] : tensor<?x3xi32>
+// CHECK: %[[SUBI_0:.*]] = arith.subi %[[DIM_1]], %[[C_0_3]] : index
// CHECK: %[[C_3:.*]] = arith.constant 3 : index
-// CHECK: %[[MASK_WRITE:.*]] = vector.create_mask %[[DIM]], %[[C_3]] : vector<8x1xi1>
-// CHECK: %[[RES:.*]] = vector.mask %[[MASK_WRITE]] { vector.transfer_write %[[READ]], %[[INIT]][%[[C_0_1]], %[[C_2]]] {in_bounds = [true, true]} : vector<8x1xi32>, tensor<?x3xi32> } : vector<8x1xi1> -> tensor<?x3xi32>
+// CHECK: %[[SUBI_1:.*]] = arith.subi %[[C_3]], %[[C_2]] : index
+// CHECK: %[[MASK_WRITE:.*]] = vector.create_mask %[[SUBI_0]], %[[SUBI_1]] : vector<8x1xi1>
+
+// CHECK: %[[RES:.*]] = vector.mask %[[MASK_WRITE]] { vector.transfer_write %[[READ]], %[[INIT]][%[[C_0_3]], %[[C_2]]] {in_bounds = [true, true]} : vector<8x1xi32>, tensor<?x3xi32> } : vector<8x1xi1> -> tensor<?x3xi32>
+
module attributes {transform.with_named_sequence} {
transform.named_sequence @__transform_main(%arg0: !transform.any_op {transform.readonly}) {
>From 1d755462df0321e8af197f8fd71baba49a670144 Mon Sep 17 00:00:00 2001
From: Andrzej Warzynski <andrzej.warzynski at arm.com>
Date: Thu, 8 Jan 2026 14:56:13 +0100
Subject: [PATCH 2/3] Replace `create` with `createOrFold`
---
.../Linalg/Transforms/Vectorization.cpp | 3 ++-
.../Linalg/vectorization/insert-slice.mlir | 27 +++++++------------
2 files changed, 12 insertions(+), 18 deletions(-)
diff --git a/mlir/lib/Dialect/Linalg/Transforms/Vectorization.cpp b/mlir/lib/Dialect/Linalg/Transforms/Vectorization.cpp
index a0ac74db4ab1a..67c9fa0e6754f 100644
--- a/mlir/lib/Dialect/Linalg/Transforms/Vectorization.cpp
+++ b/mlir/lib/Dialect/Linalg/Transforms/Vectorization.cpp
@@ -1754,7 +1754,8 @@ createWriteOrMaskedWrite(OpBuilder &builder, Location loc, Value vecToStore,
for (int64_t idx = 0; idx < vecToStoreRank; idx++) {
auto value =
getValueOrCreateConstantIndexOp(builder, loc, destSizes[diff + idx]);
- auto neg = arith::SubIOp::create(builder, loc, value, writeIndices[idx]);
+ auto neg =
+ builder.createOrFold<arith::SubIOp>(loc, value, writeIndices[idx]);
maskSizes.push_back(OpFoldResult(neg));
}
}
diff --git a/mlir/test/Dialect/Linalg/vectorization/insert-slice.mlir b/mlir/test/Dialect/Linalg/vectorization/insert-slice.mlir
index 8f96b240468ef..134f8ab3f2923 100644
--- a/mlir/test/Dialect/Linalg/vectorization/insert-slice.mlir
+++ b/mlir/test/Dialect/Linalg/vectorization/insert-slice.mlir
@@ -26,12 +26,10 @@ func.func private @insert_slice_static_sizes(%source: tensor<?x3x?x1xi32>) -> te
// CHECK-DAG: %[[C_1:.*]] = arith.constant 1 : index
// CHECK: %[[MASK_READ:.*]] = vector.create_mask %[[C_5]], %[[C_1]] : vector<8x1xi1>
// CHECK: %[[READ:.*]] = vector.mask %[[MASK_READ]] { vector.transfer_read %[[SRC_SLICE]][%[[C_0]], %[[C_0]]], %[[PAD]] {{.*}} : tensor<5x1xi32>, vector<8x1xi32> } : vector<8x1xi1> -> vector<8x1xi32>
-// CHECK: %[[C_0_2:.*]] = arith.constant 0 : index
+// CHECK-DAG: %[[C_0_2:.*]] = arith.constant 0 : index
// CHECK: %[[C_5_1:.*]] = arith.constant 5 : index
-// CHECK: %[[SUBI_0:.*]] = arith.subi %[[C_5_1]], %[[C_0_2]] : index
-// CHECK: %[[C_3:.*]] = arith.constant 3 : index
-// CHECK: %[[SUBI_1:.*]] = arith.subi %[[C_3]], %[[C_2]] : index
-// CHECK: %[[MASK_WRITE:.*]] = vector.create_mask %[[SUBI_0]], %[[SUBI_1]] : vector<8x1xi1>
+// CHECK: %[[C_1_1:.*]] = arith.constant 1 : index
+// CHECK: %[[MASK_WRITE:.*]] = vector.create_mask %[[C_5_1]], %[[C_1_1]] : vector<8x1xi1>
// CHECK: %[[RES:.*]] = vector.mask %[[MASK_WRITE]] { vector.transfer_write %[[READ]], %[[INIT]][%[[C_0_2]], %[[C_2]]] {in_bounds = [true, true]} : vector<8x1xi32>, tensor<5x3xi32> } : vector<8x1xi1> -> tensor<5x3xi32>
// CHECK: return %[[RES]] : tensor<5x3xi32>
@@ -75,10 +73,8 @@ func.func private @insert_slice_dynamic_src_dim(%source: tensor<?x3x?x1xi32>, %s
// CHECK: %[[READ:.*]] = vector.mask %[[MASK]] { vector.transfer_read %[[SRC_SLICE]][%[[C_0_1]], %[[C_0_1]]], %[[PAD]] {{.*}} : tensor<?x1xi32>, vector<8x1xi32> } : vector<8x1xi1> -> vector<8x1xi32>
// CHECK: %[[C_0_2:.*]] = arith.constant 0 : index
// CHECK: %[[C_5_2:.*]] = arith.constant 5 : index
-// CHECK: %[[SUBI_0:.*]] = arith.subi %[[C_5_2]], %[[C_0_2]] : index
-// CHECK: %[[C_3_1:.*]] = arith.constant 3 : index
-// CHECK: %[[SUBI_1:.*]] = arith.subi %[[C_3_1]], %[[C_2]] : index
-// CHECK: %[[MASK_WRITE:.*]] = vector.create_mask %[[SUBI_0]], %[[SUBI_1]] : vector<8x1xi1>
+// CHECK: %[[C_1_1:.*]] = arith.constant 1 : index
+// CHECK: %[[MASK_WRITE:.*]] = vector.create_mask %[[C_5_2]], %[[C_1_1]] : vector<8x1xi1>
// CHECK: %[[RES:.*]] = vector.mask %[[MASK_WRITE]] { vector.transfer_write %[[READ]], %[[INIT]][%[[C_0_2]], %[[C_2]]] {in_bounds = [true, true]} : vector<8x1xi32>, tensor<5x3xi32> } : vector<8x1xi1> -> tensor<5x3xi32>
// CHECK: return %[[RES]] : tensor<5x3xi32>
@@ -121,10 +117,8 @@ func.func private @insert_slice_dynamic_dest_dim(%source: tensor<?x3x?x1xi32>, %
// CHECK: %[[C_0_2:.*]] = arith.constant 0 : index
// CHECK: %[[C_0_3:.*]] = arith.constant 0 : index
// CHECK: %[[DIM_0:.*]] = tensor.dim %[[INIT]], %[[C_0_3]] : tensor<?x3xi32>
-// CHECK: %[[SUBI_0:.*]] = arith.subi %[[DIM_0]], %[[C_0_2]] : index
-// CHECK: %[[C_3:.*]] = arith.constant 3 : index
-// CHECK: %[[SUBI_1:.*]] = arith.subi %[[C_3]], %[[C_2]] : index
-// CHECK: %[[MASK_WRITE:.*]] = vector.create_mask %[[SUBI_0]], %[[SUBI_1]] : vector<8x1xi1>
+// CHECK: %[[C_1_1:.*]] = arith.constant 1 : index
+// CHECK: %[[MASK_WRITE:.*]] = vector.create_mask %[[DIM_0]], %[[C_1_1]] : vector<8x1xi1>
// CHECK: %[[RES:.*]] = vector.mask %[[MASK_WRITE]] { vector.transfer_write %[[READ]], %[[INIT]][%[[C_0_2]], %[[C_2]]] {in_bounds = [true, true]} : vector<8x1xi32>, tensor<?x3xi32> } : vector<8x1xi1> -> tensor<?x3xi32>
// CHECK: return %[[RES]] : tensor<?x3xi32>
@@ -169,10 +163,9 @@ func.func private @insert_slice_dynamic_source_and_dest_dim(%source: tensor<?x3x
// CHECK: %[[C_0_3:.*]] = arith.constant 0 : index
// CHECK: %[[C_0_4:.*]] = arith.constant 0 : index
// CHECK: %[[DIM_1:.*]] = tensor.dim %[[INIT]], %[[C_0_4]] : tensor<?x3xi32>
-// CHECK: %[[SUBI_0:.*]] = arith.subi %[[DIM_1]], %[[C_0_3]] : index
-// CHECK: %[[C_3:.*]] = arith.constant 3 : index
-// CHECK: %[[SUBI_1:.*]] = arith.subi %[[C_3]], %[[C_2]] : index
-// CHECK: %[[MASK_WRITE:.*]] = vector.create_mask %[[SUBI_0]], %[[SUBI_1]] : vector<8x1xi1>
+
+// CHECK: %[[C_1_1:.*]] = arith.constant 1 : index
+// CHECK: %[[MASK_WRITE:.*]] = vector.create_mask %[[DIM_1]], %[[C_1_1]] : vector<8x1xi1>
// CHECK: %[[RES:.*]] = vector.mask %[[MASK_WRITE]] { vector.transfer_write %[[READ]], %[[INIT]][%[[C_0_3]], %[[C_2]]] {in_bounds = [true, true]} : vector<8x1xi32>, tensor<?x3xi32> } : vector<8x1xi1> -> tensor<?x3xi32>
>From a37632ef895259db425082f96ec04fa88470ce5a Mon Sep 17 00:00:00 2001
From: Andrzej Warzynski <andrzej.warzynski at arm.com>
Date: Fri, 16 Jan 2026 15:18:35 +0000
Subject: [PATCH 3/3] Add test
---
.../Linalg/vectorization/insert-slice.mlir | 45 +++++++++++++++++++
1 file changed, 45 insertions(+)
diff --git a/mlir/test/Dialect/Linalg/vectorization/insert-slice.mlir b/mlir/test/Dialect/Linalg/vectorization/insert-slice.mlir
index 134f8ab3f2923..91ceea3b394c3 100644
--- a/mlir/test/Dialect/Linalg/vectorization/insert-slice.mlir
+++ b/mlir/test/Dialect/Linalg/vectorization/insert-slice.mlir
@@ -177,3 +177,48 @@ func.func private @insert_slice_dynamic_source_and_dest_dim(%source: tensor<?x3x
transform.yield
}
}
+
+// -----
+
+// One of the destination dimensions is dynamic and the the corresponding
+// insert index is != 0. Make sure that the mask is computed correctly (note arith.subi in the output).
+
+func.func private @insert_slice_non_zero_offset_for_dyn_dim(%source: tensor<?x3x?x1xi32>, %size: index) -> tensor<?x3xi32> {
+ %c2 = arith.constant 2 : index
+ %init = tensor.empty(%size) : tensor<?x3xi32>
+
+ %source_slice = tensor.extract_slice %source[0, %c2, 0, 0] [1, 1, 6, 1] [1, 1, 1, 1] : tensor<?x3x?x1xi32> to tensor<6x1xi32>
+ %res = tensor.insert_slice %source_slice into %init[%c2, 0] [6, 1] [1, 1] : tensor<6x1xi32> into tensor<?x3xi32>
+
+ return %res : tensor<?x3xi32>
+}
+
+// CHECK-LABEL: func.func private @insert_slice_non_zero_offset_for_dyn_dim(
+// CHECK-SAME: %[[SRC:.*]]: tensor<?x3x?x1xi32>,
+// CHECK-SAME: %[[SIZE:.*]]: index) -> tensor<?x3xi32> {
+// CHECK: %[[CST_2:.*]] = arith.constant 2 : index
+
+// CHECK: %[[INIT:.*]] = tensor.empty(%[[SIZE]]) : tensor<?x3xi32>
+// CHECK: %[[SRC_SLICE:.*]] = tensor.extract_slice %[[SRC]][0, %[[CST_2]], 0, 0]
+
+/// Read Op
+// CHECK: %[[READ:.*]] = {{.*}} vector.transfer_read %[[SRC_SLICE]]
+
+/// Mask for the write operatoin
+// CHECK: %[[CST_0:.*]] = arith.constant 0 : index
+// CHECK: %[[CST_1:.*]] = arith.constant 0 : index
+// CHECK: %[[OUT_DIM_0:.*]] = tensor.dim %[[INIT]], %[[CST_1]] : tensor<?x3xi32>
+// CHECK: %[[LB:.*]] = arith.subi %[[OUT_DIM_0]], %[[CST_2]] : index
+// CHECK: %[[CST_3:.*]] = arith.constant 3 : index
+// CHECK: %[[MASK_WRITE:.*]] = vector.create_mask %[[LB]], %[[CST_3]] : vector<8x1xi1>
+
+/// Write Op
+// CHECK: vector.mask %[[MASK_WRITE]] { vector.transfer_write %[[READ]], %[[INIT]]
+
+ 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 vector_sizes [8, 1] : !transform.any_op
+ transform.yield
+ }
+ }
More information about the Mlir-commits
mailing list