[Mlir-commits] [mlir] [mlir][vector] Account for vector.mask in transfer RAW/WAW checks (PR #228328)
llvmlistbot at llvm.org
llvmlistbot at llvm.org
Fri Oct 2 06:48:21 PDT 2026
https://github.com/Vaisman updated https://github.com/llvm/llvm-project/pull/228328
>From 12d96a9969858138c5d032d314d0ca74ff71e86b Mon Sep 17 00:00:00 2001
From: Vasili Svirski <vasili.svirski at gmail.com>
Date: Fri, 25 Sep 2026 21:42:49 +0200
Subject: [PATCH 1/2] [mlir][vector] Account for vector.mask in transfer
RAW/WAW checks
The transfer RAW/WAW checks only consider mask operands and treat an
operation nested in vector.mask as unmasked. This can remove a prior
tensor write, fold a tensor read to the stored vector, or forward a
stored vector to a memref read. In each case, the IR remains valid but
the masked-off lanes get the wrong values.
Bail out of RAW checks when either operation is region-masked, and out
of WAW when the later write is region-masked. Apply the same read check
to TransferReadAfterWriteToBroadcast.
Tests cover tensor canonicalization and memref forwarding. They also
verify that an unmasked write still eliminates an earlier masked write.
---
mlir/lib/Dialect/Vector/IR/VectorOps.cpp | 12 +++-
mlir/test/Dialect/Vector/canonicalize.mlir | 58 +++++++++++++++++++
.../Dialect/Vector/vector-transferop-opt.mlir | 34 +++++++++++
3 files changed, 102 insertions(+), 2 deletions(-)
diff --git a/mlir/lib/Dialect/Vector/IR/VectorOps.cpp b/mlir/lib/Dialect/Vector/IR/VectorOps.cpp
index d5fd5095b18b00..43ff0523e1cbca 100644
--- a/mlir/lib/Dialect/Vector/IR/VectorOps.cpp
+++ b/mlir/lib/Dialect/Vector/IR/VectorOps.cpp
@@ -239,6 +239,10 @@ static bool isSplatWriteConsistentWithMaskedRead(vector::TransferWriteOp write,
bool mlir::vector::checkSameValueRAW(vector::TransferWriteOp defWrite,
vector::TransferReadOp read) {
+ // An enclosing vector.mask leaves some lanes unwritten or padded, so the
+ // read does not see the written vector as a whole.
+ if (defWrite.isMasked() || read.isMasked())
+ return false;
return !defWrite.hasOutOfBoundsDim() &&
defWrite.getIndices() == read.getIndices() &&
defWrite.getVectorType() == read.getVectorType() &&
@@ -249,6 +253,10 @@ bool mlir::vector::checkSameValueRAW(vector::TransferWriteOp defWrite,
bool mlir::vector::checkSameValueWAW(vector::TransferWriteOp write,
vector::TransferWriteOp priorWrite) {
+ // A write under an enclosing vector.mask does not overwrite all lanes of
+ // the prior write.
+ if (write.isMasked())
+ return false;
return priorWrite.getIndices() == write.getIndices() &&
priorWrite.getMask() == write.getMask() &&
priorWrite.getVectorType() == write.getVectorType() &&
@@ -5633,8 +5641,8 @@ struct TransferReadAfterWriteToBroadcast
if (!readOp.hasPureTensorSemantics() || !defWrite.hasPureTensorSemantics())
return failure();
// Bail in the masked case (too complex atm and needed to properly account
- // for padding).
- if (readOp.getMask() || defWrite.getMask())
+ // for padding). This includes an enclosing vector.mask.
+ if (readOp.getMask() || defWrite.getMask() || readOp.isMasked())
return failure();
// If indices are not the same a shift may be required, bail.
if (readOp.getIndices() != defWrite.getIndices())
diff --git a/mlir/test/Dialect/Vector/canonicalize.mlir b/mlir/test/Dialect/Vector/canonicalize.mlir
index 6b271ee5cce19d..6650036abcf93f 100644
--- a/mlir/test/Dialect/Vector/canonicalize.mlir
+++ b/mlir/test/Dialect/Vector/canonicalize.mlir
@@ -2020,6 +2020,25 @@ func.func @negative_store_to_load_tensor(%arg0 : tensor<4x4xf32>,
// -----
+// A read under vector.mask does not see the written vector in the masked-off
+// lanes, so it must not fold to it.
+// CHECK-LABEL: func @negative_store_to_load_tensor_region_masked_read
+// CHECK: vector.transfer_write
+// CHECK: %[[R:.*]] = vector.mask %{{.*}} { vector.transfer_read
+// CHECK: return %[[R]] : vector<4xf32>
+func.func @negative_store_to_load_tensor_region_masked_read(%arg0 : tensor<4xf32>,
+ %v0 : vector<4xf32>, %mask : vector<4xi1>) -> vector<4xf32> {
+ %c0 = arith.constant 0 : index
+ %cf0 = arith.constant 0.0 : f32
+ %w0 = vector.transfer_write %v0, %arg0[%c0] {in_bounds = [true]} :
+ vector<4xf32>, tensor<4xf32>
+ %0 = vector.mask %mask { vector.transfer_read %w0[%c0], %cf0 {in_bounds = [true]} :
+ tensor<4xf32>, vector<4xf32> } : vector<4xi1> -> vector<4xf32>
+ return %0 : vector<4xf32>
+}
+
+// -----
+
// CHECK-LABEL: func @store_to_load_tensor_broadcast
// CHECK-SAME: (%[[ARG:.*]]: tensor<4x4xf32>, %[[V0:.*]]: vector<4x2xf32>)
// CHECK: %[[B:.*]] = vector.broadcast %[[V0]] : vector<4x2xf32> to vector<6x4x2xf32>
@@ -2039,6 +2058,27 @@ func.func @store_to_load_tensor_broadcast(%arg0 : tensor<4x4xf32>,
// -----
+// Same as above, but the read is under vector.mask: its masked-off lanes get
+// the padding value, so it must not become a broadcast of the stored vector.
+// CHECK-LABEL: func @negative_store_to_load_tensor_broadcast_region_masked
+// CHECK: vector.transfer_write
+// CHECK: %[[R:.*]] = vector.mask %{{.*}} { vector.transfer_read
+// CHECK-NOT: vector.broadcast
+// CHECK: return %[[R]] : vector<4x2x6xf32>
+func.func @negative_store_to_load_tensor_broadcast_region_masked(%arg0 : tensor<4x4xf32>,
+ %v0 : vector<4x2xf32>, %mask : vector<4x2xi1>) -> vector<4x2x6xf32> {
+ %c0 = arith.constant 0 : index
+ %cf0 = arith.constant 0.0 : f32
+ %w0 = vector.transfer_write %v0, %arg0[%c0, %c0] {in_bounds = [true, true]} :
+ vector<4x2xf32>, tensor<4x4xf32>
+ %0 = vector.mask %mask { vector.transfer_read %w0[%c0, %c0], %cf0 {in_bounds = [true, true, true],
+ permutation_map = affine_map<(d0, d1) -> (d0, d1, 0)>} :
+ tensor<4x4xf32>, vector<4x2x6xf32> } : vector<4x2xi1> -> vector<4x2x6xf32>
+ return %0 : vector<4x2x6xf32>
+}
+
+// -----
+
// CHECK-LABEL: func @negative_store_to_load_tensor_memref
// CHECK-NOT: vector.broadcast
// CHECK-NOT: vector.transpose
@@ -2230,6 +2270,24 @@ func.func @negative_dead_store_tensor(%arg0 : tensor<4x4xf32>,
// -----
+// A write under vector.mask does not overwrite the masked-off lanes of the
+// prior write, so the prior write must be kept.
+// CHECK-LABEL: func @negative_dead_store_tensor_region_masked
+// CHECK: %[[W0:.*]] = vector.transfer_write
+// CHECK: %[[W1:.*]] = vector.mask %{{.*}} { vector.transfer_write %{{.*}}, %[[W0]]
+// CHECK: return %[[W1]]
+func.func @negative_dead_store_tensor_region_masked(%arg0 : tensor<4xf32>,
+ %v0 : vector<4xf32>, %v1 : vector<4xf32>, %mask : vector<4xi1>) -> tensor<4xf32> {
+ %c0 = arith.constant 0 : index
+ %w0 = vector.transfer_write %v0, %arg0[%c0] {in_bounds = [true]} :
+ vector<4xf32>, tensor<4xf32>
+ %w1 = vector.mask %mask { vector.transfer_write %v1, %w0[%c0] {in_bounds = [true]} :
+ vector<4xf32>, tensor<4xf32> } : vector<4xi1> -> tensor<4xf32>
+ return %w1 : tensor<4xf32>
+}
+
+// -----
+
// CHECK: #[[$MAP:[0-9a-z]+]] = affine_map<(d0, d1) -> (d1, d0)>
// CHECK-LABEL: func @swap_extract_slice_transfer_write
diff --git a/mlir/test/Dialect/Vector/vector-transferop-opt.mlir b/mlir/test/Dialect/Vector/vector-transferop-opt.mlir
index 1c6d89d1b58dd1..a8194213ced4b1 100644
--- a/mlir/test/Dialect/Vector/vector-transferop-opt.mlir
+++ b/mlir/test/Dialect/Vector/vector-transferop-opt.mlir
@@ -607,3 +607,37 @@ func.func @forward_and_eliminate_stores_through_trivial_aliases(
vector.transfer_write %23, %subview_of_cast[%c0, %c0] {in_bounds = [true, true]} : vector<[8]x[8]xf32>, memref<?x?xf32, strided<[?, 1]>>
return
}
+
+// A read under vector.mask gets the padding value in the masked-off lanes, so
+// the stored vector must not be forwarded to it.
+// CHECK-LABEL: func @negative_forward_to_region_masked_read
+// CHECK: vector.transfer_write
+// CHECK: %[[R:.*]] = vector.mask %{{.*}} { vector.transfer_read
+// CHECK: return %[[R]]
+func.func @negative_forward_to_region_masked_read(%arg0: memref<4xf32>,
+ %v0: vector<4xf32>, %mask: vector<4xi1>) -> vector<4xf32> {
+ %c0 = arith.constant 0 : index
+ %cf0 = arith.constant 0.0 : f32
+ vector.transfer_write %v0, %arg0[%c0] {in_bounds = [true]} :
+ vector<4xf32>, memref<4xf32>
+ %0 = vector.mask %mask { vector.transfer_read %arg0[%c0], %cf0 {in_bounds = [true]} :
+ memref<4xf32>, vector<4xf32> } : vector<4xi1> -> vector<4xf32>
+ return %0 : vector<4xf32>
+}
+
+// A region-masked write is still fully overwritten by a later unmasked write
+// to the same location, so it is dead.
+// CHECK-LABEL: func @dead_region_masked_store
+// CHECK-SAME: %{{.*}}: memref<4xf32>, %[[V0:.*]]: vector<4xf32>, %[[V1:.*]]: vector<4xf32>
+// CHECK-NOT: vector.transfer_write %[[V0]]
+// CHECK: vector.transfer_write %[[V1]]
+// CHECK: return
+func.func @dead_region_masked_store(%arg0: memref<4xf32>, %v0: vector<4xf32>,
+ %v1: vector<4xf32>, %mask: vector<4xi1>) {
+ %c0 = arith.constant 0 : index
+ vector.mask %mask { vector.transfer_write %v0, %arg0[%c0] {in_bounds = [true]} :
+ vector<4xf32>, memref<4xf32> } : vector<4xi1>
+ vector.transfer_write %v1, %arg0[%c0] {in_bounds = [true]} :
+ vector<4xf32>, memref<4xf32>
+ return
+}
>From 058887f319f5e44042567405536f98282d57afc3 Mon Sep 17 00:00:00 2001
From: Vasili Svirski <vasili.svirski at gmail.com>
Date: Fri, 2 Oct 2026 15:40:30 +0200
Subject: [PATCH 2/2] [mlir][vector] Clarify the WAW comment and test a masked
write before a read
Explain why only the later write needs checking in checkSameValueWAW, and
add a test for a region-masked memref write followed by an unmasked read,
which must not be forwarded.
---
mlir/lib/Dialect/Vector/IR/VectorOps.cpp | 3 ++-
.../Dialect/Vector/vector-transferop-opt.mlir | 17 +++++++++++++++++
2 files changed, 19 insertions(+), 1 deletion(-)
diff --git a/mlir/lib/Dialect/Vector/IR/VectorOps.cpp b/mlir/lib/Dialect/Vector/IR/VectorOps.cpp
index 43ff0523e1cbca..038a2b0170896a 100644
--- a/mlir/lib/Dialect/Vector/IR/VectorOps.cpp
+++ b/mlir/lib/Dialect/Vector/IR/VectorOps.cpp
@@ -254,7 +254,8 @@ bool mlir::vector::checkSameValueRAW(vector::TransferWriteOp defWrite,
bool mlir::vector::checkSameValueWAW(vector::TransferWriteOp write,
vector::TransferWriteOp priorWrite) {
// A write under an enclosing vector.mask does not overwrite all lanes of
- // the prior write.
+ // the prior write. Only the later write needs checking: a region-masked
+ // prior write is still dead if the later write fully overwrites it.
if (write.isMasked())
return false;
return priorWrite.getIndices() == write.getIndices() &&
diff --git a/mlir/test/Dialect/Vector/vector-transferop-opt.mlir b/mlir/test/Dialect/Vector/vector-transferop-opt.mlir
index a8194213ced4b1..825e445ee66121 100644
--- a/mlir/test/Dialect/Vector/vector-transferop-opt.mlir
+++ b/mlir/test/Dialect/Vector/vector-transferop-opt.mlir
@@ -625,6 +625,23 @@ func.func @negative_forward_to_region_masked_read(%arg0: memref<4xf32>,
return %0 : vector<4xf32>
}
+// A write under vector.mask leaves its masked-off lanes unwritten, so it must
+// not be forwarded to a later unmasked read.
+// CHECK-LABEL: func @negative_forward_from_region_masked_write
+// CHECK: vector.mask %{{.*}} { vector.transfer_write
+// CHECK: %[[R:.*]] = vector.transfer_read
+// CHECK: return %[[R]]
+func.func @negative_forward_from_region_masked_write(%arg0: memref<4xf32>,
+ %v0: vector<4xf32>, %mask: vector<4xi1>) -> vector<4xf32> {
+ %c0 = arith.constant 0 : index
+ %cf0 = arith.constant 0.0 : f32
+ vector.mask %mask { vector.transfer_write %v0, %arg0[%c0] {in_bounds = [true]} :
+ vector<4xf32>, memref<4xf32> } : vector<4xi1>
+ %0 = vector.transfer_read %arg0[%c0], %cf0 {in_bounds = [true]} :
+ memref<4xf32>, vector<4xf32>
+ return %0 : vector<4xf32>
+}
+
// A region-masked write is still fully overwritten by a later unmasked write
// to the same location, so it is dead.
// CHECK-LABEL: func @dead_region_masked_store
More information about the Mlir-commits
mailing list