[fixup] Prevent the pattern from (incorrectly) applying to masked transfers

Author: momchil-velikov

mlir/lib/Dialect/ArmSVE/Transforms/LegalizeVectorStorage.cpp

@@ -309,7 +309,7 @@ struct LegalizeSVEMaskLoadConversion : public OpRewritePattern<memref::LoadOp> {
 ///   {in_bounds = [false, true, true, true]}
 ///   : memref<?x?x2x8xi8>, vector<2x[4]x2x8xi8>
 /// ```
-/// is rewriten to
+/// is rewritten to
 /// ```
 /// %collapse_shape = memref.collapse_shape %M [[0], [1, 2, 3]]
 ///   : memref<?x?x2x8xi8> into memref<?x?xi8>
@@ -324,6 +324,21 @@ struct LegalizeTransferRead : public OpRewritePattern<vector::TransferReadOp> {
   LogicalResult matchAndRewrite(vector::TransferReadOp readOp,
                                 PatternRewriter &rewriter) const override {
 
+    // Do not try to transform masked reads. For example, if we have a transfer
+    // to a `vector<[4]x4xi8>` we could have a mask like
+    //    1 1 1 0
+    //    1 1 1 0
+    //    1 1 1 0
+    //    0 0 0 0
+    // Flattening this mask would look like
+    //    1 1 1 0 1 1 1 0 1 1 1 0 0 0 0 0
+    // and we have not yet figured out an efficient way to build such a mask,
+    // neither from the mask operand, nor from the original `vector.create_mask`
+    // operation (if visible at all).
+    if (readOp.isMasked() || readOp.getMask())
+      return rewriter.notifyMatchFailure(readOp,
+                                         "masked transfers not-supported");
+
     if (!readOp.getPermutationMap().isMinorIdentity())
       return rewriter.notifyMatchFailure(readOp, "non-identity permutation");
 

mlir/test/Dialect/ArmSVE/legalize-transfer-read.mlir

@@ -224,3 +224,39 @@ func.func @negative_test_bcast_transp(%i : index, %j : index, %M : memref<?x?x?x
 
   return %A : vector<[4]x8xi8>
 }
+
+// -----
+
+// CHECK-LABEL: @negative_test_vector_mask
+// CHECK-NOT: memref.collapse
+
+func.func @negative_test_vector_mask(
+  %i : index, %j : index,
+  %M : memref<?x?x?x8xi8>, %mask : vector<[4]x8xi1>) -> vector<[4]x8xi8> {
+  
+  %c0 = arith.constant 0 : index
+  %c0_i8 = arith.constant 0 : i8
+
+  %A = vector.mask %mask {
+    vector.transfer_read %M[%i, %j, %c0, %c0], %c0_i8 {in_bounds = [true, true] } : memref<?x?x?x8xi8>, vector<[4]x8xi8>
+  } : vector<[4]x8xi1> -> vector<[4]x8xi8>
+
+  return %A : vector<[4]x8xi8>
+}
+
+// -----
+
+// CHECK-LABEL: @negative_test_mask_operand
+// CHECK-NOT: memref.collapse
+
+func.func @negative_test_mask_operand(
+  %i : index, %j : index,
+  %M : memref<?x?x?x8xi8>, %mask : vector<[4]x8xi1>) -> vector<[4]x8xi8> {
+  
+  %c0 = arith.constant 0 : index
+  %c0_i8 = arith.constant 0 : i8
+
+  %A = vector.transfer_read %M[%i, %j, %c0, %c0], %c0_i8, %mask {in_bounds = [true, true] } : memref<?x?x?x8xi8>, vector<[4]x8xi8>
+
+  return %A : vector<[4]x8xi8>
+}