[mlir][gpu] Add the `gpu-module-to-binary` pass.

[fmorac]

Summary

For an explanation of these patches see D154153.

Commit message:
This pass converts GPU modules into GPU binaries, serializing all targets present
in a GPU module by invoking the serializeToObject target attribute method.

Depends on D154147

[fmorac]

Rebasing.

[krzysz00]

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

66

Are these new #define 's I missed in the patch series?

[fmorac]

Rebasing.

[mehdi_amini]

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

74

Same as other patches: can this specify the interface instead of "Attribute"?

mlir/include/mlir/Dialect/GPU/Transforms/Passes.td

41

The restriction on ModuleOp does not seem necessary to me.

I would let it be an OperationPass and specify that this finds all the GPUModule op immediately present in the attached regions and convert them to GPUBinaryOp holding a serialization of the result of the codegen of the GPUModule based on its target information.

43

It's not useful to repeat the summary, please flesh this out.

44

If the pass runs on a GPUModuleOp present in its input, then this shouldn't be needed.

48

Document please. Do you have tests for it as well?

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

97

Nit: dyn_cast without check for nullptr

121

Pattern aren't obviously appropriate to me.

I would do something a bit more lightweight:

void GpuModuleToBinaryPass::runOnOperation() {
  for (Region &region : getOperation()->getRegion()) {
    for (Block &block : region.getBlocks()) {
      for (auto gpuModuleOP : make_early_inc(block.getOperations<GPUModuleOp>()))
          process(gpuModuleOP);
    }
  }

If we need a public API, it can be one that takes a GPUModuleOp and an OpBuilder and builds a GPUBinaryOp.

mlir/test/Dialect/GPU/module-to-binary.mlir

2

there is no diagnostics to verify here?

[fmorac]

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

74

I'll change it.

mlir/include/mlir/Dialect/GPU/Transforms/Passes.td

41

I'll change it.

43

I'll update it.

44

I'll remove it.

48

I'll add better documentation.

Tests: Currently no, I think I need to add a method for parsing attributes as options for this to work as a cmd option, it was added here for having the option on C++. I can create a patch for that (cmd parsing of attributes) and the modify this.

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

66

These are introduced in the CMakeList.txt, I added them to not reuse the ones from the existing pipeline.

97

I'll change it.

121

Yeah, the reason for doing patterns was just to expose it to be used in other passes. But I can change it,

mlir/test/Dialect/GPU/module-to-binary.mlir

2

You're right, I'll remove it.

[fmorac]

Updated tests, changed the Pass to an OperationPass and switched from patterns to looking for nested Modules.

[mehdi_amini]

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

97

Still present

[fmorac]

Added an assert check for the attribute.
Note: The verifier will also verify that the dyn_cast is valid, as it's a precondition on the elements of the array.

[fmorac]

Separated gpu-module-to-binary test into 2 tests: nvvm & rocdl.
Also updated lit.cfg.py with:

if config.run_cuda_tests:
    config.available_features.add("host-supports-nvptx")

if config.run_rocm_tests:
    config.available_features.add("host-supports-amdgpu")

Allowing the usage of // REQUIRES: host-supports-nvptx for disabling tests.

[fmorac]

Fix patch application.

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

 /// Collect all patterns to rewrite ops within the GPU dialect.
 inline void populateGpuRewritePatterns(RewritePatternSet &patterns) {
   populateGpuAllReducePatterns(patterns);
   populateGpuGlobalIdPatterns(patterns);
   populateGpuShufflePatterns(patterns);
 }
 
 namespace gpu {
+/// Searches for all GPU modules in `op` and transforms them into GPU binary
+/// operations. The resulting `gpu.binary` has `handler` as its offloading
+/// handler attribute.
+LogicalResult transformGpuModulesToBinaries(
+    Operation *op, OffloadingLLVMTranslationAttrInterface handler = nullptr,
+    const gpu::TargetOptions &options = {});
+
 /// 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;

mlir/include/mlir/Dialect/GPU/Transforms/Passes.td

 def GpuMapParallelLoopsPass
     : Pass<"gpu-map-parallel-loops", "mlir::func::FuncOp"> {
   let summary = "Greedily maps loops to GPU hardware dimensions.";
   let constructor = "mlir::createGpuMapParallelLoopsPass()";
   let description = "Greedily maps loops to GPU hardware dimensions.";
   let dependentDialects = ["mlir::gpu::GPUDialect"];
 }
 
+def GpuModuleToBinaryPass
+    : Pass<"gpu-module-to-binary", ""> {
+  let summary = "Transforms a GPU module into a GPU binary.";
+  let options = [
+    Option<"offloadingHandler", "handler", "Attribute", "nullptr",
+           "Offloading handler to be attached to the resulting binary op.">,
+    Option<"toolkitPath", "toolkit", "std::string", [{""}],
+           "Toolkit path.">,
+    ListOption<"linkFiles", "l", "std::string",
+           "Extra files to link to.">,
+    Option<"cmdOptions", "opts", "std::string", [{""}],
+           "Command line options to pass to the tools.">,
+    Option<"compilationTarget", "format", "std::string", [{"bin"}],
+           "The target representation of the compilation process.">
+  ];
+}
+
 #endif // MLIR_DIALECT_GPU_PASSES

mlir/lib/Dialect/GPU/CMakeLists.txt

add_mlir_dialect_library(MLIRGPUDialect
   )
 
 add_mlir_dialect_library(MLIRGPUTransforms
   Transforms/AllReduceLowering.cpp
   Transforms/AsyncRegionRewriter.cpp
   Transforms/GlobalIdRewriter.cpp
   Transforms/KernelOutlining.cpp
   Transforms/MemoryPromotion.cpp
+  Transforms/ModuleToBinary.cpp
   Transforms/ParallelLoopMapper.cpp
   Transforms/ShuffleRewriter.cpp
   Transforms/SerializeToBlob.cpp
   Transforms/SerializeToCubin.cpp
   Transforms/SerializeToHsaco.cpp
 
   ADDITIONAL_HEADER_DIRS
   ${MLIR_MAIN_INCLUDE_DIR}/mlir/Dialect/GPU

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

+//===- ModuleToBinary.cpp - Transforms GPU modules to GPU binaries ----------=//
+//
+// 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 the `GpuModuleToBinaryPass` pass, transforming GPU
+// modules into GPU binaries.
+//
+//===----------------------------------------------------------------------===//
+
+#include "mlir/Dialect/GPU/Transforms/Passes.h"
+
+#include "mlir/Dialect/Func/IR/FuncOps.h"
+#include "mlir/Dialect/GPU/IR/GPUDialect.h"
+#include "mlir/IR/BuiltinOps.h"
+#include "mlir/Target/LLVM/NVVM/Target.h"
+#include "mlir/Target/LLVM/ROCDL/Target.h"
+#include "mlir/Target/LLVMIR/Dialect/GPU/GPUToLLVMIRTranslation.h"
+#include "mlir/Target/LLVMIR/Dialect/LLVMIR/LLVMToLLVMIRTranslation.h"
+#include "mlir/Transforms/GreedyPatternRewriteDriver.h"
+
+#include "llvm/ADT/STLExtras.h"
+#include "llvm/ADT/StringSwitch.h"
+
+using namespace mlir;
+using namespace mlir::gpu;
+
+namespace mlir {
+#define GEN_PASS_DEF_GPUMODULETOBINARYPASS
+#include "mlir/Dialect/GPU/Transforms/Passes.h.inc"
+} // namespace mlir
+
+namespace {
+class GpuModuleToBinaryPass
+    : public impl::GpuModuleToBinaryPassBase<GpuModuleToBinaryPass> {
+public:
+  using Base::Base;
+  void getDependentDialects(DialectRegistry &registry) const override;
+  void runOnOperation() final;
+};
+} // namespace
+
+void GpuModuleToBinaryPass::getDependentDialects(
+    DialectRegistry &registry) const {
+  // Register all GPU related translations.
+  registerLLVMDialectTranslation(registry);
+  registerGPUDialectTranslation(registry);
+#if MLIR_CUDA_CONVERSIONS_ENABLED == 1
+  registerNVVMTarget(registry);
+#endif
+#if MLIR_ROCM_CONVERSIONS_ENABLED == 1
+  registerROCDLTarget(registry);
+#endif
+}
+
+void GpuModuleToBinaryPass::runOnOperation() {
+  RewritePatternSet patterns(&getContext());
+  int targetFormat = llvm::StringSwitch<int>(compilationTarget)
+                         .Cases("offloading", "llvm", TargetOptions::offload)
+                         .Cases("assembly", "isa", TargetOptions::assembly)
+                         .Cases("binary", "bin", TargetOptions::binary)
+                         .Default(-1);
+  if (targetFormat == -1)
+    getOperation()->emitError() << "Invalid format specified.";
+  TargetOptions targetOptions(
+      toolkitPath, linkFiles, cmdOptions,
+      static_cast<TargetOptions::CompilationTarget>(targetFormat));
+  if (failed(transformGpuModulesToBinaries(
+          getOperation(),
+          offloadingHandler ? dyn_cast<OffloadingLLVMTranslationAttrInterface>(
+                                  offloadingHandler.getValue())
+                            : OffloadingLLVMTranslationAttrInterface(nullptr),
+          targetOptions)))
+    return signalPassFailure();
+}
+
+namespace {
+LogicalResult moduleSerializer(GPUModuleOp op,
+                               OffloadingLLVMTranslationAttrInterface handler,
+                               const TargetOptions &targetOptions) {
+  OpBuilder builder(op->getContext());
+  SmallVector<Attribute> objects;
+  // Serialize all targets.
+  for (auto targetAttr : op.getTargetsAttr()) {
+    assert(targetAttr && "Target attribute cannot be null.");
+    auto target = dyn_cast<gpu::TargetAttrInterface>(targetAttr);
+    std::optional<SmallVector<char, 0>> object =
+        target.serializeToObject(op, targetOptions);
+
+    if (!object) {
+      op.emitError("An error happened while serializing the module.");
+      return failure();
+    }
+
+    objects.push_back(builder.getAttr<gpu::ObjectAttr>(
+        target,
+        builder.getStringAttr(StringRef(object->data(), object->size()))));
+  }
+  builder.setInsertionPointAfter(op);
+  builder.create<gpu::BinaryOp>(op.getLoc(), op.getName(), handler,
+                                builder.getArrayAttr(objects));
+  op->erase();
+  return success();
+}
+} // namespace
+
+LogicalResult mlir::gpu::transformGpuModulesToBinaries(
+    Operation *op, OffloadingLLVMTranslationAttrInterface handler,
+    const gpu::TargetOptions &targetOptions) {
+  for (Region &region : op->getRegions())
+    for (Block &block : region.getBlocks())
+      for (auto module :
+           llvm::make_early_inc_range(block.getOps<GPUModuleOp>()))
+        if (failed(moduleSerializer(module, handler, targetOptions)))
+          return failure();
+  return success();
+}

mlir/test/Dialect/GPU/module-to-binary.mlir

+// RUN: mlir-opt %s --gpu-module-to-binary="format=llvm" | FileCheck %s
+// RUN: mlir-opt %s --gpu-module-to-binary="format=isa" | FileCheck %s
+
+module attributes {gpu.container_module} {
+  // CHECK-LABEL:gpu.binary @kernel_module1
+  // CHECK:[#gpu.object<#nvvm.target<chip = "sm_70">, "{{.*}}">]
+  gpu.module @kernel_module1 [#nvvm.target<chip = "sm_70">] {
+    llvm.func @kernel(%arg0: i32, %arg1: !llvm.ptr<f32>,
+        %arg2: !llvm.ptr<f32>, %arg3: i64, %arg4: i64,
+        %arg5: i64) attributes {gpu.kernel} {
+      llvm.return
+    }
+  }
+
+  // CHECK-LABEL:gpu.binary @kernel_module2
+  // CHECK:[#gpu.object<#nvvm.target<flags = {fast}>, "{{.*}}">, #gpu.object<#nvvm.target, "{{.*}}">]
+  gpu.module @kernel_module2 [#nvvm.target<flags = {fast}>, #nvvm.target] {
+    llvm.func @kernel(%arg0: i32, %arg1: !llvm.ptr<f32>,
+        %arg2: !llvm.ptr<f32>, %arg3: i64, %arg4: i64,
+        %arg5: i64) attributes {gpu.kernel} {
+      llvm.return
+    }
+  }
+}