[mlir][gpu] Allow gpu.launch_func to be async.

[csigg]

Update documentation.

[csigg]

Rebase.

[mehdi_amini]

The semantics of async or not is *very* different from most analysis point of view. It isn't clear to me that having a single op for modeling both instead of a new operation is the right tradeoff here: I'm afraid it is too error prone.

[sanjoy.google]

In D89324#2328592, @mehdi_amini wrote:

The semantics of async or not is *very* different from most analysis point of view. It isn't clear to me that having a single op for modeling both instead of a new operation is the right tradeoff here: I'm afraid it is too error prone.

I generally agree with this, but I'm also struggling to come up with any optimizations we would actually want to do on gpu.launch_func.

[herhut]

Can you remove the printing of the extra attribute?

mlir/test/Conversion/GPUCommon/lower-launch-func-to-gpu-runtime-calls.mlir

23

It is sad that this gets printed. I assume there currently is no way around this with assemblyFormat. Does this roundtrip when reparsed? I assume operand_segment_sizes just gets overwritten with the constructed attribute.

[csigg]

Switch two test instances of gpu.launch_func from standard form to custom form. Rebase.

[csigg]

mlir/test/Conversion/GPUCommon/lower-launch-func-to-gpu-runtime-calls.mlir

23

That was the standard format. Changed to custom format now.

[csigg]

Update vulkan-runner tests.

[csigg]

Splitting 'switch to custom gpu.launch_func format' into separate change.

[csigg]

Rebase.

[mehdi_amini]

I reverted because it seems a test was crashing here: https://buildkite.com/mlir/mlir-core/builds/8962#330b9ea4-53a6-4ec9-8d10-5b40b806aa41

(I hope I got the right patch)

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

let extraClassDeclaration = [{
     LogicalResult verifyBody();
   }];
 
   // let verifier = [{ return ::verifFuncOpy(*this); }];
   let printer = [{ printGPUFuncOp(p, *this); }];
   let parser = [{ return parseGPUFuncOp(parser, result); }];
 }
 
-def GPU_LaunchFuncOp : GPU_Op<"launch_func">,
-    Arguments<(ins SymbolRefAttr:$kernel,
+def GPU_LaunchFuncOp : GPU_Op<"launch_func",
+                              [GPU_AsyncOpInterface, AttrSizedOperandSegments]>,
+    Arguments<(ins Variadic<GPU_AsyncToken>:$asyncDependencies,
+               SymbolRefAttr:$kernel,
                Index:$gridSizeX, Index:$gridSizeY, Index:$gridSizeZ,
                Index:$blockSizeX, Index:$blockSizeY, Index:$blockSizeZ,
                Variadic<AnyType>:$operands)>,
-    Results<(outs)> {
+    Results<(outs Optional<GPU_AsyncToken>:$asyncToken)> {
   let summary = "Launches a function as a GPU kernel";
 
   let description = [{
     Launch a kernel function on the specified grid of thread blocks.
     `gpu.launch` operations are lowered to `gpu.launch_func` operations by
     outlining the kernel body into a function in a dedicated module, which
     reflects the separate compilation process. The kernel function is required
     to have the `gpu.kernel` attribute. The module containing the kernel
     function is required to be a gpu.module. And finally, the module containing
     the kernel module (which thus cannot be the top-level module) is required
     to have the `gpu.container_module` attribute. The `gpu.launch_func`
-    operation has a symbol attribute named `kernel` to identify the fully 
+    operation has a symbol attribute named `kernel` to identify the fully
     specified kernel function to launch (both the gpu.module and func).
 
-    The operation takes at least six operands, with the first three operands
-    being grid sizes along x,y,z dimensions and the following three being block
-    sizes along x,y,z dimensions. When a lower-dimensional kernel is required,
-    unused sizes must be explicitly set to `1`. The remaining operands are
-    passed as arguments to the kernel function.
+    The `gpu.launch_func` supports async dependencies: the kernel does not start
+    executing until the ops producing those async dependencies have completed.
+
+    By the default, the host implicitly blocks until kernel execution has
+    completed. If the `async` keyword is present, the host does not block but
+    instead a `!gpu.async.token` is returned. Other async GPU ops can take this
+    token as dependency.
+
+    The operation requires at least the grid and block sizes along the x,y,z
+    dimensions as arguments. When a lower-dimensional kernel is required,
+    unused sizes must be explicitly set to `1`.
+
+    The remaining operands are passed as arguments to the kernel function.
 
     Example:
 
     ```mlir
     module attributes {gpu.container_module} {
 
       // This module creates a separate compilation unit for the GPU compiler.
       gpu.module @kernels {
module attributes {gpu.container_module} {
           %gDimY = "gpu.grid_dim"() {dimension = "y"} : () -> (index)
           %gDimZ = "gpu.grid_dim"() {dimension = "z"} : () -> (index)
 
           "some_op"(%bx, %tx) : (index, index) -> ()
           %42 = load %arg1[%bx] : memref<?xf32, 1>
         }
       }
 
+      %t0 = gpu.wait async
       gpu.launch_func
-          @kernels::@kernel_1                          // Kernel function.
-          blocks in (%cst, %cst, %cst)                 // Grid size.
-          threads in (%cst, %cst, %cst)                // Block size.
-          args(%arg0 : f32, %arg1 : memref<?xf32, 1>)  // Kernel arguments.
+          async                           // (Optional) Don't block host, return token.
+          [%t0]                           // (Optional) Execute only after %t0 has completed.
+          @kernels::@kernel_1             // Kernel function.
+          blocks in (%cst, %cst, %cst)    // Grid size.
+          threads in (%cst, %cst, %cst)   // Block size.
+          args(%arg0 : f32,               // (Optional) Kernel arguments.
+               %arg1 : memref<?xf32, 1>)
     }
     ```
   }];
 
   let skipDefaultBuilders = 1;
 
   let builders = [
     OpBuilder<"GPUFuncOp kernelFunc, "
friend LogicalResult GPUDialect::verifyOperationAttribute(Operation *,
                                                               NamedAttribute);
 
     /// The name of the symbol reference attribute specifying the kernel to launch.
     static StringRef getKernelAttrName() { return "kernel"; }
   }];
 
   let verifier = [{ return ::verify(*this); }];
   let assemblyFormat = [{
+      custom<AsyncDependencies>(type($asyncToken), $asyncDependencies)
       $kernel
       `blocks` `in` ` ` `(`$gridSizeX`,` $gridSizeY`,` $gridSizeZ`)`
       `threads` `in` ` ` `(`$blockSizeX`,` $blockSizeY`,` $blockSizeZ`)`
       custom<LaunchFuncOperands>($operands, type($operands))
       attr-dict
   }];
 }
 

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

void LaunchFuncOp::build(OpBuilder &builder, OperationState &result,
   // Add grid and block sizes as op operands, followed by the data operands.
   result.addOperands({gridSize.x, gridSize.y, gridSize.z, blockSize.x,
                       blockSize.y, blockSize.z});
   result.addOperands(kernelOperands);
   auto kernelModule = kernelFunc.getParentOfType<GPUModuleOp>();
   auto kernelSymbol = builder.getSymbolRefAttr(
       kernelModule.getName(), {builder.getSymbolRefAttr(kernelFunc.getName())});
   result.addAttribute(getKernelAttrName(), kernelSymbol);
+  SmallVector<int32_t, 8> segmentSizes(8, 1);
+  segmentSizes.front() = 0; // Initially no async dependencies.
+  segmentSizes.back() = static_cast<int32_t>(kernelOperands.size());
+  result.addAttribute(getOperandSegmentSizeAttr(),
+                      builder.getI32VectorAttr(segmentSizes));
 }
 
 unsigned LaunchFuncOp::getNumKernelOperands() {
-  return getNumOperands() - kNumConfigOperands;
+  return getNumOperands() - asyncDependencies().size() - kNumConfigOperands;
 }
 
 StringRef LaunchFuncOp::getKernelModuleName() {
   return kernel().getRootReference();
 }
 
 StringRef LaunchFuncOp::getKernelName() { return kernel().getLeafReference(); }
 
 Value LaunchFuncOp::getKernelOperand(unsigned i) {
-  return getOperation()->getOperand(i + kNumConfigOperands);
+  return getOperand(asyncDependencies().size() + kNumConfigOperands + i);
 }
 
 KernelDim3 LaunchFuncOp::getGridSizeOperandValues() {
-  return KernelDim3{getOperand(0), getOperand(1), getOperand(2)};
+  auto operands = getOperands().drop_front(asyncDependencies().size());
+  return KernelDim3{operands[0], operands[1], operands[2]};
 }
 
 KernelDim3 LaunchFuncOp::getBlockSizeOperandValues() {
-  return KernelDim3{getOperand(3), getOperand(4), getOperand(5)};
+  auto operands = getOperands().drop_front(asyncDependencies().size());
+  return KernelDim3{operands[3], operands[4], operands[5]};
 }
 
 static LogicalResult verify(LaunchFuncOp op) {
   auto module = op.getParentOfType<ModuleOp>();
   if (!module)
     return op.emitOpError("expected to belong to a module");
 
   if (!module.getAttrOfType<UnitAttr>(GPUDialect::getContainerModuleAttrName()))

mlir/test/Conversion/GPUCommon/lower-launch-func-to-gpu-runtime-calls.mlir

llvm.func @kernel(%arg0: !llvm.i32, %arg1: !llvm.ptr<float>,
         %arg5: !llvm.i64) attributes {gpu.kernel} {
       llvm.return
     }
   }
 
   func @foo(%buffer: memref<?xf32>) {
     %c8 = constant 8 : index
     %c32 = constant 32 : i32
-    "gpu.launch_func"(%c8, %c8, %c8, %c8, %c8, %c8, %c32, %buffer) {
-      kernel = @kernel_module::@kernel
-    } : (index, index, index, index, index, index, i32, memref<?xf32>) -> ()
+    gpu.launch_func @kernel_module::@kernel
+        blocks in (%c8, %c8, %c8)
+        threads in (%c8, %c8, %c8)
+        args(%c32 : i32, %buffer : memref<?xf32>)
     return
   }
 
   // CHECK: [[C8:%.*]] = llvm.mlir.constant(8 : index) : !llvm.i64   
   // CHECK: [[ADDRESSOF:%.*]] = llvm.mlir.addressof @[[GLOBAL]]
   // CHECK: [[C0:%.*]] = llvm.mlir.constant(0 : index)
   // CHECK: [[BINARY:%.*]] = llvm.getelementptr [[ADDRESSOF]]{{\[}}[[C0]], [[C0]]]
   // CHECK-SAME: -> !llvm.ptr<i8>

mlir/test/Conversion/GPUToSPIRV/builtins.mlir

 // RUN: mlir-opt -split-input-file -convert-gpu-to-spirv %s -o - | FileCheck %s
 
 module attributes {gpu.container_module} {
   func @builtin() {
     %c0 = constant 1 : index
-    "gpu.launch_func"(%c0, %c0, %c0, %c0, %c0, %c0) {kernel = @kernels::@builtin_workgroup_id_x} : (index, index, index, index, index, index) -> ()
+    gpu.launch_func @kernels::@builtin_workgroup_id_x
+        blocks in (%c0, %c0, %c0) threads in (%c0, %c0, %c0)
     return
   }
 
   // CHECK-LABEL:  spv.module @{{.*}} Logical GLSL450
   // CHECK: spv.globalVariable [[WORKGROUPID:@.*]] built_in("WorkgroupId")
   gpu.module @kernels {
     gpu.func @builtin_workgroup_id_x() kernel
       attributes {spv.entry_point_abi = {local_size = dense<[16, 1, 1]>: vector<3xi32>}} {
       // CHECK: [[ADDRESS:%.*]] = spv._address_of [[WORKGROUPID]]
       // CHECK-NEXT: [[VEC:%.*]] = spv.Load "Input" [[ADDRESS]]
       // CHECK-NEXT: {{%.*}} = spv.CompositeExtract [[VEC]]{{\[}}0 : i32{{\]}}
       %0 = "gpu.block_id"() {dimension = "x"} : () -> index
       gpu.return
     }
   }
 }
 
 // -----
 
 module attributes {gpu.container_module} {
   func @builtin() {
     %c0 = constant 1 : index
-    "gpu.launch_func"(%c0, %c0, %c0, %c0, %c0, %c0) {kernel = @kernels::@builtin_workgroup_id_y} : (index, index, index, index, index, index) -> ()
+    gpu.launch_func @kernels::@builtin_workgroup_id_y
+        blocks in (%c0, %c0, %c0) threads in (%c0, %c0, %c0)
     return
   }
 
   // CHECK-LABEL:  spv.module @{{.*}} Logical GLSL450
   // CHECK: spv.globalVariable [[WORKGROUPID:@.*]] built_in("WorkgroupId")
   gpu.module @kernels {
     gpu.func @builtin_workgroup_id_y() kernel
       attributes {spv.entry_point_abi = {local_size = dense<[16, 1, 1]>: vector<3xi32>}} {
       // CHECK: [[ADDRESS:%.*]] = spv._address_of [[WORKGROUPID]]
       // CHECK-NEXT: [[VEC:%.*]] = spv.Load "Input" [[ADDRESS]]
       // CHECK-NEXT: {{%.*}} = spv.CompositeExtract [[VEC]]{{\[}}1 : i32{{\]}}
       %0 = "gpu.block_id"() {dimension = "y"} : () -> index
       gpu.return
     }
   }
 }
 
 // -----
 
 module attributes {gpu.container_module} {
   func @builtin() {
     %c0 = constant 1 : index
-    "gpu.launch_func"(%c0, %c0, %c0, %c0, %c0, %c0) {kernel = @kernels::@builtin_workgroup_id_z} : (index, index, index, index, index, index) -> ()
+    gpu.launch_func @kernels::@builtin_workgroup_id_z
+        blocks in (%c0, %c0, %c0) threads in (%c0, %c0, %c0)
     return
   }
 
   // CHECK-LABEL:  spv.module @{{.*}} Logical GLSL450
   // CHECK: spv.globalVariable [[WORKGROUPID:@.*]] built_in("WorkgroupId")
   gpu.module @kernels {
     gpu.func @builtin_workgroup_id_z() kernel
       attributes {spv.entry_point_abi = {local_size = dense<[16, 1, 1]>: vector<3xi32>}} {

mlir/test/Conversion/GPUToSPIRV/if.mlir

 // RUN: mlir-opt -convert-gpu-to-spirv %s -o - | FileCheck %s
 
 module attributes {
   gpu.container_module,
   spv.target_env = #spv.target_env<
     #spv.vce<v1.0, [Shader], [SPV_KHR_storage_buffer_storage_class]>, {}>
 } {
   func @main(%arg0 : memref<10xf32>, %arg1 : i1) {
     %c0 = constant 1 : index
-    "gpu.launch_func"(%c0, %c0, %c0, %c0, %c0, %c0, %arg0, %arg1) { kernel = @kernels::@kernel_simple_selection} : (index, index, index, index, index, index, memref<10xf32>, i1) -> ()
+    gpu.launch_func @kernels::@kernel_simple_selection
+        blocks in (%c0, %c0, %c0) threads in (%c0, %c0, %c0)
+        args(%arg0 : memref<10xf32>, %arg1 : i1)
     return
   }
 
   gpu.module @kernels {
     // CHECK-LABEL: @kernel_simple_selection
     gpu.func @kernel_simple_selection(%arg2 : memref<10xf32>, %arg3 : i1) kernel
     attributes {spv.entry_point_abi = {local_size = dense<[16, 1, 1]>: vector<3xi32>}} {
       %value = constant 0.0 : f32

mlir/test/Conversion/GPUToSPIRV/load-store.mlir

func @load_store(%arg0: memref<12x4xf32>, %arg1: memref<12x4xf32>, %arg2: memref<12x4xf32>) {
     %c12 = constant 12 : index
     %0 = subi %c12, %c0 : index
     %c1 = constant 1 : index
     %c0_0 = constant 0 : index
     %c4 = constant 4 : index
     %1 = subi %c4, %c0_0 : index
     %c1_1 = constant 1 : index
     %c1_2 = constant 1 : index
-    "gpu.launch_func"(%0, %c1_2, %c1_2, %1, %c1_2, %c1_2, %arg0, %arg1, %arg2, %c0, %c0_0, %c1, %c1_1) {kernel = @kernels::@load_store_kernel} : (index, index, index, index, index, index, memref<12x4xf32>, memref<12x4xf32>, memref<12x4xf32>, index, index, index, index) -> ()
+    gpu.launch_func @kernels::@load_store_kernel
+        blocks in (%0, %c1_2, %c1_2) threads in (%1, %c1_2, %c1_2)
+        args(%arg0 : memref<12x4xf32>, %arg1 : memref<12x4xf32>, %arg2 : memref<12x4xf32>, 
+             %c0 : index, %c0_0 : index, %c1 : index, %c1_1 : index)
     return
   }
 
   // CHECK-LABEL: spv.module @{{.*}} Logical GLSL450
   gpu.module @kernels {
     // CHECK-DAG: spv.globalVariable @[[NUMWORKGROUPSVAR:.*]] built_in("NumWorkgroups") : !spv.ptr<vector<3xi32>, Input>
     // CHECK-DAG: spv.globalVariable @[[$LOCALINVOCATIONIDVAR:.*]] built_in("LocalInvocationId") : !spv.ptr<vector<3xi32>, Input>
     // CHECK-DAG: spv.globalVariable @[[$WORKGROUPIDVAR:.*]] built_in("WorkgroupId") : !spv.ptr<vector<3xi32>, Input>

mlir/test/Conversion/GPUToSPIRV/loop.mlir

 // RUN: mlir-opt -convert-gpu-to-spirv %s -o - | FileCheck %s
 
 module attributes {
   gpu.container_module,
   spv.target_env = #spv.target_env<
     #spv.vce<v1.0, [Shader], [SPV_KHR_storage_buffer_storage_class]>, {}>
 } {
   func @loop(%arg0 : memref<10xf32>, %arg1 : memref<10xf32>) {
     %c0 = constant 1 : index
-    "gpu.launch_func"(%c0, %c0, %c0, %c0, %c0, %c0, %arg0, %arg1) { kernel = @kernels::@loop_kernel} : (index, index, index, index, index, index, memref<10xf32>, memref<10xf32>) -> ()
+    gpu.launch_func @kernels::@loop_kernel
+        blocks in (%c0, %c0, %c0) threads in (%c0, %c0, %c0)
+        args(%arg0 : memref<10xf32>, %arg1 : memref<10xf32>)
     return
   }
 
   gpu.module @kernels {
     gpu.func @loop_kernel(%arg2 : memref<10xf32>, %arg3 : memref<10xf32>) kernel
     attributes {spv.entry_point_abi = {local_size = dense<[16, 1, 1]>: vector<3xi32>}} {
       // CHECK: %[[LB:.*]] = spv.constant 4 : i32
       %lb = constant 4 : index

mlir/test/Conversion/GPUToSPIRV/module-structure-opencl.mlir

gpu.func @basic_module_structure(%arg0 : f32, %arg1 : memref<12xf32, 11>) kernel
       gpu.return
     }
   }
 
   func @main() {
     %0 = "op"() : () -> (f32)
     %1 = "op"() : () -> (memref<12xf32, 11>)
     %cst = constant 1 : index
-    "gpu.launch_func"(%cst, %cst, %cst, %cst, %cst, %cst, %0, %1) { kernel = @kernels::@basic_module_structure }
-        : (index, index, index, index, index, index, f32, memref<12xf32, 11>) -> ()
+    gpu.launch_func @kernels::@basic_module_structure
+        blocks in (%cst, %cst, %cst) threads in (%cst, %cst, %cst)
+        args(%0 : f32, %1 : memref<12xf32, 11>)
     return
   }
 }

mlir/test/Conversion/GPUToSPIRV/simple.mlir

gpu.func @basic_module_structure(%arg0 : f32, %arg1 : memref<12xf32>) kernel
       gpu.return
     }
   }
 
   func @main() {
     %0 = "op"() : () -> (f32)
     %1 = "op"() : () -> (memref<12xf32>)
     %cst = constant 1 : index
-    "gpu.launch_func"(%cst, %cst, %cst, %cst, %cst, %cst, %0, %1) { kernel = @kernels::@basic_module_structure }
-        : (index, index, index, index, index, index, f32, memref<12xf32>) -> ()
+    gpu.launch_func @kernels::@basic_module_structure
+        blocks in (%cst, %cst, %cst) threads in (%cst, %cst, %cst)
+        args(%0 : f32, %1 : memref<12xf32>)
     return
   }
 }
 
 // -----
 
 module attributes {gpu.container_module} {
   gpu.module @kernels {
gpu.func @missing_entry_point_abi(%arg0 : f32, %arg1 : memref<12xf32>) kernel {
       gpu.return
     }
   }
 
   func @main() {
     %0 = "op"() : () -> (f32)
     %1 = "op"() : () -> (memref<12xf32>)
     %cst = constant 1 : index
-    "gpu.launch_func"(%cst, %cst, %cst, %cst, %cst, %cst, %0, %1) { kernel = @kernels::@missing_entry_point_abi }
-        : (index, index, index, index, index, index, f32, memref<12xf32>) -> ()
+    gpu.launch_func @kernels::@missing_entry_point_abi
+        blocks in (%cst, %cst, %cst) threads in (%cst, %cst, %cst)
+        args(%0 : f32, %1 : memref<12xf32>)
     return
   }
 }
 
 // -----
 
 module attributes {gpu.container_module} {
   gpu.module @kernels {

mlir/test/Conversion/GPUToVulkan/lower-gpu-launch-vulkan-launch.mlir

module attributes {gpu.container_module} {
   gpu.module @kernels {
     gpu.func @kernel(%arg0: memref<12xf32>) kernel {
       gpu.return
     }
   }
   func @foo() {
     %0 = alloc() : memref<12xf32>
     %c1 = constant 1 : index
-    "gpu.launch_func"(%c1, %c1, %c1, %c1, %c1, %c1, %0) {kernel = @kernels::@kernel} : (index, index, index, index, index, index, memref<12xf32>) -> ()
+    gpu.launch_func @kernels::@kernel
+        blocks in(%c1, %c1, %c1)
+        threads in(%c1, %c1, %c1)
+        args(%0 : memref<12xf32>) -> ()
     return
   }
 }

mlir/test/Dialect/GPU/invalid.mlir

func @launch_requires_gpu_return(%sz : index) {
   return
 }
 
 // -----
 
 func @launch_func_too_few_operands(%sz : index) {
   // expected-error@+1 {{expected 6 or more operands}}
   "gpu.launch_func"(%sz, %sz, %sz, %sz, %sz)
+      {operand_segment_sizes = dense<[0, 1, 1, 1, 1, 1, 0, 0]> : vector<8xi32>}
       : (index, index, index, index, index) -> ()
   return
 }
 
 // -----
 
 func @launch_func_missing_parent_module_attribute(%sz : index) {
   // expected-error@+1 {{expected the closest surrounding module to have the 'gpu.container_module' attribute}}
   gpu.launch_func @foo::@bar blocks in (%sz, %sz, %sz) threads in (%sz, %sz, %sz)
   return
 }
 
 // -----
 
 module attributes {gpu.container_module} {
   func @launch_func_missing_callee_attribute(%sz : index) {
     // expected-error@+1 {{'gpu.launch_func' op requires attribute 'kernel'}}
     "gpu.launch_func"(%sz, %sz, %sz, %sz, %sz, %sz)
+        {operand_segment_sizes = dense<[0, 1, 1, 1, 1, 1, 1, 0]> : vector<8xi32>}
         : (index, index, index, index, index, index) -> ()
     return
   }
 }
 
 // -----
 
 module attributes {gpu.container_module} {

mlir/test/Dialect/GPU/ops.mlir

gpu.module @kernels {
     }
   }
 
   func @foo() {
     %0 = "op"() : () -> (f32)
     %1 = "op"() : () -> (memref<?xf32, 1>)
     // CHECK: %{{.*}} = constant 8
     %cst = constant 8 : index
+    %t0 = gpu.wait async
 
     // CHECK: gpu.launch_func @kernels::@kernel_1 blocks in (%{{.*}}, %{{.*}}, %{{.*}}) threads in (%{{.*}}, %{{.*}}, %{{.*}}) args(%{{.*}} : f32, %{{.*}} : memref<?xf32, 1>)
     gpu.launch_func @kernels::@kernel_1 blocks in (%cst, %cst, %cst) threads in (%cst, %cst, %cst) args(%0 : f32, %1 : memref<?xf32, 1>)
 
     // CHECK: gpu.launch_func @kernels::@kernel_2 blocks in (%{{.*}}, %{{.*}}, %{{.*}}) threads in (%{{.*}}, %{{.*}}, %{{.*}})
     gpu.launch_func @kernels::@kernel_2 blocks in (%cst, %cst, %cst) threads in (%cst, %cst, %cst)
 
+    // CHECK: %{{.*}} = gpu.launch_func async [%{{.*}}] @kernels::@kernel_2 blocks in (%{{.*}}, %{{.*}}, %{{.*}}) threads in (%{{.*}}, %{{.*}}, %{{.*}})
+    %t1 = gpu.launch_func async [%t0] @kernels::@kernel_2  blocks in (%cst, %cst, %cst) threads in (%cst, %cst, %cst)
+
     return
   }
 
   gpu.module @gpu_funcs {
     // CHECK-LABEL: gpu.func @kernel_1({{.*}}: f32)
     // CHECK:       workgroup
     // CHECK:       private
     // CHECK:       attributes

mlir/test/mlir-vulkan-runner/addf.mlir

func @main() {
     %arg4 = memref_cast %arg1 : memref<8xf32> to memref<?xf32>
     %arg5 = memref_cast %arg2 : memref<8xf32> to memref<?xf32>
     call @fillResource1DFloat(%arg3, %value1) : (memref<?xf32>, f32) -> ()
     call @fillResource1DFloat(%arg4, %value2) : (memref<?xf32>, f32) -> ()
     call @fillResource1DFloat(%arg5, %value0) : (memref<?xf32>, f32) -> ()
 
     %cst1 = constant 1 : index
     %cst8 = constant 8 : index
-    "gpu.launch_func"(%cst8, %cst1, %cst1, %cst1, %cst1, %cst1, %arg0, %arg1, %arg2) { kernel = @kernels::@kernel_add }
-        : (index, index, index, index, index, index, memref<8xf32>, memref<8xf32>, memref<8xf32>) -> ()
+    gpu.launch_func @kernels::@kernel_add
+        blocks in (%cst8, %cst1, %cst1) threads in (%cst1, %cst1, %cst1)
+        args(%arg0 : memref<8xf32>, %arg1 : memref<8xf32>, %arg2 : memref<8xf32>)
     %arg6 = memref_cast %arg5 : memref<?xf32> to memref<*xf32>
     call @print_memref_f32(%arg6) : (memref<*xf32>) -> ()
     return
   }
   func @fillResource1DFloat(%0 : memref<?xf32>, %1 : f32)
   func @print_memref_f32(%ptr : memref<*xf32>)
 }
 

mlir/test/mlir-vulkan-runner/addi.mlir

func @main() {
     %arg4 = memref_cast %arg1 : memref<8x8xi32> to memref<?x?xi32>
     %arg5 = memref_cast %arg2 : memref<8x8x8xi32> to memref<?x?x?xi32>
     call @fillResource1DInt(%arg3, %value1) : (memref<?xi32>, i32) -> ()
     call @fillResource2DInt(%arg4, %value2) : (memref<?x?xi32>, i32) -> ()
     call @fillResource3DInt(%arg5, %value0) : (memref<?x?x?xi32>, i32) -> ()
 
     %cst1 = constant 1 : index
     %cst8 = constant 8 : index
-    "gpu.launch_func"(%cst8, %cst8, %cst8, %cst1, %cst1, %cst1, %arg0, %arg1, %arg2) { kernel = @kernels::@kernel_addi }
-        : (index, index, index, index, index, index, memref<8xi32>, memref<8x8xi32>, memref<8x8x8xi32>) -> ()
+    gpu.launch_func @kernels::@kernel_addi
+        blocks in (%cst8, %cst8, %cst8) threads in (%cst1, %cst1, %cst1)
+        args(%arg0 : memref<8xi32>, %arg1 : memref<8x8xi32>, %arg2 : memref<8x8x8xi32>)
     %arg6 = memref_cast %arg5 : memref<?x?x?xi32> to memref<*xi32>
     call @print_memref_i32(%arg6) : (memref<*xi32>) -> ()
     return
   }
   func @fillResource1DInt(%0 : memref<?xi32>, %1 : i32)
   func @fillResource2DInt(%0 : memref<?x?xi32>, %1 : i32)
   func @fillResource3DInt(%0 : memref<?x?x?xi32>, %1 : i32)
   func @print_memref_i32(%ptr : memref<*xi32>)
 }
 

mlir/test/mlir-vulkan-runner/addi8.mlir

func @main() {
     %arg4 = memref_cast %arg1 : memref<8x8xi8> to memref<?x?xi8>
     %arg5 = memref_cast %arg2 : memref<8x8x8xi32> to memref<?x?x?xi32>
     call @fillResource1DInt8(%arg3, %value1) : (memref<?xi8>, i8) -> ()
     call @fillResource2DInt8(%arg4, %value2) : (memref<?x?xi8>, i8) -> ()
     call @fillResource3DInt(%arg5, %value0) : (memref<?x?x?xi32>, i32) -> ()
 
     %cst1 = constant 1 : index
     %cst8 = constant 8 : index
-    "gpu.launch_func"(%cst8, %cst8, %cst8, %cst1, %cst1, %cst1, %arg0, %arg1, %arg2) { kernel = @kernels::@kernel_addi }
-        : (index, index, index, index, index, index, memref<8xi8>, memref<8x8xi8>, memref<8x8x8xi32>) -> ()
+    gpu.launch_func @kernels::@kernel_addi
+        blocks in (%cst8, %cst8, %cst8) threads in (%cst1, %cst1, %cst1)
+        args(%arg0 : memref<8xi8>, %arg1 : memref<8x8xi8>, %arg2 : memref<8x8x8xi32>)
     %arg6 = memref_cast %arg5 : memref<?x?x?xi32> to memref<*xi32>
     call @print_memref_i32(%arg6) : (memref<*xi32>) -> ()
     return
   }
   func @fillResource1DInt8(%0 : memref<?xi8>, %1 : i8)
   func @fillResource2DInt8(%0 : memref<?x?xi8>, %1 : i8)
   func @fillResource3DInt(%0 : memref<?x?x?xi32>, %1 : i32)
   func @print_memref_i32(%ptr : memref<*xi32>)
 }

mlir/test/mlir-vulkan-runner/mulf.mlir

func @main() {
     %arg4 = memref_cast %arg1 : memref<4x4xf32> to memref<?x?xf32>
     %arg5 = memref_cast %arg2 : memref<4x4xf32> to memref<?x?xf32>
     call @fillResource2DFloat(%arg3, %value1) : (memref<?x?xf32>, f32) -> ()
     call @fillResource2DFloat(%arg4, %value2) : (memref<?x?xf32>, f32) -> ()
     call @fillResource2DFloat(%arg5, %value0) : (memref<?x?xf32>, f32) -> ()
 
     %cst1 = constant 1 : index
     %cst4 = constant 4 : index
-    "gpu.launch_func"(%cst4, %cst4, %cst1, %cst1, %cst1, %cst1, %arg0, %arg1, %arg2) { kernel = @kernels::@kernel_mul }
-        : (index, index, index, index, index, index, memref<4x4xf32>, memref<4x4xf32>, memref<4x4xf32>) -> ()
+    gpu.launch_func @kernels::@kernel_mul
+        blocks in (%cst4, %cst4, %cst1) threads in(%cst1, %cst1, %cst1)
+        args(%arg0 : memref<4x4xf32>, %arg1 : memref<4x4xf32>, %arg2 : memref<4x4xf32>)
     %arg6 = memref_cast %arg5 : memref<?x?xf32> to memref<*xf32>
     call @print_memref_f32(%arg6) : (memref<*xf32>) -> ()
     return
   }
   func @fillResource2DFloat(%0 : memref<?x?xf32>, %1 : f32)
   func @print_memref_f32(%ptr : memref<*xf32>)
 }
 

mlir/test/mlir-vulkan-runner/subf.mlir

func @main() {
     %arg5 = memref_cast %arg2 : memref<8x4x4xf32> to memref<?x?x?xf32>
     call @fillResource3DFloat(%arg3, %value1) : (memref<?x?x?xf32>, f32) -> ()
     call @fillResource2DFloat(%arg4, %value2) : (memref<?x?xf32>, f32) -> ()
     call @fillResource3DFloat(%arg5, %value0) : (memref<?x?x?xf32>, f32) -> ()
 
     %cst1 = constant 1 : index
     %cst4 = constant 4 : index
     %cst8 = constant 8 : index
-    "gpu.launch_func"(%cst8, %cst4, %cst4, %cst1, %cst1, %cst1, %arg0, %arg1, %arg2) { kernel = @kernels::@kernel_sub }
-        : (index, index, index, index, index, index, memref<8x4x4xf32>, memref<4x4xf32>, memref<8x4x4xf32>) -> ()
+    gpu.launch_func @kernels::@kernel_sub
+        blocks in (%cst8, %cst4, %cst4) threads in (%cst1, %cst1, %cst1)
+        args(%arg0 : memref<8x4x4xf32>, %arg1 : memref<4x4xf32>, %arg2 : memref<8x4x4xf32>)
     %arg6 = memref_cast %arg5 : memref<?x?x?xf32> to memref<*xf32>
     call @print_memref_f32(%arg6) : (memref<*xf32>) -> ()
     return
   }
   func @fillResource2DFloat(%0 : memref<?x?xf32>, %1 : f32)
   func @fillResource3DFloat(%0 : memref<?x?x?xf32>, %1 : f32)
   func @print_memref_f32(%ptr : memref<*xf32>)
 }

mlir/test/mlir-vulkan-runner/time.mlir

func @main() {
     %arg4 = memref_cast %arg1 : memref<16384xf32> to memref<?xf32>
     %arg5 = memref_cast %arg2 : memref<16384xf32> to memref<?xf32>
     call @fillResource1DFloat(%arg3, %value1) : (memref<?xf32>, f32) -> ()
     call @fillResource1DFloat(%arg4, %value2) : (memref<?xf32>, f32) -> ()
     call @fillResource1DFloat(%arg5, %value0) : (memref<?xf32>, f32) -> ()
 
     %cst1 = constant 1 : index
     %cst128 = constant 128 : index
-    "gpu.launch_func"(%cst128, %cst1, %cst1, %cst128, %cst1, %cst1, %arg0, %arg1, %arg2) { kernel = @kernels::@kernel_add }
-        : (index, index, index, index, index, index, memref<16384xf32>, memref<16384xf32>, memref<16384xf32>) -> ()
+    gpu.launch_func @kernels::@kernel_add
+        blocks in (%cst128, %cst1, %cst1) threads in (%cst128, %cst1, %cst1)
+        args(%arg0 : memref<16384xf32>, %arg1 : memref<16384xf32>, %arg2 : memref<16384xf32>)
     %arg6 = memref_cast %arg5 : memref<?xf32> to memref<*xf32>
     return
   }
   func @fillResource1DFloat(%0 : memref<?xf32>, %1 : f32)
   func @print_memref_f32(%ptr : memref<*xf32>)
 }