diff --git a/mlir/include/mlir/Dialect/GPU/GPUOps.td b/mlir/include/mlir/Dialect/GPU/GPUOps.td --- a/mlir/include/mlir/Dialect/GPU/GPUOps.td +++ b/mlir/include/mlir/Dialect/GPU/GPUOps.td @@ -773,4 +773,25 @@ 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:$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 diff --git a/mlir/test/Dialect/GPU/ops.mlir b/mlir/test/Dialect/GPU/ops.mlir --- a/mlir/test/Dialect/GPU/ops.mlir +++ b/mlir/test/Dialect/GPU/ops.mlir @@ -160,4 +160,13 @@ %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 + } }