[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

1 ↗

(On Diff #543194)

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

1 ↗

(On Diff #543194)

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 GpuDecomposeMemrefsPass : Pass<"gpu-decompose-memrefs"> {
   let summary = "Decomposes memref index computation into explicit ops.";
   let description = [{
     This pass decomposes memref index computation into explicit computations on
     sizes/strides, obtained from `memref.extract_memref_metadata` which it tries
     to place outside of `gpu.launch` body. Memrefs are then reconstructed using
     `memref.reinterpret_cast`.
     This is needed for as some targets (SPIR-V) lower memrefs to bare pointers
     and sizes/strides for dynamically-sized memrefs are not available inside
     `gpu.launch`.
   }];
   let constructor = "mlir::createGpuDecomposeMemrefsPass()";
   let dependentDialects = [
     "mlir::gpu::GPUDialect", "mlir::memref::MemRefDialect",
     "mlir::affine::AffineDialect"
   ];
 }
 
+def GpuModuleToBinaryPass
+    : Pass<"gpu-module-to-binary", ""> {
+  let summary = "Transforms a GPU module into a GPU binary.";
+  let description = [{
+    This pass searches for all nested GPU modules and serializes the module
+    using the target attributes attached to the module, producing a GPU binary
+    with an object for every target.
+
+    The `format` argument can have the following values:
+    1. `offloading`, `llvm`: producing an offloading representation.
+    2. `assembly`, `isa`: producing assembly code.
+    3. `binary`, `bin`: producing binaries.
+  }];
+  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(MLIRGPUTransforms
   Transforms/AllReduceLowering.cpp
   Transforms/AsyncRegionRewriter.cpp
   Transforms/DecomposeMemrefs.cpp
   Transforms/GlobalIdRewriter.cpp
   Transforms/KernelOutlining.cpp
   Transforms/MemoryPromotion.cpp
+  Transforms/ModuleToBinary.cpp
   Transforms/ParallelLoopMapper.cpp
   Transforms/SerializeToBlob.cpp
   Transforms/SerializeToCubin.cpp
   Transforms/SerializeToHsaco.cpp
   Transforms/ShuffleRewriter.cpp
 
   ADDITIONAL_HEADER_DIRS
   ${MLIR_MAIN_INCLUDE_DIR}/mlir/Dialect/GPU
add_mlir_dialect_library(MLIRGPUTransforms
   MLIRExecutionEngineUtils
   MLIRGPUDialect
   MLIRIR
   MLIRIndexDialect
   MLIRLLVMDialect
   MLIRGPUToLLVMIRTranslation
   MLIRLLVMToLLVMIRTranslation
   MLIRMemRefDialect
+  MLIRNVVMTarget
   MLIRPass
   MLIRSCFDialect
   MLIRSideEffectInterfaces
   MLIRSupport
+  MLIRROCDLTarget
   MLIRTransformUtils
   )
 
 add_subdirectory(TransformOps)
 
 if(MLIR_ENABLE_CUDA_RUNNER)
   if(NOT MLIR_ENABLE_CUDA_CONVERSIONS)
     message(SEND_ERROR

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);
+    assert(target &&
+           "Target attribute doesn't implements `TargetAttrInterface`.");
+    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-nvvm.mlir

+// REQUIRES: host-supports-nvptx
+// 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
+    }
+  }
+}

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

+// REQUIRES: host-supports-amdgpu
+// 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<#rocdl.target<chip = "gfx90a">, "{{.*}}">]
+  gpu.module @kernel_module1 [#rocdl.target<chip = "gfx90a">] {
+    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<#rocdl.target<flags = {fast}>, "{{.*}}">, #gpu.object<#rocdl.target, "{{.*}}">]
+  gpu.module @kernel_module2 [#rocdl.target<flags = {fast}>, #rocdl.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
+    }
+  }
+}

mlir/test/lit.cfg.py

def have_host_jit_feature_support(feature_name):
     mlir_cpu_runner_out = mlir_cpu_runner_cmd.stdout.read().decode("ascii")
     mlir_cpu_runner_cmd.wait()
 
     return "true" in mlir_cpu_runner_out
 
 
 if have_host_jit_feature_support("jit"):
     config.available_features.add("host-supports-jit")
+
+if config.run_cuda_tests:
+    config.available_features.add("host-supports-nvptx")
+
+if config.run_rocm_tests:
+    config.available_features.add("host-supports-amdgpu")