[Mlir-commits] [mlir] d829459 - [mlir][vector] Account for vector.mask in transfer RAW/WAW checks (#228328)
llvmlistbot at llvm.org
llvmlistbot at llvm.org
Thu Oct 8 10:48:26 PDT 2026
Author: Vaisman
Date: 2026-10-08T18:48:20+01:00
New Revision: d829459f4678ef888a9f085bd26c19023c469dbb
URL: https://github.com/llvm/llvm-project/commit/d829459f4678ef888a9f085bd26c19023c469dbb
DIFF: https://github.com/llvm/llvm-project/commit/d829459f4678ef888a9f085bd26c19023c469dbb.diff
LOG: [mlir][vector] Account for vector.mask in transfer RAW/WAW checks (#228328)
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.
Added:
Modified:
mlir/lib/Dialect/Vector/IR/VectorOps.cpp
mlir/test/Dialect/Vector/canonicalize.mlir
mlir/test/Dialect/Vector/vector-transferop-opt.mlir
Removed:
################################################################################
diff --git a/mlir/lib/Dialect/Vector/IR/VectorOps.cpp b/mlir/lib/Dialect/Vector/IR/VectorOps.cpp
index d5fd5095b18b0..72bab1a37c500 100644
--- a/mlir/lib/Dialect/Vector/IR/VectorOps.cpp
+++ b/mlir/lib/Dialect/Vector/IR/VectorOps.cpp
@@ -239,6 +239,14 @@ static bool isSplatWriteConsistentWithMaskedRead(vector::TransferWriteOp write,
bool mlir::vector::checkSameValueRAW(vector::TransferWriteOp defWrite,
vector::TransferReadOp read) {
+ // An enclosing vector.mask may leave some lanes unwritten or padded, so the
+ // read may not see the written vector as a whole. This conservatively
+ // includes all-true masks. MaskOp::fold removes them during canonicalization,
+ // so this only matters when a caller such as transferOpflowOpt runs first.
+ // TODO: Account for compatible region masks instead of conservatively
+ // rejecting all region-masked operations.
+ if (defWrite.isMasked() || read.isMasked())
+ return false;
return !defWrite.hasOutOfBoundsDim() &&
defWrite.getIndices() == read.getIndices() &&
defWrite.getVectorType() == read.getVectorType() &&
@@ -249,6 +257,15 @@ bool mlir::vector::checkSameValueRAW(vector::TransferWriteOp defWrite,
bool mlir::vector::checkSameValueWAW(vector::TransferWriteOp write,
vector::TransferWriteOp priorWrite) {
+ // A write under an enclosing vector.mask may not overwrite all lanes of the
+ // prior write. This conservatively includes all-true masks, which
+ // MaskOp::fold removes during canonicalization. Only the later write needs
+ // checking: a region-masked prior write is still dead if the later write
+ // fully overwrites it.
+ // TODO: Handle a region-masked later write when it is known to fully
+ // overwrite the prior write.
+ if (write.isMasked())
+ return false;
return priorWrite.getIndices() == write.getIndices() &&
priorWrite.getMask() == write.getMask() &&
priorWrite.getVectorType() == write.getVectorType() &&
@@ -5633,8 +5650,10 @@ 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.
+ // TODO: Support mask operands and reads nested in vector.mask while
+ // properly accounting for padding.
+ 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 6b271ee5cce19..612f6dca58681 100644
--- a/mlir/test/Dialect/Vector/canonicalize.mlir
+++ b/mlir/test/Dialect/Vector/canonicalize.mlir
@@ -2020,6 +2020,31 @@ func.func @negative_store_to_load_tensor(%arg0 : tensor<4x4xf32>,
// -----
+// Same as @store_to_load_tensor, but the read is under vector.mask: its
+// masked-off lanes get the padding value, so it must not fold to the stored
+// vector.
+// CHECK-LABEL: func @negative_store_to_load_tensor_region_masked_read
+// CHECK: vector.transfer_write
+// CHECK: vector.transfer_write
+// CHECK: %[[R:.*]] = vector.mask %{{.*}} { vector.transfer_read
+// CHECK: return %[[R]] : vector<1x4xf32>
+func.func @negative_store_to_load_tensor_region_masked_read(%arg0 : tensor<4x4xf32>,
+ %v0 : vector<1x4xf32>, %v1 : vector<1x4xf32>, %mask : vector<1x4xi1>) -> vector<1x4xf32> {
+ %c1 = arith.constant 1 : index
+ %c2 = arith.constant 2 : index
+ %c0 = arith.constant 0 : index
+ %cf0 = arith.constant 0.0 : f32
+ %w0 = vector.transfer_write %v0, %arg0[%c1, %c0] {in_bounds = [true, true]} :
+ vector<1x4xf32>, tensor<4x4xf32>
+ %w1 = vector.transfer_write %v1, %w0[%c2, %c0] {in_bounds = [true, true]} :
+ vector<1x4xf32>, tensor<4x4xf32>
+ %0 = vector.mask %mask { vector.transfer_read %w1[%c1, %c0], %cf0 {in_bounds = [true, true]} :
+ tensor<4x4xf32>, vector<1x4xf32> } : vector<1x4xi1> -> vector<1x4xf32>
+ return %0 : vector<1x4xf32>
+}
+
+// -----
+
// 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 +2064,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 +2276,32 @@ func.func @negative_dead_store_tensor(%arg0 : tensor<4x4xf32>,
// -----
+// Same as @dead_store_tensor, but the last write is under vector.mask: it does
+// not overwrite the masked-off lanes, so the first write must be kept.
+// CHECK-LABEL: func @negative_dead_store_tensor_region_masked
+// CHECK-DAG: %[[C0:.*]] = arith.constant 0 : index
+// CHECK-DAG: %[[C1:.*]] = arith.constant 1 : index
+// CHECK-DAG: %[[C2:.*]] = arith.constant 2 : index
+// CHECK: vector.transfer_write {{.*}}, {{.*}}[%[[C1]], %[[C0]]
+// CHECK: vector.transfer_write {{.*}}, {{.*}}[%[[C2]], %[[C0]]
+// CHECK: %[[VTW:.*]] = vector.mask %{{.*}} { vector.transfer_write {{.*}}, {{.*}}[%[[C1]], %[[C0]]
+// CHECK: return %[[VTW]] : tensor<4x4xf32>
+func.func @negative_dead_store_tensor_region_masked(%arg0 : tensor<4x4xf32>,
+ %v0 : vector<1x4xf32>, %v1 : vector<1x4xf32>, %mask : vector<1x4xi1>) -> tensor<4x4xf32> {
+ %c1 = arith.constant 1 : index
+ %c2 = arith.constant 2 : index
+ %c0 = arith.constant 0 : index
+ %w0 = vector.transfer_write %v0, %arg0[%c1, %c0] {in_bounds = [true, true]} :
+ vector<1x4xf32>, tensor<4x4xf32>
+ %w1 = vector.transfer_write %v0, %w0[%c2, %c0] {in_bounds = [true, true]} :
+ vector<1x4xf32>, tensor<4x4xf32>
+ %w2 = vector.mask %mask { vector.transfer_write %v1, %w1[%c1, %c0] {in_bounds = [true, true]} :
+ vector<1x4xf32>, tensor<4x4xf32> } : vector<1x4xi1> -> tensor<4x4xf32>
+ return %w2 : tensor<4x4xf32>
+}
+
+// -----
+
// 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 1c6d89d1b58dd..a05c0acd88004 100644
--- a/mlir/test/Dialect/Vector/vector-transferop-opt.mlir
+++ b/mlir/test/Dialect/Vector/vector-transferop-opt.mlir
@@ -1,5 +1,99 @@
// RUN: mlir-opt %s -test-vector-transferop-opt | FileCheck %s
+//-----------------------------------------------------------------------------
+// [TransferOptimization::storeToLoadForwarding]
+//-----------------------------------------------------------------------------
+
+// CHECK-LABEL: func @forward_to_read
+// CHECK-SAME: %{{.*}}: memref<4xf32>, %[[V0:.*]]: vector<4xf32>
+// CHECK: vector.transfer_write %[[V0]]
+// CHECK-NOT: vector.transfer_read
+// CHECK: return %[[V0]]
+func.func @forward_to_read(%arg0: memref<4xf32>, %v0: vector<4xf32>)
+ -> 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.transfer_read %arg0[%c0], %cf0 {in_bounds = [true]} :
+ memref<4xf32>, vector<4xf32>
+ return %0 : vector<4xf32>
+}
+
+// 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 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>
+}
+
+//-----------------------------------------------------------------------------
+// [TransferOptimization::deadStoreOp]
+//-----------------------------------------------------------------------------
+
+// CHECK-LABEL: func @dead_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_store(%arg0: memref<4xf32>, %v0: vector<4xf32>,
+ %v1: vector<4xf32>) {
+ %c0 = arith.constant 0 : index
+ vector.transfer_write %v0, %arg0[%c0] {in_bounds = [true]} :
+ vector<4xf32>, memref<4xf32>
+ vector.transfer_write %v1, %arg0[%c0] {in_bounds = [true]} :
+ vector<4xf32>, memref<4xf32>
+ return
+}
+
+// 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
+}
+
+//-----------------------------------------------------------------------------
+// [TransferOptimization: control flow and aliasing]
+//-----------------------------------------------------------------------------
+
// CHECK-LABEL: func @forward_dead_store
// CHECK-NOT: vector.transfer_write
// CHECK-NOT: vector.transfer_read
More information about the Mlir-commits
mailing list