[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

24

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

24

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>
         }
       }
 
-      gpu.launch_func
-          @kernels::@kernel_1                      // Kernel module and function.
-          grid in (%cst, %cst, %cst)               // Grid sizes.
-          block in (%cst, %cst, %cst)              // Block sizes.
-          (%arg0 : f32, %arg1 : memref<?xf32, 1>)  // Arguments passed to the kernel.
+      %t0 = gpu.wait async
+      %t1 = gpu.launch_func
+          async                        // (Optional) Don't block host, return token.
+          [%t0]                        // (Optional) Execute only after %t0 has completed.
+          @kernels::@kernel_1          // Kernel module and function.
+          grid in (%cst, %cst, %cst)   // Grid sizes.
+          block in (%cst, %cst, %cst)  // Block sizes.
+          (%arg0 : f32,                // Arguments passed to the kernel.
+           %arg1 : memref<?xf32, 1>)
     }
     ```
   }];
 
   let skipDefaultBuilders = 1;
 
   let builders = [
     OpBuilder<"GPUFuncOp kernelFunc, "
let extraClassDeclaration = [{
 
     /// The name of the symbol reference attribute specifying the kernel to launch.
     static StringRef getKernelAttrName() { return "kernel"; }
   }];
 
   let verifier = [{ return ::verify(*this); }];
   let assemblyFormat = [{
       $kernel
+      custom<AsyncDependencies>(type($asyncToken), $asyncDependencies)
       `grid` `in` custom<Space>() `(`$gridSizeX`,` $gridSizeY`,` $gridSizeZ`)`
       `block` `in` custom<Space>() `(`$blockSizeX`,` $blockSizeY`,` $blockSizeZ`)`
       custom<LaunchFuncOperands>($operands, type($operands))
       attr-dict
   }];
 }
 
 def GPU_LaunchOp : GPU_Op<"launch">,

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>,
       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) {
+      operand_segment_sizes = dense<[0, 1, 1, 1, 1, 1, 1, 2]> : vector<8xi32>,
       kernel = @kernel_module::@kernel
     } : (index, index, index, index, index, index, i32, 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)

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 grid in (%sz, %sz, %sz) block 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 grid in (%{{.*}}, %{{.*}}, %{{.*}}) block in (%{{.*}}, %{{.*}}, %{{.*}}) (%{{.*}} : f32, %{{.*}} : memref<?xf32, 1>)
     gpu.launch_func @kernels::@kernel_1 grid in (%cst, %cst, %cst) block in (%cst, %cst, %cst) (%0 : f32, %1 : memref<?xf32, 1>)
 
     // CHECK: gpu.launch_func @kernels::@kernel_2 grid in (%{{.*}}, %{{.*}}, %{{.*}}) block in (%{{.*}}, %{{.*}}, %{{.*}}) (%{{.*}} : f32, %{{.*}} : memref<?xf32, 1>)
     gpu.launch_func @kernels::@kernel_2 grid in (%cst, %cst, %cst) block in (%cst, %cst, %cst) (%0 : f32, %1 : memref<?xf32, 1>)
 
+    // CHECK: %{{.*}} = gpu.launch_func @kernels::@kernel_1 async [%{{.*}}] grid in (%{{.*}}, %{{.*}}, %{{.*}}) block in (%{{.*}}, %{{.*}}, %{{.*}}) (%{{.*}} : f32, %{{.*}} : memref<?xf32, 1>)
+    %t1 = gpu.launch_func @kernels::@kernel_1  async [%t0] grid in (%cst, %cst, %cst) block in (%cst, %cst, %cst) (%0 : f32, %1 : memref<?xf32, 1>)
+
     return
   }
 
   gpu.module @gpu_funcs {
     // CHECK-LABEL: gpu.func @kernel_1({{.*}}: f32)
     // CHECK:       workgroup
     // CHECK:       private
     // CHECK:       attributes