[mlir][gpu] Add customer printer/parser for gpu.launch_func.

[csigg]

Add space after 'in'.

[herhut]

I like this updated syntax. @ftynse WDYT?

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

406

The custom<Space> is ugly but I see its motivation. Maybe file a bug to add support for explicit whitespace to the assembly format?

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

505

I'd prefer to expose the parsing of just the argument list instead of parsing a result type and then failing.

[csigg]

Add test for custom formatter without arguments.

[csigg]

Fix tests, expose just parsing function arguments.

[csigg]

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

406

Alternatively I could custom-parse/print the grid in or even with the operands. That would look less ugly, but it's a little less expressive.

[csigg]

Update gpu.launch_func documentation.

[csigg]

Use space literal in assemblyFormat.

[csigg]

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

406

Switched to space literal.

[ftynse]

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

21

Are we sure this isn't necessary anymore?

358

Syntax nit: I'd consider having having an prefix here too, something like args.

[csigg]

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

21

I don't think so, because we lower from launch_func to runtime calls in one go.

358

Sure, I can do that.

[csigg]

Add args keyword before argument list.
Change grid/block to blocks/threads.

[herhut]

Thanks. Just some nits.

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

358

Can you update the example, too?

mlir/include/mlir/IR/FunctionImplementation.h

61

You could add an option here whether it may have attributes, like allowAttributes or something and then parsing would just fail with a parse error.

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

487

Nit: This could be in the assemblyFormat

[csigg]

Apply reviewer comments.

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

487

The args(...) is optional, and I don't think custom<> inside an optional group is supported. I left it as is.

[csigg]

Rebase.

[herhut]

Thanks!

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

 #ifndef GPU_OPS
 #define GPU_OPS
 
 include "mlir/Dialect/GPU/GPUBase.td"
 include "mlir/Dialect/LLVMIR/LLVMOpBase.td"
 include "mlir/IR/SymbolInterfaces.td"
 include "mlir/Interfaces/SideEffectInterfaces.td"
 
-// Type constraint accepting standard integers, indices.
-def IntOrIndex : TypeConstraint<
-  Or<[AnySignlessInteger.predicate, Index.predicate]>, "integer or index">;
-
 //===----------------------------------------------------------------------===//
 // GPU Dialect operations.
 //===----------------------------------------------------------------------===//
 
 class GPU_Op<string mnemonic, list<OpTrait> traits = []> :
     Op<GPU_Dialect, mnemonic, traits>;
 
 class GPU_IndexOp<string mnemonic, list<OpTrait> traits = []> :
def GPU_GPUFuncOp : GPU_Op<"func", [HasParent<"GPUModuleOp">,
   }];
 
   // 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 IntOrIndex:$gridSizeX, IntOrIndex:$gridSizeY,
-               IntOrIndex:$gridSizeZ, IntOrIndex:$blockSizeX,
-               IntOrIndex:$blockSizeY, IntOrIndex:$blockSizeZ,
+    Arguments<(ins SymbolRefAttr:$kernel,
+               Index:$gridSizeX, Index:$gridSizeY, Index:$gridSizeZ,
+               Index:$blockSizeX, Index:$blockSizeY, Index:$blockSizeZ,
                Variadic<AnyType>:$operands)>,
     Results<(outs)> {
   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 specified
-    kernel function to launch (both the gpu.module and func).
+    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.
 
-    A custom syntax for this operation is currently not available.
-
     Example:
 
     ```mlir
     module attributes {gpu.container_module} {
 
       // This module creates a separate compilation unit for the GPU compiler.
       gpu.module @kernels {
         func @kernel_1(%arg0 : f32, %arg1 : memref<?xf32, 1>)
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"(%cst, %cst, %cst,  // Grid sizes.
-                        %cst, %cst, %cst,  // Block sizes.
-                        %arg0, %arg1)      // Arguments passed to the kernel.
-            { kernel_module = @kernels,    // Module containing the kernel.
-              kernel = "kernel_1" }        // Kernel function.
-            : (index, index, index, index, index, index, f32, 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.
     }
     ```
   }];
 
   let skipDefaultBuilders = 1;
 
   let builders = [
     OpBuilder<"GPUFuncOp kernelFunc, "
-              "Value gridSizeX, Value gridSizeY, Value gridSizeZ, "
-              "Value blockSizeX, Value blockSizeY, Value blockSizeZ, "
-              "ValueRange kernelOperands">,
-    OpBuilder<"GPUFuncOp kernelFunc, "
               "KernelDim3 gridSize, KernelDim3 blockSize, "
               "ValueRange kernelOperands">
   ];
 
   let extraClassDeclaration = [{
-    /// The kernel function specified by the operation's `kernel` attribute.
-    SymbolRefAttr kernel();
-
     /// The number of operands passed to the kernel function.
     unsigned getNumKernelOperands();
 
     /// The name of the kernel's containing module.
     StringRef getKernelModuleName();
 
     /// The name of the kernel.
     StringRef getKernelName();
let extraClassDeclaration = [{
     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 = [{
+      $kernel
+      `grid` `in` ` ` `(`$gridSizeX`,` $gridSizeY`,` $gridSizeZ`)`
+      `block` `in` ` ` `(`$blockSizeX`,` $blockSizeY`,` $blockSizeZ`)`
+      custom<LaunchFuncOperands>($operands, type($operands))
+      attr-dict
+  }];
 }
 
 def GPU_LaunchOp : GPU_Op<"launch">,
     Arguments<(ins Index:$gridSizeX, Index:$gridSizeY, Index:$gridSizeZ,
                Index:$blockSizeX, Index:$blockSizeY, Index:$blockSizeZ)>,
     Results<(outs)> {
   let summary = "GPU kernel launch operation";
 

mlir/include/mlir/IR/FunctionImplementation.h

 /// Callback type for `parseFunctionLikeOp`, the callback should produce the
 /// type that will be associated with a function-like operation from lists of
 /// function arguments and results, VariadicFlag indicates whether the function
 /// should have variadic arguments; in case of error, it may populate the last
 /// argument with a message.
 using FuncTypeBuilder = function_ref<Type(
     Builder &, ArrayRef<Type>, ArrayRef<Type>, VariadicFlag, std::string &)>;
 
+/// Parses function arguments using `parser`. The `allowVariadic` argument
+/// indicates whether functions with variadic arguments are supported. The
+/// trailing arguments are populated by this function with names, types and
+/// attributes of the arguments.
+ParseResult
+parseFunctionArgumentList(OpAsmParser &parser, bool allowVariadic,
+                          SmallVectorImpl<OpAsmParser::OperandType> &argNames,
+                          SmallVectorImpl<Type> &argTypes,
+                          SmallVectorImpl<NamedAttrList> &argAttrs,
+                          bool &isVariadic);
+
 /// Parses a function signature using `parser`. The `allowVariadic` argument
 /// indicates whether functions with variadic arguments are supported. The
 /// trailing arguments are populated by this function with names, types and
 /// attributes of the arguments and those of the results.
 ParseResult
 parseFunctionSignature(OpAsmParser &parser, bool allowVariadic,
                        SmallVectorImpl<OpAsmParser::OperandType> &argNames,
                        SmallVectorImpl<Type> &argTypes,

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

return failure(parser.parseRegion(*body, regionArgs, dataTypes) ||
                  parser.parseOptionalAttrDict(result.attributes));
 }
 
 //===----------------------------------------------------------------------===//
 // LaunchFuncOp
 //===----------------------------------------------------------------------===//
 
 void LaunchFuncOp::build(OpBuilder &builder, OperationState &result,
-                         GPUFuncOp kernelFunc, Value gridSizeX, Value gridSizeY,
-                         Value gridSizeZ, Value blockSizeX, Value blockSizeY,
-                         Value blockSizeZ, ValueRange kernelOperands) {
+                         GPUFuncOp kernelFunc, KernelDim3 gridSize,
+                         KernelDim3 blockSize, ValueRange kernelOperands) {
   // Add grid and block sizes as op operands, followed by the data operands.
-  result.addOperands(
-      {gridSizeX, gridSizeY, gridSizeZ, blockSizeX, blockSizeY, blockSizeZ});
+  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);
 }
 
-void LaunchFuncOp::build(OpBuilder &builder, OperationState &result,
-                         GPUFuncOp kernelFunc, KernelDim3 gridSize,
-                         KernelDim3 blockSize, ValueRange kernelOperands) {
-  build(builder, result, kernelFunc, gridSize.x, gridSize.y, gridSize.z,
-        blockSize.x, blockSize.y, blockSize.z, kernelOperands);
-}
-
-SymbolRefAttr LaunchFuncOp::kernel() {
-  return getAttrOfType<SymbolRefAttr>(getKernelAttrName());
-}
-
 unsigned LaunchFuncOp::getNumKernelOperands() {
   return getNumOperands() - kNumConfigOperands;
 }
 
 StringRef LaunchFuncOp::getKernelModuleName() {
   return kernel().getRootReference();
 }
 
static LogicalResult verify(LaunchFuncOp op) {
   auto kernelAttr = op.getAttrOfType<SymbolRefAttr>(op.getKernelAttrName());
   if (!kernelAttr)
     return op.emitOpError("symbol reference attribute '" +
                           op.getKernelAttrName() + "' must be specified");
 
   return success();
 }
 
+static ParseResult
+parseLaunchFuncOperands(OpAsmParser &parser,
+                        SmallVectorImpl<OpAsmParser::OperandType> &argNames,
+                        SmallVectorImpl<Type> &argTypes) {
+  SmallVector<NamedAttrList, 4> argAttrs;
+  bool isVariadic = false;
+  if (impl::parseFunctionArgumentList(parser, /*allowVariadic=*/false, argNames,
+                                      argTypes, argAttrs, isVariadic))
+    return failure();
+  for (auto pair : llvm::enumerate(argAttrs)) {
+    if (pair.value().empty())
+      continue;
+    return parser.emitError(argNames[pair.index()].location,
+                            "does not support kernel argument attributes");
+  }
+  return success();
+}
+
+static void printLaunchFuncOperands(OpAsmPrinter &printer,
+                                    OperandRange operands, TypeRange types) {
+  printer << "(";
+  llvm::interleaveComma(llvm::zip(operands, types), printer,
+                        [&](const auto &pair) {
+                          printer.printOperand(std::get<0>(pair));
+                          printer << " : ";
+                          printer.printType(std::get<1>(pair));
+                        });
+  printer << ")";
+}
+
 //===----------------------------------------------------------------------===//
 // GPUFuncOp
 //===----------------------------------------------------------------------===//
 
 /// Adds a new block argument that corresponds to buffers located in
 /// workgroup memory.
 BlockArgument GPUFuncOp::addWorkgroupAttribution(Type type) {
   auto attrName = getNumWorkgroupAttributionsAttrName();

mlir/lib/IR/FunctionImplementation.cpp

 //===- FunctionImplementation.cpp - Utilities for function-like ops -------===//
 //
 // Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
 // See https://llvm.org/LICENSE.txt for license information.
 // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
 //
 //===----------------------------------------------------------------------===//
 
 #include "mlir/IR/FunctionImplementation.h"
 #include "mlir/IR/Builders.h"
 #include "mlir/IR/FunctionSupport.h"
 #include "mlir/IR/SymbolTable.h"
 
 using namespace mlir;
 
-static ParseResult
-parseArgumentList(OpAsmParser &parser, bool allowVariadic,
-                  SmallVectorImpl<Type> &argTypes,
-                  SmallVectorImpl<OpAsmParser::OperandType> &argNames,
-                  SmallVectorImpl<NamedAttrList> &argAttrs, bool &isVariadic) {
+ParseResult mlir::impl::parseFunctionArgumentList(
+    OpAsmParser &parser, bool allowVariadic,
+    SmallVectorImpl<Type> &argTypes, SmallVectorImpl<NamedAttrList> &argAttrs,
+    bool &isVariadic) {
   if (parser.parseLParen())
     return failure();
 
   // The argument list either has to consistently have ssa-id's followed by
   // types, or just be a type list.  It isn't ok to sometimes have SSA ID's and
   // sometimes not.
   auto parseArgument = [&]() -> ParseResult {
     llvm::SMLoc loc = parser.getCurrentLocation();

 /// trailing arguments are populated by this function with names, types and
 /// attributes of the arguments and those of the results.
 ParseResult mlir::impl::parseFunctionSignature(
     OpAsmParser &parser, bool allowVariadic,
     SmallVectorImpl<OpAsmParser::OperandType> &argNames,
     SmallVectorImpl<Type> &argTypes, SmallVectorImpl<NamedAttrList> &argAttrs,
     bool &isVariadic, SmallVectorImpl<Type> &resultTypes,
     SmallVectorImpl<NamedAttrList> &resultAttrs) {
-  if (parseArgumentList(parser, allowVariadic, argTypes, argNames, argAttrs,
-                        isVariadic))
+  if (parseFunctionArgumentList(parser, allowVariadic, argNames, argTypes,
+                                argAttrs, isVariadic))
     return failure();
   if (succeeded(parser.parseOptionalArrow()))
     return parseFunctionResultList(parser, resultTypes, resultAttrs);
   return success();
 }
 
 void mlir::impl::addArgAndResultAttrs(Builder &builder, OperationState &result,
                                       ArrayRef<NamedAttrList> argAttrs,

mlir/test/Dialect/GPU/invalid.mlir

"gpu.launch_func"(%sz, %sz, %sz, %sz, %sz)
       : (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"(%sz, %sz, %sz, %sz, %sz, %sz) {foo = "bar"}
-      : (index, index, index, index, index, index) -> ()
+  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 {{symbol reference attribute 'kernel' must be specified}}
-    "gpu.launch_func"(%sz, %sz, %sz, %sz, %sz, %sz) {foo = "bar"}
+    // expected-error@+1 {{'gpu.launch_func' op requires attribute 'kernel'}}
+    "gpu.launch_func"(%sz, %sz, %sz, %sz, %sz, %sz)
         : (index, index, index, index, index, index) -> ()
     return
   }
 }
 
 // -----
 
 module attributes {gpu.container_module} {
   func @launch_func_no_function_attribute(%sz : index) {
-    // expected-error@+1 {{symbol reference attribute 'kernel' must be specified}}
-    "gpu.launch_func"(%sz, %sz, %sz, %sz, %sz, %sz) {kernel = 10}
-        : (index, index, index, index, index, index) -> ()
+    // expected-error@+1 {{custom op 'gpu.launch_func' invalid kind of attribute specified}}
+    gpu.launch_func "foo" grid in (%sz, %sz, %sz) block in (%sz, %sz, %sz) ()
     return
   }
 }
 
 // -----
 
 module attributes {gpu.container_module} {
   func @launch_func_undefined_module(%sz : index) {
     // expected-error@+1 {{kernel module 'kernels' is undefined}}
-    "gpu.launch_func"(%sz, %sz, %sz, %sz, %sz, %sz)
-    { kernel = @kernels::@kernel_1 }
-        : (index, index, index, index, index, index) -> ()
+    gpu.launch_func @kernels::@kernel_1 grid in (%sz, %sz, %sz) block in (%sz, %sz, %sz) ()
     return
   }
 }
 
 // -----
 
 module attributes {gpu.container_module} {
   module @kernels {
     // expected-error@+1 {{'gpu.func' op expects parent op 'gpu.module'}}
     gpu.func @kernel_1(%arg1 : !llvm.ptr<float>) {
       gpu.return
     }
   }
 }
 
 // -----
 
 module attributes {gpu.container_module} {
   module @kernels {
   }
 
   func @launch_func_missing_module_attribute(%sz : index) {
     // expected-error@+1 {{kernel module 'kernels' is undefined}}
-    "gpu.launch_func"(%sz, %sz, %sz, %sz, %sz, %sz)
-    { kernel = @kernels::@kernel_1 }
-        : (index, index, index, index, index, index) -> ()
+    gpu.launch_func @kernels::@kernel_1 grid in (%sz, %sz, %sz) block in (%sz, %sz, %sz) ()
     return
   }
 }
 
 // -----
 
 module attributes {gpu.container_module} {
   gpu.module @kernels { }
 
   func @launch_func_undefined_function(%sz : index) {
     // expected-error@+1 {{kernel function '@kernels::@kernel_1' is undefined}}
-    "gpu.launch_func"(%sz, %sz, %sz, %sz, %sz, %sz)
-    { kernel = @kernels::@kernel_1 }
-        : (index, index, index, index, index, index) -> ()
+    gpu.launch_func @kernels::@kernel_1 grid in (%sz, %sz, %sz) block in (%sz, %sz, %sz) ()
     return
   }
 }
 
 // -----
 
 module attributes {gpu.container_module} {
   module @kernels {
     gpu.func @kernel_1(%arg1 : !llvm.ptr<float>) kernel {
       gpu.return
     }
   }
 
   func @launch_func_missing_kernel_attr(%sz : index, %arg : !llvm.ptr<float>) {
     // expected-error@+1 {{kernel module 'kernels' is undefined}}
-    "gpu.launch_func"(%sz, %sz, %sz, %sz, %sz, %sz, %arg)
-    {kernel = @kernels::@kernel_1}
-        : (index, index, index, index, index, index, !llvm.ptr<float>) -> ()
+    gpu.launch_func @kernels::@kernel_1 grid in (%sz, %sz, %sz) block in (%sz, %sz, %sz) (%arg : !llvm.ptr<float>)
     return
   }
 }
 
 // -----
 
 module attributes {gpu.container_module} {
   gpu.module @kernels {
     gpu.func @kernel_1(%arg1 : !llvm.ptr<float>) {
       gpu.return
     }
   }
 
   func @launch_func_missing_kernel_attr(%sz : index, %arg : !llvm.ptr<float>) {
     // expected-error@+1 {{kernel function is missing the 'gpu.kernel' attribute}}
-    "gpu.launch_func"(%sz, %sz, %sz, %sz, %sz, %sz, %arg)
-    {kernel = @kernels::@kernel_1}
-        : (index, index, index, index, index, index, !llvm.ptr<float>) -> ()
+    gpu.launch_func @kernels::@kernel_1 grid in (%sz, %sz, %sz) block in (%sz, %sz, %sz) (%arg : !llvm.ptr<float>)
     return
   }
 }
 
 // -----
 
 module attributes {gpu.container_module} {
   gpu.module @kernels {
     gpu.func @kernel_1(%arg1 : !llvm.ptr<float>) kernel {
       gpu.return
     }
   }
 
   func @launch_func_kernel_operand_size(%sz : index, %arg : !llvm.ptr<float>) {
     // expected-error@+1 {{got 2 kernel operands but expected 1}}
-    "gpu.launch_func"(%sz, %sz, %sz, %sz, %sz, %sz, %arg, %arg)
-        {kernel = @kernels::@kernel_1}
-        : (index, index, index, index, index, index, !llvm.ptr<float>,
-           !llvm.ptr<float>) -> ()
+    gpu.launch_func @kernels::@kernel_1 grid in (%sz, %sz, %sz) block in (%sz, %sz, %sz) (%arg : !llvm.ptr<float>, %arg : !llvm.ptr<float>)
     return
   }
 }
 
 // -----
 
 module attributes {gpu.container_module} {
   gpu.module @kernels {
     gpu.func @kernel_1(%arg1 : f32) kernel {
       gpu.return
     }
   }
 
   func @launch_func_kernel_operand_types(%sz : index, %arg : f32) {
     // expected-err@+1 {{type of function argument 0 does not match}}
-    "gpu.launch_func"(%sz, %sz, %sz, %sz, %sz, %sz, %arg)
-        {kernel = @kernels::@kernel_1}
-        : (index, index, index, index, index, index, f32) -> ()
+    gpu.launch_func @kernels::@kernel_1 grid in (%sz, %sz, %sz) block in (%sz, %sz, %sz) (%arg : f32)
+    return
+  }
+}
+
+// -----
+
+module attributes {gpu.container_module} {
+  func @launch_func_kernel_operand_attr(%sz : index) {
+    // expected-error@+1 {{does not support kernel argument attributes}}
+    gpu.launch_func @foo::@bar grid in (%sz, %sz, %sz) block in (%sz, %sz, %sz) (%sz : index {foo})
     return
   }
 }
 
 // -----
 
 func @illegal_dimension() {
   // expected-error@+1 {{dimension "o" is invalid}}

mlir/test/Dialect/GPU/ops.mlir

module attributes {gpu.container_module} {
   }
 
   func @foo() {
     %0 = "op"() : () -> (f32)
     %1 = "op"() : () -> (memref<?xf32, 1>)
     // CHECK: %{{.*}} = constant 8
     %cst = constant 8 : index
 
-    // CHECK: "gpu.launch_func"(%{{.*}}, %{{.*}}, %{{.*}}, %{{.*}}, %{{.*}}, %{{.*}}, %{{.*}}, %{{.*}}) {kernel = @kernels::@kernel_1} : (index, index, index, index, index, index, f32, memref<?xf32, 1>) -> ()
-    "gpu.launch_func"(%cst, %cst, %cst, %cst, %cst, %cst, %0, %1)
-    { kernel = @kernels::@kernel_1}
-        : (index, index, index, index, index, index, f32, memref<?xf32, 1>) -> ()
-
-    // CHECK: "gpu.launch_func"(%{{.*}}, %{{.*}}, %{{.*}}, %{{.*}}, %{{.*}}, %{{.*}}, %{{.*}}, %{{.*}}) {kernel = @kernels::@kernel_2} : (index, index, index, index, index, index, f32, memref<?xf32, 1>) -> ()
-    "gpu.launch_func"(%cst, %cst, %cst, %cst, %cst, %cst, %0, %1)
-    { kernel = @kernels::@kernel_2}
-        : (index, index, index, index, index, index, f32, memref<?xf32, 1>) -> ()
+    // 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>)
 
     return
   }
 
   gpu.module @gpu_funcs {
     // CHECK-LABEL: gpu.func @kernel_1({{.*}}: f32)
     // CHECK:       workgroup
     // CHECK:       private

mlir/test/Dialect/GPU/outlining.mlir

func @launch() {
   %gDimZ = constant 16 : index
   // CHECK: %[[BDIMX:.*]] = constant 20
   %bDimX = constant 20 : index
   // CHECK: %[[BDIMY:.*]] = constant 24
   %bDimY = constant 24 : index
   // CHECK: %[[BDIMZ:.*]] = constant 28
   %bDimZ = constant 28 : index
 
-  // CHECK: "gpu.launch_func"(%[[GDIMX]], %[[GDIMY]], %[[GDIMZ]], %[[BDIMX]], %[[BDIMY]], %[[BDIMZ]], %[[ARG0]], %[[ARG1]]) {kernel = @launch_kernel::@launch_kernel} : (index, index, index, index, index, index, f32, memref<?xf32, 1>) -> ()
+  // CHECK: gpu.launch_func @launch_kernel::@launch_kernel grid in (%[[GDIMX]], %[[GDIMY]], %[[GDIMZ]]) block in (%[[BDIMX]], %[[BDIMY]], %[[BDIMZ]]) (%[[ARG0]] : f32, %[[ARG1]] : memref<?xf32, 1>)
   // CHECK-NOT: gpu.launch blocks
   gpu.launch blocks(%bx, %by, %bz) in (%grid_x = %gDimX, %grid_y = %gDimY,
                                        %grid_z = %gDimZ)
              threads(%tx, %ty, %tz) in (%block_x = %bDimX, %block_y = %bDimY,
                                         %block_z = %bDimZ) {
     "use"(%0): (f32) -> ()
     "some_op"(%bx, %block_x) : (index, index) -> ()
     %42 = load %1[%tx] : memref<?xf32, 1>

 
 // -----
 
 // CHECK: module attributes {gpu.container_module}
 // CHECK-LABEL: @multiple_launches
 func @multiple_launches() {
   // CHECK: %[[CST:.*]] = constant 8 : index
   %cst = constant 8 : index
-  // CHECK: "gpu.launch_func"(%[[CST]], %[[CST]], %[[CST]], %[[CST]], %[[CST]], %[[CST]]) {kernel = @multiple_launches_kernel::@multiple_launches_kernel} : (index, index, index, index, index, index) -> ()
+  // CHECK: gpu.launch_func @multiple_launches_kernel::@multiple_launches_kernel grid in (%[[CST]], %[[CST]], %[[CST]]) block in (%[[CST]], %[[CST]], %[[CST]]) ()
   gpu.launch blocks(%bx, %by, %bz) in (%grid_x = %cst, %grid_y = %cst,
                                        %grid_z = %cst)
              threads(%tx, %ty, %tz) in (%block_x = %cst, %block_y = %cst,
                                         %block_z = %cst) {
     gpu.terminator
   }
-  // CHECK: "gpu.launch_func"(%[[CST]], %[[CST]], %[[CST]], %[[CST]], %[[CST]], %[[CST]]) {kernel = @multiple_launches_kernel_0::@multiple_launches_kernel} : (index, index, index, index, index, index) -> ()
+  // CHECK: gpu.launch_func @multiple_launches_kernel_0::@multiple_launches_kernel grid in (%[[CST]], %[[CST]], %[[CST]]) block in (%[[CST]], %[[CST]], %[[CST]]) ()
   gpu.launch blocks(%bx2, %by2, %bz2) in (%grid_x2 = %cst, %grid_y2 = %cst,
                                           %grid_z2 = %cst)
              threads(%tx2, %ty2, %tz2) in (%block_x2 = %cst, %block_y2 = %cst,
                                            %block_z2 = %cst) {
     gpu.terminator
   }
   return
 }
 
 // CHECK: module @multiple_launches_kernel
 // CHECK: func @multiple_launches_kernel
 // CHECK: module @multiple_launches_kernel_0
 // CHECK: func @multiple_launches_kernel
 
 // -----
 
 // CHECK-LABEL: @extra_constants_not_inlined
 func @extra_constants_not_inlined(%arg0: memref<?xf32>) {
   // CHECK: %[[CST:.*]] = constant 8 : index
   %cst = constant 8 : index
   %cst2 = constant 2 : index
   %c0 = constant 0 : index
   %cst3 = "secret_constant"() : () -> index
-  // CHECK: "gpu.launch_func"(%[[CST]], %[[CST]], %[[CST]], %[[CST]], %[[CST]], %[[CST]], %{{.*}}, %{{.*}}) {kernel = @extra_constants_not_inlined_kernel::@extra_constants_not_inlined_kernel} : (index, index, index, index, index, index, memref<?xf32>, index) -> ()
+  // CHECK: gpu.launch_func @extra_constants_not_inlined_kernel::@extra_constants_not_inlined_kernel grid in (%[[CST]], %[[CST]], %[[CST]]) block in (%[[CST]], %[[CST]], %[[CST]]) ({{.*}} : memref<?xf32>, {{.*}} : index)
   gpu.launch blocks(%bx, %by, %bz) in (%grid_x = %cst, %grid_y = %cst,
                                        %grid_z = %cst)
              threads(%tx, %ty, %tz) in (%block_x = %cst, %block_y = %cst,
                                         %block_z = %cst) {
     "use"(%cst2, %arg0, %cst3) : (index, memref<?xf32>, index) -> ()
     gpu.terminator
   }
   return
 }
 
 // CHECK-LABEL: func @extra_constants_not_inlined_kernel(%{{.*}}: memref<?xf32>, %{{.*}}: index)
 // CHECK: constant 2
 
 // -----
 
 // CHECK-LABEL: @extra_constants
 // CHECK-SAME: %[[ARG0:.*]]: memref<?xf32>
 func @extra_constants(%arg0: memref<?xf32>) {
   // CHECK: %[[CST:.*]] = constant 8 : index
   %cst = constant 8 : index
   %cst2 = constant 2 : index
   %c0 = constant 0 : index
   %cst3 = dim %arg0, %c0 : memref<?xf32>
-  // CHECK: "gpu.launch_func"(%[[CST]], %[[CST]], %[[CST]], %[[CST]], %[[CST]], %[[CST]], %[[ARG0]]) {kernel = @extra_constants_kernel::@extra_constants_kernel} : (index, index, index, index, index, index, memref<?xf32>) -> ()
+  // CHECK: gpu.launch_func @extra_constants_kernel::@extra_constants_kernel grid in (%[[CST]], %[[CST]], %[[CST]]) block in (%[[CST]], %[[CST]], %[[CST]]) (%[[ARG0]] : memref<?xf32>)
   gpu.launch blocks(%bx, %by, %bz) in (%grid_x = %cst, %grid_y = %cst,
                                        %grid_z = %cst)
              threads(%tx, %ty, %tz) in (%block_x = %cst, %block_y = %cst,
                                         %block_z = %cst) {
     "use"(%cst2, %arg0, %cst3) : (index, memref<?xf32>, index) -> ()
     gpu.terminator
   }
   return
 }
 
-// CHECK-LABEL: func @extra_constants_kernel
+// CHECK-LABEL: func @extra_constants_kernel(
 // CHECK-SAME: %[[KARG0:.*]]: memref<?xf32>
 // CHECK: constant 2
 // CHECK: constant 0
 // CHECK: dim %[[KARG0]]
 
 // -----
 
 // CHECK-LABEL: @extra_constants_noarg
 // CHECK-SAME: %[[ARG0:.*]]: memref<?xf32>, %[[ARG1:.*]]: memref<?xf32>
 func @extra_constants_noarg(%arg0: memref<?xf32>, %arg1: memref<?xf32>) {
   // CHECK: %[[CST:.*]] = constant 8 : index
   %cst = constant 8 : index
   %cst2 = constant 2 : index
   %c0 = constant 0 : index
   // CHECK: dim %[[ARG1]]
   %cst3 = dim %arg1, %c0 : memref<?xf32>
-  // CHECK: "gpu.launch_func"(%[[CST]], %[[CST]], %[[CST]], %[[CST]], %[[CST]], %[[CST]], %[[ARG0]], %{{.*}}) {kernel = @extra_constants_noarg_kernel::@extra_constants_noarg_kernel} : (index, index, index, index, index, index, memref<?xf32>, index) -> ()
+  // CHECK: gpu.launch_func @extra_constants_noarg_kernel::@extra_constants_noarg_kernel grid in (%[[CST]], %[[CST]], %[[CST]]) block in (%[[CST]], %[[CST]], %[[CST]]) (%[[ARG0]] : memref<?xf32>, {{.*}} : index)
   gpu.launch blocks(%bx, %by, %bz) in (%grid_x = %cst, %grid_y = %cst,
                                        %grid_z = %cst)
              threads(%tx, %ty, %tz) in (%block_x = %cst, %block_y = %cst,
                                         %block_z = %cst) {
     "use"(%cst2, %arg0, %cst3) : (index, memref<?xf32>, index) -> ()
     gpu.terminator
   }
   return
 }
 
-// CHECK-LABEL: func @extra_constants_noarg_kernel
+// CHECK-LABEL: func @extra_constants_noarg_kernel(
 // CHECK-SAME: %[[KARG0:.*]]: memref<?xf32>, %[[KARG1:.*]]: index
 // CHECK: %[[KCST:.*]] = constant 2
 // CHECK: "use"(%[[KCST]], %[[KARG0]], %[[KARG1]])
 
 // -----
 
 // CHECK-LABEL: @multiple_uses
 func @multiple_uses(%arg0 : memref<?xf32>) {