[mlir][Vector] Add folding for vector.transfer ops

[nicolasvasilache]

Summary

This revision folds vector.transfer operations by updating the masked bool array attribute when more unmasked dimensions can be discovered.

[aartbik]

mlir/include/mlir/IR/AffineMap.h

102

nit: the map -> this affine map
(like above and below, now that this is instance, not static method)?

mlir/lib/Dialect/Vector/VectorOps.cpp

1731

At first reading, all these different sizes are a bit confusing (reserve, and two indices that go down).
Add a comment on the transfer format to make this more clear?

mlir/test/Dialect/Vector/canonicalize.mlir

371

I cases like this, I am always wondering on the scope of CHECK-NOT (how far does it go :-)
Perhaps add

CHECK: return

just to make the delimiters more clear?

[nicolasvasilache]

Review

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

bool isMaskedDim(unsigned dim) {
         masked()->cast<ArrayAttr>()[dim].cast<BoolAttr>().getValue();
     }
     MemRefType getMemRefType() {
       return memref().getType().cast<MemRefType>();
     }
     VectorType getVectorType() {
       return vector().getType().cast<VectorType>();
     }
+    // Number of dimensions that participate in the permutation map.
+    unsigned getTransferRank() {
+      return permutation_map().getNumResults();
+    }
+    // Number of leading dimensions that do not participate in the permutation
+    // map.
+    unsigned getLeadingMemRefRank() {
+      return getMemRefType().getRank() - permutation_map().getNumResults();
+    }
   }];
 }
 
 def Vector_TransferReadOp :
   Vector_Op<"transfer_read", [
       DeclareOpInterfaceMethods<VectorUnrollOpInterface, ["getShapeForUnroll"]>
     ]>,
     Arguments<(ins AnyMemRef:$memref, Variadic<Index>:$indices,

mlir/include/mlir/IR/AffineMap.h

public:
   bool operator==(AffineMap other) const { return other.map == map; }
   bool operator!=(AffineMap other) const { return !(other.map == map); }
 
   /// Returns true if this affine map is an identity affine map.
   /// An identity affine map corresponds to an identity affine function on the
   /// dimensional identifiers.
   bool isIdentity() const;
 
   /// Returns true if the map is a minor identity map, i.e. an identity affine
   /// map (d0, ..., dn) -> (dp, ..., dn) on the most minor dimensions.
-  static bool isMinorIdentity(AffineMap map);
+  bool isMinorIdentity() const;
 
   /// Returns true if this affine map is an empty map, i.e., () -> ().
   bool isEmpty() const;
 
   /// Returns true if this affine map is a single result constant function.
   bool isSingleConstant() const;
 
   /// Returns the constant result of this map. This methods asserts that the map

mlir/lib/Conversion/VectorToROCDL/VectorToROCDL.cpp

matchAndRewrite(Operation *op, ArrayRef<Value> operands,
                   ConversionPatternRewriter &rewriter) const override {
     auto xferOp = cast<ConcreteOp>(op);
     typename ConcreteOp::Adaptor adaptor(operands);
 
     if (xferOp.getVectorType().getRank() > 1 ||
         llvm::size(xferOp.indices()) == 0)
       return failure();
 
-    if (!AffineMap::isMinorIdentity(xferOp.permutation_map()))
+    if (!xferOp.permutation_map().isMinorIdentity())
       return failure();
 
     // Have it handled in vector->llvm conversion pass.
     if (!xferOp.isMaskedDim(0))
       return failure();
 
     auto toLLVMTy = [&](Type t) { return typeConverter.convertType(t); };
     LLVM::LLVMType vecTy =

mlir/lib/Conversion/VectorToSCF/VectorToSCF.cpp

NDTransferOpHelper(PatternRewriter &rewriter, ConcreteOp xferOp,
                      const VectorTransferToSCFOptions &options)
       : rewriter(rewriter), options(options), loc(xferOp.getLoc()),
         scope(std::make_unique<ScopedContext>(rewriter, loc)), xferOp(xferOp),
         op(xferOp.getOperation()) {
     vectorType = xferOp.getVectorType();
     // TODO: when we go to k > 1-D vectors adapt minorRank.
     minorRank = 1;
     majorRank = vectorType.getRank() - minorRank;
-    leadingRank = xferOp.getMemRefType().getRank() - (majorRank + minorRank);
+    leadingRank = xferOp.getLeadingMemRefRank();
     majorVectorType =
         VectorType::get(vectorType.getShape().take_front(majorRank),
                         vectorType.getElementType());
     minorVectorType =
         VectorType::get(vectorType.getShape().take_back(minorRank),
                         vectorType.getElementType());
     /// Memref of minor vector type is used for individual transfers.
     memRefMinorVectorType =

 
 /// Performs the rewrite.
 template <>
 LogicalResult VectorTransferRewriter<TransferReadOp>::matchAndRewrite(
     Operation *op, PatternRewriter &rewriter) const {
   using namespace mlir::edsc::op;
 
   TransferReadOp transfer = cast<TransferReadOp>(op);
-  if (AffineMap::isMinorIdentity(transfer.permutation_map())) {
+  if (transfer.permutation_map().isMinorIdentity()) {
     // If > 1D, emit a bunch of loops around 1-D vector transfers.
     if (transfer.getVectorType().getRank() > 1)
       return NDTransferOpHelper<TransferReadOp>(rewriter, transfer, options)
           .doReplace();
     // If 1-D this is now handled by the target-specific lowering.
     if (transfer.getVectorType().getRank() == 1)
       return failure();
   }

 /// TODO: implement alternatives to clipping.
 /// TODO: support non-data-parallel operations.
 template <>
 LogicalResult VectorTransferRewriter<TransferWriteOp>::matchAndRewrite(
     Operation *op, PatternRewriter &rewriter) const {
   using namespace edsc::op;
 
   TransferWriteOp transfer = cast<TransferWriteOp>(op);
-  if (AffineMap::isMinorIdentity(transfer.permutation_map())) {
+  if (transfer.permutation_map().isMinorIdentity()) {
     // If > 1D, emit a bunch of loops around 1-D vector transfers.
     if (transfer.getVectorType().getRank() > 1)
       return NDTransferOpHelper<TransferWriteOp>(rewriter, transfer, options)
           .doReplace();
     // If 1-D this is now handled by the target-specific lowering.
     if (transfer.getVectorType().getRank() == 1)
       return failure();
   }

mlir/lib/Dialect/Vector/VectorOps.cpp

static LogicalResult foldExtractOpFromTranspose(ExtractOp extractOp) {
   // If transposition permutation is larger than the ExtractOp, all minor
   // dimensions must be an identity for folding to occur. If not, individual
   // elements within the extracted value are transposed and this is not just a
   // simple folding.
   unsigned minorRank = permutation.size() - extractedPos.size();
   MLIRContext *ctx = extractOp.getContext();
   AffineMap permutationMap = AffineMap::getPermutationMap(permutation, ctx);
   AffineMap minorMap = permutationMap.getMinorSubMap(minorRank);
-  if (minorMap && !AffineMap::isMinorIdentity(minorMap))
+  if (minorMap && !minorMap.isMinorIdentity())
     return failure();
 
   //   %1 = transpose %0[x, y, z] : vector<axbxcxf32>
   //   %2 = extract %1[u, v] : vector<..xf32>
   // may turn into:
   //   %2 = extract %0[w, x] : vector<..xf32>
   // iff z == 2 and [w, x] = [x, y]^-1 o [u, v] here o denotes composition and
   // -1 denotes the inverse.
if (insertedPos.size() == extractedPos.size()) {
         }
       }
       if (fold) {
         assert(permutationMap.getNumResults() >= insertedPos.size() &&
                "expected map of rank larger than insert indexing");
         unsigned minorRank =
             permutationMap.getNumResults() - insertedPos.size();
         AffineMap minorMap = permutationMap.getMinorSubMap(minorRank);
-        if (!minorMap || AffineMap::isMinorIdentity(minorMap))
+        if (!minorMap || minorMap.isMinorIdentity())
           return insertOp.source();
       }
     }
 
     // If we haven't found a match, just continue to the next producing
     // `vector.insert` / `vector.transpose`.
     // Compute insert/transpose for the next iteration.
     insertOp = insertionDest.getDefiningOp<vector::InsertOp>();
for (OpOperand &operand : op->getOpOperands()) {
     if (castOp && canFoldIntoConsumerOp(castOp)) {
       operand.set(castOp.getOperand());
       folded = true;
     }
   }
   return success(folded);
 }
 
+template <typename TransferOp>
+static bool isInBounds(TransferOp op, int64_t resultIdx, int64_t indicesIdx) {
+  // TODO: support more aggressive createOrFold on:
+  // `op.indices()[indicesIdx] + vectorType < dim(op.memref(), indicesIdx)`
+  if (op.getMemRefType().isDynamicDim(indicesIdx))
+    return false;
+  Value index = op.indices()[indicesIdx];
+  auto cstOp = index.getDefiningOp<ConstantIndexOp>();
+  if (!cstOp)
+    return false;
+
+  int64_t memrefSize = op.getMemRefType().getDimSize(indicesIdx);
+  int64_t vectorSize = op.getVectorType().getDimSize(resultIdx);
+  return cstOp.getValue() + vectorSize <= memrefSize;
+}
+
+template <typename TransferOp>
+static LogicalResult foldTransferMaskAttribute(TransferOp op) {
+  AffineMap permutationMap = op.permutation_map();
+  if (!permutationMap.isMinorIdentity())
+    return failure();
+  bool changed = false;
+  SmallVector<bool, 4> isMasked;
+  isMasked.reserve(op.getTransferRank());
+  for (int64_t resultIdx = permutationMap.getNumResults() - 1,
+               indicesIdx = op.indices().size() - 1;
+       resultIdx >= 0; --resultIdx, --indicesIdx) {
+    // Already marked unmasked, nothing to see here.
+    if (!op.isMaskedDim(resultIdx)) {
+      isMasked.push_back(false);
+      continue;
+    }
+    // Currently masked, check whether we can statically determine it is
+    // inBounds.
+    auto inBounds = isInBounds(op, resultIdx, indicesIdx);
+    isMasked.push_back(!inBounds);
+    // We commit the pattern if it is "more inbounds".
+    changed |= inBounds;
+  }
+  if (!changed)
+    return failure();
+  // OpBuilder is only used as a helper to build an I64ArrayAttr.
+  OpBuilder b(op.getContext());
+  std::reverse(isMasked.begin(), isMasked.end());
+  op.setAttr(TransferOp::getMaskedAttrName(), b.getBoolArrayAttr(isMasked));
+  return success();
+}
+
 OpFoldResult TransferReadOp::fold(ArrayRef<Attribute>) {
   /// transfer_read(memrefcast) -> transfer_read
+  if (succeeded(foldTransferMaskAttribute(*this)))
+    return getResult();
   if (succeeded(foldMemRefCast(*this)))
     return getResult();
   return OpFoldResult();
 }
 
 Optional<SmallVector<int64_t, 4>> TransferReadOp::getShapeForUnroll() {
   auto s = getVectorType().getShape();
   return SmallVector<int64_t, 4>{s.begin(), s.end()};
if (failed(verifyTransferOp(op.getOperation(), memrefType, vectorType,
     return failure();
 
   return verifyPermutationMap(permutationMap,
                               [&op](Twine t) { return op.emitOpError(t); });
 }
 
 LogicalResult TransferWriteOp::fold(ArrayRef<Attribute>,
                                     SmallVectorImpl<OpFoldResult> &) {
+  if (succeeded(foldTransferMaskAttribute(*this)))
+    return success();
   return foldMemRefCast(*this);
 }
 
 Optional<SmallVector<int64_t, 4>> TransferWriteOp::getShapeForUnroll() {
   auto s = getVectorType().getShape();
   return SmallVector<int64_t, 4>{s.begin(), s.end()};
 }
 

mlir/lib/IR/AffineMap.cpp

 /// minor dimensions.
 AffineMap AffineMap::getMinorIdentityMap(unsigned dims, unsigned results,
                                          MLIRContext *context) {
   assert(dims >= results && "Dimension mismatch");
   auto id = AffineMap::getMultiDimIdentityMap(dims, context);
   return AffineMap::get(dims, 0, id.getResults().take_back(results), context);
 }
 
-bool AffineMap::isMinorIdentity(AffineMap map) {
-  if (!map)
-    return false;
-  return map == getMinorIdentityMap(map.getNumDims(), map.getNumResults(),
-                                    map.getContext());
+bool AffineMap::isMinorIdentity() const {
+  return *this ==
+         getMinorIdentityMap(getNumDims(), getNumResults(), getContext());
 }
 
 /// Returns an AffineMap representing a permutation.
 AffineMap AffineMap::getPermutationMap(ArrayRef<unsigned> permutation,
                                        MLIRContext *context) {
   assert(!permutation.empty() &&
          "Cannot create permutation map from empty permutation vector");
   SmallVector<AffineExpr, 4> affExprs;

mlir/test/Dialect/Vector/canonicalize.mlir

 // -----
 
 // CHECK-LABEL: cast_transfers
 func @cast_transfers(%A: memref<4x8xf32>) -> (vector<4x8xf32>) {
   %c0 = constant 0 : index
   %f0 = constant 0.0 : f32
   %0 = memref_cast %A : memref<4x8xf32> to memref<?x?xf32>
 
-  // CHECK: vector.transfer_read %{{.*}} : memref<4x8xf32>, vector<4x8xf32>
+  // CHECK: vector.transfer_read %{{.*}} {masked = [false, false]} : memref<4x8xf32>, vector<4x8xf32>
   %1 = vector.transfer_read %0[%c0, %c0], %f0 : memref<?x?xf32>, vector<4x8xf32>
 
-  // CHECK: vector.transfer_write %{{.*}} : vector<4x8xf32>, memref<4x8xf32>
+  // CHECK: vector.transfer_write %{{.*}} {masked = [false, false]} : vector<4x8xf32>, memref<4x8xf32>
   vector.transfer_write %1, %0[%c0, %c0] : vector<4x8xf32>, memref<?x?xf32>
   return %1 : vector<4x8xf32>
 }
 
 // -----
 
 // CHECK-LABEL: func @insert_extract_transpose_2d(
 //  CHECK-SAME: %[[V:[a-zA-Z0-9]*]]: vector<2x3xf32>,
func @fold_extract_transpose(
   // Not a minor identity map so intra-vector level has been permuted
   //  CHECK-NEXT: vector.transpose %[[B]], [0, 2, 3, 1]
   //  CHECK-NEXT: vector.extract %{{.*}}[0, 1, 2]
   %4 = vector.transpose %b, [0, 2, 3, 1] : vector<3x6x5x6xf32> to vector<3x5x6x6xf32>
   %5 = vector.extract %4[0, 1, 2] : vector<3x5x6x6xf32>
 
   return %1, %3, %5 : vector<6xf32>, vector<6xf32>, vector<6xf32>
 }
+
+// -----
+
+// CHECK-LABEL: fold_vector_transfers
+func @fold_vector_transfers(%A: memref<?x8xf32>) -> (vector<4x8xf32>, vector<4x9xf32>) {
+  %c0 = constant 0 : index
+  %f0 = constant 0.0 : f32
+
+  // CHECK: vector.transfer_read %{{.*}} {masked = [true, false]}
+  %1 = vector.transfer_read %A[%c0, %c0], %f0 : memref<?x8xf32>, vector<4x8xf32>
+
+  // CHECK: vector.transfer_write %{{.*}} {masked = [true, false]}
+  vector.transfer_write %1, %A[%c0, %c0] : vector<4x8xf32>, memref<?x8xf32>
+
+  // Both dims masked, attribute is elided.
+  // CHECK: vector.transfer_read %{{.*}}
+  // CHECK-NOT: masked
+  %2 = vector.transfer_read %A[%c0, %c0], %f0 : memref<?x8xf32>, vector<4x9xf32>
+
+  // Both dims masked, attribute is elided.
+  // CHECK: vector.transfer_write %{{.*}}
+  // CHECK-NOT: masked
+  vector.transfer_write %2, %A[%c0, %c0] : vector<4x9xf32>, memref<?x8xf32>
+
+  return %1, %2 : vector<4x8xf32>, vector<4x9xf32>
+}

mlir/test/Dialect/Vector/vector-transforms.mlir

 // CHECK-NEXT: %[[VTR6:.*]] = vector.transfer_read %{{.*}}[%[[C2]], %[[C0]]], %{{.*}} : memref<4x4xf32>, vector<2x2xf32>
 // CHECK-NEXT: %[[VTR7:.*]] = vector.transfer_read %{{.*}}[%[[C2]], %[[C2]]], %{{.*}} : memref<4x4xf32>, vector<2x2xf32>
 
 // CHECK-NEXT: %[[R0:.*]] = vector.contract {indexing_maps = [#map2, #map3, #map0], iterator_types = ["parallel", "reduction", "parallel"]} %[[VTR0]], %[[VTR2]], %[[VTR4]] : vector<2x2xf32>, vector<2x2xf32> into vector<2x2xf32>
 // CHECK-NEXT: %[[R1:.*]] = vector.contract {indexing_maps = [#map2, #map3, #map0], iterator_types = ["parallel", "reduction", "parallel"]} %[[VTR0]], %[[VTR3]], %[[VTR5]] : vector<2x2xf32>, vector<2x2xf32> into vector<2x2xf32>
 // CHECK-NEXT: %[[R2:.*]] = vector.contract {indexing_maps = [#map2, #map3, #map0], iterator_types = ["parallel", "reduction", "parallel"]} %[[VTR1]], %[[VTR2]], %[[VTR6]] : vector<2x2xf32>, vector<2x2xf32> into vector<2x2xf32>
 // CHECK-NEXT: %[[R3:.*]] = vector.contract {indexing_maps = [#map2, #map3, #map0], iterator_types = ["parallel", "reduction", "parallel"]} %[[VTR1]], %[[VTR3]], %[[VTR7]] : vector<2x2xf32>, vector<2x2xf32> into vector<2x2xf32>
 
-// CHECK-NEXT: vector.transfer_write %[[R0]], %{{.*}}[%[[C0]], %[[C0]]] : vector<2x2xf32>, memref<4x4xf32>
-// CHECK-NEXT: vector.transfer_write %[[R1]], %{{.*}}[%[[C0]], %[[C2]]] : vector<2x2xf32>, memref<4x4xf32>
-// CHECK-NEXT: vector.transfer_write %[[R2]], %{{.*}}[%[[C2]], %[[C0]]] : vector<2x2xf32>, memref<4x4xf32>
-// CHECK-NEXT: vector.transfer_write %[[R3]], %{{.*}}[%[[C2]], %[[C2]]] : vector<2x2xf32>, memref<4x4xf32>
+// CHECK-NEXT: vector.transfer_write %[[R0]], %{{.*}}[%[[C0]], %[[C0]]] {masked = [false, false]} : vector<2x2xf32>, memref<4x4xf32>
+// CHECK-NEXT: vector.transfer_write %[[R1]], %{{.*}}[%[[C0]], %[[C2]]] {masked = [false, false]} : vector<2x2xf32>, memref<4x4xf32>
+// CHECK-NEXT: vector.transfer_write %[[R2]], %{{.*}}[%[[C2]], %[[C0]]] {masked = [false, false]} : vector<2x2xf32>, memref<4x4xf32>
+// CHECK-NEXT: vector.transfer_write %[[R3]], %{{.*}}[%[[C2]], %[[C2]]] {masked = [false, false]} : vector<2x2xf32>, memref<4x4xf32>
 // CHECK-NEXT: return
 
 func @contraction4x4_ikj_xfer_read(%arg0 : memref<4x2xf32>,
                                    %arg1 : memref<2x4xf32>,
                                    %arg2 : memref<4x4xf32>) {
   %c0 = constant 0 : index
   %cf0 = constant 0.0 : f32
 
func @vector_transfers(%arg0: index, %arg1: index) {
   %cst = constant 0.000000e+00 : f32
   %0 = alloc(%arg0, %arg1) : memref<?x?xf32>
   %1 = alloc(%arg0, %arg1) : memref<?x?xf32>
   %2 = alloc(%arg0, %arg1) : memref<?x?xf32>
   %cst_0 = constant 1.000000e+00 : f32
   %cst_1 = constant 2.000000e+00 : f32
   affine.for %arg2 = 0 to %arg0 step 4 {
     affine.for %arg3 = 0 to %arg1 step 4 {
-      %4 = vector.transfer_read %0[%arg2, %arg3], %cst  {permutation_map = affine_map<(d0, d1) -> (d0, d1)>} : memref<?x?xf32>, vector<4x4xf32>
-      %5 = vector.transfer_read %1[%arg2, %arg3], %cst  {permutation_map = affine_map<(d0, d1) -> (d0, d1)>} : memref<?x?xf32>, vector<4x4xf32>
+      %4 = vector.transfer_read %0[%arg2, %arg3], %cst {permutation_map = affine_map<(d0, d1) -> (d0, d1)>} : memref<?x?xf32>, vector<4x4xf32>
+      %5 = vector.transfer_read %1[%arg2, %arg3], %cst {permutation_map = affine_map<(d0, d1) -> (d0, d1)>} : memref<?x?xf32>, vector<4x4xf32>
       %6 = addf %4, %5 : vector<4x4xf32>
       vector.transfer_write %6, %2[%arg2, %arg3] {permutation_map = affine_map<(d0, d1) -> (d0, d1)>} : vector<4x4xf32>, memref<?x?xf32>
     }
   }
   return
 }
 
 // CHECK-LABEL: func @tuple_get(%arg0: vector<4xf32>, %arg1: vector<8xf32>)
func @cancelling_shape_cast_ops(%arg0 : vector<2x4xf32>) -> vector<2x4xf32> {
   %0 = vector.shape_cast %arg0 : vector<2x4xf32> to vector<8xf32>
   %1 = vector.shape_cast %0 : vector<8xf32> to vector<2x4xf32>
   return %1 : vector<2x4xf32>
 }
 
 // CHECK-LABEL: func @vector_transfers_vector_element_type
 //      CHECK: %[[C0:.*]] = constant 0 : index
 //      CHECK: %[[C1:.*]] = constant 1 : index
-//      CHECK: %[[VTR0:.*]] = vector.transfer_read %{{.*}}[%[[C0]], %[[C0]], %[[C0]]], %{{.*}} : memref<6x2x1xvector<2x4xf32>>, vector<1x1x2x4xf32>
-// CHECK-NEXT: %[[VTR1:.*]] = vector.transfer_read %{{.*}}[%[[C0]], %[[C1]], %[[C0]]], %{{.*}} : memref<6x2x1xvector<2x4xf32>>, vector<1x1x2x4xf32>
-// CHECK-NEXT: vector.transfer_write %[[VTR0]], %{{.*}}[%[[C0]], %[[C0]], %[[C0]]] : vector<1x1x2x4xf32>, memref<6x2x1xvector<2x4xf32>>
-// CHECK-NEXT: vector.transfer_write %[[VTR1]], %{{.*}}[%[[C0]], %[[C1]], %[[C0]]] : vector<1x1x2x4xf32>, memref<6x2x1xvector<2x4xf32>>
+//      CHECK: %[[VTR0:.*]] = vector.transfer_read %{{.*}}[%[[C0]], %[[C0]], %[[C0]]], %{{.*}} {masked = [false, false]} : memref<6x2x1xvector<2x4xf32>>, vector<1x1x2x4xf32>
+// CHECK-NEXT: %[[VTR1:.*]] = vector.transfer_read %{{.*}}[%[[C0]], %[[C1]], %[[C0]]], %{{.*}} {masked = [false, false]} : memref<6x2x1xvector<2x4xf32>>, vector<1x1x2x4xf32>
+// CHECK-NEXT: vector.transfer_write %[[VTR0]], %{{.*}}[%[[C0]], %[[C0]], %[[C0]]] {masked = [false, false]} : vector<1x1x2x4xf32>, memref<6x2x1xvector<2x4xf32>>
+// CHECK-NEXT: vector.transfer_write %[[VTR1]], %{{.*}}[%[[C0]], %[[C1]], %[[C0]]] {masked = [false, false]} : vector<1x1x2x4xf32>, memref<6x2x1xvector<2x4xf32>>
 
 func @vector_transfers_vector_element_type() {
   %c0 = constant 0 : index
   %cf0 = constant 0.000000e+00 : f32
   %vf0 = splat %cf0 : vector<2x4xf32>
 
   %0 = alloc() : memref<6x2x1xvector<2x4xf32>>