Skip to content

Commit

Permalink
[mlir][memref] Add memref.copy operation
Browse files Browse the repository at this point in the history
As the name suggests, it copies from one memref to another.

Differential Revision: https://reviews.llvm.org/D104657
  • Loading branch information
Stephan Herhut committed Jun 22, 2021
1 parent d177988 commit bb6afc6
Show file tree
Hide file tree
Showing 3 changed files with 63 additions and 0 deletions.
37 changes: 37 additions & 0 deletions mlir/include/mlir/Dialect/MemRef/IR/MemRefOps.td
Original file line number Diff line number Diff line change
Expand Up @@ -444,6 +444,43 @@ def CloneOp : MemRef_Op<"clone", [
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
//===----------------------------------------------------------------------===//
Expand Down
16 changes: 16 additions & 0 deletions mlir/test/Dialect/MemRef/invalid.mlir
Original file line number Diff line number Diff line change
Expand Up @@ -215,3 +215,19 @@ func @mismatched_types() {
%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
}
10 changes: 10 additions & 0 deletions mlir/test/Dialect/MemRef/ops.mlir
Original file line number Diff line number Diff line change
Expand Up @@ -69,6 +69,16 @@ func @memref_clone() {
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>
Expand Down

0 comments on commit bb6afc6

Please sign in to comment.