[mlir][memref] Add memref.copy operation

[herhut]

Summary

As the name suggests, it copies from one memref to another.

[pifon2a]

mlir/include/mlir/Dialect/MemRef/IR/MemRefOps.td

452

does it mean SameOperandsType, although we don't have this trait in mlir yet? Are different layouts allowed?

[herhut]

Added more detail in comment

[herhut]

mlir/include/mlir/Dialect/MemRef/IR/MemRefOps.td

452

I was not sure whether SameOperandsType would imply that the types have to be identical, which they don't. They may differ in the used maps and I have extended the comment a little.

mlir/include/mlir/Dialect/MemRef/IR/MemRefOps.td

def CloneOp : MemRef_Op<"clone", [
   let assemblyFormat = "$input attr-dict `:` type($input) `to` type($output)";
 
   let hasFolder = 1;
   let verifier = ?;
   let hasCanonicalizer = 1;
 }
 
 //===----------------------------------------------------------------------===//
+// CopyOp
+//===----------------------------------------------------------------------===//
+
+def CopyOp : MemRef_Op<"copy",
+    [CopyOpInterface, SameOperandsElementType, SameOperandsShape]> {
+
+  let description = [{
+    Copies the data from the source to the destination memref.
+
+    Usage:
+
+    ```mlir
+    memref.copy %arg0, %arg1 : memref<?xf32> to memref<?xf32>
+    ```
+
+    Source and destination are expected to have the same element type and shape.
+    Otherwise, the result is undefined. They may have different layouts.
+  }];
+
+  let arguments = (ins Arg<AnyRankedOrUnrankedMemRef, "the memref to copy from",
+                           [MemRead]>:$source,
+                       Arg<AnyRankedOrUnrankedMemRef, "the memref to copy to",
+                           [MemWrite]>:$target);
+
+  let extraClassDeclaration = [{
+    Value getSource() { return source();}
+    Value getTarget() { return target(); }
+  }];
+
+  let assemblyFormat = [{
+    $source `,` $target attr-dict `:` type($source) `to` type($target)
+  }];
+
+  let verifier = ?;
+}
+
+//===----------------------------------------------------------------------===//
 // DeallocOp
 //===----------------------------------------------------------------------===//
 
 def MemRef_DeallocOp : MemRef_Op<"dealloc", [MemRefsNormalizable]> {
   let summary = "memory deallocation operation";
   let description = [{
     The `dealloc` operation frees the region of memory referenced by a memref
     which was originally created by the `alloc` operation.

mlir/test/Dialect/MemRef/invalid.mlir

 
 memref.global @gv : memref<3xi32>
 
 func @mismatched_types() {
   // expected-error @+1 {{result type 'memref<3xf32>' does not match type 'memref<3xi32>' of the global memref @gv}}
   %0 = memref.get_global @gv : memref<3xf32>
   return
 }
+
+// -----
+
+func @copy_different_shape(%arg0: memref<2xf32>, %arg1: memref<3xf32>) {
+  // expected-error @+1 {{op requires the same shape for all operands}}
+  memref.copy %arg0, %arg1 : memref<2xf32> to memref<3xf32>
+  return
+}
+
+// -----
+
+func @copy_different_eltype(%arg0: memref<2xf32>, %arg1: memref<2xf16>) {
+  // expected-error @+1 {{op requires the same element type for all operands}}
+  memref.copy %arg0, %arg1 : memref<2xf32> to memref<2xf16>
+  return
+}

mlir/test/Dialect/MemRef/ops.mlir

 // CHECK-LABEL: func @memref_clone
 func @memref_clone() {
   %0 = memref.alloc() : memref<2xf32>
   %1 = memref.cast %0 : memref<2xf32> to memref<*xf32>
   %2 = memref.clone %1 : memref<*xf32> to memref<*xf32>
   return
 }
 
+// CHECK-LABEL: func @memref_copy
+func @memref_copy() {
+  %0 = memref.alloc() : memref<2xf32>
+  %1 = memref.cast %0 : memref<2xf32> to memref<*xf32>
+  %2 = memref.alloc() : memref<2xf32>
+  %3 = memref.cast %0 : memref<2xf32> to memref<*xf32>
+  memref.copy %1, %3 : memref<*xf32> to memref<*xf32>
+  return
+}
+
 // CHECK-LABEL: func @memref_dealloc
 func @memref_dealloc() {
   %0 = memref.alloc() : memref<2xf32>
   %1 = memref.cast %0 : memref<2xf32> to memref<*xf32>
   memref.dealloc %1 : memref<*xf32>
   return
 }
 
 
 // CHECK-LABEL: func @memref_alloca_scope
 func @memref_alloca_scope() {
   memref.alloca_scope {
     memref.alloca_scope.return
   }
   return
 }