[mlir][gpu] Add gpu.host_wait op.

[csigg]

Summary

This op could naturally be combined with gpu.create_token, by making a token result optional: If it returns a token, it behaves like create_token. Otherwise, it behaves like host_wait.
I do like the idea of combining them, but I don't have a good name for it, and the two use cases are quite different (create_token is created by lowering within GPU only, host_wait is added upstream).

[herhut]

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

776

Hmm, is this op really asynchronous? I find it strange that it implements this interface. It also breaks the assumption that all asynchronous operations return a token to sync on their effects.

Also, as all operands of this op are tokens, we already naturally have a way to add extra operands.

787

Is this actually needed?

794

I like the syntax with [] for tokens. But as this is not an async operation, maybe it should have the tokens as proper operands with ()?

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

def GPU_CreateTokenOp : GPU_Op<"create_token", [GPU_AsyncOpInterface]> {
   }];
 
   let arguments = (ins Variadic<GPU_AsyncToken>:$asyncDependencies);
   let results = (outs GPU_AsyncToken:$asyncToken);
 
   let assemblyFormat = "(`[` $asyncDependencies^ `]`)? attr-dict";
 }
 
+def GPU_WaitOp : GPU_Op<"host_wait", [GPU_AsyncOpInterface]> {
+  let summary = "GPU async blocking wait operation";
+  let description = [{
+    The `host_wait` operation blocks until all ops producing the async
+    dependency tokens finished execution.
+
+    All dependent memory operations are visible to the host once this op
+    completes.
+  }];
+
+  let arguments = (ins Variadic<GPU_AsyncToken>:$asyncDependencies);
+  let results = (outs);
+
+  let extraClassDeclaration = [{
+      // This op never returns an async token.
+      Value asyncToken() const { return Value(); }
+  }];
+
+  let assemblyFormat = "`[` $asyncDependencies `]` attr-dict";
+}
+
 #endif // GPU_OPS

mlir/test/Dialect/GPU/ops.mlir

func @async_create_token() {
     // CHECK: %[[t0:.*]] = gpu.create_token
     %0 = gpu.create_token
     // CHECK: %[[t1:.*]] = gpu.create_token[%[[t0]]]
     %1 = gpu.create_token[%0]
     // CHECK: gpu.create_token[%[[t0]], %[[t1]]]
     %2 = gpu.create_token[%0, %1]
     return
   }
+
+  func @async_host_wait() {
+    // CHECK-LABEL: func @async_host_wait
+    // CHECK: %[[t0:.*]] = gpu.create_token
+    %0 = gpu.create_token
+    // CHECK: gpu.host_wait[%[[t0]]]
+    gpu.host_wait[%0]
+    return
+  }
 }