[MLIR] Make gpu.launch_func rewrite pattern part of the LLVM lowering pass.

[herhut]

Nice, some nits.

mlir/include/mlir/Conversion/Passes.td

86

nit: formatting

mlir/lib/Conversion/GPUCommon/ConvertLaunchFuncToRuntimeCalls.cpp

233

nit: formatting

294

Maybe return rewriter.notifyMatchfailure(op, "Cannot convert...")?

302

SymbolTable::lookupNearestSymbolFrom ?

[csigg]

Rebase, apply reviewer comments.

[csigg]

Fix.

[csigg]

Fix.

[csigg]

Fix.

[csigg]

Fix.

[jurahul]

mlir/lib/Conversion/GPUCommon/ConvertLaunchFuncToRuntimeCalls.cpp

302

Would it make sense to make this a member function of gpu::LaunchFuncOp if this lookup is done in multiple places?

[csigg]

mlir/lib/Conversion/GPUCommon/ConvertLaunchFuncToRuntimeCalls.cpp

302

As far as I can tell, this is only used here.

[herhut]

mlir/lib/Conversion/GPUCommon/ConvertLaunchFuncToRuntimeCalls.cpp

291–295

Maybe put this lambda in a local and give it a name?

[csigg]

Add isLLVMType function.

mlir/include/mlir/Conversion/GPUCommon/GPUCommonPass.h

 
 /// Creates a pass to convert a gpu.launch_func operation into a sequence of
 /// GPU runtime calls.
 ///
 /// This pass does not generate code to call GPU runtime APIs directly but
 /// instead uses a small wrapper library that exports a stable and conveniently
 /// typed ABI on top of GPU runtimes such as CUDA or ROCm (HIP).
 std::unique_ptr<OperationPass<ModuleOp>>
-createConvertGpuLaunchFuncToGpuRuntimeCallsPass(
-    StringRef gpuBinaryAnnotation = "");
+createGpuToLLVMConversionPass(StringRef gpuBinaryAnnotation = "");
 
 /// Collect a set of patterns to convert from the GPU dialect to LLVM.
 void populateGpuToLLVMConversionPatterns(LLVMTypeConverter &converter,
                                          OwningRewritePatternList &patterns,
                                          StringRef gpuBinaryAnnotation);
 
 /// Creates a pass to convert kernel functions into GPU target object blobs.
 ///

mlir/include/mlir/Conversion/Passes.td

let summary = "Convert the operations from the avx512 dialect into the LLVM "
                 "dialect";
   let constructor = "mlir::createConvertAVX512ToLLVMPass()";
 }
 
 //===----------------------------------------------------------------------===//
 // GPUCommon
 //===----------------------------------------------------------------------===//
 
-def ConvertGpuLaunchFuncToGpuRuntimeCalls : Pass<"launch-func-to-gpu-runtime",
+def GpuToLLVMConversionPass : Pass<"gpu-to-llvm",
                                                  "ModuleOp"> {
-  let summary = "Convert all launch_func ops to GPU runtime calls";
-  let constructor = "mlir::createConvertGpuLaunchFuncToGpuRuntimeCallsPass()";
+  let summary = "Convert GPU dialect to LLVM dialect with GPU runtime calls";
+  let constructor = "mlir::createGpuToLLVMConversionPass()";
   let options = [
     Option<"gpuBinaryAnnotation", "gpu-binary-annotation", "std::string",
-           "\"nvvm.cubin\"",
-           "Annotation attribute string for GPU binary">,
+           "", "Annotation attribute string for GPU binary">,
   ];
 }
 
 //===----------------------------------------------------------------------===//
 // GPUToNVVM
 //===----------------------------------------------------------------------===//
 
 def ConvertGpuOpsToNVVMOps : Pass<"convert-gpu-to-nvvm", "gpu::GPUModuleOp"> {

mlir/lib/Conversion/GPUCommon/ConvertLaunchFuncToRuntimeCalls.cpp

 #include "llvm/Support/FormatVariadic.h"
 
 using namespace mlir;
 
 static constexpr const char *kGpuBinaryStorageSuffix = "_gpubin_cst";
 
 namespace {
 
-class GpuLaunchFuncToGpuRuntimeCallsPass
-    : public ConvertGpuLaunchFuncToGpuRuntimeCallsBase<
-          GpuLaunchFuncToGpuRuntimeCallsPass> {
+class GpuToLLVMConversionPass
+    : public GpuToLLVMConversionPassBase<GpuToLLVMConversionPass> {
 public:
-  GpuLaunchFuncToGpuRuntimeCallsPass(StringRef gpuBinaryAnnotation) {
+  GpuToLLVMConversionPass(StringRef gpuBinaryAnnotation) {
     if (!gpuBinaryAnnotation.empty())
       this->gpuBinaryAnnotation = gpuBinaryAnnotation.str();
   }
 
   // Run the dialect converter on the module.
   void runOnOperation() override;
 };
 
class ConvertLaunchFuncOpToGpuRuntimeCallPattern
     : public ConvertOpToGpuRuntimeCallPattern<gpu::LaunchFuncOp> {
 public:
   ConvertLaunchFuncOpToGpuRuntimeCallPattern(LLVMTypeConverter &typeConverter,
                                              StringRef gpuBinaryAnnotation)
       : ConvertOpToGpuRuntimeCallPattern<gpu::LaunchFuncOp>(typeConverter),
         gpuBinaryAnnotation(gpuBinaryAnnotation) {}
 
 private:
-  void addParamToArray(OpBuilder &builder, Location loc, Value param,
-                       Value array, unsigned pos, Value one) const;
-  Value generateParamsArray(gpu::LaunchFuncOp launchOp, unsigned numArguments,
-                            OpBuilder &builder) const;
+  Value generateParamsArray(gpu::LaunchFuncOp launchOp,
+                            ArrayRef<Value> operands, OpBuilder &builder) const;
   Value generateKernelNameConstant(StringRef moduleName, StringRef name,
                                    Location loc, OpBuilder &builder) const;
 
   LogicalResult
   matchAndRewrite(Operation *op, ArrayRef<Value> operands,
                   ConversionPatternRewriter &rewriter) const override;
 
   llvm::SmallString<32> gpuBinaryAnnotation;
 };
 
 class EraseGpuModuleOpPattern : public OpRewritePattern<gpu::GPUModuleOp> {
   using OpRewritePattern<gpu::GPUModuleOp>::OpRewritePattern;
 
   LogicalResult matchAndRewrite(gpu::GPUModuleOp op,
                                 PatternRewriter &rewriter) const override {
     // GPU kernel modules are no longer necessary since we have a global
     // constant with the CUBIN, or HSACO data.
     rewriter.eraseOp(op);
     return success();
   }
 };
 
 } // namespace
 
-void GpuLaunchFuncToGpuRuntimeCallsPass::runOnOperation() {
+void GpuToLLVMConversionPass::runOnOperation() {
   LLVMTypeConverter converter(&getContext());
   OwningRewritePatternList patterns;
+  populateStdToLLVMConversionPatterns(converter, patterns);
   populateGpuToLLVMConversionPatterns(converter, patterns, gpuBinaryAnnotation);
 
   LLVMConversionTarget target(getContext());
   if (failed(applyPartialConversion(getOperation(), target, patterns)))
     signalPassFailure();
 }
 
 LLVM::CallOp FunctionCallBuilder::create(Location loc, OpBuilder &builder,
                                          ArrayRef<Value> arguments) const {
   auto module = builder.getBlock()->getParent()->getParentOfType<ModuleOp>();
   auto function = [&] {
     if (auto function = module.lookupSymbol<LLVM::LLVMFuncOp>(functionName))
       return function;
     return OpBuilder(module.getBody()->getTerminator())
         .create<LLVM::LLVMFuncOp>(loc, functionName, functionType);
   }();
   return builder.create<LLVM::CallOp>(
       loc, const_cast<LLVM::LLVMType &>(functionType).getFunctionResultType(),
       builder.getSymbolRefAttr(function), arguments);
 }
 
-/// Emits the IR with the following structure:
-///
-///   %data = llvm.alloca 1 x type-of(<param>)
-///   llvm.store <param>, %data
-///   %typeErased = llvm.bitcast %data to !llvm<"i8*">
-///   %addr = llvm.getelementptr <array>[<pos>]
-///   llvm.store %typeErased, %addr
-///
-/// This is necessary to construct the array of arguments passed to the kernel
-/// function as accepted by cuLaunchKernel, i.e. as a void** that points to
-/// array of stack-allocated type-erased pointers to the actual arguments.
-void ConvertLaunchFuncOpToGpuRuntimeCallPattern::addParamToArray(
-    OpBuilder &builder, Location loc, Value param, Value array, unsigned pos,
-    Value one) const {
-  auto memLocation = builder.create<LLVM::AllocaOp>(
-      loc, param.getType().cast<LLVM::LLVMType>().getPointerTo(), one,
-      /*alignment=*/0);
-  builder.create<LLVM::StoreOp>(loc, param, memLocation);
-  auto casted =
-      builder.create<LLVM::BitcastOp>(loc, llvmPointerType, memLocation);
-
-  auto index = builder.create<LLVM::ConstantOp>(loc, llvmInt32Type,
-                                                builder.getI32IntegerAttr(pos));
-  auto gep = builder.create<LLVM::GEPOp>(loc, llvmPointerPointerType, array,
-                                         index.getResult());
-  builder.create<LLVM::StoreOp>(loc, casted, gep);
-}
-
-// Generates a parameters array to be used with a CUDA / ROCm (HIP) kernel
-// launch call. The arguments are extracted from the launchOp.
+// Creates a struct containing all kernel parameters on the stack and returns
+// an array of type-erased pointers to the fields of the struct. The array can
+// then be passed to the CUDA / ROCm (HIP) kernel launch calls.
 // The generated code is essentially as follows:
 //
-// %array = alloca(numparams * sizeof(void *))
-// for (i : [0, NumKernelOperands))
-//   %array[i] = cast<void*>(KernelOperand[i])
+// %struct = alloca(sizeof(struct { Parameters... }))
+// %array = alloca(NumParameters * sizeof(void *))
+// for (i : [0, NumParameters))
+//   %fieldPtr = llvm.getelementptr %struct[0, i]
+//   llvm.store parameters[i], %fieldPtr
+//   %elementPtr = llvm.getelementptr %array[i]
+//   llvm.store %fieldPtr, %elementPtr
 // return %array
 Value ConvertLaunchFuncOpToGpuRuntimeCallPattern::generateParamsArray(
-    gpu::LaunchFuncOp launchOp, unsigned numArguments,
+    gpu::LaunchFuncOp launchOp, ArrayRef<Value> operands,
     OpBuilder &builder) const {
+  auto loc = launchOp.getLoc();
   auto numKernelOperands = launchOp.getNumKernelOperands();
-  Location loc = launchOp.getLoc();
+  auto arguments = typeConverter.promoteMemRefDescriptors(
+      loc, launchOp.getOperands().take_back(numKernelOperands),
+      operands.take_back(numKernelOperands), builder);
+  auto numArguments = arguments.size();
+  SmallVector<LLVM::LLVMType, 4> argumentTypes;
+  argumentTypes.reserve(numArguments);
+  for (auto argument : arguments) {
+    // Convert non-memref types. For LLVM types it's just a no-op.
+    auto llvmType = typeConverter.convertType(argument.getType());
+    argumentTypes.push_back(llvmType.cast<LLVM::LLVMType>());
+  }
+  auto structType = LLVM::LLVMType::createStructTy(argumentTypes, llvm::None);
   auto one = builder.create<LLVM::ConstantOp>(loc, llvmInt32Type,
                                               builder.getI32IntegerAttr(1));
+  auto structPtr = builder.create<LLVM::AllocaOp>(
+      loc, structType.getPointerTo(), one, /*alignment=*/0);
   auto arraySize = builder.create<LLVM::ConstantOp>(
       loc, llvmInt32Type, builder.getI32IntegerAttr(numArguments));
-  auto array = builder.create<LLVM::AllocaOp>(loc, llvmPointerPointerType,
+  auto arrayPtr = builder.create<LLVM::AllocaOp>(loc, llvmPointerPointerType,
                                                  arraySize, /*alignment=*/0);
-
-  unsigned pos = 0;
-  for (unsigned idx = 0; idx < numKernelOperands; ++idx) {
-    auto operand = launchOp.getKernelOperand(idx);
-    auto llvmType = operand.getType().cast<LLVM::LLVMType>();
-
-    // Assume all struct arguments come from MemRef. If this assumption does not
-    // hold anymore then we `launchOp` to lower from MemRefType and not after
-    // LLVMConversion has taken place and the MemRef information is lost.
-    if (!llvmType.isStructTy()) {
-      addParamToArray(builder, loc, operand, array, pos++, one);
-      continue;
-    }
-
-    // Put individual components of a memref descriptor into the flat argument
-    // list. We cannot use unpackMemref from LLVM lowering here because we have
-    // no access to MemRefType that had been lowered away.
-    for (int32_t j = 0, ej = llvmType.getStructNumElements(); j < ej; ++j) {
-      auto elemType = llvmType.getStructElementType(j);
-      if (elemType.isArrayTy()) {
-        for (int32_t k = 0, ek = elemType.getArrayNumElements(); k < ek; ++k) {
-          Value elem = builder.create<LLVM::ExtractValueOp>(
-              loc, elemType.getArrayElementType(), operand,
-              builder.getI32ArrayAttr({j, k}));
-          addParamToArray(builder, loc, elem, array, pos++, one);
-        }
-      } else {
-        assert((elemType.isIntegerTy() || elemType.isFloatTy() ||
-                elemType.isDoubleTy() || elemType.isPointerTy()) &&
-               "expected scalar type");
-        Value strct = builder.create<LLVM::ExtractValueOp>(
-            loc, elemType, operand, builder.getI32ArrayAttr(j));
-        addParamToArray(builder, loc, strct, array, pos++, one);
-      }
-    }
+  auto zero = builder.create<LLVM::ConstantOp>(loc, llvmInt32Type,
+                                               builder.getI32IntegerAttr(0));
+  for (auto en : llvm::enumerate(arguments)) {
+    auto index = builder.create<LLVM::ConstantOp>(
+        loc, llvmInt32Type, builder.getI32IntegerAttr(en.index()));
+    auto fieldPtr =
+        builder.create<LLVM::GEPOp>(loc, structType.getPointerTo(), structPtr,
+                                    ArrayRef<Value>{zero, index.getResult()});
+    builder.create<LLVM::StoreOp>(loc, en.value(), fieldPtr);
+    auto elementPtr = builder.create<LLVM::GEPOp>(loc, llvmPointerPointerType,
+                                                  arrayPtr, index.getResult());
+    auto casted =
+        builder.create<LLVM::BitcastOp>(loc, llvmPointerType, fieldPtr);
+    builder.create<LLVM::StoreOp>(loc, casted, elementPtr);
   }
-
-  return array;
+  return arrayPtr;
 }
 
 // Generates an LLVM IR dialect global that contains the name of the given
 // kernel function as a C string, and returns a pointer to its beginning.
 // The code is essentially:
 //
 // llvm.global constant @kernel_name("function_name\00")
 // func(...) {

 // %3 = call %moduleGetFunction(%1, %2)
 // %4 = call %streamCreate()
 // %5 = <see generateParamsArray>
 // call %launchKernel(%3, <launchOp operands 0..5>, 0, %4, %5, nullptr)
 // call %streamSynchronize(%4)
 LogicalResult ConvertLaunchFuncOpToGpuRuntimeCallPattern::matchAndRewrite(
     Operation *op, ArrayRef<Value> operands,
     ConversionPatternRewriter &rewriter) const {
-  Location loc = op->getLoc();
+  if (!llvm::all_of(operands, [](Value operand) {
+        return operand.getType().isa<LLVM::LLVMType>();
+      }))
+    return failure(); // Cannot convert if operands aren't of LLVM type.
+
   auto launchOp = cast<gpu::LaunchFuncOp>(op);
-  auto moduleOp = op->getParentOfType<ModuleOp>();
+  Location loc = launchOp.getLoc();
 
   // Create an LLVM global with CUBIN extracted from the kernel annotation and
   // obtain a pointer to the first byte in it.
   auto kernelModule =
-      moduleOp.lookupSymbol<gpu::GPUModuleOp>(launchOp.getKernelModuleName());
+      launchOp.getParentOfType<ModuleOp>().lookupSymbol<gpu::GPUModuleOp>(
+          launchOp.getKernelModuleName());
   assert(kernelModule && "expected a kernel module");
 
   auto binaryAttr = kernelModule.getAttrOfType<StringAttr>(gpuBinaryAnnotation);
   if (!binaryAttr) {
     kernelModule.emitOpError()
         << "missing " << gpuBinaryAnnotation << " attribute";
     return failure();
   }
 
   SmallString<128> nameBuffer(kernelModule.getName());
   nameBuffer.append(kGpuBinaryStorageSuffix);
   Value data = LLVM::createGlobalString(loc, rewriter, nameBuffer.str(),
                                         binaryAttr.getValue(),
                                         LLVM::Linkage::Internal, llvmDialect);
 
   auto module = moduleLoadCallBuilder.create(loc, rewriter, data);
   // Get the function from the module. The name corresponds to the name of
   // the kernel function.
   auto kernelName = generateKernelNameConstant(
       launchOp.getKernelModuleName(), launchOp.getKernelName(), loc, rewriter);
   auto function = moduleGetFunctionCallBuilder.create(
       loc, rewriter, {module.getResult(0), kernelName});
-  // Grab the global stream needed for execution.
-  auto stream = streamCreateCallBuilder.create(loc, rewriter, {});
-
-  // Get the launch target.
-  auto gpuFuncOp = SymbolTable::lookupNearestSymbolFrom<LLVM::LLVMFuncOp>(
-      launchOp, launchOp.kernel());
-  if (!gpuFuncOp) {
-    launchOp.emitOpError() << "corresponding kernel function not found";
-    return failure();
-  }
-  // Build array of kernel parameters.
-  auto kernelParams =
-      generateParamsArray(launchOp, gpuFuncOp.getNumArguments(), rewriter);
-
-  // Invoke the function with required arguments.
   auto zero = rewriter.create<LLVM::ConstantOp>(loc, llvmInt32Type,
                                                 rewriter.getI32IntegerAttr(0));
-  auto nullpointer =
-      rewriter.create<LLVM::IntToPtrOp>(loc, llvmPointerPointerType, zero);
+  // Grab the global stream needed for execution.
+  auto stream = streamCreateCallBuilder.create(loc, rewriter, {});
+  // Create array of pointers to kernel arguments.
+  auto kernelParams = generateParamsArray(launchOp, operands, rewriter);
+  auto nullpointer = rewriter.create<LLVM::NullOp>(loc, llvmPointerPointerType);
   launchKernelCallBuilder.create(
       loc, rewriter,
       {function.getResult(0), launchOp.gridSizeX(), launchOp.gridSizeY(),
        launchOp.gridSizeZ(), launchOp.blockSizeX(), launchOp.blockSizeY(),
        launchOp.blockSizeZ(), zero, /* sharedMemBytes */
        stream.getResult(0),         /* stream */
        kernelParams,                /* kernel params */
        nullpointer /* extra */});
   streamSynchronizeCallBuilder.create(loc, rewriter, stream.getResult(0));
 
   rewriter.eraseOp(op);
   return success();
 }
 
 std::unique_ptr<mlir::OperationPass<mlir::ModuleOp>>
-mlir::createConvertGpuLaunchFuncToGpuRuntimeCallsPass(
-    StringRef gpuBinaryAnnotation) {
-  return std::make_unique<GpuLaunchFuncToGpuRuntimeCallsPass>(
-      gpuBinaryAnnotation);
+mlir::createGpuToLLVMConversionPass(StringRef gpuBinaryAnnotation) {
+  return std::make_unique<GpuToLLVMConversionPass>(gpuBinaryAnnotation);
 }
 
 void mlir::populateGpuToLLVMConversionPatterns(
     LLVMTypeConverter &converter, OwningRewritePatternList &patterns,
     StringRef gpuBinaryAnnotation) {
   patterns.insert<ConvertLaunchFuncOpToGpuRuntimeCallPattern>(
       converter, gpuBinaryAnnotation);
   patterns.insert<EraseGpuModuleOpPattern>(&converter.getContext());
 }

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

-// RUN: mlir-opt -allow-unregistered-dialect %s --launch-func-to-gpu-runtime="gpu-binary-annotation=nvvm.cubin" | FileCheck %s
-// RUN: mlir-opt -allow-unregistered-dialect %s --launch-func-to-gpu-runtime="gpu-binary-annotation=rocdl.hsaco" | FileCheck %s --check-prefix=ROCDL
+// RUN: mlir-opt -allow-unregistered-dialect %s --gpu-to-llvm="gpu-binary-annotation=nvvm.cubin" | FileCheck %s
+// RUN: mlir-opt -allow-unregistered-dialect %s --gpu-to-llvm="gpu-binary-annotation=rocdl.hsaco" | FileCheck %s --check-prefix=ROCDL
 
 module attributes {gpu.container_module} {
 
   // CHECK: llvm.mlir.global internal constant @[[kernel_name:.*]]("kernel\00")
   // CHECK: llvm.mlir.global internal constant @[[global:.*]]("CUBIN")
   // ROCDL: llvm.mlir.global internal constant @[[global:.*]]("HSACO")
 
   gpu.module @kernel_module attributes {nvvm.cubin = "CUBIN", rocdl.hsaco = "HSACO"} {

mlir/tools/mlir-cuda-runner/mlir-cuda-runner.cpp

static LogicalResult runMLIRPasses(ModuleOp m) {
   const char gpuBinaryAnnotation[] = "nvvm.cubin";
   pm.addPass(createGpuKernelOutliningPass());
   auto &kernelPm = pm.nest<gpu::GPUModuleOp>();
   kernelPm.addPass(createStripDebugInfoPass());
   kernelPm.addPass(createLowerGpuOpsToNVVMOpsPass());
   kernelPm.addPass(createConvertGPUKernelToBlobPass(
       translateModuleToNVVMIR, compilePtxToCubin, "nvptx64-nvidia-cuda",
       "sm_35", "+ptx60", gpuBinaryAnnotation));
-  pm.addPass(createLowerToLLVMPass());
-  pm.addPass(
-      createConvertGpuLaunchFuncToGpuRuntimeCallsPass(gpuBinaryAnnotation));
+  pm.addPass(createGpuToLLVMConversionPass(gpuBinaryAnnotation));
 
   return pm.run(m);
 }
 
 int main(int argc, char **argv) {
   registerPassManagerCLOptions();
   mlir::registerAllDialects();
   llvm::InitLLVM y(argc, argv);

mlir/tools/mlir-rocm-runner/mlir-rocm-runner.cpp

static LogicalResult runMLIRPasses(ModuleOp m) {
   const char gpuBinaryAnnotation[] = "rocdl.hsaco";
   pm.addPass(createGpuKernelOutliningPass());
   auto &kernelPm = pm.nest<gpu::GPUModuleOp>();
   kernelPm.addPass(createStripDebugInfoPass());
   kernelPm.addPass(createLowerGpuOpsToROCDLOpsPass());
   kernelPm.addPass(createConvertGPUKernelToBlobPass(
       compileModuleToROCDLIR, compileISAToHsaco, tripleName, targetChip,
       features, gpuBinaryAnnotation));
-  pm.addPass(createLowerToLLVMPass());
-  pm.addPass(
-      createConvertGpuLaunchFuncToGpuRuntimeCallsPass(gpuBinaryAnnotation));
+  pm.addPass(createGpuToLLVMConversionPass(gpuBinaryAnnotation));
 
   return pm.run(m);
 }
 
 int main(int argc, char **argv) {
   registerPassManagerCLOptions();
   mlir::registerAllDialects();
   llvm::InitLLVM y(argc, argv);