[mlir][vector] Prevent masked transfer read/write identity folding#192966
Merged
Conversation
zackc6
requested review from
Groverkss,
banach-space,
dcaballe and
nicolasvasilache
as code owners
Member
|
@llvm/pr-subscribers-mlir Author: zackc6 ChangesAvoid folding transfer_read + transfer_write to the base tensor when either op is masked, since masked transfer semantics are not identity. Add regression coverage for masked read-only, masked write-only, and masked read+write cases. Full diff: https://github.com/llvm/llvm-project/pull/192966.diff 2 Files Affected:
diff --git a/mlir/lib/Dialect/Vector/IR/VectorOps.cpp b/mlir/lib/Dialect/Vector/IR/VectorOps.cpp
index 3d3e49134363f..f78b3c37832c2 100644
--- a/mlir/lib/Dialect/Vector/IR/VectorOps.cpp
+++ b/mlir/lib/Dialect/Vector/IR/VectorOps.cpp
@@ -5850,6 +5850,9 @@ static LogicalResult foldReadInitWrite(TransferWriteOp write,
// Bail on potential out-of-bounds accesses.
if (read.hasOutOfBoundsDim() || write.hasOutOfBoundsDim())
return failure();
+ // Masked transfers have padding/select semantics and are not identity folds.
+ if (read.getMask() || write.getMask())
+ return failure();
// Tensor types must be the same.
if (read.getBase().getType() != rankedTensorType)
return failure();
diff --git a/mlir/test/Dialect/Vector/canonicalize.mlir b/mlir/test/Dialect/Vector/canonicalize.mlir
index de32ab45aa7ea..6aa92ab79a0dd 100644
--- a/mlir/test/Dialect/Vector/canonicalize.mlir
+++ b/mlir/test/Dialect/Vector/canonicalize.mlir
@@ -1816,6 +1816,61 @@ func.func @transfer_folding_1(%t0: tensor<2x3x4xf32>, %t1: tensor<2x3x4xf32>)
// -----
+// CHECK-LABEL: func @negative_transfer_folding_masked_read
+// CHECK: vector.transfer_read {{.*}}, {{.*}}, %[[MASK:.*]]
+// CHECK: %[[R:.*]] = vector.transfer_write
+// CHECK: return %[[R]]
+func.func @negative_transfer_folding_masked_read(
+ %t0: tensor<2x3x4xf32>, %t1: tensor<2x3x4xf32>,
+ %mask: vector<2x3x4xi1>) -> tensor<2x3x4xf32> {
+ %c0 = arith.constant 0 : index
+ %pad = arith.constant 0.0 : f32
+ %v = vector.transfer_read %t0[%c0, %c0, %c0], %pad, %mask {in_bounds = [true, true, true]} :
+ tensor<2x3x4xf32>, vector<2x3x4xf32>
+ %r = vector.transfer_write %v, %t1[%c0, %c0, %c0] {in_bounds = [true, true, true]} :
+ vector<2x3x4xf32>, tensor<2x3x4xf32>
+ return %r : tensor<2x3x4xf32>
+}
+
+// -----
+
+// CHECK-LABEL: func @negative_transfer_folding_masked_write
+// CHECK: vector.transfer_read
+// CHECK: %[[R:.*]] = vector.transfer_write {{.*}}, {{.*}}, %[[MASK:.*]]
+// CHECK: return %[[R]]
+func.func @negative_transfer_folding_masked_write(
+ %t0: tensor<2x3x4xf32>, %t1: tensor<2x3x4xf32>,
+ %mask: vector<2x3x4xi1>) -> tensor<2x3x4xf32> {
+ %c0 = arith.constant 0 : index
+ %pad = arith.constant 0.0 : f32
+ %v = vector.transfer_read %t0[%c0, %c0, %c0], %pad {in_bounds = [true, true, true]} :
+ tensor<2x3x4xf32>, vector<2x3x4xf32>
+ %r = vector.transfer_write %v, %t1[%c0, %c0, %c0], %mask {in_bounds = [true, true, true]} :
+ vector<2x3x4xf32>, tensor<2x3x4xf32>
+ return %r : tensor<2x3x4xf32>
+}
+
+// -----
+
+// CHECK-LABEL: func @negative_transfer_folding_masked_read_and_write
+// CHECK-SAME: %[[MASK:[0-9a-zA-Z_]+]]: vector<2x3x4xi1>
+// CHECK: %[[V:.*]] = vector.transfer_read {{.*}}, {{.*}}, %[[MASK]]
+// CHECK: %[[R:.*]] = vector.transfer_write %[[V]], {{.*}}, %[[MASK]]
+// CHECK: return %[[R]]
+func.func @negative_transfer_folding_masked_read_and_write(
+ %t0: tensor<2x3x4xf32>, %t1: tensor<2x3x4xf32>,
+ %mask: vector<2x3x4xi1>) -> tensor<2x3x4xf32> {
+ %c0 = arith.constant 0 : index
+ %pad = arith.constant 0.0 : f32
+ %v = vector.transfer_read %t0[%c0, %c0, %c0], %pad, %mask {in_bounds = [true, true, true]} :
+ tensor<2x3x4xf32>, vector<2x3x4xf32>
+ %r = vector.transfer_write %v, %t1[%c0, %c0, %c0], %mask {in_bounds = [true, true, true]} :
+ vector<2x3x4xf32>, tensor<2x3x4xf32>
+ return %r : tensor<2x3x4xf32>
+}
+
+// -----
+
// CHECK-LABEL: func @store_after_load_tensor
// CHECK-SAME: (%[[ARG:.*]]: tensor<4x4xf32>)
// CHECK-NOT: vector.transfer_read
|
Member
|
@llvm/pr-subscribers-mlir-vector Author: zackc6 ChangesAvoid folding transfer_read + transfer_write to the base tensor when either op is masked, since masked transfer semantics are not identity. Add regression coverage for masked read-only, masked write-only, and masked read+write cases. Full diff: https://github.com/llvm/llvm-project/pull/192966.diff 2 Files Affected:
diff --git a/mlir/lib/Dialect/Vector/IR/VectorOps.cpp b/mlir/lib/Dialect/Vector/IR/VectorOps.cpp
index 3d3e49134363f..f78b3c37832c2 100644
--- a/mlir/lib/Dialect/Vector/IR/VectorOps.cpp
+++ b/mlir/lib/Dialect/Vector/IR/VectorOps.cpp
@@ -5850,6 +5850,9 @@ static LogicalResult foldReadInitWrite(TransferWriteOp write,
// Bail on potential out-of-bounds accesses.
if (read.hasOutOfBoundsDim() || write.hasOutOfBoundsDim())
return failure();
+ // Masked transfers have padding/select semantics and are not identity folds.
+ if (read.getMask() || write.getMask())
+ return failure();
// Tensor types must be the same.
if (read.getBase().getType() != rankedTensorType)
return failure();
diff --git a/mlir/test/Dialect/Vector/canonicalize.mlir b/mlir/test/Dialect/Vector/canonicalize.mlir
index de32ab45aa7ea..6aa92ab79a0dd 100644
--- a/mlir/test/Dialect/Vector/canonicalize.mlir
+++ b/mlir/test/Dialect/Vector/canonicalize.mlir
@@ -1816,6 +1816,61 @@ func.func @transfer_folding_1(%t0: tensor<2x3x4xf32>, %t1: tensor<2x3x4xf32>)
// -----
+// CHECK-LABEL: func @negative_transfer_folding_masked_read
+// CHECK: vector.transfer_read {{.*}}, {{.*}}, %[[MASK:.*]]
+// CHECK: %[[R:.*]] = vector.transfer_write
+// CHECK: return %[[R]]
+func.func @negative_transfer_folding_masked_read(
+ %t0: tensor<2x3x4xf32>, %t1: tensor<2x3x4xf32>,
+ %mask: vector<2x3x4xi1>) -> tensor<2x3x4xf32> {
+ %c0 = arith.constant 0 : index
+ %pad = arith.constant 0.0 : f32
+ %v = vector.transfer_read %t0[%c0, %c0, %c0], %pad, %mask {in_bounds = [true, true, true]} :
+ tensor<2x3x4xf32>, vector<2x3x4xf32>
+ %r = vector.transfer_write %v, %t1[%c0, %c0, %c0] {in_bounds = [true, true, true]} :
+ vector<2x3x4xf32>, tensor<2x3x4xf32>
+ return %r : tensor<2x3x4xf32>
+}
+
+// -----
+
+// CHECK-LABEL: func @negative_transfer_folding_masked_write
+// CHECK: vector.transfer_read
+// CHECK: %[[R:.*]] = vector.transfer_write {{.*}}, {{.*}}, %[[MASK:.*]]
+// CHECK: return %[[R]]
+func.func @negative_transfer_folding_masked_write(
+ %t0: tensor<2x3x4xf32>, %t1: tensor<2x3x4xf32>,
+ %mask: vector<2x3x4xi1>) -> tensor<2x3x4xf32> {
+ %c0 = arith.constant 0 : index
+ %pad = arith.constant 0.0 : f32
+ %v = vector.transfer_read %t0[%c0, %c0, %c0], %pad {in_bounds = [true, true, true]} :
+ tensor<2x3x4xf32>, vector<2x3x4xf32>
+ %r = vector.transfer_write %v, %t1[%c0, %c0, %c0], %mask {in_bounds = [true, true, true]} :
+ vector<2x3x4xf32>, tensor<2x3x4xf32>
+ return %r : tensor<2x3x4xf32>
+}
+
+// -----
+
+// CHECK-LABEL: func @negative_transfer_folding_masked_read_and_write
+// CHECK-SAME: %[[MASK:[0-9a-zA-Z_]+]]: vector<2x3x4xi1>
+// CHECK: %[[V:.*]] = vector.transfer_read {{.*}}, {{.*}}, %[[MASK]]
+// CHECK: %[[R:.*]] = vector.transfer_write %[[V]], {{.*}}, %[[MASK]]
+// CHECK: return %[[R]]
+func.func @negative_transfer_folding_masked_read_and_write(
+ %t0: tensor<2x3x4xf32>, %t1: tensor<2x3x4xf32>,
+ %mask: vector<2x3x4xi1>) -> tensor<2x3x4xf32> {
+ %c0 = arith.constant 0 : index
+ %pad = arith.constant 0.0 : f32
+ %v = vector.transfer_read %t0[%c0, %c0, %c0], %pad, %mask {in_bounds = [true, true, true]} :
+ tensor<2x3x4xf32>, vector<2x3x4xf32>
+ %r = vector.transfer_write %v, %t1[%c0, %c0, %c0], %mask {in_bounds = [true, true, true]} :
+ vector<2x3x4xf32>, tensor<2x3x4xf32>
+ return %r : tensor<2x3x4xf32>
+}
+
+// -----
+
// CHECK-LABEL: func @store_after_load_tensor
// CHECK-SAME: (%[[ARG:.*]]: tensor<4x4xf32>)
// CHECK-NOT: vector.transfer_read
|
Groverkss
reviewed
Apr 20, 2026
Comment on lines
+5853
to
+5855
| // Masked transfers have padding/select semantics and are not identity folds. | ||
| if (read.getMask() || write.getMask()) | ||
| return failure(); |
Member
There was a problem hiding this comment.
If the padding is poison this is okay. Can you add the case?
Member
Author
There was a problem hiding this comment.
Sure, added, please check :)
zackc6
force-pushed
the
fix/mlir-vector-masked-transfer-fold
branch
from
1316df3 to
2b441ba
Compare
Avoid folding transfer_read + transfer_write to the base tensor when either op is masked, since masked transfer semantics are not identity. Add regression coverage for masked read-only, masked write-only, and masked read+write cases.
zackc6
force-pushed
the
fix/mlir-vector-masked-transfer-fold
branch
from
2b441ba to
b7f120c
Compare
Groverkss
approved these changes
Apr 22, 2026
Member
Author
|
@Groverkss thanks a lot, could you help me to merge it? |
linuxlonelyeagle
pushed a commit
to linuxlonelyeagle/llvm-project
that referenced
this pull request
Apr 23, 2026
…lvm#192966) Avoid folding transfer_read + transfer_write to the base tensor when either op is masked, since masked transfer semantics are not identity. Add regression coverage for masked read-only, masked write-only, and masked read+write cases.
yingopq
pushed a commit
to yingopq/llvm-project
that referenced
this pull request
Apr 29, 2026
…lvm#192966) Avoid folding transfer_read + transfer_write to the base tensor when either op is masked, since masked transfer semantics are not identity. Add regression coverage for masked read-only, masked write-only, and masked read+write cases.
KHicketts
pushed a commit
to KHicketts/llvm-project
that referenced
this pull request
Apr 30, 2026
…lvm#192966) Avoid folding transfer_read + transfer_write to the base tensor when either op is masked, since masked transfer semantics are not identity. Add regression coverage for masked read-only, masked write-only, and masked read+write cases.
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
3 participants
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Avoid folding transfer_read + transfer_write to the base tensor when either op is masked, since masked transfer semantics are not identity. Add regression coverage for masked read-only, masked write-only, and masked read+write cases.