[mlir] Add gpu.memcpy op.

[csigg]

Swap %src and %dst, doh.

[herhut]

mlir/include/mlir/Dialect/GPU/GPUOps.td

883

SameOperandsAndResultShape and SameOperandsAndResultElementType should make the custom verifier redundant.

901

This syntax is somewhat unprecedented, to have types intermixed with operands. I think I would prefer something like

gpu.memcpy %dst %src : type, type

We could be extra verbose and do

gpu.memcopy %src to %dst : type, type

[csigg]

Change assembly syntax.
Improve verifier.

[csigg]

mlir/include/mlir/Dialect/GPU/GPUOps.td

883

That doesn't work because of the !gpu.async.token operands. I've stole the implementations from those verifiers though.

901

I changed it to gpu.memcpy %dst %src : type, type.

[herhut]

mlir/include/mlir/Dialect/GPU/GPUOps.td

901

Hyper nit: Remove extra space in front of comma.

[csigg]

Remove space.

mlir/include/mlir/Dialect/GPU/GPUOps.td

def GPU_DeallocOp : GPU_Op<"dealloc", [
   let results = (outs Optional<GPU_AsyncToken>:$asyncToken);
 
   let assemblyFormat = [{
     custom<AsyncDependencies>(type($asyncToken), $asyncDependencies)
     $memref attr-dict `:` type($memref)
   }];
 }
 
+def GPU_MemcpyOp : GPU_Op<"memcpy", [
+    GPU_AsyncOpInterface, MemoryEffects<[MemRead, MemWrite]>
+  ]> {
+
+  let summary = "GPU memcpy operation";
+
+  let description = [{
+    The `gpu.memcpy` operation copies the content of one memref to another.
+
+    The op does not execute before all async dependencies have finished
+    executing.
+
+    If the `async` keyword is present, the op is executed asynchronously (i.e.
+    it does not block until the execution has finished on the device). In
+    that case, it returns a !gpu.async.token.
+
+    Example:
+
+    ```mlir
+    %token = gpu.memcpy async [%dep] %src : memref<?xf32>, %dst : memref<?xf32, 1>
+    ```
+  }];
+
+  let arguments = (ins Variadic<GPU_AsyncToken>:$asyncDependencies,
+                   Arg<AnyMemRef, "", [MemRead]>:$src,
+                   Arg<AnyMemRef, "", [MemWrite]>:$dst);
+  let results = (outs Optional<GPU_AsyncToken>:$asyncToken);
+
+  let assemblyFormat = [{
+    custom<AsyncDependencies>(type($asyncToken), $asyncDependencies)
+    $src `:` type($src)`,` $dst `:` type($dst) attr-dict
+  }];
+  let verifier = [{ return ::verify(*this); }];
+}
+
 #endif // GPU_OPS

mlir/lib/Dialect/GPU/IR/GPUDialect.cpp

static void print(OpAsmPrinter &p, GPUModuleOp op) {
   p << op.getOperationName() << ' ';
   p.printSymbolName(op.getName());
   p.printOptionalAttrDictWithKeyword(op.getAttrs(),
                                      {SymbolTable::getSymbolAttrName()});
   p.printRegion(op->getRegion(0), /*printEntryBlockArgs=*/false,
                 /*printBlockTerminators=*/false);
 }
 
+//===----------------------------------------------------------------------===//
+// GPUMemcpyOp
+//===----------------------------------------------------------------------===//
+
+static LogicalResult verify(MemcpyOp op) {
+  auto srcType = op.src().getType().cast<MemRefType>();
+  auto dstType = op.dst().getType().cast<MemRefType>();
+
+  if (srcType.getElementType() != dstType.getElementType())
+    return op.emitOpError("arguments have incompatible element type");
+
+  if (srcType.getShape() != dstType.getShape())
+    return op.emitOpError("arguments have incompatible shape");
+
+  return success();
+}
+
 static ParseResult parseAsyncDependencies(
     OpAsmParser &parser, Type &asyncTokenType,
     SmallVectorImpl<OpAsmParser::OperandType> &asyncDependencies) {
   auto loc = parser.getCurrentLocation();
   if (succeeded(parser.parseOptionalKeyword("async"))) {
     if (parser.getNumResults() == 0)
       return parser.emitError(loc, "needs to be named when marked 'async'");
     asyncTokenType = parser.getBuilder().getType<AsyncTokenType>();

mlir/test/Dialect/GPU/invalid.mlir

 }
 
 // -----
 
 func @async_wait_without_result() {
   // expected-error @+1 {{custom op 'gpu.wait' needs to be named when marked 'async'}}
   gpu.wait async
 }
+
+// -----
+
+func @memcpy_incompatible_type(%src : memref<?xf32>, %dst : memref<?xi32>) {
+  // expected-error @+1 {{'gpu.memcpy' op arguments have incompatible element type}}
+  gpu.memcpy %src : memref<?xf32>, %dst : memref<?xi32>
+}
+
+// -----
+
+func @memcpy_incompatible_shape(%src : memref<7xf32>, %dst : memref<9xf32>) {
+  // expected-error @+1 {{'gpu.memcpy' op arguments have incompatible shape}}
+  gpu.memcpy %src : memref<7xf32>, %dst : memref<9xf32>
+}

mlir/test/Dialect/GPU/ops.mlir

func @async_wait() {
     // CHECK: gpu.wait [%[[t0]], %[[t1]]]
     // CHECK-NOT: async
     gpu.wait [%0, %1]
     // CHECK: gpu.wait
     // CHECK-NOT: async
     gpu.wait // Valid, but a no-op.
     return
   }
+
+  func @memcpy(%src : memref<3x7xf32>, %dst : memref<3x7xf32, 1>) {
+    // CHECK-LABEL: func @memcpy
+    // CHECK: gpu.memcpy {{.*}} : memref<3x7xf32>, {{.*}} : memref<3x7xf32, 1>
+    gpu.memcpy %src : memref<3x7xf32>, %dst : memref<3x7xf32, 1>
+    // CHECK: %[[t0:.*]] = gpu.wait async
+    %0 = gpu.wait async
+    // CHECK: {{.*}} = gpu.memcpy async [%[[t0]]] {{.*}} : memref<3x7xf32>, {{.*}} : memref<3x7xf32, 1>
+    %1 = gpu.memcpy async [%0] %src : memref<3x7xf32>, %dst : memref<3x7xf32, 1>
+    return
+  }
 }