[mlir][memref] Implement fast lowering of memref.copy

[herhut]

Summary

In the absence of maps, we can lower memref.copy to a memcpy.

[bondhugula]

Do you want to instead create one pattern with two functions in it? It'll lead to less overhead in the greedy rewrite driver and perhaps also easier to later choose between the two.

mlir/lib/Conversion/MemRefToLLVM/MemRefToLLVM.cpp

781

Doc comment here - along the lines of your revision summary and title.

[mehdi_amini]

mlir/lib/Conversion/MemRefToLLVM/MemRefToLLVM.cpp

727

This should likely be documented, I think notifyFailure would do it here.

[herhut]

Add comments and failure reason.

[herhut]

In D116099#3205178, @bondhugula wrote:

Do you want to instead create one pattern with two functions in it? It'll lead to less overhead in the greedy rewrite driver and perhaps also easier to later choose between the two.

I though about that, too. Another way would be to give the memcpy based pattern higher benefit and actually make then non-exclusive. That way one could optionally blend in the memcpy based pattern.

One issue I currently see with this entire thing is that it relies on maps being present as a way to understand whether the memref has identity strides and offset. I am not sure whether this assumption actually holds for all users.

[bondhugula]

In D116099#3206432, @herhut wrote:

In D116099#3205178, @bondhugula wrote:

Do you want to instead create one pattern with two functions in it? It'll lead to less overhead in the greedy rewrite driver and perhaps also easier to later choose between the two.

I though about that, too. Another way would be to give the memcpy based pattern higher benefit and actually make then non-exclusive. That way one could optionally blend in the memcpy based pattern.

All of this still leads to two patterns.

One issue I currently see with this entire thing is that it relies on maps being present as a way to understand whether the memref has identity strides and offset. I am not sure whether this assumption actually holds for all users.

You are just using isIdentity() -- the API doesn't name map outside. In the inside, it's still a map.

mlir/lib/Conversion/MemRefToLLVM/MemRefToLLVM.cpp

710–711

non-identity maps -> non-identity layout

There is no reference to map in the API and in your code below and there isn't a need to. The MemRefLayoutAttrInterface hides it.

[herhut]

In D116099#3206683, @bondhugula wrote:

In D116099#3206432, @herhut wrote:

In D116099#3205178, @bondhugula wrote:

Do you want to instead create one pattern with two functions in it? It'll lead to less overhead in the greedy rewrite driver and perhaps also easier to later choose between the two.

I though about that, too. Another way would be to give the memcpy based pattern higher benefit and actually make then non-exclusive. That way one could optionally blend in the memcpy based pattern.

All of this still leads to two patterns.

Yes, my question is more whether that is useful, to have it as two patterns vs. one. The benefit modelling would be one reason. If everybody thinks that this is not useful enough for the cost, I can merge them.

One issue I currently see with this entire thing is that it relies on maps being present as a way to understand whether the memref has identity strides and offset. I am not sure whether this assumption actually holds for all users.

You are just using isIdentity() -- the API doesn't name map outside. In the inside, it's still a map.

I was not clear. I was wondering about the case where the static type of the memref does not have a map but the descriptor at runtime still does not have identity strides. Currently, one can create this setting using memref.reinterpret_cast by omitting the map from the target type. Maybe that should be illegal.

[bondhugula]

You are just using isIdentity() -- the API doesn't name map outside. In the inside, it's still a map.

I was not clear. I was wondering about the case where the static type of the memref does not have a map but the descriptor at runtime still does not have identity strides. Currently, one can create this setting using memref.reinterpret_cast by omitting the map from the target type. Maybe that should be illegal.

This would be an invalid op and should have failed the verifier in the first place. (Note that the memref always has a map -- it's not printed if it's identity as you know. Also, identity strides (all ones) don't correspond to an identity map -- N^2, N, 1 for example would correspond to an identity map for a 3-d memref for example.)

[herhut]

Combined the two patterns.

[herhut]

In D116099#3225198, @bondhugula wrote:

You are just using isIdentity() -- the API doesn't name map outside. In the inside, it's still a map.

I was not clear. I was wondering about the case where the static type of the memref does not have a map but the descriptor at runtime still does not have identity strides. Currently, one can create this setting using memref.reinterpret_cast by omitting the map from the target type. Maybe that should be illegal.

This would be an invalid op and should have failed the verifier in the first place. (Note that the memref always has a map -- it's not printed if it's identity as you know. Also, identity strides (all ones) don't correspond to an identity map -- N^2, N, 1 for example would correspond to an identity map for a 3-d memref for example.)

I have sent https://reviews.llvm.org/D116601 to improve verification. One can still create this case dynamically but as we consider that illegal, the result of memref.copy is undefined in such cases, so it is fine to just use the intrinsics.

[bondhugula]

LGTM - thanks!

[herhut]

rebase

[herhut]

fix test

mlir/lib/Conversion/MemRefToLLVM/MemRefToLLVM.cpp

if (srcType.isa<MemRefType>() && dstType.isa<UnrankedMemRefType>()) {
       auto loadOp = rewriter.create<LLVM::LoadOp>(loc, castPtr);
       rewriter.replaceOp(memRefCastOp, loadOp.getResult());
     } else {
       llvm_unreachable("Unsupported unranked memref to unranked memref cast");
     }
   }
 };
 
-struct MemRefCopyOpLowering : public ConvertOpToLLVMPattern<memref::CopyOp> {
+/// Pattern to lower a `memref.copy` to the generic `MemrefCopyFn`.
+///
+/// This pattern only applies in the presence of non-identity maps on the
+/// operands.
+struct MemRefCopyOpWithMapsLowering
+    : public ConvertOpToLLVMPattern<memref::CopyOp> {
   using ConvertOpToLLVMPattern<memref::CopyOp>::ConvertOpToLLVMPattern;
 
   LogicalResult
   matchAndRewrite(memref::CopyOp op, OpAdaptor adaptor,
                   ConversionPatternRewriter &rewriter) const override {
     auto loc = op.getLoc();
     auto srcType = op.source().getType().cast<BaseMemRefType>();
     auto targetType = op.target().getType().cast<BaseMemRefType>();
 
+    if (srcType.hasRank() &&
+        srcType.cast<MemRefType>().getLayout().isIdentity() &&
+        targetType.hasRank() &&
+        targetType.cast<MemRefType>().getLayout().isIdentity())
+      return rewriter.notifyMatchFailure(
+          op, "identity layouts should be lowered to memcpy");
+
     // First make sure we have an unranked memref descriptor representation.
     auto makeUnranked = [&, this](Value ranked, BaseMemRefType type) {
       auto rank = rewriter.create<LLVM::ConstantOp>(
           loc, getIndexType(), rewriter.getIndexAttr(type.getRank()));
       auto *typeConverter = getTypeConverter();
       auto ptr =
           typeConverter->promoteOneMemRefDescriptor(loc, ranked, rewriter);
       auto voidPtr =
matchAndRewrite(memref::CopyOp op, OpAdaptor adaptor,
     rewriter.create<LLVM::CallOp>(loc, copyFn,
                                   ValueRange{elemSize, sourcePtr, targetPtr});
     rewriter.eraseOp(op);
 
     return success();
   }
 };
 
+/// Pattern to lower a `memref.copy` to the llvm `memcpy` intrinsic.
+///
+/// This pattern is a fast-path for copies that operate on memrefs with
+/// identity maps, i.e., where the data is contiguous in memory.
+struct MemRefCopyOpLowering : public ConvertOpToLLVMPattern<memref::CopyOp> {
+  using ConvertOpToLLVMPattern<memref::CopyOp>::ConvertOpToLLVMPattern;
+
+  LogicalResult
+  matchAndRewrite(memref::CopyOp op, OpAdaptor adaptor,
+                  ConversionPatternRewriter &rewriter) const override {
+    auto loc = op.getLoc();
+    auto srcType = op.source().getType().dyn_cast<MemRefType>();
+    auto targetType = op.target().getType().dyn_cast<MemRefType>();
+
+    // Only ranked memrefs with identity layout are supported.
+    if (!srcType || !targetType || !srcType.getLayout().isIdentity() ||
+        !targetType.getLayout().isIdentity())
+      return rewriter.notifyMatchFailure(
+          op, "non-identity layout cannot uses memcpy");
+
+    MemRefDescriptor srcDesc(adaptor.source());
+
+    // Compute number of elements.
+    Value numElements;
+    for (int pos = 0; pos < srcType.getRank(); ++pos) {
+      auto size = srcDesc.size(rewriter, loc, pos);
+      numElements = numElements
+                        ? rewriter.create<LLVM::MulOp>(loc, numElements, size)
+                        : size;
+    }
+    // Get element size.
+    auto sizeInBytes = getSizeInBytes(loc, srcType.getElementType(), rewriter);
+    // Compute total.
+    Value totalSize =
+        rewriter.create<LLVM::MulOp>(loc, numElements, sizeInBytes);
+
+    Value srcBasePtr = srcDesc.alignedPtr(rewriter, loc);
+    MemRefDescriptor targetDesc(adaptor.target());
+    Value targetBasePtr = targetDesc.alignedPtr(rewriter, loc);
+    Value isVolatile = rewriter.create<LLVM::ConstantOp>(
+        loc, typeConverter->convertType(rewriter.getI1Type()),
+        rewriter.getBoolAttr(false));
+    rewriter.create<LLVM::MemcpyOp>(loc, targetBasePtr, srcBasePtr, totalSize,
+                                    isVolatile);
+    rewriter.eraseOp(op);
+
+    return success();
+  }
+};
+
 /// Extracts allocated, aligned pointers and offset from a ranked or unranked
 /// memref type. In unranked case, the fields are extracted from the underlying
 /// ranked descriptor.
 static void extractPointersAndOffset(Location loc,
                                      ConversionPatternRewriter &rewriter,
                                      LLVMTypeConverter &typeConverter,
                                      Value originalOperand,
                                      Value convertedOperand,
patterns.add<
       AllocaScopeOpLowering,
       AssumeAlignmentOpLowering,
       DimOpLowering,
       GlobalMemrefOpLowering,
       GetGlobalMemrefOpLowering,
       LoadOpLowering,
       MemRefCastOpLowering,
       MemRefCopyOpLowering,
+      MemRefCopyOpWithMapsLowering,
       MemRefReinterpretCastOpLowering,
       MemRefReshapeOpLowering,
       PrefetchOpLowering,
       RankOpLowering,
       ReassociatingReshapeOpConversion<memref::ExpandShapeOp>,
       ReassociatingReshapeOpConversion<memref::CollapseShapeOp>,
       StoreOpLowering,
       SubViewOpLowering,

mlir/test/mlir-cpu-runner/copy.mlir

func @main() -> () {
   memref.copy %input, %copy : memref<2x3xf32> to memref<2x3xf32>
   %unranked_copy = memref.cast %copy : memref<2x3xf32> to memref<*xf32>
   call @print_memref_f32(%unranked_copy) : (memref<*xf32>) -> ()
   // CHECK: rank = 2 offset = 0 sizes = [2, 3] strides = [3, 1]
   // CHECK-NEXT: [0,   1,   2]
   // CHECK-NEXT: [3,   4,   5]
 
   %copy_two = memref.alloc() : memref<3x2xf32>
-  %copy_two_casted = memref.reinterpret_cast %copy_two to offset: [0], sizes: [2,3], strides:[1, 2]
-    : memref<3x2xf32> to memref<2x3xf32>
-  memref.copy %input, %copy_two_casted : memref<2x3xf32> to memref<2x3xf32>
+  %copy_two_casted = memref.reinterpret_cast %copy_two to offset: [0], sizes: [2, 3], strides:[1, 2]
+    : memref<3x2xf32> to memref<2x3xf32, offset: 0, strides: [1, 2]>
+  memref.copy %input, %copy_two_casted : memref<2x3xf32> to memref<2x3xf32, offset: 0, strides: [1, 2]>
   %unranked_copy_two = memref.cast %copy_two : memref<3x2xf32> to memref<*xf32>
   call @print_memref_f32(%unranked_copy_two) : (memref<*xf32>) -> ()
   // CHECK: rank = 2 offset = 0 sizes = [3, 2] strides = [2, 1]
   // CHECK-NEXT: [0,   3]
   // CHECK-NEXT: [1,   4]
   // CHECK-NEXT: [2,   5]
 
   %input_empty = memref.alloc() : memref<3x0x1xf32>
   %copy_empty = memref.alloc() : memref<3x0x1xf32>
   // Copying an empty shape should do nothing (and should not crash).
   memref.copy %input_empty, %copy_empty : memref<3x0x1xf32> to memref<3x0x1xf32>
+
+  %input_empty_casted = memref.reinterpret_cast %input_empty to offset: [0], sizes: [0, 3, 1], strides:[1, 3, 1]
+    : memref<3x0x1xf32> to memref<0x3x1xf32>
+  %copy_empty_casted = memref.alloc() : memref<0x3x1xf32>
+  // Copying a casted empty shape should do nothing (and should not crash).
+  memref.copy %input_empty_casted, %copy_empty_casted : memref<0x3x1xf32> to memref<0x3x1xf32>
+
   memref.dealloc %copy_empty : memref<3x0x1xf32>
   memref.dealloc %input_empty : memref<3x0x1xf32>
   memref.dealloc %copy_two : memref<3x2xf32>
   memref.dealloc %copy : memref<2x3xf32>
   memref.dealloc %input : memref<2x3xf32>
   return
 }