[mlir][Vector] Add an optional "masked" boolean array attribute to vector transfer operations

[nicolasvasilache]

Summary

Vector transfer ops semantic is extended to allow specifying a per-dimension masked
attribute. When the attribute is false on a particular dimension, lowering to LLVM emits
unmasked load and store operations.

[ftynse]

mlir/include/mlir/Dialect/Vector/VectorOps.td

1124

The comment looks wrong given that the builder takes a permutationMap

mlir/lib/Conversion/VectorToLLVM/ConvertVectorToLLVM.cpp

761–762

I wonder if we actually need template specialization here, or just overloading the function on the operation type would have sufficed?

mlir/lib/Dialect/Vector/VectorOps.cpp

1305

It's already an ArrayAttr, no need to cast

1308

Nit: "requires optional" sounds weird to me. "expects the optional 'masked' attr to be of <...> rank when present"?

1356

Isn't it already an ArrayAttr? In any case, I think you should just cast because the verifier ensures it's an array attribute.

[nicolasvasilache]

Address review.

[nicolasvasilache]

Rebase

[nicolasvasilache]

Messed the rebase, will revisit.

[nicolasvasilache]

Rebase

[nicolasvasilache]

Rebased properly.

[dcaballe]

Thanks, Nicolas! It looks good! Just a minor comment below.
This patch introduces an attribute to set if a vector load/store is masked but how is that mask being computed/passed to the vector load/store?
IIUC, the current implementation assumes that there is no divergent control flow and therefore the loop exit condition can be used as mask for all the vector loads/stores, is that correct?

mlir/test/Conversion/VectorToLLVM/vector-to-llvm.mlir

923

I think we should make padding optional and avoid it when masked is false. It doesn't make sense, right?

mlir/include/mlir/Dialect/Vector/VectorOps.td

let extraClassDeclaration = [{
     void getOffsets(SmallVectorImpl<int64_t> &results);
   }];
   let hasCanonicalizer = 1;
   let assemblyFormat = "$vector attr-dict `:` type($vector) `to` type(results)";
 }
 
 class Vector_TransferOpBase<int dummy> {
   code extraTransferDeclaration = [{
+    static StringRef getMaskedAttrName() { return "masked"; }
     static StringRef getPermutationMapAttrName() { return "permutation_map"; }
     MemRefType getMemRefType() {
       return memref().getType().cast<MemRefType>();
     }
     VectorType getVectorType() {
       return vector().getType().cast<VectorType>();
     }
+    bool isMaskedDim(unsigned dim) {
+      return !masked() ||
+        masked()->cast<ArrayAttr>()[dim].cast<BoolAttr>().getValue();
+    }
     /// Build the default minor identity map suitable for a vector transfer.
     /// This also handles the case memref<... x vector<...>> -> vector<...> in
     /// which the rank of tte identity map must take the vector element type
     /// into account.
     static AffineMap getTransferMinorIdentityMap(
       MemRefType memRefType, VectorType vectorType);
   }];
 }
 
 def Vector_TransferReadOp :
   Vector_Op<"transfer_read">,
   Vector_TransferOpBase<0>,
     Arguments<(ins AnyMemRef:$memref, Variadic<Index>:$indices,
-               AffineMapAttr:$permutation_map, AnyType:$padding)>,
+               AffineMapAttr:$permutation_map, AnyType:$padding,
+               OptionalAttr<BoolArrayAttr>:$masked)>,
     Results<(outs AnyVector:$vector)> {
 
   let summary = "Reads a supervector from memory into an SSA vector value.";
 
   let description = [{
-    The `vector.transfer_read` op performs a blocking read from a slice within
-    a [MemRef](../LangRef.md#memref-type) supplied as its first operand
+    The `vector.transfer_read` op performs a read from a slice within a
+    [MemRef](../LangRef.md#memref-type) supplied as its first operand
     into a [vector](../LangRef.md#vector-type) of the same base elemental type.
 
     A memref operand with vector element type, must have its vector element
     type match a suffix (shape and element type) of the vector (e.g.
     memref<3x2x6x4x3xf32>, vector<1x1x4x3xf32>).
 
     The slice is further defined by a full-rank index within the MemRef,
-    supplied as the operands `2 .. 1 + rank(memref)`. The permutation_map
-    [attribute](../LangRef.md#attributes) is an
+    supplied as the operands `2 .. 1 + rank(memref)`.
+
+    The permutation_map [attribute](../LangRef.md#attributes) is an
     [affine-map](Affine.md#affine-maps) which specifies the transposition on the
     slice to match the vector shape. The permutation map may be implicit and
     ommitted from parsing and printing if it is the canonical minor identity map
     (i.e. if it does not permute or broadcast any dimension).
 
     The size of the slice is specified by the size of the vector, given as the
     return type.
 
     An `ssa-value` of the same elemental type as the MemRef is provided as the
     last operand to specify padding in the case of out-of-bounds accesses.
 
+    An optional boolean array attribute is provided to specify which dimensions
+    of the transfer need masking. When a dimension is specified as not requiring
+    masking, the `vector.transfer_read` may be lowered to simple loads. The
+    absence of this `masked` attribute signifies that all dimensions of the
+    transfer need to be masked.
+
     This operation is called 'read' by opposition to 'load' because the
     super-vector granularity is generally not representable with a single
     hardware register. A `vector.transfer_read` is thus a mid-level abstraction
     that supports super-vectorization with non-effecting padding for full-tile
     only operations.
 
     More precisely, let's dive deeper into the permutation_map for the following
     MLIR:
%4 = vector.transfer_read %arg1[%c3, %c3], %vf0
       {permutation_map = (d0, d1)->(d0, d1)}
         : memref<?x?xvector<4x3xf32>>, vector<1x1x4x3xf32>
     ```
   }];
 
   let builders = [
     // Builder that sets padding to zero.
     OpBuilder<"OpBuilder &builder, OperationState &result, VectorType vector, "
-              "Value memref, ValueRange indices, AffineMap permutationMap">,
+              "Value memref, ValueRange indices, AffineMap permutationMap, "
+              "ArrayRef<bool> maybeMasked = {}">,
     // Builder that sets permutation map (resp. padding) to
     // 'getMinorIdentityMap' (resp. zero).
     OpBuilder<"OpBuilder &builder, OperationState &result, VectorType vector, "
-              "Value memref, ValueRange indices">
+              "Value memref, ValueRange indices, "
+              "ArrayRef<bool> maybeMasked = {}">
   ];
 
   let extraClassDeclaration = extraTransferDeclaration # [{
   }];
 }
 
 def Vector_TransferWriteOp :
   Vector_Op<"transfer_write">,
   Vector_TransferOpBase<0>,
     Arguments<(ins AnyVector:$vector, AnyMemRef:$memref,
                Variadic<Index>:$indices,
-               AffineMapAttr:$permutation_map)> {
+               AffineMapAttr:$permutation_map,
+               OptionalAttr<BoolArrayAttr>:$masked)> {
 
   let summary = "The vector.transfer_write op writes a supervector to memory.";
 
   let description = [{
-    The `vector.transfer_write` performs a blocking write from a
+    The `vector.transfer_write` op performs a write from a
     [vector](../LangRef.md#vector-type), supplied as its first operand, into a
     slice within a [MemRef](../LangRef.md#memref-type) of the same base
     elemental type, supplied as its second operand.
 
     A vector memref operand must have its vector element type match a suffix
     (shape and element type) of the vector (e.g. memref<3x2x6x4x3xf32>,
     vector<1x1x4x3xf32>).
 
     The slice is further defined by a full-rank index within the MemRef,
     supplied as the operands `3 .. 2 + rank(memref)`.
+
     The permutation_map [attribute](../LangRef.md#attributes) is an
     [affine-map](Affine.md#affine-maps) which specifies the transposition on the
     slice to match the vector shape. The permutation map may be implicit and
     ommitted from parsing and printing if it is the canonical minor identity map
     (i.e. if it does not permute or broadcast any dimension).
 
     The size of the slice is specified by the size of the vector.
 
+    An optional boolean array attribute is provided to specify which dimensions
+    of the transfer need masking. When a dimension is specified as not requiring
+    masking, the `vector.transfer_write` may be lowered to simple stores. The
+    absence of this `mask` attribute signifies that all dimensions of the
+    transfer need to be masked.
+
     This operation is called 'write' by opposition to 'store' because the
     super-vector granularity is generally not representable with a single
     hardware register. A `vector.transfer_write` is thus a
     mid-level abstraction that supports super-vectorization with non-effecting
     padding for full-tile-only code. It is the responsibility of
     `vector.transfer_write`'s implementation to ensure the memory writes are
     valid. Different lowerings may be pertinent depending on the hardware
     support.
vector.transfer_write %4, %arg1[%c3, %c3]
       {permutation_map = (d0, d1)->(d0, d1)}
         : vector<1x1x4x3xf32>, memref<?x?xvector<4x3xf32>>
     ```
   }];
 
   let builders = [
     // Builder that sets permutation map to 'getMinorIdentityMap'.
     OpBuilder<"OpBuilder &builder, OperationState &result, Value vector, "
-              "Value memref, ValueRange indices">
+              "Value memref, ValueRange indices, "
+              "ArrayRef<bool> maybeMasked = {}">,
+    OpBuilder<"OpBuilder &builder, OperationState &result, Value vector, "
+              "Value memref, ValueRange indices, AffineMap permutationMap">,
   ];
 
   let extraClassDeclaration = extraTransferDeclaration # [{
   }];
 }
 
 def Vector_ShapeCastOp :
   Vector_Op<"shape_cast", [NoSideEffect]>,

mlir/lib/Conversion/VectorToLLVM/ConvertVectorToLLVM.cpp

for (auto indexedSize : llvm::enumerate(targetMemRefType.getShape())) {
       desc.setStride(rewriter, loc, index, stride);
     }
 
     rewriter.replaceOp(op, {desc});
     return success();
   }
 };
 
-template <typename ConcreteOp>
-LogicalResult replaceTransferOp(ConversionPatternRewriter &rewriter,
-                                LLVMTypeConverter &typeConverter, Location loc,
-                                Operation *op, ArrayRef<Value> operands,
-                                Value dataPtr, Value mask);
-
 LogicalResult getLLVMTypeAndAlignment(LLVMTypeConverter &typeConverter,
                                       Type type, LLVM::LLVMType &llvmType,
                                       unsigned &align) {
   auto convertedType = typeConverter.convertType(type);
   if (!convertedType)
     return failure();
 
   llvmType = convertedType.template cast<LLVM::LLVMType>();
   auto dataLayout = typeConverter.getDialect()->getLLVMModule().getDataLayout();
   align = dataLayout.getPrefTypeAlignment(llvmType.getUnderlyingType());
   return success();
 }
 
-template <>
-LogicalResult replaceTransferOp<TransferReadOp>(
-    ConversionPatternRewriter &rewriter, LLVMTypeConverter &typeConverter,
-    Location loc, Operation *op, ArrayRef<Value> operands, Value dataPtr,
-    Value mask) {
-  auto xferOp = cast<TransferReadOp>(op);
+LogicalResult
+replaceTransferOpWithLoadOrStore(ConversionPatternRewriter &rewriter,
+                                 LLVMTypeConverter &typeConverter, Location loc,
+                                 TransferReadOp xferOp,
+                                 ArrayRef<Value> operands, Value dataPtr) {
+  LLVM::LLVMType vecTy;
+  unsigned align;
+  if (failed(getLLVMTypeAndAlignment(typeConverter, xferOp.getVectorType(),
+                                     vecTy, align)))
+    return failure();
+  rewriter.replaceOpWithNewOp<LLVM::LoadOp>(xferOp, dataPtr);
+  return success();
+}
+
+LogicalResult replaceTransferOpWithMasked(ConversionPatternRewriter &rewriter,
+                                          LLVMTypeConverter &typeConverter,
+                                          Location loc, TransferReadOp xferOp,
+                                          ArrayRef<Value> operands,
+                                          Value dataPtr, Value mask) {
   auto toLLVMTy = [&](Type t) { return typeConverter.convertType(t); };
   VectorType fillType = xferOp.getVectorType();
   Value fill = rewriter.create<SplatOp>(loc, fillType, xferOp.padding());
   fill = rewriter.create<LLVM::DialectCastOp>(loc, toLLVMTy(fillType), fill);
 
   LLVM::LLVMType vecTy;
   unsigned align;
   if (failed(getLLVMTypeAndAlignment(typeConverter, xferOp.getVectorType(),
                                      vecTy, align)))
     return failure();
 
   rewriter.replaceOpWithNewOp<LLVM::MaskedLoadOp>(
-      op, vecTy, dataPtr, mask, ValueRange{fill},
+      xferOp, vecTy, dataPtr, mask, ValueRange{fill},
       rewriter.getI32IntegerAttr(align));
   return success();
 }
 
-template <>
-LogicalResult replaceTransferOp<TransferWriteOp>(
-    ConversionPatternRewriter &rewriter, LLVMTypeConverter &typeConverter,
-    Location loc, Operation *op, ArrayRef<Value> operands, Value dataPtr,
-    Value mask) {
+LogicalResult
+replaceTransferOpWithLoadOrStore(ConversionPatternRewriter &rewriter,
+                                 LLVMTypeConverter &typeConverter, Location loc,
+                                 TransferWriteOp xferOp,
+                                 ArrayRef<Value> operands, Value dataPtr) {
   auto adaptor = TransferWriteOpOperandAdaptor(operands);
+  LLVM::LLVMType vecTy;
+  unsigned align;
+  if (failed(getLLVMTypeAndAlignment(typeConverter, xferOp.getVectorType(),
+                                     vecTy, align)))
+    return failure();
+  rewriter.replaceOpWithNewOp<LLVM::StoreOp>(xferOp, adaptor.vector(), dataPtr);
+  return success();
+}
 
-  auto xferOp = cast<TransferWriteOp>(op);
+LogicalResult replaceTransferOpWithMasked(ConversionPatternRewriter &rewriter,
+                                          LLVMTypeConverter &typeConverter,
+                                          Location loc, TransferWriteOp xferOp,
+                                          ArrayRef<Value> operands,
+                                          Value dataPtr, Value mask) {
+  auto adaptor = TransferWriteOpOperandAdaptor(operands);
   LLVM::LLVMType vecTy;
   unsigned align;
   if (failed(getLLVMTypeAndAlignment(typeConverter, xferOp.getVectorType(),
                                      vecTy, align)))
     return failure();
 
   rewriter.replaceOpWithNewOp<LLVM::MaskedStoreOp>(
-      op, adaptor.vector(), dataPtr, mask, rewriter.getI32IntegerAttr(align));
+      xferOp, adaptor.vector(), dataPtr, mask,
+      rewriter.getI32IntegerAttr(align));
   return success();
 }
 
 static TransferReadOpOperandAdaptor
 getTransferOpAdapter(TransferReadOp xferOp, ArrayRef<Value> operands) {
   return TransferReadOpOperandAdaptor(operands);
 }
 
matchAndRewrite(Operation *op, ArrayRef<Value> operands,
     Value vectorDataPtr;
     if (memRefType.getMemorySpace() == 0)
       vectorDataPtr =
           rewriter.create<LLVM::BitcastOp>(loc, vecTy.getPointerTo(), dataPtr);
     else
       vectorDataPtr = rewriter.create<LLVM::AddrSpaceCastOp>(
           loc, vecTy.getPointerTo(), dataPtr);
 
+    if (!xferOp.isMaskedDim(0))
+      return replaceTransferOpWithLoadOrStore(rewriter, typeConverter, loc,
+                                              xferOp, operands, vectorDataPtr);
+
     // 2. Create a vector with linear indices [ 0 .. vector_length - 1 ].
     unsigned vecWidth = vecTy.getVectorNumElements();
     VectorType vectorCmpType = VectorType::get(vecWidth, i64Type);
     SmallVector<int64_t, 8> indices;
     indices.reserve(vecWidth);
     for (unsigned i = 0; i < vecWidth; ++i)
       indices.push_back(i);
     Value linearIndices = rewriter.create<ConstantOp>(
matchAndRewrite(Operation *op, ArrayRef<Value> operands,
     dim = rewriter.create<IndexCastOp>(loc, i64Type, dim);
     dim = rewriter.create<SplatOp>(loc, vectorCmpType, dim);
     Value mask =
         rewriter.create<CmpIOp>(loc, CmpIPredicate::slt, offsetVector, dim);
     mask = rewriter.create<LLVM::DialectCastOp>(loc, toLLVMTy(mask.getType()),
                                                 mask);
 
     // 5. Rewrite as a masked read / write.
-    return replaceTransferOp<ConcreteOp>(rewriter, typeConverter, loc, op,
-                                         operands, vectorDataPtr, mask);
+    return replaceTransferOpWithMasked(rewriter, typeConverter, loc, xferOp,
   }
 };
 
 class VectorPrintOpConversion : public ConvertToLLVMPattern {
 public:
   explicit VectorPrintOpConversion(MLIRContext *context,
                                    LLVMTypeConverter &typeConverter)
       : ConvertToLLVMPattern(vector::PrintOp::getOperationName(), context,

mlir/lib/Conversion/VectorToSCF/VectorToSCF.cpp

 }
 
 template <typename ConcreteOp>
 template <typename LambdaThen, typename LambdaElse>
 void NDTransferOpHelper<ConcreteOp>::emitInBounds(
     ValueRange majorIvs, ValueRange majorOffsets,
     MemRefBoundsCapture &memrefBounds, LambdaThen thenBlockBuilder,
     LambdaElse elseBlockBuilder) {
-  Value inBounds = std_constant_int(/*value=*/1, /*width=*/1);
+  Value inBounds;
   SmallVector<Value, 4> majorIvsPlusOffsets;
   majorIvsPlusOffsets.reserve(majorIvs.size());
+  unsigned idx = 0;
   for (auto it : llvm::zip(majorIvs, majorOffsets, memrefBounds.getUbs())) {
     Value iv = std::get<0>(it), off = std::get<1>(it), ub = std::get<2>(it);
     using namespace mlir::edsc::op;
     majorIvsPlusOffsets.push_back(iv + off);
+    if (xferOp.isMaskedDim(leadingRank + idx)) {
       Value inBounds2 = majorIvsPlusOffsets.back() < ub;
-    inBounds = inBounds && inBounds2;
+      inBounds = (inBounds) ? (inBounds && inBounds2) : inBounds2;
+    }
+    ++idx;
   }
 
+  if (inBounds) {
     auto ifOp = ScopedContext::getBuilderRef().create<scf::IfOp>(
         ScopedContext::getLocation(), TypeRange{}, inBounds,
         /*withElseRegion=*/std::is_same<ConcreteOp, TransferReadOp>());
     BlockBuilder(&ifOp.thenRegion().front(),
                  Append())([&] { thenBlockBuilder(majorIvsPlusOffsets); });
     if (std::is_same<ConcreteOp, TransferReadOp>())
       BlockBuilder(&ifOp.elseRegion().front(),
                    Append())([&] { elseBlockBuilder(majorIvsPlusOffsets); });
+  } else {
+    // Just build the body of the then block right here.
+    thenBlockBuilder(majorIvsPlusOffsets);
+  }
 }
 
 template <> LogicalResult NDTransferOpHelper<TransferReadOp>::doReplace() {
   Value alloc = std_alloc(memRefMinorVectorType);
 
   emitLoops([&](ValueRange majorIvs, ValueRange leadingOffsets,
                 ValueRange majorOffsets, ValueRange minorOffsets,
                 MemRefBoundsCapture &memrefBounds) {
     // If in-bounds, index into memref and lower to 1-D transfer read.
     auto thenBlockBuilder = [&](ValueRange majorIvsPlusOffsets) {
       SmallVector<Value, 8> indexing;
       indexing.reserve(leadingRank + majorRank + minorRank);
       indexing.append(leadingOffsets.begin(), leadingOffsets.end());
       indexing.append(majorIvsPlusOffsets.begin(), majorIvsPlusOffsets.end());
       indexing.append(minorOffsets.begin(), minorOffsets.end());
-      // Lower to 1-D vector_transfer_read and let recursion handle it.
+
       Value memref = xferOp.memref();
       auto map = TransferReadOp::getTransferMinorIdentityMap(
           xferOp.getMemRefType(), minorVectorType);
-      auto loaded1D =
-          vector_transfer_read(minorVectorType, memref, indexing,
-                               AffineMapAttr::get(map), xferOp.padding());
+      ArrayAttr masked;
+      if (xferOp.isMaskedDim(xferOp.getVectorType().getRank() - 1)) {
+        OpBuilder &b = ScopedContext::getBuilderRef();
+        masked = b.getBoolArrayAttr({true});
+      }
+      auto loaded1D = vector_transfer_read(minorVectorType, memref, indexing,
+                                           AffineMapAttr::get(map),
+                                           xferOp.padding(), masked);
       // Store the 1-D vector.
       std_store(loaded1D, alloc, majorIvs);
     };
     // If out-of-bounds, just store a splatted vector.
     auto elseBlockBuilder = [&](ValueRange majorIvsPlusOffsets) {
       auto vector = std_splat(minorVectorType, xferOp.padding());
       std_store(vector, alloc, majorIvs);
     };
template <> LogicalResult NDTransferOpHelper<TransferWriteOp>::doReplace() {
 
   std_store(xferOp.vector(),
             vector_type_cast(MemRefType::get({}, vectorType), alloc));
 
   emitLoops([&](ValueRange majorIvs, ValueRange leadingOffsets,
                 ValueRange majorOffsets, ValueRange minorOffsets,
                 MemRefBoundsCapture &memrefBounds) {
     auto thenBlockBuilder = [&](ValueRange majorIvsPlusOffsets) {
-      // Lower to 1-D vector_transfer_write and let recursion handle it.
       SmallVector<Value, 8> indexing;
       indexing.reserve(leadingRank + majorRank + minorRank);
       indexing.append(leadingOffsets.begin(), leadingOffsets.end());
       indexing.append(majorIvsPlusOffsets.begin(), majorIvsPlusOffsets.end());
       indexing.append(minorOffsets.begin(), minorOffsets.end());
       // Lower to 1-D vector_transfer_write and let recursion handle it.
       Value loaded1D = std_load(alloc, majorIvs);
       auto map = TransferWriteOp::getTransferMinorIdentityMap(
           xferOp.getMemRefType(), minorVectorType);
+      ArrayAttr masked;
+      if (xferOp.isMaskedDim(xferOp.getVectorType().getRank() - 1)) {
+        OpBuilder &b = ScopedContext::getBuilderRef();
+        masked = b.getBoolArrayAttr({true});
+      }
       vector_transfer_write(loaded1D, xferOp.memref(), indexing,
-                            AffineMapAttr::get(map));
+                            AffineMapAttr::get(map), masked);
     };
     // Don't write anything when out of bounds.
     auto elseBlockBuilder = [&](ValueRange majorIvsPlusOffsets) {};
     emitInBounds(majorIvs, majorOffsets, memrefBounds, thenBlockBuilder,
                  elseBlockBuilder);
   });
 
   rewriter.eraseOp(op);

mlir/lib/Dialect/Affine/Transforms/SuperVectorize.cpp

if (auto store = dyn_cast<AffineStoreOp>(opInst)) {
 
     auto permutationMap =
         makePermutationMap(opInst, indices, state->strategy->loopToVectorDim);
     if (!permutationMap)
       return nullptr;
     LLVM_DEBUG(dbgs() << "\n[early-vect]+++++ permutationMap: ");
     LLVM_DEBUG(permutationMap.print(dbgs()));
     auto transfer = b.create<vector::TransferWriteOp>(
-        opInst->getLoc(), vectorValue, memRef, indices,
-        AffineMapAttr::get(permutationMap));
+        opInst->getLoc(), vectorValue, memRef, indices, permutationMap);
     auto *res = transfer.getOperation();
     LLVM_DEBUG(dbgs() << "\n[early-vect]+++++ vectorized store: " << *res);
     // "Terminals" (i.e. AffineStoreOps) are erased on the spot.
     opInst->erase();
     return res;
   }
   if (opInst->getNumRegions() != 0)
     return nullptr;

mlir/lib/Dialect/Vector/VectorOps.cpp

void ExtractStridedSliceOp::getCanonicalizationPatterns(
   // Pattern to rewrite a ExtractStridedSliceOp(ConstantMaskOp) ->
   // ConstantMaskOp.
   results.insert<StridedSliceConstantMaskFolder>(context);
 }
 
 //===----------------------------------------------------------------------===//
 // TransferReadOp
 //===----------------------------------------------------------------------===//
+
+namespace impl {
+/// Build the default minor identity map suitable for a vector transfer. This
+/// also handles the case memref<... x vector<...>> -> vector<...> in which the
+/// rank of the identity map must take the vector element type into account.
+static AffineMap getTransferMinorIdentityMap(MemRefType memRefType,
+                                             VectorType vectorType) {
+  int64_t elementVectorRank = 0;
+  VectorType elementVectorType =
+      memRefType.getElementType().dyn_cast<VectorType>();
+  if (elementVectorType)
+    elementVectorRank += elementVectorType.getRank();
+  return AffineMap::getMinorIdentityMap(
+      memRefType.getRank(), vectorType.getRank() - elementVectorRank,
+      memRefType.getContext());
+}
+} // namespace impl
+
 template <typename EmitFun>
 static LogicalResult verifyPermutationMap(AffineMap permutationMap,
                                           EmitFun emitOpError) {
   SmallVector<bool, 8> seen(permutationMap.getNumInputs(), false);
   for (auto expr : permutationMap.getResults()) {
     auto dim = expr.dyn_cast<AffineDimExpr>();
     auto zero = expr.dyn_cast<AffineConstantExpr>();
     if (zero) {
for (auto expr : permutationMap.getResults()) {
     }
     seen[dim.getPosition()] = true;
   }
   return success();
 }
 
 static LogicalResult verifyTransferOp(Operation *op, MemRefType memrefType,
                                       VectorType vectorType,
-                                      AffineMap permutationMap) {
+                                      AffineMap permutationMap,
+                                      ArrayAttr optionalMasked) {
   auto memrefElementType = memrefType.getElementType();
   if (auto memrefVectorElementType = memrefElementType.dyn_cast<VectorType>()) {
     // Memref has vector element type.
 
     // Check that 'memrefVectorElementType' and vector element types match.
     if (memrefVectorElementType.getElementType() != vectorType.getElementType())
       return op->emitOpError(
           "requires memref and vector types of the same elemental type");
static LogicalResult verifyTransferOp(Operation *op, MemRefType memrefType,
   }
 
   if (permutationMap.getNumSymbols() != 0)
     return op->emitOpError("requires permutation_map without symbols");
   if (permutationMap.getNumInputs() != memrefType.getRank())
     return op->emitOpError("requires a permutation_map with input dims of the "
                            "same rank as the memref type");
 
-  return success();
+  if (optionalMasked) {
+    if (permutationMap.getNumResults() !=
+        static_cast<int64_t>(optionalMasked.size()))
+      return op->emitOpError("expects the optional masked attr of same rank as "
+                             "permutation_map results: ")
+             << AffineMapAttr::get(permutationMap);
   }
 
-namespace impl {
-/// Build the default minor identity map suitable for a vector transfer. This
-/// also handles the case memref<... x vector<...>> -> vector<...> in which the
-/// rank of the identity map must take the vector element type into account.
-static AffineMap getTransferMinorIdentityMap(MemRefType memRefType,
-                                             VectorType vectorType) {
-  int64_t elementVectorRank = 0;
-  VectorType elementVectorType =
-      memRefType.getElementType().dyn_cast<VectorType>();
-  if (elementVectorType)
-    elementVectorRank += elementVectorType.getRank();
-  return AffineMap::getMinorIdentityMap(
-      memRefType.getRank(), vectorType.getRank() - elementVectorRank,
-      memRefType.getContext());
+  return success();
 }
-} // namespace impl
 
-/// Builder that sets permutation map and padding to 'getMinorIdentityMap' and
-/// zero, respectively, by default.
+/// Builder that sets padding to zero.
 void TransferReadOp::build(OpBuilder &builder, OperationState &result,
                            VectorType vector, Value memref, ValueRange indices,
-                           AffineMap permutationMap) {
+                           AffineMap permutationMap,
+                           ArrayRef<bool> maybeMasked) {
   Type elemType = vector.cast<VectorType>().getElementType();
   Value padding = builder.create<ConstantOp>(result.location, elemType,
                                              builder.getZeroAttr(elemType));
-  build(builder, result, vector, memref, indices, permutationMap, padding);
+  if (maybeMasked.empty())
+    return build(builder, result, vector, memref, indices, permutationMap,
+                 padding, ArrayAttr());
+  ArrayAttr maskedArrayAttr = builder.getBoolArrayAttr(maybeMasked);
+  build(builder, result, vector, memref, indices, permutationMap, padding,
+        maskedArrayAttr);
 }
 
 /// Builder that sets permutation map (resp. padding) to 'getMinorIdentityMap'
 /// (resp. zero).
 void TransferReadOp::build(OpBuilder &builder, OperationState &result,
                            VectorType vectorType, Value memref,
-                           ValueRange indices) {
-  build(builder, result, vectorType, memref, indices,
-        getTransferMinorIdentityMap(memref.getType().cast<MemRefType>(),
-                                    vectorType));
+                           ValueRange indices, ArrayRef<bool> maybeMasked) {
+  auto permMap = getTransferMinorIdentityMap(
+      memref.getType().cast<MemRefType>(), vectorType);
+  build(builder, result, vectorType, memref, indices, permMap, maybeMasked);
 }
 
 /// Build the default minor identity map suitable for a vector transfer. This
 /// also handles the case memref<... x vector<...>> -> vector<...> in which the
 /// rank of the identity map must take the vector element type into account.
 AffineMap TransferReadOp::getTransferMinorIdentityMap(MemRefType memRefType,
                                                       VectorType vectorType) {
   return ::impl::getTransferMinorIdentityMap(memRefType, vectorType);
 }
 
 template <typename TransferOp>
 void printTransferAttrs(OpAsmPrinter &p, TransferOp op) {
-  SmallVector<StringRef, 1> elidedAttrs;
+  SmallVector<StringRef, 2> elidedAttrs;
   if (AffineMap::isMinorIdentity(op.permutation_map()))
     elidedAttrs.push_back(op.getPermutationMapAttrName());
+  bool elideMasked = true;
+  if (auto maybeMasked = op.masked()) {
+    for (auto attr : *maybeMasked) {
+      if (!attr.template cast<BoolAttr>().getValue()) {
+        elideMasked = false;
+        break;
+      }
+    }
+  }
+  if (elideMasked)
+    elidedAttrs.push_back(op.getMaskedAttrName());
   p.printOptionalAttrDict(op.getAttrs(), elidedAttrs);
 }
 
 static void print(OpAsmPrinter &p, TransferReadOp op) {
   p << op.getOperationName() << " " << op.memref() << "[" << op.indices()
     << "], " << op.padding();
   printTransferAttrs(p, op);
   p << " : " << op.getMemRefType() << ", " << op.getVectorType();
static LogicalResult verify(TransferReadOp op) {
   auto paddingType = op.padding().getType();
   auto permutationMap = op.permutation_map();
   auto memrefElementType = memrefType.getElementType();
 
   if (static_cast<int64_t>(op.indices().size()) != memrefType.getRank())
     return op.emitOpError("requires ") << memrefType.getRank() << " indices";
 
   if (failed(verifyTransferOp(op.getOperation(), memrefType, vectorType,
-                              permutationMap)))
+                              permutationMap,
+                              op.masked() ? *op.masked() : ArrayAttr())))
     return failure();
 
   if (auto memrefVectorElementType = memrefElementType.dyn_cast<VectorType>()) {
     // Memref has vector element type.
     // Check that 'memrefVectorElementType' and 'paddingType' types match.
     if (memrefVectorElementType != paddingType)
       return op.emitOpError(
           "requires memref element type and padding type to match.");

 /// rank of the identity map must take the vector element type into account.
 AffineMap TransferWriteOp::getTransferMinorIdentityMap(MemRefType memRefType,
                                                        VectorType vectorType) {
   return ::impl::getTransferMinorIdentityMap(memRefType, vectorType);
 }
 
 /// Builder that sets permutation map to 'getMinorIdentityMap'.
 void TransferWriteOp::build(OpBuilder &builder, OperationState &result,
-                            Value vector, Value memref, ValueRange indices) {
+                            Value vector, Value memref, ValueRange indices,
+                            ArrayRef<bool> maybeMasked) {
   auto vectorType = vector.getType().cast<VectorType>();
   auto permMap = getTransferMinorIdentityMap(
       memref.getType().cast<MemRefType>(), vectorType);
-  build(builder, result, vector, memref, indices, permMap);
+  if (maybeMasked.empty())
+    return build(builder, result, vector, memref, indices, permMap,
+                 ArrayAttr());
+  ArrayAttr maskedArrayAttr = builder.getBoolArrayAttr(maybeMasked);
+  build(builder, result, vector, memref, indices, permMap, maskedArrayAttr);
+}
+
+/// Builder that sets permutation map to 'getMinorIdentityMap'.
+void TransferWriteOp::build(OpBuilder &builder, OperationState &result,
+                            Value vector, Value memref, ValueRange indices,
+                            AffineMap permutationMap) {
+  build(builder, result, vector, memref, indices,
+        /*maybeMasked=*/ArrayRef<bool>{});
 }
 
 static ParseResult parseTransferWriteOp(OpAsmParser &parser,
                                         OperationState &result) {
   llvm::SMLoc typesLoc;
   OpAsmParser::OperandType vectorInfo, memrefInfo;
   SmallVector<OpAsmParser::OperandType, 8> indexInfo;
   SmallVector<Type, 2> types;
static LogicalResult verify(TransferWriteOp op) {
   MemRefType memrefType = op.getMemRefType();
   VectorType vectorType = op.getVectorType();
   auto permutationMap = op.permutation_map();
 
   if (llvm::size(op.indices()) != memrefType.getRank())
     return op.emitOpError("requires ") << memrefType.getRank() << " indices";
 
   if (failed(verifyTransferOp(op.getOperation(), memrefType, vectorType,
-                              permutationMap)))
+                              permutationMap,
+                              op.masked() ? *op.masked() : ArrayAttr())))
     return failure();
 
   return verifyPermutationMap(permutationMap,
                               [&op](Twine t) { return op.emitOpError(t); });
 }
 
 //===----------------------------------------------------------------------===//
 // ShapeCastOp

mlir/lib/Dialect/Vector/VectorTransforms.cpp

LogicalResult matchAndRewrite(vector::TransferReadOp xferReadOp,
     int64_t numSlices = resultTupleType.size();
     SmallVector<Value, 4> vectorTupleValues(numSlices);
     SmallVector<Value, 4> indices(xferReadOp.indices().begin(),
                                   xferReadOp.indices().end());
     auto createSlice = [&](unsigned index, ArrayRef<Value> sliceIndices) {
       // Get VectorType for slice 'i'.
       auto sliceVectorType = resultTupleType.getType(index);
       // Create split TransferReadOp for 'sliceUser'.
+      // `masked` attribute propagates conservatively: if the coarse op didn't
+      // need masking, the fine op doesn't either.
       vectorTupleValues[index] = rewriter.create<vector::TransferReadOp>(
           loc, sliceVectorType, xferReadOp.memref(), sliceIndices,
-          xferReadOp.permutation_map(), xferReadOp.padding());
+          xferReadOp.permutation_map(), xferReadOp.padding(),
+          xferReadOp.masked() ? *xferReadOp.masked() : ArrayAttr());
     };
     generateTransferOpSlices(memrefElementType, sourceVectorType,
                              resultTupleType, sizes, strides, indices, rewriter,
                              createSlice);
 
     // Create tuple of splice xfer read operations.
     Value tupleOp = rewriter.create<vector::TupleOp>(loc, resultTupleType,
                                                      vectorTupleValues);
LogicalResult matchAndRewrite(vector::TransferWriteOp xferWriteOp,
 
     Location loc = xferWriteOp.getLoc();
     auto memrefElementType =
         xferWriteOp.memref().getType().cast<MemRefType>().getElementType();
     SmallVector<Value, 4> indices(xferWriteOp.indices().begin(),
                                   xferWriteOp.indices().end());
     auto createSlice = [&](unsigned index, ArrayRef<Value> sliceIndices) {
       // Create split TransferWriteOp for source vector 'tupleOp.operand[i]'.
+      // `masked` attribute propagates conservatively: if the coarse op didn't
+      // need masking, the fine op doesn't either.
       rewriter.create<vector::TransferWriteOp>(
           loc, tupleOp.getOperand(index), xferWriteOp.memref(), sliceIndices,
-          xferWriteOp.permutation_map());
+          xferWriteOp.permutation_map(),
+          xferWriteOp.masked() ? *xferWriteOp.masked() : ArrayAttr());
     };
     generateTransferOpSlices(memrefElementType, resultVectorType,
                              sourceTupleType, sizes, strides, indices, rewriter,
                              createSlice);
 
     // Erase old 'xferWriteOp'.
     rewriter.eraseOp(xferWriteOp);
     return success();

mlir/test/Conversion/VectorToLLVM/vector-to-llvm.mlir

 //  CHECK-SAME: !llvm<"{ float addrspace(3)*, float addrspace(3)*, i64, [1 x i64], [1 x i64] }">
 //
 // 3. Check address apce for GEP is correct.
 //       CHECK: %[[gep_b:.*]] = llvm.getelementptr {{.*}} :
 //  CHECK-SAME: (!llvm<"float addrspace(3)*">, !llvm.i64) -> !llvm<"float addrspace(3)*">
 //       CHECK: %[[vecPtr_b:.*]] = llvm.addrspacecast %[[gep_b]] :
 //  CHECK-SAME: !llvm<"float addrspace(3)*"> to !llvm<"<17 x float>*">
 
+func @transfer_read_1d_not_masked(%A : memref<?xf32>, %base: index) -> vector<17xf32> {
+  %f7 = constant 7.0: f32
+  %f = vector.transfer_read %A[%base], %f7 {masked = [false]} :
+    memref<?xf32>, vector<17xf32>
+  return %f: vector<17xf32>
+}
+// CHECK-LABEL: func @transfer_read_1d_not_masked
+//  CHECK-SAME: %[[BASE:[a-zA-Z0-9]*]]: !llvm.i64) -> !llvm<"<17 x float>">
+//
+// 1. Bitcast to vector form.
+//       CHECK: %[[gep:.*]] = llvm.getelementptr {{.*}} :
+//  CHECK-SAME: (!llvm<"float*">, !llvm.i64) -> !llvm<"float*">
+//       CHECK: %[[vecPtr:.*]] = llvm.bitcast %[[gep]] :
+//  CHECK-SAME: !llvm<"float*"> to !llvm<"<17 x float>*">
+//
+// 2. Rewrite as a load.
+//       CHECK: %[[loaded:.*]] = llvm.load %[[vecPtr]] : !llvm<"<17 x float>*">
+
 func @genbool_1d() -> vector<8xi1> {
   %0 = vector.constant_mask [4] : vector<8xi1>
   return %0 : vector<8xi1>
 }
 // CHECK-LABEL: func @genbool_1d
 // CHECK: %[[T0:.*]] = llvm.mlir.constant(1 : i1) : !llvm.i1
 // CHECK: %[[T1:.*]] = llvm.mlir.constant(dense<false> : vector<8xi1>) : !llvm<"<8 x i1>">
 // CHECK: %[[T2:.*]] = llvm.mlir.constant(0 : i64) : !llvm.i64
 // CHECK: %[[T3:.*]] = llvm.insertelement %[[T0]], %[[T1]][%[[T2]] : !llvm.i64] : !llvm<"<8 x i1>">
 // CHECK: %[[T4:.*]] = llvm.mlir.constant(1 : i64) : !llvm.i64
 // CHECK: %[[T5:.*]] = llvm.insertelement %[[T0]], %[[T3]][%[[T4]] : !llvm.i64] : !llvm<"<8 x i1>">
 // CHECK: %[[T6:.*]] = llvm.mlir.constant(2 : i64) : !llvm.i64
 // CHECK: %[[T7:.*]] = llvm.insertelement %[[T0]], %[[T5]][%[[T6]] : !llvm.i64] : !llvm<"<8 x i1>">
 // CHECK: %[[T8:.*]] = llvm.mlir.constant(3 : i64) : !llvm.i64
 // CHECK: %[[T9:.*]] = llvm.insertelement %[[T0]], %[[T7]][%[[T8]] : !llvm.i64] : !llvm<"<8 x i1>">
 // CHECK: llvm.return %9 : !llvm<"<8 x i1>">

mlir/test/Conversion/VectorToLoops/vector-to-loops.mlir

 
 // CHECK-LABEL: transfer_read_progressive(
 //  CHECK-SAME:   %[[A:[a-zA-Z0-9]+]]: memref<?x?xf32>,
 //  CHECK-SAME:   %[[base:[a-zA-Z0-9]+]]: index
 func @transfer_read_progressive(%A : memref<?x?xf32>, %base: index) -> vector<17x15xf32> {
   // CHECK: %[[cst:.*]] = constant 7.000000e+00 : f32
   %f7 = constant 7.0: f32
 
-  // CHECK-DAG: %[[cond0:.*]] = constant 1 : i1
   // CHECK-DAG: %[[splat:.*]] = constant dense<7.000000e+00> : vector<15xf32>
   // CHECK-DAG: %[[alloc:.*]] = alloc() : memref<17xvector<15xf32>>
   // CHECK-DAG: %[[dim:.*]] = dim %[[A]], 0 : memref<?x?xf32>
   // CHECK: affine.for %[[I:.*]] = 0 to 17 {
   // CHECK:   %[[add:.*]] = affine.apply #[[MAP0]](%[[I]])[%[[base]]]
-  // CHECK:   %[[cmp:.*]] = cmpi "slt", %[[add]], %[[dim]] : index
-  // CHECK:   %[[cond1:.*]] = and %[[cmp]], %[[cond0]] : i1
+  // CHECK:   %[[cond1:.*]] = cmpi "slt", %[[add]], %[[dim]] : index
   // CHECK:   scf.if %[[cond1]] {
   // CHECK:     %[[vec_1d:.*]] = vector.transfer_read %[[A]][%[[add]], %[[base]]], %[[cst]] : memref<?x?xf32>, vector<15xf32>
   // CHECK:     store %[[vec_1d]], %[[alloc]][%[[I]]] : memref<17xvector<15xf32>>
   // CHECK:   } else {
   // CHECK:     store %[[splat]], %[[alloc]][%[[I]]] : memref<17xvector<15xf32>>
   // CHECK:   }
   // CHECK: %[[vmemref:.*]] = vector.type_cast %[[alloc]] : memref<17xvector<15xf32>> to memref<vector<17x15xf32>>
   // CHECK: %[[cst:.*]] = load %[[vmemref]][] : memref<vector<17x15xf32>>

 // CHECK-DAG: #[[MAP0:.*]] = affine_map<(d0)[s0] -> (d0 + s0)>
 // CHECK-DAG: #[[MAP1:.*]] = affine_map<(d0, d1) -> (d1)>
 
 // CHECK-LABEL: transfer_write_progressive(
 //  CHECK-SAME:   %[[A:[a-zA-Z0-9]+]]: memref<?x?xf32>,
 //  CHECK-SAME:   %[[base:[a-zA-Z0-9]+]]: index,
 //  CHECK-SAME:   %[[vec:[a-zA-Z0-9]+]]: vector<17x15xf32>
 func @transfer_write_progressive(%A : memref<?x?xf32>, %base: index, %vec: vector<17x15xf32>) {
-  // CHECK: %[[cond0:.*]] = constant 1 : i1
   // CHECK: %[[alloc:.*]] = alloc() : memref<17xvector<15xf32>>
   // CHECK: %[[vmemref:.*]] = vector.type_cast %[[alloc]] : memref<17xvector<15xf32>> to memref<vector<17x15xf32>>
   // CHECK: store %[[vec]], %[[vmemref]][] : memref<vector<17x15xf32>>
   // CHECK: %[[dim:.*]] = dim %[[A]], 0 : memref<?x?xf32>
   // CHECK: affine.for %[[I:.*]] = 0 to 17 {
   // CHECK:   %[[add:.*]] = affine.apply #[[MAP0]](%[[I]])[%[[base]]]
   // CHECK:   %[[cmp:.*]] = cmpi "slt", %[[add]], %[[dim]] : index
-  // CHECK:   %[[cond1:.*]] = and %[[cmp]], %[[cond0]] : i1
-  // CHECK:   scf.if %[[cond1]] {
+  // CHECK:   scf.if %[[cmp]] {
   // CHECK:     %[[vec_1d:.*]] = load %0[%[[I]]] : memref<17xvector<15xf32>>
   // CHECK:     vector.transfer_write %[[vec_1d]], %[[A]][%[[add]], %[[base]]] : vector<15xf32>, memref<?x?xf32>
   // CHECK:   }
   vector.transfer_write %vec, %A[%base, %base]
       {permutation_map = affine_map<(d0, d1) -> (d0, d1)>} :
     vector<17x15xf32>, memref<?x?xf32>
   return
 }
+
+// -----
+
+// CHECK-DAG: #[[MAP0:.*]] = affine_map<(d0)[s0] -> (d0 + s0)>
+// CHECK-DAG: #[[MAP1:.*]] = affine_map<(d0, d1) -> (d1)>
+
+// CHECK-LABEL: transfer_write_progressive_not_masked(
+//  CHECK-SAME:   %[[A:[a-zA-Z0-9]+]]: memref<?x?xf32>,
+//  CHECK-SAME:   %[[base:[a-zA-Z0-9]+]]: index,
+//  CHECK-SAME:   %[[vec:[a-zA-Z0-9]+]]: vector<17x15xf32>
+func @transfer_write_progressive_not_masked(%A : memref<?x?xf32>, %base: index, %vec: vector<17x15xf32>) {
+  // CHECK-NOT:    scf.if
+  // CHECK-NEXT: %[[alloc:.*]] = alloc() : memref<17xvector<15xf32>>
+  // CHECK-NEXT: %[[vmemref:.*]] = vector.type_cast %[[alloc]] : memref<17xvector<15xf32>> to memref<vector<17x15xf32>>
+  // CHECK-NEXT: store %[[vec]], %[[vmemref]][] : memref<vector<17x15xf32>>
+  // CHECK-NEXT: affine.for %[[I:.*]] = 0 to 17 {
+  // CHECK-NEXT:   %[[add:.*]] = affine.apply #[[MAP0]](%[[I]])[%[[base]]]
+  // CHECK-NEXT:   %[[vec_1d:.*]] = load %0[%[[I]]] : memref<17xvector<15xf32>>
+  // CHECK-NEXT:   vector.transfer_write %[[vec_1d]], %[[A]][%[[add]], %[[base]]] : vector<15xf32>, memref<?x?xf32>
+  vector.transfer_write %vec, %A[%base, %base] {masked = [false, false]} :
+    vector<17x15xf32>, memref<?x?xf32>
+  return
+}

mlir/test/Dialect/Vector/invalid.mlir

func @test_vector.transfer_read(%arg0: memref<?x?xvector<4x3xf32>>) {
   %f0 = constant 0.0 : f32
   %vf0 = splat %f0 : vector<4x3xf32>
   // expected-error@+1 {{ requires memref vector element shape to match suffix of vector result shape}}
   %0 = vector.transfer_read %arg0[%c3, %c3], %vf0 {permutation_map = affine_map<(d0, d1)->(d0, d1)>} : memref<?x?xvector<4x3xf32>>, vector<1x1x2x3xf32>
 }
 
 // -----
 
+func @test_vector.transfer_read(%arg0: memref<?x?xvector<2x3xf32>>) {
+  %c3 = constant 3 : index
+  %f0 = constant 0.0 : f32
+  %vf0 = splat %f0 : vector<2x3xf32>
+  // expected-error@+1 {{ expects the optional masked attr of same rank as permutation_map results: affine_map<(d0, d1) -> (d0, d1)>}}
+  %0 = vector.transfer_read %arg0[%c3, %c3], %vf0 {masked = [false], permutation_map = affine_map<(d0, d1)->(d0, d1)>} : memref<?x?xvector<2x3xf32>>, vector<1x1x2x3xf32>
+}
+
+// -----
+
 func @test_vector.transfer_write(%arg0: memref<?x?xf32>) {
   %c3 = constant 3 : index
   %cst = constant dense<3.0> : vector<128 x f32>
   // expected-error@+1 {{expected 5 operand types but had 4}}
   %0 = "vector.transfer_write"(%cst, %arg0, %c3, %c3, %c3) {permutation_map = affine_map<()->(0)>} : (vector<128xf32>, memref<?x?xf32>, index, index) -> ()
 }
 
 // -----

mlir/test/Dialect/Vector/ops.mlir

func @vector_transfer_ops(%arg0: memref<?x?xf32>,
   // CHECK: vector.transfer_read
   %1 = vector.transfer_read %arg0[%c3, %c3], %f0 {permutation_map = affine_map<(d0, d1)->(d1, d0)>} : memref<?x?xf32>, vector<3x7xf32>
   // CHECK: vector.transfer_read
   %2 = vector.transfer_read %arg0[%c3, %c3], %cst {permutation_map = affine_map<(d0, d1)->(d0)>} : memref<?x?xf32>,  vector<128xf32>
   // CHECK: vector.transfer_read
   %3 = vector.transfer_read %arg0[%c3, %c3], %cst {permutation_map = affine_map<(d0, d1)->(d1)>} : memref<?x?xf32>,  vector<128xf32>
   // CHECK: vector.transfer_read %{{.*}}[%[[C3]], %[[C3]]], %{{.*}} : memref<?x?xvector<4x3xf32>>, vector<1x1x4x3xf32>
   %4 = vector.transfer_read %arg1[%c3, %c3], %vf0 {permutation_map = affine_map<(d0, d1)->(d0, d1)>} : memref<?x?xvector<4x3xf32>>, vector<1x1x4x3xf32>
+  // CHECK: vector.transfer_read %{{.*}}[%[[C3]], %[[C3]]], %{{.*}} {masked = [true, false]} : memref<?x?xvector<4x3xf32>>, vector<1x1x4x3xf32>
+  %5 = vector.transfer_read %arg1[%c3, %c3], %vf0 {masked = [true, false]} : memref<?x?xvector<4x3xf32>>, vector<1x1x4x3xf32>
 
   // CHECK: vector.transfer_write
   vector.transfer_write %0, %arg0[%c3, %c3] {permutation_map = affine_map<(d0, d1)->(d0)>} : vector<128xf32>, memref<?x?xf32>
   // CHECK: vector.transfer_write
   vector.transfer_write %1, %arg0[%c3, %c3] {permutation_map = affine_map<(d0, d1)->(d1, d0)>} : vector<3x7xf32>, memref<?x?xf32>
   // CHECK: vector.transfer_write %{{.*}}, %{{.*}}[%[[C3]], %[[C3]]] : vector<1x1x4x3xf32>, memref<?x?xvector<4x3xf32>>
   vector.transfer_write %4, %arg1[%c3, %c3] {permutation_map = affine_map<(d0, d1)->(d0, d1)>} : vector<1x1x4x3xf32>, memref<?x?xvector<4x3xf32>>
+  // CHECK: vector.transfer_write %{{.*}}, %{{.*}}[%[[C3]], %[[C3]]] : vector<1x1x4x3xf32>, memref<?x?xvector<4x3xf32>>
+  vector.transfer_write %5, %arg1[%c3, %c3] {masked = [true, true]} : vector<1x1x4x3xf32>, memref<?x?xvector<4x3xf32>>
 
   return
 }
 
 // CHECK-LABEL: @vector_broadcast
 func @vector_broadcast(%a: f32, %b: vector<16xf32>, %c: vector<1x16xf32>, %d: vector<8x1xf32>) -> vector<8x16xf32> {
   // CHECK: vector.broadcast %{{.*}} : f32 to vector<16xf32>
   %0 = vector.broadcast %a : f32 to vector<16xf32>