[mlir] Add base class for GpuKernelToBlobPass

[csigg]

Summary

Instead of configuring kernel-to-cubin/rocdl lowering through callbacks, introduce a base class that target-specific passes can derive from.

Put the base class in GPU/Transforms, according to the discussion in D98203.

The mlir-cuda-runner will go away shortly, and the mlir-rocdl-runner as well at some point. I therefore kept the existing code path working and will remove it in a separate step.

Depends On D98168

[csigg]

Rebase

[csigg]

Add options for triple, chip, features.

[csigg]

Remove accidental mlir-cuda-runner change.

[herhut]

mlir/include/mlir/Dialect/GPU/Passes.h

49

This is not public API but rather an implementation detail. Could it move to an implementation header in lib?

mlir/lib/Dialect/GPU/Transforms/SerializeToBlob.cpp

12

Not sure this comment is correct. It translates the entire module, no?

33

I have never seen a pass copy options before. How is this normally handled?

[herhut]

After off-line discussion that this will be morphed into the final state in multiple changes, this looks like a good first step. Please fix the comment, though.

[rriddle]

mlir/lib/Conversion/GPUCommon/ConvertKernelFuncToBlob.cpp

76

This is copying the string, why not StringRef?

mlir/lib/Dialect/GPU/Transforms/SerializeToBlob.cpp

33

Pass options are copied implicitly as part of clone, i.e. you shouldn't need to do anything.

[csigg]

Fix file comment.

[csigg]

mlir/include/mlir/Dialect/GPU/Passes.h

49

At the moment it is public API because it's used outside.

Once we cleaned up the use in Conversion/GPUCommon/ConvertKernelFuncToBlob.cpp (which I will do in a separate CL), we can move it.

mlir/lib/Dialect/GPU/Transforms/SerializeToBlob.cpp

33

Pass::Option is currently missing a copy c'tor. I didn't want to delay this change (or make unrelated changes) and will clean it up separately.

[csigg]

Change isa string to reference.

[csigg]

mlir/lib/Conversion/GPUCommon/ConvertKernelFuncToBlob.cpp

76

Good point. I changed it to a string const-reference because it is used as null-terminated char pointer.

[csigg]

Reduce public #includes.

[csigg]

Move LLVM dependencies.

[csigg]

Fix.

[rriddle]

mlir/lib/Dialect/GPU/Transforms/SerializeToBlob.cpp

33

(Did you miss this comment?)

[csigg]

mlir/lib/Dialect/GPU/Transforms/SerializeToBlob.cpp

33

I did miss your comment, sorry. I will send a follow-up change.

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

 class Operation;
 class OwningRewritePatternList;
 
 template <typename T>
 class OperationPass;
 
 namespace gpu {
 class GPUModuleOp;
-
-/// Returns the default annotation name for GPU binary blobs.
-std::string getDefaultGpuBinaryAnnotation();
 } // namespace gpu
 
 namespace LLVM {
 class LLVMDialect;
 } // namespace LLVM
 
 using OwnedBlob = std::unique_ptr<std::vector<char>>;
 using BlobGenerator =

mlir/include/mlir/Dialect/GPU/Passes.h

 //===- Passes.h - Pass Entrypoints ------------------------------*- C++ -*-===//
 //
 // 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
 //
 //===----------------------------------------------------------------------===//
 //
 // This header file defines prototypes that expose pass constructors.
 //
 //===----------------------------------------------------------------------===//
 
 #ifndef MLIR_DIALECT_GPU_PASSES_H_
 #define MLIR_DIALECT_GPU_PASSES_H_
 
+#include "mlir/Dialect/GPU/GPUDialect.h"
 #include "mlir/Pass/Pass.h"
 
+namespace llvm {
+class TargetMachine;
+class LLVMContext;
+class Module;
+} // namespace llvm
+
 namespace mlir {
 /// Replaces `gpu.launch` with `gpu.launch_func` by moving the region into
 /// a separate kernel function.
 std::unique_ptr<OperationPass<ModuleOp>> createGpuKernelOutliningPass();
 
 /// Rewrites a function region so that GPU ops execute asynchronously.
 std::unique_ptr<OperationPass<FuncOp>> createGpuAsyncRegionPass();
 
 /// Collect a set of patterns to rewrite all-reduce ops within the GPU dialect.
 void populateGpuAllReducePatterns(MLIRContext *context,
                                   OwningRewritePatternList &patterns);
 
 /// Collect all patterns to rewrite ops within the GPU dialect.
 inline void populateGpuRewritePatterns(MLIRContext *context,
                                        OwningRewritePatternList &patterns) {
   populateGpuAllReducePatterns(context, patterns);
 }
 
+namespace gpu {
+/// Returns the default annotation name for GPU binary blobs.
+std::string getDefaultGpuBinaryAnnotation();
+
+/// Base pass class to serialize kernel functions through LLVM into
+/// user-specified IR and add the resulting blob as module attribute.
+class SerializeToBlobPass : public OperationPass<gpu::GPUModuleOp> {
+public:
+  SerializeToBlobPass(TypeID passID);
+  SerializeToBlobPass(const SerializeToBlobPass &other);
+
+  void runOnOperation() final;
+
+protected:
+  // Creates an LLVM target machine from given arguments.
+  static std::unique_ptr<llvm::TargetMachine>
+  createTargetMachine(Location loc, StringRef triple, StringRef chip,
+                      StringRef features);
+
+private:
+  // Creates the LLVM target machine to generate the ISA.
+  virtual std::unique_ptr<llvm::TargetMachine> createTargetMachine() = 0;
+
+  // Translates the 'getOperation()' result to an LLVM module.
+  virtual std::unique_ptr<llvm::Module>
+  translateToLLVMIR(llvm::LLVMContext &llvmContext) = 0;
+
+  // Serializes the target ISA to binary form.
+  virtual std::unique_ptr<std::vector<char>>
+  serializeISA(const std::string isa) = 0;
+
+protected:
+  Option<std::string> gpuBinaryAnnotation{
+      *this, "gpu-binary-annotation",
+      llvm::cl::desc("Annotation attribute string for GPU binary"),
+      llvm::cl::init(getDefaultGpuBinaryAnnotation())};
+};
+} // namespace gpu
+
 //===----------------------------------------------------------------------===//
 // Registration
 //===----------------------------------------------------------------------===//
 
 /// Generate the code for registering passes.
 #define GEN_PASS_REGISTRATION
 #include "mlir/Dialect/GPU/Passes.h.inc"
 
 } // namespace mlir
 
 #endif // MLIR_DIALECT_GPU_PASSES_H_

mlir/lib/Conversion/GPUCommon/ConvertKernelFuncToBlob.cpp

 // corresponding binary blob that can be executed on a GPU. Currently
 // only translates the function itself but no dependencies.
 //
 //===----------------------------------------------------------------------===//
 
 #include "mlir/Conversion/GPUCommon/GPUCommonPass.h"
 
 #include "mlir/Dialect/GPU/GPUDialect.h"
+#include "mlir/Dialect/GPU/Passes.h"
 #include "mlir/Dialect/LLVMIR/LLVMDialect.h"
 #include "mlir/IR/Attributes.h"
 #include "mlir/IR/Builders.h"
 #include "mlir/IR/BuiltinOps.h"
 #include "mlir/Pass/Pass.h"
 #include "mlir/Pass/PassRegistry.h"
 #include "mlir/Support/LogicalResult.h"
 

 
 /// A pass converting tagged kernel modules to a blob with target instructions.
 ///
 /// If tagged as a kernel module, each contained function is translated to
 /// user-specified IR. A user provided BlobGenerator then compiles the IR to
 /// GPU binary code, which is then attached as an attribute to the function.
 /// The function body is erased.
 class GpuKernelToBlobPass
-    : public PassWrapper<GpuKernelToBlobPass, OperationPass<gpu::GPUModuleOp>> {
+    : public PassWrapper<GpuKernelToBlobPass, gpu::SerializeToBlobPass> {
 public:
   GpuKernelToBlobPass(LoweringCallback loweringCallback,
                       BlobGenerator blobGenerator, StringRef triple,
                       StringRef targetChip, StringRef features,
                       StringRef gpuBinaryAnnotation)
       : loweringCallback(loweringCallback), blobGenerator(blobGenerator),
         triple(triple), targetChip(targetChip), features(features) {
     if (!gpuBinaryAnnotation.empty())
       this->gpuBinaryAnnotation = gpuBinaryAnnotation.str();
   }
 
-  GpuKernelToBlobPass(const GpuKernelToBlobPass &other)
-      : loweringCallback(other.loweringCallback),
-        blobGenerator(other.blobGenerator), triple(other.triple),
-        targetChip(other.targetChip), features(other.features) {}
-
-  void runOnOperation() override {
-    gpu::GPUModuleOp module = getOperation();
-
-    // Lower the module to an LLVM IR module using a separate context to enable
-    // multi-threaded processing.
-    llvm::LLVMContext llvmContext;
-    std::unique_ptr<llvm::Module> llvmModule =
-        loweringCallback(module, llvmContext, "LLVMDialectModule");
-    if (!llvmModule)
-      return signalPassFailure();
-
-    // Translate the llvm module to a target blob and attach the result as
-    // attribute to the module.
-    if (auto blobAttr = translateGPUModuleToBinaryAnnotation(
-            *llvmModule, module.getLoc(), module.getName()))
-      module->setAttr(gpuBinaryAnnotation, blobAttr);
-    else
-      signalPassFailure();
-  }
-
 private:
-  std::string translateModuleToISA(llvm::Module &module,
-                                   llvm::TargetMachine &targetMachine);
-
-  /// Converts llvmModule to a blob with target instructions using the
-  /// user-provided generator. Location is used for error reporting and name is
-  /// forwarded to the blob generator to use in its logging mechanisms.
-  OwnedBlob convertModuleToBlob(llvm::Module &llvmModule, Location loc,
-                                StringRef name);
-
-  /// Translates llvmModule to a blob with target instructions and returns the
-  /// result as attribute.
-  StringAttr translateGPUModuleToBinaryAnnotation(llvm::Module &llvmModule,
-                                                  Location loc, StringRef name);
+  // Creates the LLVM target machine to generate the ISA.
+  std::unique_ptr<llvm::TargetMachine> createTargetMachine() override {
+    return gpu::SerializeToBlobPass::createTargetMachine(
+        getOperation()->getLoc(), triple, targetChip, features);
+  }
+
+  // Translates the 'getOperation()' result to an LLVM module.
+  std::unique_ptr<llvm::Module>
+  translateToLLVMIR(llvm::LLVMContext &llvmContext) override {
+    return loweringCallback(getOperation(), llvmContext, "LLVMDialectModule");
+  }
+
+  // Serializes the target ISA to binary form.
+  std::unique_ptr<std::vector<char>>
+  serializeISA(const std::string isa) override {
+    return blobGenerator(isa, getOperation()->getLoc(),
+                         getOperation().getName());
+  }
 
   LoweringCallback loweringCallback;
   BlobGenerator blobGenerator;
 
-  llvm::Triple triple;
+  std::string triple;
   std::string targetChip;
   std::string features;
-
-  Option<std::string> gpuBinaryAnnotation{
-      *this, "gpu-binary-annotation",
-      llvm::cl::desc("Annotation attribute string for GPU binary"),
-      llvm::cl::init(gpu::getDefaultGpuBinaryAnnotation())};
 };
 
 } // anonymous namespace
 
-std::string gpu::getDefaultGpuBinaryAnnotation() { return "gpu.binary"; }
-
-std::string
-GpuKernelToBlobPass::translateModuleToISA(llvm::Module &module,
-                                          llvm::TargetMachine &targetMachine) {
-  std::string targetISA;
-  {
-    llvm::raw_string_ostream stream(targetISA);
-    llvm::buffer_ostream pstream(stream);
-    llvm::legacy::PassManager codegenPasses;
-    targetMachine.addPassesToEmitFile(codegenPasses, pstream, nullptr,
-                                      llvm::CGFT_AssemblyFile);
-    codegenPasses.run(module);
-  }
-
-  return targetISA;
-}
-
-OwnedBlob GpuKernelToBlobPass::convertModuleToBlob(llvm::Module &llvmModule,
-                                                   Location loc,
-                                                   StringRef name) {
-  std::unique_ptr<llvm::TargetMachine> targetMachine;
-  {
-    std::string error;
-    const llvm::Target *target =
-        llvm::TargetRegistry::lookupTarget("", triple, error);
-    if (target == nullptr) {
-      emitError(loc, "cannot initialize target triple");
-      return {};
-    }
-    targetMachine.reset(target->createTargetMachine(triple.str(), targetChip,
-                                                    features, {}, {}));
-    if (targetMachine == nullptr) {
-      emitError(loc, "cannot initialize target machine");
-      return {};
-    }
-  }
-
-  llvmModule.setDataLayout(targetMachine->createDataLayout());
-
-  auto targetISA = translateModuleToISA(llvmModule, *targetMachine);
-
-  return blobGenerator(targetISA, loc, name);
-}
-
-StringAttr GpuKernelToBlobPass::translateGPUModuleToBinaryAnnotation(
-    llvm::Module &llvmModule, Location loc, StringRef name) {
-  auto blob = convertModuleToBlob(llvmModule, loc, name);
-  if (!blob)
-    return {};
-  return StringAttr::get(loc->getContext(), {blob->data(), blob->size()});
-}
-
 std::unique_ptr<OperationPass<gpu::GPUModuleOp>>
 mlir::createConvertGPUKernelToBlobPass(LoweringCallback loweringCallback,
                                        BlobGenerator blobGenerator,
                                        StringRef triple, StringRef targetChip,
                                        StringRef features,
                                        StringRef gpuBinaryAnnotation) {
   return std::make_unique<GpuKernelToBlobPass>(loweringCallback, blobGenerator,
                                                triple, targetChip, features,
                                                gpuBinaryAnnotation);
 }

mlir/lib/Conversion/GPUCommon/ConvertLaunchFuncToRuntimeCalls.cpp

 
 #include "mlir/Conversion/GPUCommon/GPUCommonPass.h"
 
 #include "../PassDetail.h"
 #include "mlir/Conversion/AsyncToLLVM/AsyncToLLVM.h"
 #include "mlir/Conversion/StandardToLLVM/ConvertStandardToLLVM.h"
 #include "mlir/Dialect/Async/IR/Async.h"
 #include "mlir/Dialect/GPU/GPUDialect.h"
+#include "mlir/Dialect/GPU/Passes.h"
 #include "mlir/Dialect/LLVMIR/LLVMDialect.h"
 #include "mlir/IR/Attributes.h"
 #include "mlir/IR/Builders.h"
 #include "mlir/IR/BuiltinOps.h"
 #include "mlir/IR/BuiltinTypes.h"
 
 #include "llvm/ADT/STLExtras.h"
 #include "llvm/IR/DataLayout.h"

mlir/lib/Dialect/GPU/Transforms/SerializeToBlob.cpp

+//===- SerializeToBlob.cpp - MLIR GPU lowering pass -----------------------===//
+//
+// 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
+//
+//===----------------------------------------------------------------------===//
+//
+// This file implements a base class for a pass to convert gpu kernel functions
+// into a corresponding binary blob that can be executed on a GPU. Currently
+// only translates the function itself but no dependencies.
+//
+//===----------------------------------------------------------------------===//
+
+#include "mlir/Dialect/GPU/Passes.h"
+#include "mlir/Pass/Pass.h"
+#include "llvm/IR/LegacyPassManager.h"
+#include "llvm/Support/TargetRegistry.h"
+#include "llvm/Support/TargetSelect.h"
+#include "llvm/Target/TargetMachine.h"
+
+using namespace mlir;
+
+std::string gpu::getDefaultGpuBinaryAnnotation() { return "gpu.binary"; }
+
+gpu::SerializeToBlobPass::SerializeToBlobPass(TypeID passID)
+    : OperationPass<gpu::GPUModuleOp>(passID) {}
+
+gpu::SerializeToBlobPass::SerializeToBlobPass(const SerializeToBlobPass &other)
+    : OperationPass<gpu::GPUModuleOp>(other) {}
+
+static std::string translateToISA(llvm::Module &llvmModule,
+                                  llvm::TargetMachine &targetMachine) {
+  llvmModule.setDataLayout(targetMachine.createDataLayout());
+
+  std::string targetISA;
+  llvm::raw_string_ostream stream(targetISA);
+  llvm::buffer_ostream pstream(stream);
+  llvm::legacy::PassManager codegenPasses;
+  targetMachine.addPassesToEmitFile(codegenPasses, pstream, nullptr,
+                                    llvm::CGFT_AssemblyFile);
+  codegenPasses.run(llvmModule);
+  return targetISA;
+}
+
+void gpu::SerializeToBlobPass::runOnOperation() {
+  // Lower the module to an LLVM IR module using a separate context to enable
+  // multi-threaded processing.
+  llvm::LLVMContext llvmContext;
+  std::unique_ptr<llvm::Module> llvmModule = translateToLLVMIR(llvmContext);
+  if (!llvmModule)
+    return signalPassFailure();
+
+  // Lower the LLVM IR module to target ISA.
+  std::unique_ptr<llvm::TargetMachine> targetMachine = createTargetMachine();
+  if (!targetMachine)
+    return signalPassFailure();
+
+  std::string targetISA = translateToISA(*llvmModule, *targetMachine);
+
+  // Serialize the target ISA.
+  std::unique_ptr<std::vector<char>> blob = serializeISA(targetISA);
+  if (!blob)
+    return signalPassFailure();
+
+  // Add the blob as module attribute.
+  auto attr = StringAttr::get(&getContext(), {blob->data(), blob->size()});
+  getOperation()->setAttr(gpuBinaryAnnotation, attr);
+}
+
+std::unique_ptr<llvm::TargetMachine>
+gpu::SerializeToBlobPass::createTargetMachine(Location loc, StringRef triple,
+                                              StringRef chip,
+                                              StringRef features) {
+  std::string error;
+  const llvm::Target *target =
+      llvm::TargetRegistry::lookupTarget(triple.str(), error);
+  if (!target) {
+    emitError(loc, Twine("failed to lookup target: ") + error);
+    return {};
+  }
+  llvm::TargetMachine *machine =
+      target->createTargetMachine(triple.str(), chip, features, {}, {});
+  if (!machine) {
+    emitError(loc, "failed to create target machine");
+    return {};
+  }
+
+  return std::unique_ptr<llvm::TargetMachine>{machine};
+}

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

OwnedBlob compilePtxToCubin(const std::string ptx, Location loc,
 
   return result;
 }
 
 struct GpuToCubinPipelineOptions
     : public mlir::PassPipelineOptions<GpuToCubinPipelineOptions> {
   Option<std::string> gpuBinaryAnnotation{
       *this, "gpu-binary-annotation",
-      llvm::cl::desc("Annotation attribute string for GPU binary"),
-      llvm::cl::init(gpu::getDefaultGpuBinaryAnnotation())};
+      llvm::cl::desc("Annotation attribute string for GPU binary")};
 };
 
 // Register cuda-runner specific passes.
 static void registerCudaRunnerPasses() {
   PassPipelineRegistration<GpuToCubinPipelineOptions> registerGpuToCubin(
       "gpu-to-cubin", "Generate CUBIN from gpu.launch regions",
       [&](OpPassManager &pm, const GpuToCubinPipelineOptions &options) {
         pm.addPass(createGpuKernelOutliningPass());