Add `gpu.create_token` op.

[mehdi_amini]

mlir/include/mlir/Dialect/GPU/GPUBase.td

60 ↗

(On Diff #296694)

Is this interface more generic than the GPU dialect? Can you include GPU in the name to avoid confusion with the async dialect.

62 ↗

(On Diff #296694)

Please extend the documentation.

83 ↗

(On Diff #296694)

if constexpr is a c++17 feature, LLVM is using C++14.

92 ↗

(On Diff #296694)

Nit: this is a non-trivial block of C++, can this go inside an implementation file?

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

759

Is ForkOp intended here?

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

327 ↗

(On Diff #296694)

This hunk seems unrelated to this revision, please split and commit ahead.

[mehdi_amini]

Actually seems like you generated a mixed diff that includes changes from https://reviews.llvm.org/D88972 ?

[csigg]

In D88972#2317357, @mehdi_amini wrote:

Actually seems like you generated a mixed diff that includes changes from https://reviews.llvm.org/D88972 ?

This is a follow-up change. Is there a better way than creating a new revision and adding the parent revision in phabricator?

[mehdi_amini]

Phab works with "diff", somehow you create this revision as a diff against master? In general I do an interactive rebase, edit individually each commit and arc diff HEAD~

[csigg]

Remove base revision, address comments.

[csigg]

mlir/include/mlir/Dialect/GPU/GPUBase.td

60 ↗

(On Diff #296694)

Sorry, this was meant to be in the ::mlir::gpu namespace. I left the name unchanged.

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

759

Oops. Thanks for catching that.

[csigg]

In D88972#2317773, @mehdi_amini wrote:

Phab works with "diff", somehow you create this revision as a diff against master? In general I do an interactive rebase, edit individually each commit and arc diff HEAD~

Got it, thanks.

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

let description = [{
     launched afterwards. Writes from the device are guaranteed to be visible on
     the host after synchronizing with the device kernel completion.
   }];
 
   let assemblyFormat = "$value attr-dict `:` type($value)";
   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<GPU_AsyncToken>:$asyncDependencies);
+  let results = (outs GPU_AsyncToken:$asyncToken);
+
+  let assemblyFormat = "(`[` $asyncDependencies^ `]`)? attr-dict";
+}
+
 #endif // GPU_OPS

mlir/test/Dialect/GPU/ops.mlir

module attributes {gpu.container_module} {
 
   gpu.module @explicit_attributions {
     // CHECK-LABEL: gpu.func @kernel_1({{.*}}: f32, {{.*}}: memref<?xf32>) workgroup({{.*}}: memref<5xf32, 3>) private({{.*}}: memref<5xf32, 5>)
     "gpu.func"() ( {
     ^bb0(%arg0: f32, %arg1: memref<?xf32>, %arg2: memref<5xf32, 3>, %arg3: memref<5xf32, 5>):
       "gpu.return"() : () -> ()
     } ) {gpu.kernel, sym_name = "kernel_1", type = (f32, memref<?xf32>) -> (), 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
+  }
 }