[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

96

Nit: dyn_cast without check for nullptr

120

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

1

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.

96

I'll change it.

120

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

1

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

96

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);
 }
 
+/// Collect a set of patterns to rewrite GPU modules into GPU binary operations.
+void populateGpuModuleToBinaryPatterns(RewritePatternSet &patterns,
+                                       Attribute objectManager = nullptr,
+                                       const gpu::TargetOptions &options = {});
+
 namespace gpu {
 /// 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);
 

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", "mlir::ModuleOp"> {
+  let summary = "Transforms a GPU module into a GPU binary.";
+  let description = "Transforms a GPU module into a GPU binary.";
+  let dependentDialects = ["mlir::gpu::GPUDialect"];
+  let options = [
+    Option<"objectManager", "object-manager", "Attribute",
+           /*default=*/"nullptr",
+           "Object manager attribute.">,
+    Option<"toolkitPath", "toolkit", "std::string",
+           /*default=*/"\"\"",
+           "Toolkit path.">,
+    ListOption<"bitcodeFiles", "l", "std::string",
+           "Extra bitcode files to link to.">,
+  ];
+}
+
 #endif // MLIR_DIALECT_GPU_PASSES

mlir/lib/Dialect/GPU/CMakeLists.txt

add_mlir_dialect_library(MLIRGPUTargets
   )
 
 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
add_mlir_dialect_library(MLIRGPUTransforms
   MLIRGPUToLLVMIRTranslation
   MLIRLLVMToLLVMIRTranslation
   MLIRMemRefDialect
   MLIRPass
   MLIRSCFDialect
   MLIRSideEffectInterfaces
   MLIRSupport
   MLIRTransformUtils
+
+  PRIVATE
+  MLIRGPUTargets
   )
 
 add_subdirectory(TransformOps)
 
 if(MLIR_ENABLE_CUDA_RUNNER)
   if(NOT MLIR_ENABLE_CUDA_CONVERSIONS)
     message(SEND_ERROR
       "Building mlir with cuda support requires the NVPTX backend")

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/LLVMIR/Dialect/GPU/GPUToLLVMIRTranslation.h"
+#include "mlir/Target/LLVMIR/Dialect/LLVMIR/LLVMToLLVMIRTranslation.h"
+#include "mlir/Target/LLVMIR/Dialect/NVVM/NVVMToLLVMIRTranslation.h"
+#include "mlir/Target/LLVMIR/Dialect/ROCDL/ROCDLToLLVMIRTranslation.h"
+#include "mlir/Transforms/GreedyPatternRewriteDriver.h"
+
+#include "llvm/ADT/STLExtras.h"
+
+using namespace mlir;
+
+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;
+};
+
+// Rewriter for transforming GPU modules to GPU binaries.
+class GPUModuleToBinaryRewriter : public OpRewritePattern<gpu::GPUModuleOp> {
+public:
+  GPUModuleToBinaryRewriter(Attribute objectManager,
+                            const gpu::TargetOptions &targetOptions,
+                            MLIRContext *context, PatternBenefit benefit = 1,
+                            ArrayRef<StringRef> generatedNames = {});
+
+  LogicalResult matchAndRewrite(gpu::GPUModuleOp op,
+                                PatternRewriter &rewriter) const override;
+
+private:
+  Attribute objectManager;
+  const gpu::TargetOptions &targetOptions;
+};
+} // namespace
+
+void GpuModuleToBinaryPass::getDependentDialects(
+    DialectRegistry &registry) const {
+  // Register all GPU related translations.
+  registerLLVMDialectTranslation(registry);
+  registerGPUDialectTranslation(registry);
+#ifdef MLIR_GPU_NVPTX_TARGET_ENABLED
+  registerNVVMDialectTranslation(registry);
+#endif
+#ifdef MLIR_GPU_AMDGPU_TARGET_ENABLED
+  registerROCDLDialectTranslation(registry);
+#endif
+}
+
+void GpuModuleToBinaryPass::runOnOperation() {
+  RewritePatternSet patterns(&getContext());
+  gpu::TargetOptions targetOptions(toolkitPath, bitcodeFiles);
+  populateGpuModuleToBinaryPatterns(patterns, objectManager, targetOptions);
+  FrozenRewritePatternSet patternSet(std::move(patterns));
+  if (failed(applyPatternsAndFoldGreedily(getOperation(), patternSet)))
+    return signalPassFailure();
+}
+
+GPUModuleToBinaryRewriter::GPUModuleToBinaryRewriter(
+    Attribute objectManager, const gpu::TargetOptions &targetOptions,
+    MLIRContext *context, PatternBenefit benefit,
+    ArrayRef<StringRef> generatedNames)
+    : OpRewritePattern(context, benefit, generatedNames),
+      objectManager(objectManager), targetOptions(targetOptions) {}
+
+LogicalResult
+GPUModuleToBinaryRewriter::matchAndRewrite(gpu::GPUModuleOp op,
+                                           PatternRewriter &rewriter) const {
+  SmallVector<Attribute> objects;
+
+  // Serialize all targets.
+  for (auto targetAttr : op.getTargetsAttr()) {
+    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(rewriter.getAttr<gpu::ObjectAttr>(
+        targetAttr,
+        rewriter.getStringAttr(StringRef(object->data(), object->size()))));
+  }
+
+  rewriter.replaceOpWithNewOp<gpu::BinaryOp>(op, op.getName(), objectManager,
+                                             rewriter.getArrayAttr(objects));
+  return success();
+}
+
+void ::mlir::populateGpuModuleToBinaryPatterns(
+    RewritePatternSet &patterns, Attribute objectManager,
+    const gpu::TargetOptions &options) {
+  patterns.add<GPUModuleToBinaryRewriter>(objectManager, options,
+                                          patterns.getContext());
+}

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

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