[Mlir-commits] [mlir] [mlir][vector] Do not flatten or scalarize region-masked transfers (PR #231799)
llvmlistbot at llvm.org
llvmlistbot at llvm.org
Sat Oct 17 08:39:03 PDT 2026
https://github.com/Vaisman updated https://github.com/llvm/llvm-project/pull/231799
>From bb3369610af1a3c976a486a0c5a483004a7e53eb Mon Sep 17 00:00:00 2001
From: Vasili Svirski <vasili.svirski at gmail.com>
Date: Tue, 29 Sep 2026 10:35:54 +0200
Subject: [PATCH 1/3] [mlir][vector] Do not flatten or scalarize region-masked
transfers
The transfer flattening patterns and RewriteScalarWrite already bail out
on transfers with a mask operand, but do not detect an enclosing
vector.mask. Rewriting such a transfer creates multiple operations
inside the mask region, causing verification to fail.
Also bail out on region-masked transfers, consistent with the existing
mask-operand behavior.
Regression tests cover region-masked reads and writes in the flattening
patterns and a region-masked scalar write. They fail without the fix.
---
.../Transforms/VectorTransferOpTransforms.cpp | 6 +--
.../scalar-vector-transfer-to-memref.mlir | 23 ++++++++++
.../Vector/vector-transfer-flatten.mlir | 43 +++++++++++++++++++
3 files changed, 69 insertions(+), 3 deletions(-)
diff --git a/mlir/lib/Dialect/Vector/Transforms/VectorTransferOpTransforms.cpp b/mlir/lib/Dialect/Vector/Transforms/VectorTransferOpTransforms.cpp
index 74b1ec98afccae..b8c2f9ab45e426 100644
--- a/mlir/lib/Dialect/Vector/Transforms/VectorTransferOpTransforms.cpp
+++ b/mlir/lib/Dialect/Vector/Transforms/VectorTransferOpTransforms.cpp
@@ -848,7 +848,7 @@ class FlattenContiguousRowMajorTransferReadPattern
LDBG() << " -> Not minor identity permutation map, skipping";
return failure();
}
- if (transferReadOp.getMask()) {
+ if (transferReadOp.getMask() || transferReadOp.isMasked()) {
LDBG() << " -> Has mask, skipping";
return failure();
}
@@ -962,7 +962,7 @@ class FlattenContiguousRowMajorTransferWritePattern
return failure();
if (!transferWriteOp.getPermutationMap().isMinorIdentity())
return failure();
- if (transferWriteOp.getMask())
+ if (transferWriteOp.getMask() || transferWriteOp.isMasked())
return failure();
// Determine vector dimensions to collapse.
@@ -1124,7 +1124,7 @@ class RewriteScalarWrite : public OpRewritePattern<vector::TransferWriteOp> {
if (!llvm::all_of(vecType.getShape(), [](int64_t sz) { return sz == 1; }))
return failure();
// Mask not supported.
- if (xferOp.getMask())
+ if (xferOp.getMask() || xferOp.isMasked())
return failure();
// Map not supported.
if (!xferOp.getPermutationMap().isMinorIdentity())
diff --git a/mlir/test/Dialect/Vector/scalar-vector-transfer-to-memref.mlir b/mlir/test/Dialect/Vector/scalar-vector-transfer-to-memref.mlir
index 7fec1c6ba56420..7b451131664b9f 100644
--- a/mlir/test/Dialect/Vector/scalar-vector-transfer-to-memref.mlir
+++ b/mlir/test/Dialect/Vector/scalar-vector-transfer-to-memref.mlir
@@ -178,3 +178,26 @@ func.func @transfer_read_2d_extract_dynamic(%m: memref<?x?xf32>, %row_idx: index
%elem = vector.extract %vec[%row_offset, %col_offset] : f32 from vector<10x5xf32>
return %elem : f32
}
+
+// -----
+
+// A region-masked transfer cannot be scalarized because the rewrite would
+// create multiple operations inside the vector.mask region.
+
+// CHECK-LABEL: func.func @region_masked_transfer_write
+// CHECK: vector.mask %{{.+}} { vector.transfer_write
+// CHECK: return
+
+// MULTIUSE-LABEL: func.func @region_masked_transfer_write
+// MULTIUSE: vector.mask %{{.+}} { vector.transfer_write
+// MULTIUSE: return
+func.func @region_masked_transfer_write(
+ %mem: memref<1x1xf32>, %value: vector<1x1xf32>,
+ %mask: vector<1x1xi1>) {
+ %c0 = arith.constant 0 : index
+ vector.mask %mask {
+ vector.transfer_write %value, %mem[%c0, %c0] {in_bounds = [true, true]} :
+ vector<1x1xf32>, memref<1x1xf32>
+ } : vector<1x1xi1>
+ return
+}
diff --git a/mlir/test/Dialect/Vector/vector-transfer-flatten.mlir b/mlir/test/Dialect/Vector/vector-transfer-flatten.mlir
index b048af24acfcd1..554493f220fb9b 100644
--- a/mlir/test/Dialect/Vector/vector-transfer-flatten.mlir
+++ b/mlir/test/Dialect/Vector/vector-transfer-flatten.mlir
@@ -826,3 +826,46 @@ func.func @transfer_write_multi_dim_unit_vector(
// CHECK-128B-LABEL: func @transfer_write_multi_dim_unit_vector
// CHECK-128B: vector.shape_cast {{.*}}: vector<1x1x1xi8> to vector<1xi8>
// CHECK-128B: vector.transfer_write {{.*}}: vector<1xi8>, memref<5x4x3x2xi8>
+
+// -----
+
+// Region-masked transfers cannot be flattened because the rewrite would
+// create multiple operations inside the vector.mask region.
+
+// CHECK-LABEL: func.func @region_masked_transfer_read
+// CHECK: %[[READ:.+]] = vector.mask %{{.+}} { vector.transfer_read
+// CHECK: return %[[READ]]
+
+// CHECK-128B-LABEL: func.func @region_masked_transfer_read
+// CHECK-128B: %[[READ:.+]] = vector.mask %{{.+}} { vector.transfer_read
+// CHECK-128B: return %[[READ]]
+func.func @region_masked_transfer_read(
+ %mem: memref<2x2xf32>, %mask: vector<2x2xi1>) -> vector<2x2xf32> {
+ %c0 = arith.constant 0 : index
+ %cst = arith.constant 0.0 : f32
+ %read = vector.mask %mask {
+ vector.transfer_read %mem[%c0, %c0], %cst {in_bounds = [true, true]} :
+ memref<2x2xf32>, vector<2x2xf32>
+ } : vector<2x2xi1> -> vector<2x2xf32>
+ return %read : vector<2x2xf32>
+}
+
+// -----
+
+// CHECK-LABEL: func.func @region_masked_transfer_write
+// CHECK: vector.mask %{{.+}} { vector.transfer_write
+// CHECK: return
+
+// CHECK-128B-LABEL: func.func @region_masked_transfer_write
+// CHECK-128B: vector.mask %{{.+}} { vector.transfer_write
+// CHECK-128B: return
+func.func @region_masked_transfer_write(
+ %mem: memref<2x2xf32>, %value: vector<2x2xf32>,
+ %mask: vector<2x2xi1>) {
+ %c0 = arith.constant 0 : index
+ vector.mask %mask {
+ vector.transfer_write %value, %mem[%c0, %c0] {in_bounds = [true, true]} :
+ vector<2x2xf32>, memref<2x2xf32>
+ } : vector<2x2xi1>
+ return
+}
>From 16756b1e8ddad923e83d88850860f333c31bf030 Mon Sep 17 00:00:00 2001
From: Vasili Svirski <vasili.svirski at gmail.com>
Date: Thu, 15 Oct 2026 09:43:52 +0200
Subject: [PATCH 2/3] [mlir][vector] Add region-masked scalar read test
---
.../scalar-vector-transfer-to-memref.mlir | 24 +++++++++++++++++++
1 file changed, 24 insertions(+)
diff --git a/mlir/test/Dialect/Vector/scalar-vector-transfer-to-memref.mlir b/mlir/test/Dialect/Vector/scalar-vector-transfer-to-memref.mlir
index 7b451131664b9f..7e330142c48941 100644
--- a/mlir/test/Dialect/Vector/scalar-vector-transfer-to-memref.mlir
+++ b/mlir/test/Dialect/Vector/scalar-vector-transfer-to-memref.mlir
@@ -201,3 +201,27 @@ func.func @region_masked_transfer_write(
} : vector<1x1xi1>
return
}
+
+// -----
+
+// A region-masked read is not scalarized: the vector.extract source is the
+// vector.mask result, not the transfer_read.
+
+// CHECK-LABEL: func.func @region_masked_transfer_read
+// CHECK: %[[READ:.+]] = vector.mask %{{.+}} { vector.transfer_read
+// CHECK: vector.extract %[[READ]][0, 1]
+
+// MULTIUSE-LABEL: func.func @region_masked_transfer_read
+// MULTIUSE: %[[READ:.+]] = vector.mask %{{.+}} { vector.transfer_read
+// MULTIUSE: vector.extract %[[READ]][0, 1]
+func.func @region_masked_transfer_read(
+ %mem: memref<2x2xf32>, %mask: vector<2x2xi1>) -> f32 {
+ %c0 = arith.constant 0 : index
+ %cst = arith.constant 0.0 : f32
+ %read = vector.mask %mask {
+ vector.transfer_read %mem[%c0, %c0], %cst {in_bounds = [true, true]} :
+ memref<2x2xf32>, vector<2x2xf32>
+ } : vector<2x2xi1> -> vector<2x2xf32>
+ %elem = vector.extract %read[0, 1] : f32 from vector<2x2xf32>
+ return %elem : f32
+}
>From 4f29904288aebe52b957e21337f363ea37cdcfdb Mon Sep 17 00:00:00 2001
From: Vasili Svirski <vasili.svirski at gmail.com>
Date: Sat, 17 Oct 2026 17:38:28 +0200
Subject: [PATCH 3/3] [mlir][vector] Rework masked transfer tests
---
.../scalar-vector-transfer-to-memref.mlir | 103 ++++++++-------
.../Vector/vector-transfer-flatten.mlir | 125 ++++++++++++------
2 files changed, 138 insertions(+), 90 deletions(-)
diff --git a/mlir/test/Dialect/Vector/scalar-vector-transfer-to-memref.mlir b/mlir/test/Dialect/Vector/scalar-vector-transfer-to-memref.mlir
index 7e330142c48941..1323495023ee8d 100644
--- a/mlir/test/Dialect/Vector/scalar-vector-transfer-to-memref.mlir
+++ b/mlir/test/Dialect/Vector/scalar-vector-transfer-to-memref.mlir
@@ -30,6 +30,36 @@ func.func @transfer_read_1d(%m: memref<?x?x?xf32>, %idx: index, %idx2: index) ->
// -----
+// Scalarizing masked transfers is not supported yet.
+
+// CHECK-LABEL: func @negative_transfer_read_1d_with_mask(
+// CHECK: %[[r:.*]] = vector.transfer_read
+// CHECK: vector.extract %[[r]]
+// CHECK-NOT: memref.load
+func.func @negative_transfer_read_1d_with_mask(%m: memref<?x?x?xf32>, %idx: index, %idx2: index, %mask: vector<5xi1>) -> f32 {
+ %cst = arith.constant 0.0 : f32
+ %0 = vector.transfer_read %m[%idx, %idx, %idx], %cst, %mask {in_bounds = [true]} : memref<?x?x?xf32>, vector<5xf32>
+ %1 = vector.extract %0[%idx2] : f32 from vector<5xf32>
+ return %1 : f32
+}
+
+// -----
+
+// Scalarizing masked transfers is not supported yet.
+
+// CHECK-LABEL: func @negative_transfer_read_1d_masked(
+// CHECK: %[[r:.*]] = vector.mask %{{.*}} { vector.transfer_read
+// CHECK: vector.extract %[[r]]
+// CHECK-NOT: memref.load
+func.func @negative_transfer_read_1d_masked(%m: memref<?x?x?xf32>, %idx: index, %idx2: index, %mask: vector<5xi1>) -> f32 {
+ %cst = arith.constant 0.0 : f32
+ %0 = vector.mask %mask { vector.transfer_read %m[%idx, %idx, %idx], %cst {in_bounds = [true]} : memref<?x?x?xf32>, vector<5xf32> } : vector<5xi1> -> vector<5xf32>
+ %1 = vector.extract %0[%idx2] : f32 from vector<5xf32>
+ return %1 : f32
+}
+
+// -----
+
// CHECK-LABEL: func @tensor_transfer_read_0d(
// CHECK-SAME: %[[t:.*]]: tensor<?x?x?xf32>, %[[idx:.*]]: index
// CHECK: %[[r:.*]] = tensor.extract %[[t]][%[[idx]], %[[idx]], %[[idx]]]
@@ -65,6 +95,32 @@ func.func @transfer_write_1d(%m: memref<?x?x?xf32>, %idx: index, %f: f32) {
// -----
+// Scalarizing masked transfers is not supported yet.
+
+// CHECK-LABEL: func @negative_transfer_write_1d_with_mask(
+// CHECK: vector.transfer_write
+// CHECK-NOT: memref.store
+func.func @negative_transfer_write_1d_with_mask(%m: memref<?x?x?xf32>, %idx: index, %f: f32, %mask: vector<1xi1>) {
+ %0 = vector.broadcast %f : f32 to vector<1xf32>
+ vector.transfer_write %0, %m[%idx, %idx, %idx], %mask : vector<1xf32>, memref<?x?x?xf32>
+ return
+}
+
+// -----
+
+// Scalarizing masked transfers is not supported yet.
+
+// CHECK-LABEL: func @negative_transfer_write_1d_masked(
+// CHECK: vector.mask %{{.*}} { vector.transfer_write
+// CHECK-NOT: memref.store
+func.func @negative_transfer_write_1d_masked(%m: memref<?x?x?xf32>, %idx: index, %f: f32, %mask: vector<1xi1>) {
+ %0 = vector.broadcast %f : f32 to vector<1xf32>
+ vector.mask %mask { vector.transfer_write %0, %m[%idx, %idx, %idx] : vector<1xf32>, memref<?x?x?xf32> } : vector<1xi1>
+ return
+}
+
+// -----
+
// CHECK-LABEL: func @tensor_transfer_write_0d(
// CHECK-SAME: %[[t:.*]]: tensor<?x?x?xf32>, %[[idx:.*]]: index, %[[f:.*]]: f32
// CHECK: %[[r:.*]] = tensor.insert %[[f]] into %[[t]][%[[idx]], %[[idx]], %[[idx]]]
@@ -178,50 +234,3 @@ func.func @transfer_read_2d_extract_dynamic(%m: memref<?x?xf32>, %row_idx: index
%elem = vector.extract %vec[%row_offset, %col_offset] : f32 from vector<10x5xf32>
return %elem : f32
}
-
-// -----
-
-// A region-masked transfer cannot be scalarized because the rewrite would
-// create multiple operations inside the vector.mask region.
-
-// CHECK-LABEL: func.func @region_masked_transfer_write
-// CHECK: vector.mask %{{.+}} { vector.transfer_write
-// CHECK: return
-
-// MULTIUSE-LABEL: func.func @region_masked_transfer_write
-// MULTIUSE: vector.mask %{{.+}} { vector.transfer_write
-// MULTIUSE: return
-func.func @region_masked_transfer_write(
- %mem: memref<1x1xf32>, %value: vector<1x1xf32>,
- %mask: vector<1x1xi1>) {
- %c0 = arith.constant 0 : index
- vector.mask %mask {
- vector.transfer_write %value, %mem[%c0, %c0] {in_bounds = [true, true]} :
- vector<1x1xf32>, memref<1x1xf32>
- } : vector<1x1xi1>
- return
-}
-
-// -----
-
-// A region-masked read is not scalarized: the vector.extract source is the
-// vector.mask result, not the transfer_read.
-
-// CHECK-LABEL: func.func @region_masked_transfer_read
-// CHECK: %[[READ:.+]] = vector.mask %{{.+}} { vector.transfer_read
-// CHECK: vector.extract %[[READ]][0, 1]
-
-// MULTIUSE-LABEL: func.func @region_masked_transfer_read
-// MULTIUSE: %[[READ:.+]] = vector.mask %{{.+}} { vector.transfer_read
-// MULTIUSE: vector.extract %[[READ]][0, 1]
-func.func @region_masked_transfer_read(
- %mem: memref<2x2xf32>, %mask: vector<2x2xi1>) -> f32 {
- %c0 = arith.constant 0 : index
- %cst = arith.constant 0.0 : f32
- %read = vector.mask %mask {
- vector.transfer_read %mem[%c0, %c0], %cst {in_bounds = [true, true]} :
- memref<2x2xf32>, vector<2x2xf32>
- } : vector<2x2xi1> -> vector<2x2xf32>
- %elem = vector.extract %read[0, 1] : f32 from vector<2x2xf32>
- return %elem : f32
-}
diff --git a/mlir/test/Dialect/Vector/vector-transfer-flatten.mlir b/mlir/test/Dialect/Vector/vector-transfer-flatten.mlir
index 554493f220fb9b..8287812ac9e253 100644
--- a/mlir/test/Dialect/Vector/vector-transfer-flatten.mlir
+++ b/mlir/test/Dialect/Vector/vector-transfer-flatten.mlir
@@ -33,6 +33,48 @@ func.func @transfer_read_dims_match_contiguous(
// CHECK-128B-LABEL: func @transfer_read_dims_match_contiguous
// CHECK-128B: memref.collapse_shape
+// Flattening masked transfers is not supported yet.
+
+func.func @negative_transfer_read_dims_match_contiguous_with_mask(
+ %mem : memref<5x4x3x2xi8, strided<[24, 6, 2, 1], offset: ?>>,
+ %mask : vector<5x4x3x2xi1>) -> vector<5x4x3x2xi8> {
+
+ %c0 = arith.constant 0 : index
+ %cst = arith.constant 0 : i8
+ %res = vector.transfer_read %mem[%c0, %c0, %c0, %c0], %cst, %mask :
+ memref<5x4x3x2xi8, strided<[24, 6, 2, 1], offset: ?>>, vector<5x4x3x2xi8>
+ return %res : vector<5x4x3x2xi8>
+}
+
+// CHECK-LABEL: func @negative_transfer_read_dims_match_contiguous_with_mask(
+// CHECK-NOT: memref.collapse_shape
+// CHECK-NOT: vector.shape_cast
+
+// CHECK-128B-LABEL: func @negative_transfer_read_dims_match_contiguous_with_mask(
+// CHECK-128B-NOT: memref.collapse_shape
+
+// Flattening masked transfers is not supported yet.
+
+func.func @negative_transfer_read_dims_match_contiguous_masked(
+ %mem : memref<5x4x3x2xi8, strided<[24, 6, 2, 1], offset: ?>>,
+ %mask : vector<5x4x3x2xi1>) -> vector<5x4x3x2xi8> {
+
+ %c0 = arith.constant 0 : index
+ %cst = arith.constant 0 : i8
+ %res = vector.mask %mask {
+ vector.transfer_read %mem[%c0, %c0, %c0, %c0], %cst :
+ memref<5x4x3x2xi8, strided<[24, 6, 2, 1], offset: ?>>, vector<5x4x3x2xi8>
+ } : vector<5x4x3x2xi1> -> vector<5x4x3x2xi8>
+ return %res : vector<5x4x3x2xi8>
+}
+
+// CHECK-LABEL: func @negative_transfer_read_dims_match_contiguous_masked(
+// CHECK-NOT: memref.collapse_shape
+// CHECK-NOT: vector.shape_cast
+
+// CHECK-128B-LABEL: func @negative_transfer_read_dims_match_contiguous_masked(
+// CHECK-128B-NOT: memref.collapse_shape
+
func.func @transfer_read_dims_match_contiguous_scalable(
%mem : memref<5x4x3x2xi8, strided<[24, 6, 2, 1], offset: ?>>) -> vector<5x4x3x[2]xi8> {
@@ -435,6 +477,46 @@ func.func @transfer_write_dims_match_contiguous(
// CHECK-128B-LABEL: func @transfer_write_dims_match_contiguous(
// CHECK-128B: memref.collapse_shape
+// Flattening masked transfers is not supported yet.
+
+func.func @negative_transfer_write_dims_match_contiguous_with_mask(
+ %mem : memref<5x4x3x2xi8, strided<[24, 6, 2, 1], offset: ?>>,
+ %vec : vector<5x4x3x2xi8>, %mask : vector<5x4x3x2xi1>) {
+
+ %c0 = arith.constant 0 : index
+ vector.transfer_write %vec, %mem [%c0, %c0, %c0, %c0], %mask :
+ vector<5x4x3x2xi8>, memref<5x4x3x2xi8, strided<[24, 6, 2, 1], offset: ?>>
+ return
+}
+
+// CHECK-LABEL: func @negative_transfer_write_dims_match_contiguous_with_mask(
+// CHECK-NOT: memref.collapse_shape
+// CHECK-NOT: vector.shape_cast
+
+// CHECK-128B-LABEL: func @negative_transfer_write_dims_match_contiguous_with_mask(
+// CHECK-128B-NOT: memref.collapse_shape
+
+// Flattening masked transfers is not supported yet.
+
+func.func @negative_transfer_write_dims_match_contiguous_masked(
+ %mem : memref<5x4x3x2xi8, strided<[24, 6, 2, 1], offset: ?>>,
+ %vec : vector<5x4x3x2xi8>, %mask : vector<5x4x3x2xi1>) {
+
+ %c0 = arith.constant 0 : index
+ vector.mask %mask {
+ vector.transfer_write %vec, %mem [%c0, %c0, %c0, %c0] :
+ vector<5x4x3x2xi8>, memref<5x4x3x2xi8, strided<[24, 6, 2, 1], offset: ?>>
+ } : vector<5x4x3x2xi1>
+ return
+}
+
+// CHECK-LABEL: func @negative_transfer_write_dims_match_contiguous_masked(
+// CHECK-NOT: memref.collapse_shape
+// CHECK-NOT: vector.shape_cast
+
+// CHECK-128B-LABEL: func @negative_transfer_write_dims_match_contiguous_masked(
+// CHECK-128B-NOT: memref.collapse_shape
+
func.func @transfer_write_dims_match_contiguous_scalable(
%mem : memref<5x4x3x2xi8, strided<[24, 6, 2, 1], offset: ?>>,
%vec : vector<5x4x3x[2]xi8>) {
@@ -826,46 +908,3 @@ func.func @transfer_write_multi_dim_unit_vector(
// CHECK-128B-LABEL: func @transfer_write_multi_dim_unit_vector
// CHECK-128B: vector.shape_cast {{.*}}: vector<1x1x1xi8> to vector<1xi8>
// CHECK-128B: vector.transfer_write {{.*}}: vector<1xi8>, memref<5x4x3x2xi8>
-
-// -----
-
-// Region-masked transfers cannot be flattened because the rewrite would
-// create multiple operations inside the vector.mask region.
-
-// CHECK-LABEL: func.func @region_masked_transfer_read
-// CHECK: %[[READ:.+]] = vector.mask %{{.+}} { vector.transfer_read
-// CHECK: return %[[READ]]
-
-// CHECK-128B-LABEL: func.func @region_masked_transfer_read
-// CHECK-128B: %[[READ:.+]] = vector.mask %{{.+}} { vector.transfer_read
-// CHECK-128B: return %[[READ]]
-func.func @region_masked_transfer_read(
- %mem: memref<2x2xf32>, %mask: vector<2x2xi1>) -> vector<2x2xf32> {
- %c0 = arith.constant 0 : index
- %cst = arith.constant 0.0 : f32
- %read = vector.mask %mask {
- vector.transfer_read %mem[%c0, %c0], %cst {in_bounds = [true, true]} :
- memref<2x2xf32>, vector<2x2xf32>
- } : vector<2x2xi1> -> vector<2x2xf32>
- return %read : vector<2x2xf32>
-}
-
-// -----
-
-// CHECK-LABEL: func.func @region_masked_transfer_write
-// CHECK: vector.mask %{{.+}} { vector.transfer_write
-// CHECK: return
-
-// CHECK-128B-LABEL: func.func @region_masked_transfer_write
-// CHECK-128B: vector.mask %{{.+}} { vector.transfer_write
-// CHECK-128B: return
-func.func @region_masked_transfer_write(
- %mem: memref<2x2xf32>, %value: vector<2x2xf32>,
- %mask: vector<2x2xi1>) {
- %c0 = arith.constant 0 : index
- vector.mask %mask {
- vector.transfer_write %value, %mem[%c0, %c0] {in_bounds = [true, true]} :
- vector<2x2xf32>, memref<2x2xf32>
- } : vector<2x2xi1>
- return
-}
More information about the Mlir-commits
mailing list