[mlir][Vector] Add folding of memref_cast into vector_transfer ops

[nicolasvasilache]

Summary

This revision adds a common folding pattern that starts appearing on
vector_transfer ops.

[ThomasRaoux]

LGTM

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

[{
     /// 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) {
         return impl::getTransferMinorIdentityMap(memRefType, vectorType);
       }
   }];
+
+  let hasFolder = 1;
 }
 
 def Vector_TransferWriteOp :
   Vector_Op<"transfer_write">,
     Arguments<(ins AnyVector:$vector, AnyMemRef:$memref,
                Variadic<Index>:$indices,
                AffineMapAttr:$permutation_map,
                OptionalAttr<BoolArrayAttr>:$masked)> {
[{
     /// 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) {
         return impl::getTransferMinorIdentityMap(memRefType, vectorType);
       }
   }];
+
+  let hasFolder = 1;
 }
 
 def Vector_ShapeCastOp :
   Vector_Op<"shape_cast", [NoSideEffect]>,
     Arguments<(ins AnyTypeOf<[AnyVector, TupleOf<[AnyVector]>]>:$source)>,
     Results<(outs AnyTypeOf<[AnyVector, TupleOf<[AnyVector]>]>:$result)> {
   let summary = "shape_cast casts between vector shapes";
   let description = [{

mlir/lib/Dialect/Vector/VectorOps.cpp

if (paddingType != vectorType.getElementType())
       return op.emitOpError(
           "requires formal padding and vector of the same elemental type");
   }
 
   return verifyPermutationMap(permutationMap,
                               [&op](Twine t) { return op.emitOpError(t); });
 }
 
+/// This is a common class used for patterns of the form
+/// ```
+///    someop(memrefcast) -> someop
+/// ```
+/// It folds the source of the memref_cast into the root operation directly.
+static LogicalResult foldMemRefCast(Operation *op) {
+  bool folded = false;
+  for (OpOperand &operand : op->getOpOperands()) {
+    auto castOp = operand.get().getDefiningOp<MemRefCastOp>();
+    if (castOp && canFoldIntoConsumerOp(castOp)) {
+      operand.set(castOp.getOperand());
+      folded = true;
+    }
+  }
+  return success(folded);
+}
+
+OpFoldResult TransferReadOp::fold(ArrayRef<Attribute>) {
+  /// transfer_read(memrefcast) -> transfer_read
+  if (succeeded(foldMemRefCast(*this)))
+    return getResult();
+  return OpFoldResult();
+}
+
 //===----------------------------------------------------------------------===//
 // TransferWriteOp
 //===----------------------------------------------------------------------===//
 
 /// Builder that sets permutation map to 'getMinorIdentityMap'.
 void TransferWriteOp::build(OpBuilder &builder, OperationState &result,
                             Value vector, Value memref, ValueRange indices,
                             ArrayRef<bool> maybeMasked) {
if (failed(verifyTransferOp(op.getOperation(), memrefType, vectorType,
                               permutationMap,
                               op.masked() ? *op.masked() : ArrayAttr())))
     return failure();
 
   return verifyPermutationMap(permutationMap,
                               [&op](Twine t) { return op.emitOpError(t); });
 }
 
+LogicalResult TransferWriteOp::fold(ArrayRef<Attribute>,
+                                    SmallVectorImpl<OpFoldResult> &) {
+  return foldMemRefCast(*this);
+}
+
 //===----------------------------------------------------------------------===//
 // ShapeCastOp
 //===----------------------------------------------------------------------===//
 
 /// Returns true if each element of 'a' is equal to the product of a contiguous
 /// sequence of the elements of 'b'. Returns false otherwise.
 static bool isValidShapeCast(ArrayRef<int64_t> a, ArrayRef<int64_t> b) {
   unsigned rankA = a.size();

mlir/test/Dialect/Vector/canonicalize.mlir

func @transpose_3D_sequence(%arg : vector<4x3x2xf32>) -> vector<4x3x2xf32> {
   %5 = vector.transpose %4, [2, 1, 0] : vector<2x3x4xf32> to vector<4x3x2xf32>
   // CHECK-NOT: transpose
   %6 = vector.transpose %3, [2, 1, 0] : vector<2x3x4xf32> to vector<4x3x2xf32>
   // CHECK: [[ADD:%.*]] = addf [[T5]], [[ARG]]
   %7 = addf %5, %6 : vector<4x3x2xf32>
   // CHECK-NEXT: return [[ADD]]
   return %7 : vector<4x3x2xf32>
 }
+
+// -----
+
+// 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>
+  %1 = vector.transfer_read %0[%c0, %c0], %f0 : memref<?x?xf32>, vector<4x8xf32>
+
+  // CHECK: vector.transfer_write %{{.*}} : vector<4x8xf32>, memref<4x8xf32>
+  vector.transfer_write %1, %0[%c0, %c0] : vector<4x8xf32>, memref<?x?xf32>
+  return %1 : vector<4x8xf32>
+}