[mlir][gpu] Modifies `gpu.launch_func` to allow lowering it after gpu-to-llvm.

[fmorac]

Summary

For an explanation of these patches see D154153.

Commit message:
In order to lower gpu.launch_func after running gpu-to-llvm it must be
able to handle lowered types -eg. index -> i64. This patch also allows the op
to refer to GPU binaries and not only GPU modules.

Depends on D154132.

[fmorac]

Correct dependencies.

[fmorac]

Update dependencies.

[fmorac]

Rebasing.

[krzysz00]

Comments wrt assembly format

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

1023

This is just AllTypesMatch in the tablegen

1034

I don't think this'll be needed with an AllTypesMatch

[fmorac]

Added the AllTypesMatch trait.

[mehdi_amini]

Please add tests here

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

277–278

Update the comment?

981

Since we're touching the build methods, can you update these to use properties instead?

should look like:

Properties &prop = result.getOrAddProperties<Properties>();
prop.operand_segment_sizes = builder.getDenseI32ArrayAttr(segmentSizes);

[fmorac]

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

277–278

I'll update it.

981

Yes.

[mehdi_amini]

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

981

Actually prop.odsOperandSegmentSizes, I just changed it ;)

[fmorac]

Change builders to use properties.

[mehdi_amini]

Can you add tests for the new syntax and new types?

[fmorac]

In D154137#4556082, @mehdi_amini wrote:

Can you add tests for the new syntax and new types?

Sure, I can add syntax & type tests here, however tests for verifying the compatibility with gpu.binary have to be added after the introduction of #gpu.select_object.

[fmorac]

Added tests for the syntax.

[mehdi_amini]

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

1019

Sorry, I removed the ods prefix...

[fmorac]

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

1019

Is that change in trunk? Because I just did a build and it's still using odsOperandSegmentSizes.

[mehdi_amini]

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

1019

I forgot to push... now it is there :)

[fmorac]

Removing Ods prefix.

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

let extraClassDeclaration = [{
     /// Verifies the body of the function.
     LogicalResult verifyBody();
   }];
   let hasCustomAssemblyFormat = 1;
 
   let hasVerifier = 1;
 }
 
+def LaunchIndx : AnyTypeOf<[Index, I32, I64]>;
+
 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,
+               LaunchIndx:$gridSizeX, LaunchIndx:$gridSizeY, LaunchIndx:$gridSizeZ,
+               LaunchIndx:$blockSizeX, LaunchIndx:$blockSizeY, LaunchIndx:$blockSizeZ,
                Optional<I32>:$dynamicSharedMemorySize,
-               Variadic<AnyType>:$kernelOperands)>,
+               Variadic<AnyType>:$kernelOperands,
+               Optional<AnyType>:$asyncObject)>,
     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
def GPU_LaunchFuncOp : GPU_Op<"launch_func",
 
   let skipDefaultBuilders = 1;
 
   let builders = [
     OpBuilder<(ins "GPUFuncOp":$kernelFunc, "KernelDim3":$gridSize,
       "KernelDim3":$blockSize, "Value":$dynamicSharedMemorySize,
       "ValueRange":$kernelOperands,
       CArg<"Type", "nullptr">:$asyncTokenType,
-      CArg<"ValueRange", "{}">:$asyncDependencies)>
+      CArg<"ValueRange", "{}">:$asyncDependencies)>,
+    OpBuilder<(ins "SymbolRefAttr":$kernel, "KernelDim3":$gridSize,
+      "KernelDim3":$blockSize, "Value":$dynamicSharedMemorySize,
+      "ValueRange":$kernelOperands,
+      CArg<"Value", "nullptr">:$asyncObject)>
   ];
 
   let extraClassDeclaration = [{
     /// The name of the kernel's containing module.
     StringAttr getKernelModuleName();
 
     /// The name of the kernel.
     StringAttr getKernelName();
let extraClassDeclaration = [{
     // This needs to quietly verify if attributes with names defined below are
     // present since it is run before the verifier of this op.
     friend LogicalResult GPUDialect::verifyOperationAttribute(Operation *,
                                                               NamedAttribute);
   }];
 
   let assemblyFormat = [{
       custom<AsyncDependencies>(type($asyncToken), $asyncDependencies)
+      (`<` $asyncObject^ type($asyncObject) `>`)?
       $kernel
-      `blocks` `in` ` ` `(`$gridSizeX`,` $gridSizeY`,` $gridSizeZ`)`
-      `threads` `in` ` ` `(`$blockSizeX`,` $blockSizeY`,` $blockSizeZ`)`
+      `blocks` `in` custom<Dim3>($gridSizeX, type($gridSizeX), $gridSizeY, type($gridSizeY), $gridSizeZ, type($gridSizeZ))
+      `threads` `in` custom<Dim3>($blockSizeX, type($blockSizeX), $blockSizeY, type($blockSizeY), $blockSizeZ, type($blockSizeZ))
       (`dynamic_shared_memory_size` $dynamicSharedMemorySize^)?
       custom<LaunchFuncOperands>($kernelOperands, type($kernelOperands)) attr-dict
   }];
   let hasVerifier = 1;
 }
 
 def GPU_LaunchOp : GPU_Op<"launch", [
       AutomaticAllocationScope, AttrSizedOperandSegments, GPU_AsyncOpInterface,

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

if (!launchOp->getParentOp() ||
       return success();
 
     // Ignore launch ops with missing attributes here. The errors will be
     // reported by the verifiers of those ops.
     if (!launchOp->getAttrOfType<SymbolRefAttr>(
             LaunchFuncOp::getKernelAttrName(launchOp->getName())))
       return success();
 
     // Check that `launch_func` refers to a well-formed GPU kernel module.
-    StringAttr kernelModuleName = launchOp.getKernelModuleName();
-    auto kernelModule = module.lookupSymbol<GPUModuleOp>(kernelModuleName);
+    StringAttr kernelContainerName = launchOp.getKernelModuleName();
+    Operation *kernelContainer = module.lookupSymbol(kernelContainerName);
+    if (!kernelContainer)
+      return launchOp.emitOpError()
+             << "kernel container '" << kernelContainerName.getValue()
+             << "' is undefined";
+
+    // If the container is a GPU binary op return success.
+    if (isa<BinaryOp>(kernelContainer))
+      return success();
+
+    auto kernelModule = dyn_cast<GPUModuleOp>(kernelContainer);
     if (!kernelModule)
       return launchOp.emitOpError()
-             << "kernel module '" << kernelModuleName.getValue()
+             << "kernel module '" << kernelContainerName.getValue()
              << "' is undefined";
 
     // Check that `launch_func` refers to a well-formed kernel function.
     Operation *kernelFunc = module.lookupSymbol(launchOp.getKernelAttr());
     if (!kernelFunc)
       return launchOp.emitOpError("kernel function '")
              << launchOp.getKernel() << "' is undefined";
     auto kernelConvertedFunction = dyn_cast<FunctionOpInterface>(kernelFunc);
void LaunchFuncOp::build(OpBuilder &builder, OperationState &result,
   if (dynamicSharedMemorySize)
     result.addOperands(dynamicSharedMemorySize);
   result.addOperands(kernelOperands);
   auto kernelModule = kernelFunc->getParentOfType<GPUModuleOp>();
   auto kernelSymbol =
       SymbolRefAttr::get(kernelModule.getNameAttr(),
                          {SymbolRefAttr::get(kernelFunc.getNameAttr())});
   result.addAttribute(getKernelAttrName(result.name), kernelSymbol);
-  SmallVector<int32_t, 9> segmentSizes(9, 1);
+  SmallVector<int32_t, 9> segmentSizes(10, 1);
   segmentSizes.front() = asyncDependencies.size();
-  segmentSizes[segmentSizes.size() - 2] = dynamicSharedMemorySize ? 1 : 0;
-  segmentSizes.back() = static_cast<int32_t>(kernelOperands.size());
+  segmentSizes[segmentSizes.size() - 3] = dynamicSharedMemorySize ? 1 : 0;
+  segmentSizes[segmentSizes.size() - 2] =
+      static_cast<int32_t>(kernelOperands.size());
+  segmentSizes.back() = 0;
+  result.addAttribute(getOperandSegmentSizeAttr(),
+                      builder.getDenseI32ArrayAttr(segmentSizes));
+}
+
+void LaunchFuncOp::build(OpBuilder &builder, OperationState &result,
+                         SymbolRefAttr kernel, KernelDim3 gridSize,
+                         KernelDim3 getBlockSize, Value dynamicSharedMemorySize,
+                         ValueRange kernelOperands, Value asyncObject) {
+  // Add grid and block sizes as op operands, followed by the data operands.
+  result.addOperands({gridSize.x, gridSize.y, gridSize.z, getBlockSize.x,
+                      getBlockSize.y, getBlockSize.z});
+  if (dynamicSharedMemorySize)
+    result.addOperands(dynamicSharedMemorySize);
+  result.addOperands(kernelOperands);
+  if (asyncObject)
+    result.addOperands(asyncObject);
+  result.addAttribute(getKernelAttrName(result.name), kernel);
+  SmallVector<int32_t, 9> segmentSizes(10, 1);
+  segmentSizes.front() = 0;
+  segmentSizes[segmentSizes.size() - 3] = dynamicSharedMemorySize ? 1 : 0;
+  segmentSizes[segmentSizes.size() - 2] =
+      static_cast<int32_t>(kernelOperands.size());
+
+  segmentSizes.back() = asyncObject ? 1 : 0;
   result.addAttribute(getOperandSegmentSizeAttr(),
                       builder.getDenseI32ArrayAttr(segmentSizes));
 }
 
 StringAttr LaunchFuncOp::getKernelModuleName() {
   return getKernel().getRootReference();
 }
 
 StringAttr LaunchFuncOp::getKernelName() {
   return getKernel().getLeafReference();
LogicalResult LaunchFuncOp::verify() {
   auto module = (*this)->getParentOfType<ModuleOp>();
   if (!module)
     return emitOpError("expected to belong to a module");
 
   if (!module->getAttrOfType<UnitAttr>(
           GPUDialect::getContainerModuleAttrName()))
     return emitOpError("expected the closest surrounding module to have the '" +
                        GPUDialect::getContainerModuleAttrName() +
                        "' attribute");
 
+  KernelDim3 grid = getGridSizeOperandValues();
+  KernelDim3 block = getBlockSizeOperandValues();
+  if (grid.x.getType() != grid.y.getType() ||
+      grid.x.getType() != grid.z.getType() ||
+      grid.x.getType() != block.x.getType() ||
+      grid.x.getType() != block.y.getType() ||
+      grid.x.getType() != block.z.getType())
+    return emitOpError(
+        "expected the grid and block sizes all having the same type");
   return success();
 }
 
+static ParseResult
+parseDim3(OpAsmParser &parser, OpAsmParser::UnresolvedOperand &sizeX,
+          Type &sizeXTy, OpAsmParser::UnresolvedOperand &sizeY, Type &sizeYTy,
+          OpAsmParser::UnresolvedOperand &sizeZ, Type &sizeZTy) {
+  if (parser.parseLParen() || parser.parseOperand(sizeX) ||
+      parser.parseComma() || parser.parseOperand(sizeY) ||
+      parser.parseComma() || parser.parseOperand(sizeZ) || parser.parseRParen())
+    return failure();
+  SmallVector<Type, 1> types;
+  if (failed(parser.parseOptionalColonTypeList(types)))
+    return failure();
+  if (types.size()) {
+    sizeXTy = types[0];
+    sizeYTy = types[0];
+    sizeZTy = types[0];
+  } else {
+    types.push_back(IndexType::get(parser.getContext()));
+    sizeXTy = types[0];
+    sizeYTy = types[0];
+    sizeZTy = types[0];
+  }
+  return success();
+}
+
+static void printDim3(OpAsmPrinter &printer, Operation *op, Value sizeX,
+                      Type sizeXTy, Value sizeY, Type sizeYTy, Value sizeZ,
+                      Type sizeZTy) {
+  printer << '(' << sizeX << ", " << sizeY << ", " << sizeZ << ')';
+  IndexType indexType = IndexType::get(op->getContext());
+  if (indexType != sizeXTy)
+    printer << " : " << sizeXTy;
+}
+
 static ParseResult parseLaunchFuncOperands(
     OpAsmParser &parser,
     SmallVectorImpl<OpAsmParser::UnresolvedOperand> &argNames,
     SmallVectorImpl<Type> &argTypes) {
   if (parser.parseOptionalKeyword("args"))
     return success();
 
   auto parseElement = [&]() -> ParseResult {

mlir/test/Dialect/GPU/invalid.mlir

func.func @launch_func_no_function_attribute(%sz : index) {
     return
   }
 }
 
 // -----
 
 module attributes {gpu.container_module} {
   func.func @launch_func_undefined_module(%sz : index) {
-    // expected-error@+1 {{kernel module 'kernels' is undefined}}
+    // expected-error@+1 {{kernel container 'kernels' is undefined}}
     gpu.launch_func @kernels::@kernel_1 blocks in (%sz, %sz, %sz) threads in (%sz, %sz, %sz)
     return
   }
 }
 
 // -----
 
 module attributes {gpu.container_module} {