[Mlir-commits] [mlir] [mlir][vector] Do not flatten or scalarize region-masked transfers (PR #231799)

[llvmlistbot at llvm.org]

llvmorg-github-actions[bot] wrote:


<!--LLVM PR SUMMARY COMMENT-->

@llvm/pr-subscribers-mlir-vector

Author: Vaisman

<details>
<summary>Changes</summary>

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. 

---
Full diff: https://github.com/llvm/llvm-project/pull/231799.diff


3 Files Affected:

- (modified) mlir/lib/Dialect/Vector/Transforms/VectorTransferOpTransforms.cpp (+3-3) 
- (modified) mlir/test/Dialect/Vector/scalar-vector-transfer-to-memref.mlir (+23) 
- (modified) mlir/test/Dialect/Vector/vector-transfer-flatten.mlir (+43) 


``````````diff
diff --git a/mlir/lib/Dialect/Vector/Transforms/VectorTransferOpTransforms.cpp b/mlir/lib/Dialect/Vector/Transforms/VectorTransferOpTransforms.cpp
index 74b1ec98afcca..b8c2f9ab45e42 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 7fec1c6ba5642..7b451131664b9 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 b048af24acfcd..554493f220fb9 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
+}

``````````

</details>


https://github.com/llvm/llvm-project/pull/231799