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 @@ -756,4 +756,21 @@ let verifier = [{ return success(); }]; } +def GPU_CreateTokenOp : GPU_Op<"create_token", [GPU_AsyncOpInterface]> { + let summary = "Creates a GPU async token."; + let description = [{ + This op creates a new async token from a list of async dependencies. + + The op is inserted during lowering so that each async token is used only + once. This makes forks in async execution explicit. + + The op does not imply any host synchronization. + }]; + + let arguments = (ins Variadic:$asyncDependencies); + let results = (outs GPU_AsyncToken:$asyncToken); + + 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 @@ -143,4 +143,15 @@ "gpu.return"() : () -> () } ) {gpu.kernel, sym_name = "kernel_1", type = (f32, memref) -> (), workgroup_attributions = 1: i64} : () -> () } + + func @async() { + // CHECK-LABEL: func @async + // 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 + } }