diff --git a/mlir/include/mlir/Dialect/GPU/Transforms/Passes.h b/mlir/include/mlir/Dialect/GPU/Transforms/Passes.h --- a/mlir/include/mlir/Dialect/GPU/Transforms/Passes.h +++ b/mlir/include/mlir/Dialect/GPU/Transforms/Passes.h @@ -69,6 +69,11 @@ 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. diff --git a/mlir/include/mlir/Dialect/GPU/Transforms/Passes.td b/mlir/include/mlir/Dialect/GPU/Transforms/Passes.td --- a/mlir/include/mlir/Dialect/GPU/Transforms/Passes.td +++ b/mlir/include/mlir/Dialect/GPU/Transforms/Passes.td @@ -37,4 +37,21 @@ 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 diff --git a/mlir/lib/Dialect/GPU/CMakeLists.txt b/mlir/lib/Dialect/GPU/CMakeLists.txt --- a/mlir/lib/Dialect/GPU/CMakeLists.txt +++ b/mlir/lib/Dialect/GPU/CMakeLists.txt @@ -79,6 +79,7 @@ Transforms/GlobalIdRewriter.cpp Transforms/KernelOutlining.cpp Transforms/MemoryPromotion.cpp + Transforms/ModuleToBinary.cpp Transforms/ParallelLoopMapper.cpp Transforms/ShuffleRewriter.cpp Transforms/SerializeToBlob.cpp @@ -118,6 +119,9 @@ MLIRSideEffectInterfaces MLIRSupport MLIRTransformUtils + + PRIVATE + MLIRGPUTargets ) add_subdirectory(TransformOps) diff --git a/mlir/lib/Dialect/GPU/Transforms/ModuleToBinary.cpp b/mlir/lib/Dialect/GPU/Transforms/ModuleToBinary.cpp new file mode 100644 --- /dev/null +++ b/mlir/lib/Dialect/GPU/Transforms/ModuleToBinary.cpp @@ -0,0 +1,120 @@ +//===- 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 { +public: + using Base::Base; + void getDependentDialects(DialectRegistry ®istry) const override; + void runOnOperation() final; +}; + +// Rewriter for transforming GPU modules to GPU binaries. +class GPUModuleToBinaryRewriter : public OpRewritePattern { +public: + GPUModuleToBinaryRewriter(Attribute objectManager, + const gpu::TargetOptions &targetOptions, + MLIRContext *context, PatternBenefit benefit = 1, + ArrayRef generatedNames = {}); + + LogicalResult matchAndRewrite(gpu::GPUModuleOp op, + PatternRewriter &rewriter) const override; + +private: + Attribute objectManager; + const gpu::TargetOptions &targetOptions; +}; +} // namespace + +void GpuModuleToBinaryPass::getDependentDialects( + DialectRegistry ®istry) 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 generatedNames) + : OpRewritePattern(context, benefit, generatedNames), + objectManager(objectManager), targetOptions(targetOptions) {} + +LogicalResult +GPUModuleToBinaryRewriter::matchAndRewrite(gpu::GPUModuleOp op, + PatternRewriter &rewriter) const { + SmallVector objects; + + // Serialize all targets. + for (auto targetAttr : op.getTargetsAttr()) { + auto target = dyn_cast(targetAttr); + std::optional> object = + target.serializeToObject(op, targetOptions); + + if (!object) { + op.emitError("An error happened while serializing the module."); + return failure(); + } + + objects.push_back(rewriter.getAttr( + targetAttr, + rewriter.getStringAttr(StringRef(object->data(), object->size())))); + } + + rewriter.replaceOpWithNewOp(op, op.getName(), objectManager, + rewriter.getArrayAttr(objects)); + return success(); +} + +void ::mlir::populateGpuModuleToBinaryPatterns( + RewritePatternSet &patterns, Attribute objectManager, + const gpu::TargetOptions &options) { + patterns.add(objectManager, options, + patterns.getContext()); +} diff --git a/mlir/test/Dialect/GPU/module-to-binary.mlir b/mlir/test/Dialect/GPU/module-to-binary.mlir new file mode 100644 --- /dev/null +++ b/mlir/test/Dialect/GPU/module-to-binary.mlir @@ -0,0 +1,23 @@ +// 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, "{{.*}}">] + gpu.module @kernel_module1 [#gpu.nvptx] { + llvm.func @kernel(%arg0: i32, %arg1: !llvm.ptr, + %arg2: !llvm.ptr, %arg3: i64, %arg4: i64, + %arg5: i64) attributes {gpu.kernel} { + llvm.return + } + } + + // CHECK-LABEL:gpu.binary @kernel_module2 + // CHECK:[#gpu.object<#gpu.nvptx, "{{.*}}">, #gpu.object<#gpu.nvptx, "{{.*}}">] + gpu.module @kernel_module2 [#gpu.nvptx, #gpu.nvptx] { + llvm.func @kernel(%arg0: i32, %arg1: !llvm.ptr, + %arg2: !llvm.ptr, %arg3: i64, %arg4: i64, + %arg5: i64) attributes {gpu.kernel} { + llvm.return + } + } +}