[mlir] Fix masked vector transfer ops with broadcasts

[springerm]

Summary

Broadcast dimensions of a vector transfer op have no corresponding dimension in the mask vector. E.g., a 2-D TransferReadOp, where one dimension is a broadcast, can have a 1-D mask attribute.

This commit also adds a few additional transfer op integration tests for various combinations of broadcasts, masking, dim transposes, etc.

[springerm]

Fix build

[aartbik]

mlir/lib/Conversion/VectorToSCF/ProgressiveVectorToSCF.cpp

604

typo: broadcasted

mlir/lib/Interfaces/VectorInterfaces.cpp

17

I believe, in C++ files we are supposed to use the style rather than adding namespaces around code

mlir::vector::detail::transfermaskType

mlir/test/Integration/Dialect/Vector/CPU/test-transfer-read-1d.mlir

25

Very nitty, here and below, period at end of comment

[springerm]

rebase

[springerm]

no change

[springerm]

no change

[nicolasvasilache]

mlir/include/mlir/Interfaces/VectorInterfaces.td

173

empty vector type usually does it, is there something special about the Optional here ?

mlir/lib/Conversion/VectorToSCF/ProgressiveVectorToSCF.cpp

147–151

Does not return -> Returns null ?

Also, can we enumerate the 3 cases:

1. X
2. Y
3. Z

152–153

maybe is usually associated with optional in my mind, null Value is fine, can we drop the "maybe"?

[springerm]

address review comments

[springerm]

mlir/test/Integration/Dialect/Vector/CPU/test-transfer-read-1d.mlir

25

I usually put dots only at the end of sentences, not headlines etc.

[aartbik]

mlir/test/Integration/Dialect/Vector/CPU/test-transfer-read-1d.mlir

25

Yes, but you are part of a larger suite now, and all other files here follow that style.

[springerm]

address comments

[springerm]

no change

[springerm]

address comments

[springerm]

no change

[springerm]

no change

[springerm]

no change

[springerm]

no change

[springerm]

no change

[ThomasRaoux]

mlir/lib/Dialect/Vector/VectorOps.cpp

2546

Should we add something in verifyTransferOp to check that the mask type is consistent?

[springerm]

mlir/lib/Dialect/Vector/VectorOps.cpp

2546

Yes, I think there should be a check here, because the mask attribute could be modified.

But in general, parser.resolveOperand(maskInfo, maskType, result.operands) already raises an error if a type different from maskType is found during parsing if I remember correctly. Do you know why this is the case? Why does resolveOperand do type checking (maybe in addition to the verifier)?

[ThomasRaoux]

mlir/lib/Dialect/Vector/VectorOps.cpp

2546

Yes the parser will check that the types are consistent. Since the parser infer the mask type it will check that the value has the correct type - the verifier will check that a value is used with the same type as it is defined.

That mean you won't be able to write a simple test for this case but I still think it may be valuable to have the verifier so that when ops are created with the builder there is also a check.

[springerm]

address comments

[springerm]

no change

[springerm]

no change

[springerm]

masks not supported with vector element type of memref

[springerm]

no change

[springerm]

no change

[nicolasvasilache]

Same thing as in the previous PR, would be good to have some minimal IR check too.
Approving conditioned on that.

Thanks!

mlir/lib/Conversion/VectorToSCF/ProgressiveVectorToSCF.cpp

82–84

Value maskBuffer = ... ?

91

looks like a good helper to add to vector.transfer ops : isBroadcast(int dim) ?

[springerm]

address comments

[springerm]

no change

[springerm]

no change

[springerm]

upload newest version

mlir/include/mlir/Interfaces/VectorInterfaces.h

 
 #ifndef MLIR_INTERFACES_VECTORINTERFACES_H
 #define MLIR_INTERFACES_VECTORINTERFACES_H
 
 #include "mlir/IR/AffineMap.h"
 #include "mlir/IR/BuiltinTypes.h"
 #include "mlir/IR/OpDefinition.h"
 
+namespace mlir {
+namespace vector {
+namespace detail {
+
+/// Given the vector type and the permutation map of a vector transfer op,
+/// compute the expected mask type.
+VectorType transferMaskType(VectorType vecType, AffineMap map);
+
+} // namespace detail
+} // namespace vector
+} // namespace mlir
+
 /// Include the generated interface declarations.
 #include "mlir/Interfaces/VectorInterfaces.h.inc"
 
 #endif // MLIR_INTERFACES_VECTORINTERFACES_H

mlir/include/mlir/Interfaces/VectorInterfaces.td

InterfaceMethod<
       /*desc=*/"Return the permutation map.",
       /*retTy=*/"AffineMap",
       /*methodName=*/"permutation_map",
       /*args=*/(ins),
       /*methodBody=*/"return $_op.permutation_map();"
       /*defaultImplementation=*/
     >,
     InterfaceMethod<
+      /*desc=*/[{ Returns true if the specified dimension is a broadcast. }],
+      /*retTy=*/"bool",
+      /*methodName=*/"isBroadcastDim",
+      /*args=*/(ins "unsigned":$idx),
+      /*methodBody=*/"",
+      /*defaultImplementation=*/[{
+        auto expr = $_op.permutation_map().getResult(idx);
+        return expr.template isa<AffineConstantExpr>() &&
+               expr.template dyn_cast<AffineConstantExpr>().getValue() == 0;
+      }]
+    >,
+    InterfaceMethod<
       /*desc=*/[{ Returns true if at least one of the dimensions in the
                   permutation map is a broadcast.}],
       /*retTy=*/"bool",
       /*methodName=*/"hasBroadcastDim",
       /*args=*/(ins),
       /*methodBody=*/"",
       /*defaultImplementation=*/[{
-        return llvm::any_of(
-            $_op.permutation_map().getResults(),
-            [](AffineExpr e) {
-                return e.isa<AffineConstantExpr>() &&
-                       e.dyn_cast<AffineConstantExpr>().getValue() == 0; });
+        for (unsigned i = 0; i < $_op.permutation_map().getNumResults(); ++i) {
+          if ($_op.isBroadcastDim(i))
+            return true;
+        }
+        return false;
       }]
     >,
     InterfaceMethod<
       /*desc=*/"Return the `in_bounds` boolean ArrayAttr.",
       /*retTy=*/"Optional<ArrayAttr>",
       /*methodName=*/"in_bounds",
       /*args=*/(ins),
       /*methodBody=*/"return $_op.in_bounds();"
InterfaceMethod<
       /*methodName=*/"getVectorType",
       /*args=*/(ins),
       /*methodBody=*/"",
       /*defaultImplementation=*/[{
         return $_op.vector().getType().template dyn_cast<VectorType>();
         }]
     >,
     InterfaceMethod<
+      /*desc=*/"Return the mask type if the op has a mask.",
+      /*retTy=*/"VectorType",
+      /*methodName=*/"getMaskType",
+      /*args=*/(ins),
+      /*methodBody=*/"",
+      /*defaultImplementation=*/[{
+        return $_op.mask()
+            ? mlir::vector::detail::transferMaskType(
+                $_op.getVectorType(), $_op.permutation_map())
+            : VectorType();
+      }]
+    >,
+    InterfaceMethod<
       /*desc=*/[{ Return the number of dimensions that participate in the
                   permutation map.}],
       /*retTy=*/"unsigned",
       /*methodName=*/"getTransferRank",
       /*args=*/(ins),
       /*methodBody=*/"",
       /*defaultImplementation=*/
         "return $_op.permutation_map().getNumResults();"

mlir/lib/Conversion/VectorToSCF/ProgressiveVectorToSCF.cpp

static BufferAllocs allocBuffers(OpTy xferOp) {
   b.setInsertionPointToStart(&scope->getRegion(0).front());
 
   BufferAllocs result;
   auto bufferType = MemRefType::get({}, xferOp.getVectorType());
   result.dataBuffer = memref_alloca(bufferType).value;
 
   if (xferOp.mask()) {
     auto maskType = MemRefType::get({}, xferOp.mask().getType());
-    result.maskBuffer = memref_alloca(maskType).value;
-    memref_store(xferOp.mask(), result.maskBuffer);
+    Value maskBuffer = memref_alloca(maskType);
+    memref_store(xferOp.mask(), maskBuffer);
+    result.maskBuffer = memref_load(maskBuffer);
   }
 
   return result;
 }
 
 /// Given a vector transfer op, calculate which dimension of the `source`
 /// memref should be unpacked in the next application of TransferOpConversion.
 /// A return value of None indicates a broadcast.
 template <typename OpTy>
 static Optional<int64_t> unpackedDim(OpTy xferOp) {
   auto map = xferOp.permutation_map();
   if (auto expr = map.getResult(0).template dyn_cast<AffineDimExpr>()) {
     return expr.getPosition();
   }
-  assert(map.getResult(0).template isa<AffineConstantExpr>() &&
+  assert(xferOp.isBroadcastDim(0) &&
          "Expected AffineDimExpr or AffineConstantExpr");
   return None;
 }
 
 /// Compute the permutation map for the new (N-1)-D vector transfer op. This
 /// map is identical to the current permutation map, but the first result is
 /// omitted.
 template <typename OpTy>
static void maybeYieldValue(
   if (hasRetVal) {
     builder.create<scf::YieldOp>(loc, value);
   } else {
     builder.create<scf::YieldOp>(loc);
   }
 }
 
 /// Generates a boolean Value that is true if the iv-th bit in xferOp's mask
-/// is set to true. Does not return a Value if the transfer op is not 1D or
-/// if the transfer op does not have a mask.
+/// is set to true. No such check is generated under following circumstances:
+/// * xferOp does not have a mask.
+/// * xferOp's mask is not 1D. (In case of (N>1)-D, a subvector of the mask is
+///   computed and attached to the new transfer op in the pattern.)
+/// * The to-be-unpacked dim of xferOp is a broadcast.
 template <typename OpTy>
-static Value maybeGenerateMaskCheck(OpBuilder &builder, OpTy xferOp, Value iv) {
-  if (xferOp.getVectorType().getRank() != 1)
-    return Value();
+static Value generateMaskCheck(OpBuilder &builder, OpTy xferOp, Value iv) {
   if (!xferOp.mask())
     return Value();
+  if (xferOp.getMaskType().getRank() != 1)
+    return Value();
+  if (xferOp.isBroadcastDim(0))
+    return Value();
 
   auto ivI32 = std_index_cast(IntegerType::get(builder.getContext(), 32), iv);
   return vector_extract_element(xferOp.mask(), ivI32).value;
 }
 
 /// Helper function TransferOpConversion and TransferOp1dConversion.
 /// Generate an in-bounds check if the transfer op may go out-of-bounds on the
 /// specified dimension `dim` with the loop iteration variable `iv`.
if (!xferOp.isDimInBounds(0) && !isBroadcast) {
     auto memrefDim =
         memref_dim(xferOp.source(), std_constant_index(dim.getValue()));
     using edsc::op::operator+;
     auto memrefIdx = xferOp.indices()[dim.getValue()] + iv;
     cond = std_cmpi_sgt(memrefDim.value, memrefIdx);
   }
 
   // Condition check 2: Masked in?
-  if (auto maskCond = maybeGenerateMaskCheck(builder, xferOp, iv)) {
+  if (auto maskCond = generateMaskCheck(builder, xferOp, iv)) {
     if (cond) {
       cond = builder.create<AndOp>(xferOp.getLoc(), cond, maskCond);
     } else {
       cond = maskCond;
     }
   }
 
   // If the condition is non-empty, generate an SCF::IfOp.
LogicalResult matchAndRewrite(TransferReadOp xferOp,
     if (checkPrepareXferOp(xferOp).failed())
       return failure();
 
     ScopedContext scope(rewriter, xferOp.getLoc());
     auto buffers = allocBuffers(xferOp);
     auto *newXfer = rewriter.clone(*xferOp.getOperation());
     newXfer->setAttr(kPassLabel, rewriter.getUnitAttr());
     if (xferOp.mask()) {
-      auto loadedMask = memref_load(buffers.maskBuffer);
-      dyn_cast<TransferReadOp>(newXfer).maskMutable().assign(loadedMask);
+      dyn_cast<TransferReadOp>(newXfer).maskMutable().assign(
+          buffers.maskBuffer);
     }
 
     memref_store(newXfer->getResult(0), buffers.dataBuffer);
     rewriter.replaceOpWithNewOp<memref::LoadOp>(xferOp, buffers.dataBuffer);
 
     return success();
   }
 };
LogicalResult matchAndRewrite(TransferWriteOp xferOp,
     memref_store(xferOp.vector(), buffers.dataBuffer);
     auto loadedVec = memref_load(buffers.dataBuffer);
     rewriter.updateRootInPlace(xferOp, [&]() {
       xferOp.vectorMutable().assign(loadedVec);
       xferOp->setAttr(kPassLabel, rewriter.getUnitAttr());
     });
 
     if (xferOp.mask()) {
-      auto loadedMask = memref_load(buffers.maskBuffer);
       rewriter.updateRootInPlace(
-          xferOp, [&]() { xferOp.maskMutable().assign(loadedMask); });
+          xferOp, [&]() { xferOp.maskMutable().assign(buffers.maskBuffer); });
     }
 
     return success();
   }
 };
 
 /// Progressive lowering of vector transfer ops: Unpack one dimension.
 ///
LogicalResult matchAndRewrite(OpTy xferOp,
     auto castedDataBuffer = vector_type_cast(castedDataType, dataBuffer);
 
     // If the xferOp has a mask: Find and cast mask buffer.
     Value castedMaskBuffer;
     if (xferOp.mask()) {
       auto maskBuffer = getMaskBuffer(xferOp);
       auto maskBufferType =
           maskBuffer.getType().template dyn_cast<MemRefType>();
+      if (xferOp.isBroadcastDim(0) || xferOp.getMaskType().getRank() == 1) {
+        // Do not unpack a dimension of the mask, if:
+        // * To-be-unpacked transfer op dimension is a broadcast.
+        // * Mask is 1D, i.e., the mask cannot be further unpacked.
+        //   (That means that all remaining dimensions of the transfer op must
+        //   be broadcasted.)
+        castedMaskBuffer = maskBuffer;
+      } else {
         auto castedMaskType = unpackOneDim(maskBufferType);
         castedMaskBuffer = vector_type_cast(castedMaskType, maskBuffer);
       }
+    }
 
     // Loop bounds and step.
     auto lb = std_constant_index(0).value;
     auto ub = std_constant_index(
                   castedDataType.getDimSize(castedDataType.getRank() - 1))
                   .value;
     auto step = std_constant_index(1).value;
 
     // Generate for loop.
     rewriter.create<scf::ForOp>(
         xferOp.getLoc(), lb, ub, step, ValueRange(),
         [&](OpBuilder &b, Location loc, Value iv,
             ValueRange /*loopState*/) {
       ScopedContext scope(b, loc);
       generateInBoundsCheck(
           xferOp, iv, b, unpackedDim(xferOp),
           /*inBoundsCase=*/
           [&](OpBuilder &b, Location /*loc*/) {
             // Create new transfer op.
             OpTy newXfer =
                 Strategy<OpTy>::rewriteOp(b, xferOp, castedDataBuffer, iv);
 
             // If old transfer op has a mask: Set mask on new transfer op.
-            if (xferOp.mask()) {
+            // Special case: If the mask of the old transfer op is 1D and the
+            //               unpacked dim is not a broadcast, no mask is needed
+            //               on the new transfer op.
+            if (xferOp.mask() && (xferOp.isBroadcastDim(0) ||
+                                  xferOp.getMaskType().getRank() > 1)) {
               OpBuilder::InsertionGuard guard(b);
               b.setInsertionPoint(newXfer); // Insert load before newXfer.
 
               SmallVector<Value, 8> loadIndices;
               Strategy<OpTy>::getBufferIndices(xferOp, loadIndices);
+              // In case of broadcast: Use same indices to load from memref as
+              // before.
+              if (!xferOp.isBroadcastDim(0))
                 loadIndices.push_back(iv);
 
               auto mask = memref_load(castedMaskBuffer, loadIndices);
               rewriter.updateRootInPlace(
                   newXfer, [&]() { newXfer.maskMutable().assign(mask); });
             }
           },
           /*outOfBoundsCase=*/
           [&](OpBuilder &b, Location /*loc*/) {
assert(map.getNumResults() == 1 &&
          "Expected 1 permutation map result for 1D transfer");
   if (auto expr = map.getResult(0).template dyn_cast<AffineDimExpr>()) {
     auto dim = expr.getPosition();
     using edsc::op::operator+;
     memrefIndices[dim] = memrefIndices[dim] + iv;
     return dim;
   }
 
-  assert(map.getResult(0).template isa<AffineConstantExpr>() &&
+  assert(xferOp.isBroadcastDim(0) &&
          "Expected AffineDimExpr or AffineConstantExpr");
   return None;
 }
 
 /// Codegen strategy for TransferOp1dConversion, depending on the
 /// operation.
 template <typename OpTy>
 struct Strategy1d;

mlir/lib/Dialect/Vector/VectorOps.cpp

for (auto expr : permutationMap.getResults()) {
     }
     seen[dim.getPosition()] = true;
   }
   return success();
 }
 
 static LogicalResult verifyTransferOp(Operation *op, ShapedType shapedType,
                                       VectorType vectorType,
+                                      VectorType maskType,
                                       AffineMap permutationMap,
                                       ArrayAttr inBounds) {
   if (op->hasAttr("masked")) {
     return op->emitOpError("masked attribute has been removed. "
                            "Use in_bounds instead.");
   }
 
   if (!shapedType.isa<MemRefType, RankedTensorType>())
if (auto vectorElementType = elementType.dyn_cast<VectorType>()) {
     if (sourceVecEltRank > resultVecRank)
       return op->emitOpError(
           "requires source vector element and vector result ranks to match.");
     unsigned rankOffset = resultVecRank - sourceVecEltRank;
     // Check that permutation map results match 'rankOffset' of vector type.
     if (permutationMap.getNumResults() != rankOffset)
       return op->emitOpError("requires a permutation_map with result dims of "
                              "the same rank as the vector type");
+
+    if (maskType)
+      return op->emitOpError("does not support masks with vector element type");
   } else {
     // Memref or tensor has scalar element type.
     unsigned resultVecSize =
         dataLayout.getTypeSizeInBits(vectorType.getElementType()) *
         vectorType.getShape().back();
     if (resultVecSize % dataLayout.getTypeSizeInBits(elementType) != 0)
       return op->emitOpError(
           "requires the bitwidth of the minor 1-D vector to be an integral "
           "multiple of the bitwidth of the source element type");
 
     // Check that permutation map results match rank of vector type.
     if (permutationMap.getNumResults() != vectorType.getRank())
       return op->emitOpError("requires a permutation_map with result dims of "
                              "the same rank as the vector type");
+
+    VectorType expectedMaskType =
+        vector::detail::transferMaskType(vectorType, permutationMap);
+    if (maskType && expectedMaskType != maskType)
+      return op->emitOpError("expects mask type consistent with permutation "
+                             "map: ")
+             << maskType;
   }
 
   if (permutationMap.getNumSymbols() != 0)
     return op->emitOpError("requires permutation_map without symbols");
   if (permutationMap.getNumInputs() != shapedType.getRank())
     return op->emitOpError("requires a permutation_map with input dims of the "
                            "same rank as the source type");
 
static ParseResult parseTransferReadOp(OpAsmParser &parser,
   auto indexType = builder.getIndexType();
   auto shapedType = types[0].dyn_cast<ShapedType>();
   if (!shapedType || !shapedType.isa<MemRefType, RankedTensorType>())
     return parser.emitError(typesLoc, "requires memref or ranked tensor type");
   VectorType vectorType = types[1].dyn_cast<VectorType>();
   if (!vectorType)
     return parser.emitError(typesLoc, "requires vector type");
   auto permutationAttrName = TransferReadOp::getPermutationMapAttrName();
-  auto attr = result.attributes.get(permutationAttrName);
-  if (!attr) {
+  Attribute mapAttr = result.attributes.get(permutationAttrName);
+  if (!mapAttr) {
     auto permMap = getTransferMinorIdentityMap(shapedType, vectorType);
-    result.attributes.set(permutationAttrName, AffineMapAttr::get(permMap));
+    mapAttr = AffineMapAttr::get(permMap);
+    result.attributes.set(permutationAttrName, mapAttr);
   }
   if (parser.resolveOperand(sourceInfo, shapedType, result.operands) ||
       parser.resolveOperands(indexInfo, indexType, result.operands) ||
       parser.resolveOperand(paddingInfo, shapedType.getElementType(),
                             result.operands))
     return failure();
   if (hasMask.succeeded()) {
-    auto maskType = VectorType::get(vectorType.getShape(), builder.getI1Type());
+    if (shapedType.getElementType().dyn_cast<VectorType>())
+      return parser.emitError(
+          maskInfo.location, "does not support masks with vector element type");
+    auto map = mapAttr.dyn_cast<AffineMapAttr>().getValue();
+    // Instead of adding the mask type as an op type, compute it based on the
+    // vector type and the permutation map (to keep the type signature small).
+    auto maskType = mlir::vector::detail::transferMaskType(vectorType, map);
     if (parser.resolveOperand(maskInfo, maskType, result.operands))
       return failure();
   }
   result.addAttribute(
       TransferReadOp::getOperandSegmentSizeAttr(),
       builder.getI32VectorAttr({1, static_cast<int32_t>(indexInfo.size()), 1,
                                 static_cast<int32_t>(hasMask.succeeded())}));
   return parser.addTypeToList(vectorType, result.types);
 }
 
 static LogicalResult verify(TransferReadOp op) {
   // Consistency of elemental types in source and vector.
   ShapedType shapedType = op.getShapedType();
   VectorType vectorType = op.getVectorType();
+  VectorType maskType = op.getMaskType();
   auto paddingType = op.padding().getType();
   auto permutationMap = op.permutation_map();
   auto sourceElementType = shapedType.getElementType();
 
   if (static_cast<int64_t>(op.indices().size()) != shapedType.getRank())
     return op.emitOpError("requires ") << shapedType.getRank() << " indices";
 
   if (failed(verifyTransferOp(op.getOperation(), shapedType, vectorType,
-                              permutationMap,
+                              maskType, permutationMap,
                               op.in_bounds() ? *op.in_bounds() : ArrayAttr())))
     return failure();
 
   if (auto sourceVectorElementType = sourceElementType.dyn_cast<VectorType>()) {
     // Source has vector element type.
     // Check that 'sourceVectorElementType' and 'paddingType' types match.
     if (sourceVectorElementType != paddingType)
       return op.emitOpError(
if (!attr) {
     auto permMap = getTransferMinorIdentityMap(shapedType, vectorType);
     result.attributes.set(permutationAttrName, AffineMapAttr::get(permMap));
   }
   if (parser.resolveOperand(vectorInfo, vectorType, result.operands) ||
       parser.resolveOperand(sourceInfo, shapedType, result.operands) ||
       parser.resolveOperands(indexInfo, indexType, result.operands))
     return failure();
   if (hasMask.succeeded()) {
+    if (shapedType.getElementType().dyn_cast<VectorType>())
+      return parser.emitError(
+          maskInfo.location, "does not support masks with vector element type");
     auto maskType = VectorType::get(vectorType.getShape(), builder.getI1Type());
     if (parser.resolveOperand(maskInfo, maskType, result.operands))
       return failure();
   }
   result.addAttribute(
       TransferWriteOp::getOperandSegmentSizeAttr(),
       builder.getI32VectorAttr({1, 1, static_cast<int32_t>(indexInfo.size()),
                                 static_cast<int32_t>(hasMask.succeeded())}));
static void print(OpAsmPrinter &p, TransferWriteOp op) {
   printTransferAttrs(p, cast<VectorTransferOpInterface>(op.getOperation()));
   p << " : " << op.getVectorType() << ", " << op.getShapedType();
 }
 
 static LogicalResult verify(TransferWriteOp op) {
   // Consistency of elemental types in shape and vector.
   ShapedType shapedType = op.getShapedType();
   VectorType vectorType = op.getVectorType();
+  VectorType maskType = op.getMaskType();
   auto permutationMap = op.permutation_map();
 
   if (llvm::size(op.indices()) != shapedType.getRank())
     return op.emitOpError("requires ") << shapedType.getRank() << " indices";
 
   // We do not allow broadcast dimensions on TransferWriteOps for the moment,
   // as the semantics is unclear. This can be revisited later if necessary.
   if (op.hasBroadcastDim())
     return op.emitOpError("should not have broadcast dimensions");
 
   if (failed(verifyTransferOp(op.getOperation(), shapedType, vectorType,
-                              permutationMap,
+                              maskType, permutationMap,
                               op.in_bounds() ? *op.in_bounds() : ArrayAttr())))
     return failure();
 
   return verifyPermutationMap(permutationMap,
                               [&op](Twine t) { return op.emitOpError(t); });
 }
 
 /// Fold:

mlir/lib/Interfaces/VectorInterfaces.cpp

 //===- VectorInterfaces.cpp - Unrollable vector operations -*- C++ -*-===//
 //
 // Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
 // See https://llvm.org/LICENSE.txt for license information.
 // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
 //
 //===----------------------------------------------------------------------===//
 
 #include "mlir/Interfaces/VectorInterfaces.h"
 
 using namespace mlir;
 
+VectorType mlir::vector::detail::transferMaskType(VectorType vecType,
+                                                  AffineMap map) {
+  auto i1Type = IntegerType::get(map.getContext(), 1);
+  SmallVector<int64_t, 8> shape;
+  for (int64_t i = 0; i < vecType.getRank(); ++i) {
+    // Only result dims have a corresponding dim in the mask.
+    if (auto expr = map.getResult(i).template isa<AffineDimExpr>()) {
+      shape.push_back(vecType.getDimSize(i));
+    }
+  }
+  return shape.empty() ? VectorType() : VectorType::get(shape, i1Type);
+}
+
 //===----------------------------------------------------------------------===//
 // VectorUnroll Interfaces
 //===----------------------------------------------------------------------===//
 
 /// Include the definitions of the VectorUnroll interfaces.
 #include "mlir/Interfaces/VectorInterfaces.cpp.inc"

mlir/test/Dialect/Vector/invalid.mlir

func @test_vector.transfer_read(%arg0: memref<?x?x?xf32>) {
   %c3 = constant 3 : index
   %cst = constant 3.0 : f32
   // expected-error@+1 {{requires a permutation_map that is a permutation (found one dim used more than once)}}
   %0 = vector.transfer_read %arg0[%c3, %c3, %c3], %cst {permutation_map = affine_map<(d0, d1, d2)->(d0, d0)>} : memref<?x?x?xf32>, vector<3x7xf32>
 }
 
 // -----
 
+func @test_vector.transfer_read(%arg0: memref<?x?x?xf32>) {
+  %c1 = constant 1 : i1
+  %c3 = constant 3 : index
+  %cst = constant 3.0 : f32
+  // expected-note@+1 {{prior use here}}
+  %mask = splat %c1 : vector<3x8x7xi1>
+  // expected-error@+1 {{expects different type than prior uses: 'vector<3x7xi1>' vs 'vector<3x8x7xi1>'}}
+  %0 = vector.transfer_read %arg0[%c3, %c3, %c3], %cst, %mask {permutation_map = affine_map<(d0, d1, d2)->(d0, 0, d2)>} : memref<?x?x?xf32>, vector<3x8x7xf32>
+}
+
+// -----
+
 func @test_vector.transfer_read(%arg0: memref<?x?xvector<4x3xf32>>) {
   %c3 = constant 3 : index
   %f0 = constant 0.0 : f32
   %vf0 = splat %f0 : vector<4x3xf32>
   // expected-error@+1 {{requires source vector element and vector result ranks to match}}
   %0 = vector.transfer_read %arg0[%c3, %c3], %vf0 {permutation_map = affine_map<(d0, d1)->(d0, d1)>} : memref<?x?xvector<4x3xf32>>, vector<3xf32>
 }
 
func @test_vector.transfer_read(%arg0: memref<?x?xvector<2x3xf32>>) {
   %f0 = constant 0.0 : f32
   %vf0 = splat %f0 : vector<2x3xf32>
   // expected-error@+1 {{ expects the optional in_bounds attr of same rank as permutation_map results: affine_map<(d0, d1) -> (d0, d1)>}}
   %0 = vector.transfer_read %arg0[%c3, %c3], %vf0 {in_bounds = [true], permutation_map = affine_map<(d0, d1)->(d0, d1)>} : memref<?x?xvector<2x3xf32>>, 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>
+  %mask = splat %c1 : vector<2x3xi1>
+  // expected-error@+1 {{does not support masks with vector element type}}
+  %0 = vector.transfer_read %arg0[%c3, %c3], %vf0, %mask {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 3.0 : f32
   // expected-error@+1 {{requires two types}}
   vector.transfer_write %arg0, %arg0[%c3, %c3] : memref<?x?xf32>
 }
 
 // -----

mlir/test/Dialect/Vector/ops.mlir

 // RUN: mlir-opt %s | mlir-opt | FileCheck %s
 
 // CHECK-LABEL: func @vector_transfer_ops(
 func @vector_transfer_ops(%arg0: memref<?x?xf32>,
                           %arg1 : memref<?x?xvector<4x3xf32>>,
                           %arg2 : memref<?x?xvector<4x3xi32>>,
-                          %arg3 : memref<?x?xvector<4x3xindex>>) {
+                          %arg3 : memref<?x?xvector<4x3xindex>>,
+                          %arg4 : memref<?x?x?xf32>) {
   // CHECK: %[[C3:.*]] = constant 3 : index
   %c3 = constant 3 : index
   %cst = constant 3.0 : f32
   %f0 = constant 0.0 : f32
   %c0 = constant 0 : i32
   %i0 = constant 0 : index
+  %i1 = constant 1 : i1
 
   %vf0 = splat %f0 : vector<4x3xf32>
   %v0 = splat %c0 : vector<4x3xi32>
   %vi0 = splat %i0 : vector<4x3xindex>
   %m = constant dense<[0, 0, 1, 0, 1]> : vector<5xi1>
-
+  %m2 = splat %i1 : vector<5x4xi1>
   //
   // CHECK: vector.transfer_read
   %0 = vector.transfer_read %arg0[%c3, %c3], %f0 {permutation_map = affine_map<(d0, d1)->(d0)>} : memref<?x?xf32>, vector<128xf32>
   // 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]]], %{{.*}} {in_bounds = [false, true]} : memref<?x?xvector<4x3xf32>>, vector<1x1x4x3xf32>
   %5 = vector.transfer_read %arg1[%c3, %c3], %vf0 {in_bounds = [false, true]} : memref<?x?xvector<4x3xf32>>, vector<1x1x4x3xf32>
   // CHECK: vector.transfer_read %{{.*}}[%[[C3]], %[[C3]]], %{{.*}} : memref<?x?xvector<4x3xi32>>, vector<5x24xi8>
   %6 = vector.transfer_read %arg2[%c3, %c3], %v0 : memref<?x?xvector<4x3xi32>>, vector<5x24xi8>
   // CHECK: vector.transfer_read %{{.*}}[%[[C3]], %[[C3]]], %{{.*}} : memref<?x?xvector<4x3xindex>>, vector<5x48xi8>
   %7 = vector.transfer_read %arg3[%c3, %c3], %vi0 : memref<?x?xvector<4x3xindex>>, vector<5x48xi8>
   // CHECK: vector.transfer_read %{{.*}}[%[[C3]], %[[C3]]], %{{.*}}, %{{.*}} : memref<?x?xf32>, vector<5xf32>
   %8 = vector.transfer_read %arg0[%c3, %c3], %f0, %m : memref<?x?xf32>, vector<5xf32>
+  // CHECK: vector.transfer_read %{{.*}}[%[[C3]], %[[C3]], %[[C3]]], %{{.*}}, %{{.*}} : memref<?x?x?xf32>, vector<5x4x8xf32>
+  %9 = vector.transfer_read %arg4[%c3, %c3, %c3], %f0, %m2 {permutation_map = affine_map<(d0, d1, d2)->(d1, d0, 0)>} : memref<?x?x?xf32>, vector<5x4x8xf32>
 
   // 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>>

mlir/test/Integration/Dialect/Vector/CPU/test-transfer-read-1d.mlir

 // RUN: mlir-opt %s -test-progressive-convert-vector-to-scf -lower-affine -convert-scf-to-std -convert-vector-to-llvm -convert-std-to-llvm | \
 // RUN: mlir-cpu-runner -e entry -entry-point-result=void \
 // RUN:   -shared-libs=%mlir_integration_test_dir/libmlir_c_runner_utils%shlibext | \
 // RUN: FileCheck %s
 
 // Test for special cases of 1D vector transfer ops.
 
+memref.global "private" @gv : memref<5x6xf32> =
+    dense<[[0. , 1. , 2. , 3. , 4. , 5. ],
+           [10., 11., 12., 13., 14., 15.],
+           [20., 21., 22., 23., 24., 25.],
+           [30., 31., 32., 33., 34., 35.],
+           [40., 41., 42., 43., 44., 45.]]>
+
+// Non-contiguous, strided load.
 func @transfer_read_1d(%A : memref<?x?xf32>, %base1 : index, %base2 : index) {
   %fm42 = constant -42.0: f32
   %f = vector.transfer_read %A[%base1, %base2], %fm42
       {permutation_map = affine_map<(d0, d1) -> (d0)>}
       : memref<?x?xf32>, vector<9xf32>
   vector.print %f: vector<9xf32>
   return
 }
 
+// Broadcast.
 func @transfer_read_1d_broadcast(
     %A : memref<?x?xf32>, %base1 : index, %base2 : index) {
   %fm42 = constant -42.0: f32
   %f = vector.transfer_read %A[%base1, %base2], %fm42
       {permutation_map = affine_map<(d0, d1) -> (0)>}
       : memref<?x?xf32>, vector<9xf32>
   vector.print %f: vector<9xf32>
   return
 }
 
+// Non-contiguous, strided load.
 func @transfer_read_1d_in_bounds(
     %A : memref<?x?xf32>, %base1 : index, %base2 : index) {
   %fm42 = constant -42.0: f32
   %f = vector.transfer_read %A[%base1, %base2], %fm42
       {permutation_map = affine_map<(d0, d1) -> (d0)>, in_bounds = [true]}
       : memref<?x?xf32>, vector<3xf32>
   vector.print %f: vector<3xf32>
   return
 }
 
+// Non-contiguous, strided load.
 func @transfer_read_1d_mask(
     %A : memref<?x?xf32>, %base1 : index, %base2 : index) {
   %fm42 = constant -42.0: f32
   %mask = constant dense<[1, 0, 1, 0, 1, 1, 1, 0, 1]> : vector<9xi1>
   %f = vector.transfer_read %A[%base1, %base2], %fm42, %mask
       {permutation_map = affine_map<(d0, d1) -> (d0)>}
       : memref<?x?xf32>, vector<9xf32>
   vector.print %f: vector<9xf32>
   return
 }
 
+// Non-contiguous, strided load.
 func @transfer_read_1d_mask_in_bounds(
     %A : memref<?x?xf32>, %base1 : index, %base2 : index) {
   %fm42 = constant -42.0: f32
   %mask = constant dense<[1, 0, 1]> : vector<3xi1>
   %f = vector.transfer_read %A[%base1, %base2], %fm42, %mask
       {permutation_map = affine_map<(d0, d1) -> (d0)>, in_bounds = [true]}
       : memref<?x?xf32>, vector<3xf32>
   vector.print %f: vector<3xf32>
   return
 }
 
+// Non-contiguous, strided store.
 func @transfer_write_1d(%A : memref<?x?xf32>, %base1 : index, %base2 : index) {
   %fn1 = constant -1.0 : f32
   %vf0 = splat %fn1 : vector<7xf32>
   vector.transfer_write %vf0, %A[%base1, %base2]
     {permutation_map = affine_map<(d0, d1) -> (d0)>}
     : vector<7xf32>, memref<?x?xf32>
   return
 }
 
+// Non-contiguous, strided store.
+func @transfer_write_1d_mask(%A : memref<?x?xf32>, %base1 : index, %base2 : index) {
+  %fn1 = constant -2.0 : f32
+  %vf0 = splat %fn1 : vector<7xf32>
+  %mask = constant dense<[1, 0, 1, 0, 1, 1, 1]> : vector<7xi1>
+  vector.transfer_write %vf0, %A[%base1, %base2], %mask
+    {permutation_map = affine_map<(d0, d1) -> (d0)>}
+    : vector<7xf32>, memref<?x?xf32>
+  return
+}
+
 func @entry() {
   %c0 = constant 0: index
   %c1 = constant 1: index
   %c2 = constant 2: index
   %c3 = constant 3: index
-  %f10 = constant 10.0: f32
-  // work with dims of 4, not of 3
-  %first = constant 5: index
-  %second = constant 6: index
-  %A = memref.alloc(%first, %second) : memref<?x?xf32>
-  scf.for %i = %c0 to %first step %c1 {
-    %i32 = index_cast %i : index to i32
-    %fi = sitofp %i32 : i32 to f32
-    %fi10 = mulf %fi, %f10 : f32
-    scf.for %j = %c0 to %second step %c1 {
-        %j32 = index_cast %j : index to i32
-        %fj = sitofp %j32 : i32 to f32
-        %fres = addf %fi10, %fj : f32
-        memref.store %fres, %A[%i, %j] : memref<?x?xf32>
-    }
-  }
+  %0 = memref.get_global @gv : memref<5x6xf32>
+  %A = memref.cast %0 : memref<5x6xf32> to memref<?x?xf32>
 
-  // Read from 2D memref on first dimension. Cannot be lowered to an LLVM
-  // vector load. Instead, generates scalar loads.
+  // 1. Read from 2D memref on first dimension. Cannot be lowered to an LLVM
+  //    vector load. Instead, generates scalar loads.
   call @transfer_read_1d(%A, %c1, %c2) : (memref<?x?xf32>, index, index) -> ()
-  // Write to 2D memref on first dimension. Cannot be lowered to an LLVM
-  // vector store. Instead, generates scalar stores.
+  // CHECK: ( 12, 22, 32, 42, -42, -42, -42, -42, -42 )
+
+  // 2. Write to 2D memref on first dimension. Cannot be lowered to an LLVM
+  //    vector store. Instead, generates scalar stores.
   call @transfer_write_1d(%A, %c3, %c2) : (memref<?x?xf32>, index, index) -> ()
-  // (Same as above.)
+
+  // 3. (Same as 1. To check if 2 works correctly.)
   call @transfer_read_1d(%A, %c0, %c2) : (memref<?x?xf32>, index, index) -> ()
-  // Read a scalar from a 2D memref and broadcast the value to a 1D vector.
-  // Generates a loop with vector.insertelement.
+  // CHECK: ( 2, 12, 22, -1, -1, -42, -42, -42, -42 )
+
+  // 4. Read a scalar from a 2D memref and broadcast the value to a 1D vector.
+  //    Generates a loop with vector.insertelement.
   call @transfer_read_1d_broadcast(%A, %c1, %c2)
       : (memref<?x?xf32>, index, index) -> ()
-  //  Read from 2D memref on first dimension. Accesses are in-bounds, so no
-  // if-check is generated inside the generated loop.
+  // CHECK: ( 12, 12, 12, 12, 12, 12, 12, 12, 12 )
+
+  // 5. Read from 2D memref on first dimension. Accesses are in-bounds, so no
+  //    if-check is generated inside the generated loop.
   call @transfer_read_1d_in_bounds(%A, %c1, %c2)
       : (memref<?x?xf32>, index, index) -> ()
-  // Optional mask attribute is specified and, in addition, there may be
-  // out-of-bounds accesses.
+  // CHECK: ( 12, 22, -1 )
+
+  // 6. Optional mask attribute is specified and, in addition, there may be
+  //    out-of-bounds accesses.
   call @transfer_read_1d_mask(%A, %c1, %c2)
       : (memref<?x?xf32>, index, index) -> ()
-  // Same as above, but accesses are in-bounds.
+  // CHECK: ( 12, -42, -1, -42, -42, -42, -42, -42, -42 )
+
+  // 7. Same as 6, but accesses are in-bounds.
   call @transfer_read_1d_mask_in_bounds(%A, %c1, %c2)
       : (memref<?x?xf32>, index, index) -> ()
+  // CHECK: ( 12, -42, -1 )
+
+  // 8. Write to 2D memref on first dimension with a mask.
+  call @transfer_write_1d_mask(%A, %c1, %c0)
+      : (memref<?x?xf32>, index, index) -> ()
+
+  // 9. (Same as 1. To check if 8 works correctly.)
+  call @transfer_read_1d(%A, %c0, %c0) : (memref<?x?xf32>, index, index) -> ()
+  // CHECK: ( 0, -2, 20, -2, 40, -42, -42, -42, -42 )
+
   return
 }
-
-// CHECK: ( 12, 22, 32, 42, -42, -42, -42, -42, -42 )
-// CHECK: ( 2, 12, 22, -1, -1, -42, -42, -42, -42 )
-// CHECK: ( 12, 12, 12, 12, 12, 12, 12, 12, 12 )
-// CHECK: ( 12, 22, -1 )
-// CHECK: ( 12, -42, -1, -42, -42, -42, -42, -42, -42 )
-// CHECK: ( 12, -42, -1 )

mlir/test/Integration/Dialect/Vector/CPU/test-transfer-read-2d.mlir

 // RUN: mlir-opt %s -test-progressive-convert-vector-to-scf -lower-affine -convert-scf-to-std -convert-vector-to-llvm -convert-std-to-llvm | \
 // RUN: mlir-cpu-runner -e entry -entry-point-result=void  \
 // RUN:   -shared-libs=%mlir_integration_test_dir/libmlir_c_runner_utils%shlibext | \
 // RUN: FileCheck %s
 
+memref.global "private" @gv : memref<3x4xf32> = dense<[[0. , 1. , 2. , 3. ],
+                                                       [10., 11., 12., 13.],
+                                                       [20., 21., 22., 23.]]>
+
+// Vector load.
 func @transfer_read_2d(%A : memref<?x?xf32>, %base1: index, %base2: index) {
   %fm42 = constant -42.0: f32
   %f = vector.transfer_read %A[%base1, %base2], %fm42
       {permutation_map = affine_map<(d0, d1) -> (d0, d1)>} :
     memref<?x?xf32>, vector<4x9xf32>
   vector.print %f: vector<4x9xf32>
   return
 }
 
+// Vector load with mask.
 func @transfer_read_2d_mask(%A : memref<?x?xf32>, %base1: index, %base2: index) {
   %fm42 = constant -42.0: f32
   %mask = constant dense<[[1, 0, 1, 0, 1, 1, 1, 0, 1],
                           [0, 0, 1, 1, 1, 1, 1, 0, 1],
                           [1, 1, 1, 1, 1, 1, 1, 0, 1],
                           [0, 0, 1, 0, 1, 1, 1, 0, 1]]> : vector<4x9xi1>
   %f = vector.transfer_read %A[%base1, %base2], %fm42, %mask
       {permutation_map = affine_map<(d0, d1) -> (d0, d1)>} :
     memref<?x?xf32>, vector<4x9xf32>
   vector.print %f: vector<4x9xf32>
   return
 }
 
+// Vector load with mask + transpose.
+func @transfer_read_2d_mask_transposed(
+    %A : memref<?x?xf32>, %base1: index, %base2: index) {
+  %fm42 = constant -42.0: f32
+  %mask = constant dense<[[1, 0, 1, 0], [0, 0, 1, 0],
+                          [1, 1, 1, 1], [0, 1, 1, 0],
+                          [1, 1, 1, 1], [1, 1, 1, 1],
+                          [1, 1, 1, 1], [0, 0, 0, 0],
+                          [1, 1, 1, 1]]> : vector<9x4xi1>
+  %f = vector.transfer_read %A[%base1, %base2], %fm42, %mask
+      {permutation_map = affine_map<(d0, d1) -> (d1, d0)>} :
+    memref<?x?xf32>, vector<9x4xf32>
+  vector.print %f: vector<9x4xf32>
+  return
+}
+
+// Vector load with mask + broadcast.
+func @transfer_read_2d_mask_broadcast(
+    %A : memref<?x?xf32>, %base1: index, %base2: index) {
+  %fm42 = constant -42.0: f32
+  %mask = constant dense<[1, 0, 1, 0, 1, 1, 1, 0, 1]> : vector<9xi1>
+  %f = vector.transfer_read %A[%base1, %base2], %fm42, %mask
+      {permutation_map = affine_map<(d0, d1) -> (0, d1)>} :
+    memref<?x?xf32>, vector<4x9xf32>
+  vector.print %f: vector<4x9xf32>
+  return
+}
+
+// Transpose + vector load with mask + broadcast.
+func @transfer_read_2d_mask_transpose_broadcast_last_dim(
+    %A : memref<?x?xf32>, %base1: index, %base2: index) {
+  %fm42 = constant -42.0: f32
+  %mask = constant dense<[1, 0, 1, 1]> : vector<4xi1>
+  %f = vector.transfer_read %A[%base1, %base2], %fm42, %mask
+      {permutation_map = affine_map<(d0, d1) -> (d1, 0)>} :
+    memref<?x?xf32>, vector<4x9xf32>
+  vector.print %f: vector<4x9xf32>
+  return
+}
+
+// Load + transpose.
 func @transfer_read_2d_transposed(
     %A : memref<?x?xf32>, %base1: index, %base2: index) {
   %fm42 = constant -42.0: f32
   %f = vector.transfer_read %A[%base1, %base2], %fm42
       {permutation_map = affine_map<(d0, d1) -> (d1, d0)>} :
     memref<?x?xf32>, vector<4x9xf32>
   vector.print %f: vector<4x9xf32>
   return
 }
 
+// Load 1D + broadcast to 2D.
 func @transfer_read_2d_broadcast(
     %A : memref<?x?xf32>, %base1: index, %base2: index) {
   %fm42 = constant -42.0: f32
   %f = vector.transfer_read %A[%base1, %base2], %fm42
       {permutation_map = affine_map<(d0, d1) -> (d1, 0)>} :
     memref<?x?xf32>, vector<4x9xf32>
   vector.print %f: vector<4x9xf32>
   return
 }
 
+// Vector store.
 func @transfer_write_2d(%A : memref<?x?xf32>, %base1: index, %base2: index) {
   %fn1 = constant -1.0 : f32
   %vf0 = splat %fn1 : vector<1x4xf32>
   vector.transfer_write %vf0, %A[%base1, %base2]
     {permutation_map = affine_map<(d0, d1) -> (d0, d1)>} :
     vector<1x4xf32>, memref<?x?xf32>
   return
 }
 
+// Vector store with mask.
+func @transfer_write_2d_mask(%A : memref<?x?xf32>, %base1: index, %base2: index) {
+  %fn1 = constant -2.0 : f32
+  %mask = constant dense<[[1, 0, 1, 0]]> : vector<1x4xi1>
+  %vf0 = splat %fn1 : vector<1x4xf32>
+  vector.transfer_write %vf0, %A[%base1, %base2], %mask
+    {permutation_map = affine_map<(d0, d1) -> (d0, d1)>} :
+    vector<1x4xf32>, memref<?x?xf32>
+  return
+}
+
 func @entry() {
   %c0 = constant 0: index
   %c1 = constant 1: index
   %c2 = constant 2: index
   %c3 = constant 3: index
-  %c4 = constant 4: index
-  %c5 = constant 5: index
-  %c8 = constant 5: index
-  %f10 = constant 10.0: f32
-  // work with dims of 4, not of 3
-  %first = constant 3: index
-  %second = constant 4: index
-  %A = memref.alloc(%first, %second) : memref<?x?xf32>
-  scf.for %i = %c0 to %first step %c1 {
-    %i32 = index_cast %i : index to i32
-    %fi = sitofp %i32 : i32 to f32
-    %fi10 = mulf %fi, %f10 : f32
-    scf.for %j = %c0 to %second step %c1 {
-        %j32 = index_cast %j : index to i32
-        %fj = sitofp %j32 : i32 to f32
-        %fres = addf %fi10, %fj : f32
-        memref.store %fres, %A[%i, %j] : memref<?x?xf32>
-    }
-  }
-  // On input, memory contains [[ 0, 1, 2, ...], [10, 11, 12, ...], ...]
-  // Read shifted by 2 and pad with -42:
+  %0 = memref.get_global @gv : memref<3x4xf32>
+  %A = memref.cast %0 : memref<3x4xf32> to memref<?x?xf32>
+
+  // 1. Read 2D vector from 2D memref.
   call @transfer_read_2d(%A, %c1, %c2) : (memref<?x?xf32>, index, index) -> ()
-  // Same as above, but transposed
+  // CHECK: ( ( 12, 13, -42, -42, -42, -42, -42, -42, -42 ), ( 22, 23, -42, -42, -42, -42, -42, -42, -42 ), ( -42, -42, -42, -42, -42, -42, -42, -42, -42 ), ( -42, -42, -42, -42, -42, -42, -42, -42, -42 ) )
+
+  // 2. Read 2D vector from 2D memref at specified location and transpose the
+  //    result.
   call @transfer_read_2d_transposed(%A, %c1, %c2)
       : (memref<?x?xf32>, index, index) -> ()
-  // Write into memory shifted by 3
-  call @transfer_write_2d(%A, %c3, %c1) : (memref<?x?xf32>, index, index) -> ()
-  // Read shifted by 0 and pad with -42:
-  call @transfer_read_2d(%A, %c0, %c0) : (memref<?x?xf32>, index, index) -> ()
-  // Same as above, but apply a mask
+  // CHECK: ( ( 12, 22, -42, -42, -42, -42, -42, -42, -42 ), ( 13, 23, -42, -42, -42, -42, -42, -42, -42 ), ( -42, -42, -42, -42, -42, -42, -42, -42, -42 ), ( -42, -42, -42, -42, -42, -42, -42, -42, -42 ) )
+
+  // 3. Read 2D vector from 2D memref with a 2D mask. In addition, some
+  //    accesses are out-of-bounds.
   call @transfer_read_2d_mask(%A, %c0, %c0)
       : (memref<?x?xf32>, index, index) -> ()
-  // Same as above, but without mask and transposed
-  call @transfer_read_2d_transposed(%A, %c0, %c0)
+  // CHECK: ( ( 0, -42, 2, -42, -42, -42, -42, -42, -42 ), ( -42, -42, 12, 13, -42, -42, -42, -42, -42 ), ( 20, 21, 22, 23, -42, -42, -42, -42, -42 ), ( -42, -42, -42, -42, -42, -42, -42, -42, -42 ) )
+
+  // 4. Same as 3, but transpose the result.
+  call @transfer_read_2d_mask_transposed(%A, %c0, %c0)
       : (memref<?x?xf32>, index, index) -> ()
-  // Second vector dimension is a broadcast
+  // CHECK: ( ( 0, -42, 20, -42 ), ( -42, -42, 21, -42 ), ( 2, 12, 22, -42 ), ( -42, 13, 23, -42 ), ( -42, -42, -42, -42 ), ( -42, -42, -42, -42 ), ( -42, -42, -42, -42 ), ( -42, -42, -42, -42 ), ( -42, -42, -42, -42 ) )
+
+  // 5. Read 1D vector from 2D memref at specified location and broadcast the
+  //    result to 2D.
   call @transfer_read_2d_broadcast(%A, %c1, %c2)
       : (memref<?x?xf32>, index, index) -> ()
+  // CHECK: ( ( 12, 12, 12, 12, 12, 12, 12, 12, 12 ), ( 13, 13, 13, 13, 13, 13, 13, 13, 13 ), ( -42, -42, -42, -42, -42, -42, -42, -42, -42 ), ( -42, -42, -42, -42, -42, -42, -42, -42, -42 ) )
+
+  // 6. Read 1D vector from 2D memref at specified location with mask and
+  //    broadcast the result to 2D.
+  call @transfer_read_2d_mask_broadcast(%A, %c2, %c1)
+      : (memref<?x?xf32>, index, index) -> ()
+  // CHECK: ( ( 21, -42, 23, -42, -42, -42, -42, -42, -42 ), ( 21, -42, 23, -42, -42, -42, -42, -42, -42 ), ( 21, -42, 23, -42, -42, -42, -42, -42, -42 ), ( 21, -42, 23, -42, -42, -42, -42, -42, -42 ) )
+
+  // 7. Read 1D vector from 2D memref (second dimension) at specified location
+  //    with mask and broadcast the result to 2D. In this test case, mask
+  //    elements must be evaluated before lowering to an (N>1)-D transfer.
+  call @transfer_read_2d_mask_transpose_broadcast_last_dim(%A, %c0, %c1)
+      : (memref<?x?xf32>, index, index) -> ()
+  // CHECK: ( ( 1, 1, 1, 1, 1, 1, 1, 1, 1 ), ( -42, -42, -42, -42, -42, -42, -42, -42, -42 ), ( 3, 3, 3, 3, 3, 3, 3, 3, 3 ), ( -42, -42, -42, -42, -42, -42, -42, -42, -42 ) )
+
+  // 8. Write 2D vector into 2D memref at specified location.
+  call @transfer_write_2d(%A, %c1, %c2) : (memref<?x?xf32>, index, index) -> ()
+
+  // 9. Read memref to verify step 8.
+  call @transfer_read_2d(%A, %c0, %c0) : (memref<?x?xf32>, index, index) -> ()
+  // CHECK: ( ( 0, 1, 2, 3, -42, -42, -42, -42, -42 ), ( 10, 11, -1, -1, -42, -42, -42, -42, -42 ), ( 20, 21, 22, 23, -42, -42, -42, -42, -42 ), ( -42, -42, -42, -42, -42, -42, -42, -42, -42 ) )
+
+  // 10. Write 2D vector into 2D memref at specified location with mask.
+  call @transfer_write_2d_mask(%A, %c0, %c2) : (memref<?x?xf32>, index, index) -> ()
+
+  // 11. Read memref to verify step 10.
+  call @transfer_read_2d(%A, %c0, %c0) : (memref<?x?xf32>, index, index) -> ()
+  // CHECK: ( ( 0, 1, -2, 3, -42, -42, -42, -42, -42 ), ( 10, 11, -1, -1, -42, -42, -42, -42, -42 ), ( 20, 21, 22, 23, -42, -42, -42, -42, -42 ), ( -42, -42, -42, -42, -42, -42, -42, -42, -42 ) )
+
   return
 }
 
-// CHECK: ( ( 12, 13, -42, -42, -42, -42, -42, -42, -42 ), ( 22, 23, -42, -42, -42, -42, -42, -42, -42 ), ( -42, -42, -42, -42, -42, -42, -42, -42, -42 ), ( -42, -42, -42, -42, -42, -42, -42, -42, -42 ) )
-// CHECK: ( ( 12, 22, -42, -42, -42, -42, -42, -42, -42 ), ( 13, 23, -42, -42, -42, -42, -42, -42, -42 ), ( -42, -42, -42, -42, -42, -42, -42, -42, -42 ), ( -42, -42, -42, -42, -42, -42, -42, -42, -42 ) )
-// CHECK: ( ( 0, 1, 2, 3, -42, -42, -42, -42, -42 ), ( 10, 11, 12, 13, -42, -42, -42, -42, -42 ), ( 20, 21, 22, 23, -42, -42, -42, -42, -42 ), ( -42, -42, -42, -42, -42, -42, -42, -42, -42 ) )
-// CHECK: ( ( 0, -42, 2, -42, -42, -42, -42, -42, -42 ), ( -42, -42, 12, 13, -42, -42, -42, -42, -42 ), ( 20, 21, 22, 23, -42, -42, -42, -42, -42 ), ( -42, -42, -42, -42, -42, -42, -42, -42, -42 ) )
-// CHECK: ( ( 0, 10, 20, -42, -42, -42, -42, -42, -42 ), ( 1, 11, 21, -42, -42, -42, -42, -42, -42 ), ( 2, 12, 22, -42, -42, -42, -42, -42, -42 ), ( 3, 13, 23, -42, -42, -42, -42, -42, -42 ) )
-// CHECK: ( ( 12, 12, 12, 12, 12, 12, 12, 12, 12 ), ( 13, 13, 13, 13, 13, 13, 13, 13, 13 ), ( -42, -42, -42, -42, -42, -42, -42, -42, -42 ), ( -42, -42, -42, -42, -42, -42, -42, -42, -42 ) )

mlir/test/Integration/Dialect/Vector/CPU/test-transfer-read-3d.mlir

-// RUN: mlir-opt %s -convert-vector-to-scf -lower-affine -convert-scf-to-std -convert-vector-to-llvm -convert-std-to-llvm | \
-// RUN: mlir-cpu-runner -e entry -entry-point-result=void  \
-// RUN:   -shared-libs=%mlir_integration_test_dir/libmlir_c_runner_utils%shlibext | \
-// RUN: FileCheck %s
-
 // RUN: mlir-opt %s -test-progressive-convert-vector-to-scf -lower-affine -convert-scf-to-std -convert-vector-to-llvm -convert-std-to-llvm | \
 // RUN: mlir-cpu-runner -e entry -entry-point-result=void  \
 // RUN:   -shared-libs=%mlir_integration_test_dir/libmlir_c_runner_utils%shlibext | \
 // RUN: FileCheck %s
 
-// Test case is based on test-transfer-read-2d.
-
 func @transfer_read_3d(%A : memref<?x?x?x?xf32>,
                        %o: index, %a: index, %b: index, %c: index) {
   %fm42 = constant -42.0: f32
   %f = vector.transfer_read %A[%o, %a, %b, %c], %fm42
       : memref<?x?x?x?xf32>, vector<2x5x3xf32>
   vector.print %f: vector<2x5x3xf32>
   return
 }
 
 func @transfer_read_3d_broadcast(%A : memref<?x?x?x?xf32>,
                                  %o: index, %a: index, %b: index, %c: index) {
   %fm42 = constant -42.0: f32
   %f = vector.transfer_read %A[%o, %a, %b, %c], %fm42
       {permutation_map = affine_map<(d0, d1, d2, d3) -> (d1, 0, d3)>}
       : memref<?x?x?x?xf32>, vector<2x5x3xf32>
   vector.print %f: vector<2x5x3xf32>
   return
 }
 
+func @transfer_read_3d_mask_broadcast(
+    %A : memref<?x?x?x?xf32>, %o: index, %a: index, %b: index, %c: index) {
+  %fm42 = constant -42.0: f32
+  %mask = constant dense<[0, 1]> : vector<2xi1>
+  %f = vector.transfer_read %A[%o, %a, %b, %c], %fm42, %mask
+      {permutation_map = affine_map<(d0, d1, d2, d3) -> (d1, 0, 0)>}
+      : memref<?x?x?x?xf32>, vector<2x5x3xf32>
+  vector.print %f: vector<2x5x3xf32>
+  return
+}
+
 func @transfer_read_3d_transposed(%A : memref<?x?x?x?xf32>,
                                   %o: index, %a: index, %b: index, %c: index) {
   %fm42 = constant -42.0: f32
   %f = vector.transfer_read %A[%o, %a, %b, %c], %fm42
       {permutation_map = affine_map<(d0, d1, d2, d3) -> (d3, d0, d1)>}
       : memref<?x?x?x?xf32>, vector<3x5x3xf32>
   vector.print %f: vector<3x5x3xf32>
   return
scf.for %i = %c0 to %first step %c1 {
           %fk1 = addf %f2, %fk : f32
           %fmul = mulf %fadded, %fk1 : f32
           memref.store %fmul, %A[%o, %i, %j, %k] : memref<?x?x?x?xf32>
         }
       }
     }
   }
 
+  // 1. Read 3D vector from 4D memref.
   call @transfer_read_3d(%A, %c0, %c0, %c0, %c0)
       : (memref<?x?x?x?xf32>, index, index, index, index) -> ()
+  // CHECK: ( ( ( 0, 0, -42 ), ( 2, 3, -42 ), ( 4, 6, -42 ), ( 6, 9, -42 ), ( -42, -42, -42 ) ), ( ( 20, 30, -42 ), ( 22, 33, -42 ), ( 24, 36, -42 ), ( 26, 39, -42 ), ( -42, -42, -42 ) ) )
+
+  // 2. Write 3D vector to 4D memref.
   call @transfer_write_3d(%A, %c0, %c0, %c1, %c1)
       : (memref<?x?x?x?xf32>, index, index, index, index) -> ()
+
+  // 3. Read memref to verify step 2.
   call @transfer_read_3d(%A, %c0, %c0, %c0, %c0)
       : (memref<?x?x?x?xf32>, index, index, index, index) -> ()
+  // CHECK: ( ( ( 0, 0, -42 ), ( 2, -1, -42 ), ( 4, -1, -42 ), ( 6, -1, -42 ), ( -42, -42, -42 ) ), ( ( 20, 30, -42 ), ( 22, -1, -42 ), ( 24, -1, -42 ), ( 26, -1, -42 ), ( -42, -42, -42 ) ) )
+
+  // 4. Read 3D vector from 4D memref and transpose vector.
   call @transfer_read_3d_transposed(%A, %c0, %c0, %c0, %c0)
       : (memref<?x?x?x?xf32>, index, index, index, index) -> ()
+  // CHECK: ( ( ( 0, 20, 40 ), ( 0, 20, 40 ), ( 0, 20, 40 ), ( 0, 20, 40 ), ( 0, 20, 40 ) ), ( ( 0, 30, 60 ), ( 0, 30, 60 ), ( 0, 30, 60 ), ( 0, 30, 60 ), ( 0, 30, 60 ) ), ( ( -42, -42, -42 ), ( -42, -42, -42 ), ( -42, -42, -42 ), ( -42, -42, -42 ), ( -42, -42, -42 ) ) )
+
+  // 5. Read 1D vector from 4D memref and broadcast vector to 3D.
   call @transfer_read_3d_broadcast(%A, %c0, %c0, %c0, %c0)
       : (memref<?x?x?x?xf32>, index, index, index, index) -> ()
+  // CHECK: ( ( ( 0, 0, -42 ), ( 0, 0, -42 ), ( 0, 0, -42 ), ( 0, 0, -42 ), ( 0, 0, -42 ) ), ( ( 20, 30, -42 ), ( 20, 30, -42 ), ( 20, 30, -42 ), ( 20, 30, -42 ), ( 20, 30, -42 ) ) )
+
+  // 6. Read 1D vector from 4D memref with mask and broadcast vector to 3D.
+  call @transfer_read_3d_mask_broadcast(%A, %c0, %c0, %c0, %c0)
+      : (memref<?x?x?x?xf32>, index, index, index, index) -> ()
+  // CHECK: ( ( ( -42, -42, -42 ), ( -42, -42, -42 ), ( -42, -42, -42 ), ( -42, -42, -42 ), ( -42, -42, -42 ) ), ( ( 20, 20, 20 ), ( 20, 20, 20 ), ( 20, 20, 20 ), ( 20, 20, 20 ), ( 20, 20, 20 ) ) )
+
   return
 }
-
-// CHECK: ( ( ( 0, 0, -42 ), ( 2, 3, -42 ), ( 4, 6, -42 ), ( 6, 9, -42 ), ( -42, -42, -42 ) ), ( ( 20, 30, -42 ), ( 22, 33, -42 ), ( 24, 36, -42 ), ( 26, 39, -42 ), ( -42, -42, -42 ) ) )
-// CHECK: ( ( ( 0, 0, -42 ), ( 2, -1, -42 ), ( 4, -1, -42 ), ( 6, -1, -42 ), ( -42, -42, -42 ) ), ( ( 20, 30, -42 ), ( 22, -1, -42 ), ( 24, -1, -42 ), ( 26, -1, -42 ), ( -42, -42, -42 ) ) )
-// CHECK: ( ( ( 0, 20, 40 ), ( 0, 20, 40 ), ( 0, 20, 40 ), ( 0, 20, 40 ), ( 0, 20, 40 ) ), ( ( 0, 30, 60 ), ( 0, 30, 60 ), ( 0, 30, 60 ), ( 0, 30, 60 ), ( 0, 30, 60 ) ), ( ( -42, -42, -42 ), ( -42, -42, -42 ), ( -42, -42, -42 ), ( -42, -42, -42 ), ( -42, -42, -42 ) ) )
-// CHECK: ( ( ( 0, 0, -42 ), ( 0, 0, -42 ), ( 0, 0, -42 ), ( 0, 0, -42 ), ( 0, 0, -42 ) ), ( ( 20, 30, -42 ), ( 20, 30, -42 ), ( 20, 30, -42 ), ( 20, 30, -42 ), ( 20, 30, -42 ) ) )