[mlir] Add gpu.memzero operation

[bondhugula]

LGTM.

Do you plan to add a lowering via CUDA memset? If not, this could be a useful starter issue to list on LLVM Github issues.

mlir/test/Dialect/GPU/ops.mlir

261

Nit: Indent CHECK-LABEL on par with the line above.

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

def GPU_MemsetOp : GPU_Op<"memset",
 
   let assemblyFormat = [{
     custom<AsyncDependencies>(type($asyncToken), $asyncDependencies)
     $dst`,` $value `:` type($dst)`,` type($value) attr-dict
   }];
   let hasFolder = 1;
 }
 
+def GPU_MemzeroOp : GPU_Op<"memzero", [GPU_AsyncOpInterface]> {
+
+  let summary = "GPU memzero operation";
+
+  let description = [{
+    The `gpu.memzero` operation sets the content of memref to zero value.
+
+    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.memzero async [%dep] %dst : memref<?xf32, 1>
+    ```
+  }];
+
+  let arguments = (ins Variadic<GPU_AsyncToken>:$asyncDependencies,
+                   Arg<AnyMemRef, "", [MemWrite]>:$dst);
+  let results = (outs Optional<GPU_AsyncToken>:$asyncToken);
+
+  let assemblyFormat = [{
+    custom<AsyncDependencies>(type($asyncToken), $asyncDependencies)
+    $dst `:` type($dst) attr-dict
+  }];
+  let hasFolder = 1;
+}
+
 def GPU_SetDefaultDeviceOp : GPU_Op<"set_default_device",
                                     [MemoryEffects<[MemWrite]>]>,
     Arguments<(ins I32:$devIndex)> {
   let summary = "Set default GPU for operations after this by index";
   let description = [{
     Operation that sets the current default GPU, using a zero-based index
     into the set of GPUs on the system. The default GPU setting may be
     thread-local.

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

LogicalResult MemcpyOp::fold(ArrayRef<Attribute> operands,
   return foldMemRefCast(*this);
 }
 
 LogicalResult MemsetOp::fold(ArrayRef<Attribute> operands,
                              SmallVectorImpl<::mlir::OpFoldResult> &results) {
   return foldMemRefCast(*this);
 }
 
+LogicalResult MemzeroOp::fold(ArrayRef<Attribute> operands,
+                              SmallVectorImpl<::mlir::OpFoldResult> &results) {
+  return foldMemRefCast(*this);
+}
+
 //===----------------------------------------------------------------------===//
 // GPU_WaitOp
 //===----------------------------------------------------------------------===//
 
 namespace {
 
 /// Remove gpu.wait op use of gpu.wait op def without async dependencies.
 /// %t = gpu.wait async []       // No async dependencies.

mlir/test/Dialect/GPU/canonicalize.mlir

 func.func @memset_after_cast(%arg0: memref<10xf32>, %arg1: f32) {
   // CHECK-NOT: memref.cast
   // CHECK: gpu.memset
   %0 = memref.cast %arg0 : memref<10xf32> to memref<?xf32>
   gpu.memset %0, %arg1 : memref<?xf32>, f32
   return
 }
 
+// CHECK-LABEL: @memzero_after_cast
+func.func @memzero_after_cast(%arg0: memref<10xf32>) {
+  // CHECK-NOT: memref.cast
+  // CHECK: gpu.memzero
+  %0 = memref.cast %arg0 : memref<10xf32> to memref<?xf32>
+  gpu.memzero %0 : memref<?xf32>
+  return
+}
+
 // -----
 
 // Test case: Folding of memref.dim(gpu.alloc(%size), %idx) -> %size
 // CHECK-LABEL: func @gpu_dim_of_alloc(
 //  CHECK-SAME:     %[[SIZE:[0-9a-z]+]]: index
 //  CHECK-NEXT:   return %[[SIZE]] : index
 func.func @gpu_dim_of_alloc(%size: index) -> index {
   %0 = gpu.alloc(%size) : memref<?xindex>

mlir/test/Dialect/GPU/ops.mlir

func.func @memset(%dst : memref<3x7xf32>, %value : f32) {
     gpu.memset %dst, %value : memref<3x7xf32>, f32
     // CHECK: %[[t0:.*]] = gpu.wait async
     %0 = gpu.wait async
     // CHECK: {{.*}} = gpu.memset async [%[[t0]]] {{.*}}, {{.*}} : memref<3x7xf32>, f32
     %1 = gpu.memset async [%0] %dst, %value : memref<3x7xf32>, f32
     return
   }
 
+  func.func @memzero(%dst : memref<3x7xf32>) {
+    // CHECK-LABEL: func @memzero
+    // CHECK: gpu.memzero {{.*}} : memref<3x7xf32>
+    gpu.memzero %dst : memref<3x7xf32>
+    // CHECK: %[[t0:.*]] = gpu.wait async
+    %0 = gpu.wait async
+    // CHECK: {{.*}} = gpu.memzero async [%[[t0]]] {{.*}} : memref<3x7xf32>
+    %1 = gpu.memzero async [%0] %dst : memref<3x7xf32>
+    return
+  }
+
   func.func @mmamatrix_valid_element_type(%src : memref<32x32xf16, affine_map<(d0, d1) -> (d0 * 64 + d1)>>){
     // CHECK-LABEL: func @mmamatrix_valid_element_type
     %wg = memref.alloca() {alignment = 32} : memref<32x32xf16, 3>
     // CHECK: %[[wg:.*]] = memref.alloca()
     %i = arith.constant 16 : index
     // CHECK: %[[i:.*]] = arith.constant 16 : index
      %cst = arith.constant 1.000000e+00 : f32
     // CHECK: %[[cst:.*]] = arith.constant 1.000000e+00 : f32