[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

410

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

410

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

410

Switched to space literal.

[ftynse]

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

21

Are we sure this isn't necessary anymore?

366

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.

366

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

366

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
module attributes {gpu.container_module} {
                         %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>)
               -> ()
     }
     ```
   }];
 
   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` 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">,
     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 parseSpace(OpAsmParser &parser) { return success(); }
+
+static void printSpace(OpAsmPrinter &printer) { printer << " "; }
+
+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>) {

mlir/test/mlir-tblgen/op-format-spec.td

def DirectiveCustomValidA : TestFormat_Op<"custom_valid_a", [{
   custom<MyDirective>($operand) attr-dict
 }]>, Arguments<(ins Optional<I64>:$operand)>;
 def DirectiveCustomValidB : TestFormat_Op<"custom_valid_b", [{
   custom<MyDirective>($operand, type($operand), type($result)) attr-dict
 }]>, Arguments<(ins I64:$operand)>, Results<(outs I64:$result)>;
 def DirectiveCustomValidC : TestFormat_Op<"custom_valid_c", [{
   custom<MyDirective>($attr) attr-dict
 }]>, Arguments<(ins I64Attr:$attr)>;
+def DirectiveCustomValidD : TestFormat_Op<"custom_valid_d", [{
+  custom<MyDirective>() attr-dict
+}]>, Arguments<(ins)>;
 
 //===----------------------------------------------------------------------===//
 // functional-type
 
 // CHECK: error: 'functional-type' is only valid as a top-level directive
 def DirectiveFunctionalTypeInvalidA : TestFormat_Op<"functype_invalid_a", [{
   functional-type(functional-type)
 }]>;

mlir/tools/mlir-tblgen/OpFormatGen.cpp

FormatParser::parseCustomDirective(std::unique_ptr<Element> &element,
   if (failed(parseToken(Token::identifier,
                         "expected custom directive name identifier")) ||
       failed(parseToken(Token::greater,
                         "expected '>' after custom directive name")) ||
       failed(parseToken(Token::l_paren,
                         "expected '(' before custom directive parameters")))
     return ::mlir::failure();
 
-  // Parse the child elements for this optional group.=
+  // Parse the child elements for this optional group.
   std::vector<std::unique_ptr<Element>> elements;
+
+  if (curToken.getKind() != Token::r_paren) {
     do {
       if (failed(parseCustomDirectiveParameter(elements)))
         return ::mlir::failure();
       if (curToken.getKind() != Token::comma)
         break;
       consumeToken();
     } while (true);
+  }
 
   if (failed(parseToken(Token::r_paren,
                         "expected ')' after custom directive parameters")))
     return ::mlir::failure();
 
   // After parsing all of the elements, ensure that all type directives refer
   // only to variables.
   for (auto &ele : elements) {