diff --git a/mlir/include/mlir/Dialect/GPU/IR/CompilationInterfaces.h b/mlir/include/mlir/Dialect/GPU/IR/CompilationInterfaces.h index d4b16a1de8edd..8cf8a7e55517d 100644 --- a/mlir/include/mlir/Dialect/GPU/IR/CompilationInterfaces.h +++ b/mlir/include/mlir/Dialect/GPU/IR/CompilationInterfaces.h @@ -51,7 +51,7 @@ class TargetOptions { /// `Fatbin`. TargetOptions( StringRef toolkitPath = {}, ArrayRef linkFiles = {}, - StringRef cmdOptions = {}, + StringRef cmdOptions = {}, StringRef elfSection = {}, CompilationTarget compilationTarget = getDefaultCompilationTarget(), function_ref getSymbolTableCallback = {}, function_ref initialLlvmIRCallback = {}, @@ -71,6 +71,9 @@ class TargetOptions { /// Returns the command line options. StringRef getCmdOptions() const; + /// Returns the ELF section. + StringRef getELFSection() const; + /// Returns a tokenization of the command line options. std::pair> tokenizeCmdOptions() const; @@ -110,6 +113,7 @@ class TargetOptions { TargetOptions( TypeID typeID, StringRef toolkitPath = {}, ArrayRef linkFiles = {}, StringRef cmdOptions = {}, + StringRef elfSection = {}, CompilationTarget compilationTarget = getDefaultCompilationTarget(), function_ref getSymbolTableCallback = {}, function_ref initialLlvmIRCallback = {}, @@ -127,6 +131,9 @@ class TargetOptions { /// process. std::string cmdOptions; + /// ELF Section where the binary needs to be located + std::string elfSection; + /// Compilation process target format. CompilationTarget compilationTarget; diff --git a/mlir/include/mlir/Dialect/GPU/Transforms/Passes.td b/mlir/include/mlir/Dialect/GPU/Transforms/Passes.td index 4a9ddafdd177d..e9e5dade84c98 100644 --- a/mlir/include/mlir/Dialect/GPU/Transforms/Passes.td +++ b/mlir/include/mlir/Dialect/GPU/Transforms/Passes.td @@ -87,16 +87,15 @@ def GpuModuleToBinaryPass 3. `binary`, `bin`: produces binaries. 4. `fatbinary`, `fatbin`: produces fatbinaries. }]; - let options = [ - 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", [{"fatbin"}], - "The target representation of the compilation process."> - ]; + let options = + [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", [{"fatbin"}], + "The target representation of the compilation process.">, + Option<"elfSection", "section", "std::string", [{""}], + "ELF section where binary is to be located.">]; } def GpuNVVMAttachTarget: Pass<"nvvm-attach-target", ""> { diff --git a/mlir/lib/Dialect/GPU/IR/GPUDialect.cpp b/mlir/lib/Dialect/GPU/IR/GPUDialect.cpp index ee00fbeb28b61..1fad251b2f79e 100644 --- a/mlir/lib/Dialect/GPU/IR/GPUDialect.cpp +++ b/mlir/lib/Dialect/GPU/IR/GPUDialect.cpp @@ -2484,27 +2484,31 @@ KernelMetadataAttr KernelTableAttr::lookup(StringAttr key) const { TargetOptions::TargetOptions( StringRef toolkitPath, ArrayRef linkFiles, - StringRef cmdOptions, CompilationTarget compilationTarget, + StringRef cmdOptions, StringRef elfSection, + CompilationTarget compilationTarget, function_ref getSymbolTableCallback, function_ref initialLlvmIRCallback, function_ref linkedLlvmIRCallback, function_ref optimizedLlvmIRCallback, function_ref isaCallback) : TargetOptions(TypeID::get(), toolkitPath, linkFiles, - cmdOptions, compilationTarget, getSymbolTableCallback, - initialLlvmIRCallback, linkedLlvmIRCallback, - optimizedLlvmIRCallback, isaCallback) {} + cmdOptions, elfSection, compilationTarget, + getSymbolTableCallback, initialLlvmIRCallback, + linkedLlvmIRCallback, optimizedLlvmIRCallback, + isaCallback) {} TargetOptions::TargetOptions( TypeID typeID, StringRef toolkitPath, ArrayRef linkFiles, - StringRef cmdOptions, CompilationTarget compilationTarget, + StringRef cmdOptions, StringRef elfSection, + CompilationTarget compilationTarget, function_ref getSymbolTableCallback, function_ref initialLlvmIRCallback, function_ref linkedLlvmIRCallback, function_ref optimizedLlvmIRCallback, function_ref isaCallback) : toolkitPath(toolkitPath.str()), linkFiles(linkFiles), - cmdOptions(cmdOptions.str()), compilationTarget(compilationTarget), + cmdOptions(cmdOptions.str()), elfSection(elfSection.str()), + compilationTarget(compilationTarget), getSymbolTableCallback(getSymbolTableCallback), initialLlvmIRCallback(initialLlvmIRCallback), linkedLlvmIRCallback(linkedLlvmIRCallback), @@ -2519,6 +2523,8 @@ ArrayRef TargetOptions::getLinkFiles() const { return linkFiles; } StringRef TargetOptions::getCmdOptions() const { return cmdOptions; } +StringRef TargetOptions::getELFSection() const { return elfSection; } + SymbolTable *TargetOptions::getSymbolTable() const { return getSymbolTableCallback ? getSymbolTableCallback() : nullptr; } diff --git a/mlir/lib/Dialect/GPU/Transforms/ModuleToBinary.cpp b/mlir/lib/Dialect/GPU/Transforms/ModuleToBinary.cpp index 86a3b4780e88c..295ece4782fdb 100644 --- a/mlir/lib/Dialect/GPU/Transforms/ModuleToBinary.cpp +++ b/mlir/lib/Dialect/GPU/Transforms/ModuleToBinary.cpp @@ -69,8 +69,8 @@ void GpuModuleToBinaryPass::runOnOperation() { return &parentTable.value(); }; - TargetOptions targetOptions(toolkitPath, linkFiles, cmdOptions, *targetFormat, - lazyTableBuilder); + TargetOptions targetOptions(toolkitPath, linkFiles, cmdOptions, elfSection, + *targetFormat, lazyTableBuilder); if (failed(transformGpuModulesToBinaries( getOperation(), OffloadingLLVMTranslationAttrInterface(nullptr), targetOptions))) diff --git a/mlir/lib/Target/LLVM/NVVM/Target.cpp b/mlir/lib/Target/LLVM/NVVM/Target.cpp index 3c92359915ded..e39c63d2b422b 100644 --- a/mlir/lib/Target/LLVM/NVVM/Target.cpp +++ b/mlir/lib/Target/LLVM/NVVM/Target.cpp @@ -664,9 +664,18 @@ NVVMTargetAttrImpl::createObject(Attribute attribute, Operation *module, gpu::CompilationTarget format = options.getCompilationTarget(); DictionaryAttr objectProps; Builder builder(attribute.getContext()); + SmallVector properties; if (format == gpu::CompilationTarget::Assembly) - objectProps = builder.getDictionaryAttr( - {builder.getNamedAttr("O", builder.getI32IntegerAttr(target.getO()))}); + properties.push_back( + builder.getNamedAttr("O", builder.getI32IntegerAttr(target.getO()))); + + if (auto section = options.getELFSection(); !section.empty()) + properties.push_back( + builder.getNamedAttr("section", builder.getStringAttr(section))); + + if (!properties.empty()) + objectProps = builder.getDictionaryAttr(properties); + return builder.getAttr( attribute, format, builder.getStringAttr(StringRef(object.data(), object.size())), diff --git a/mlir/test/Dialect/GPU/module-to-binary-nvvm.mlir b/mlir/test/Dialect/GPU/module-to-binary-nvvm.mlir index c286c8bc9042f..8b93e7ef983a6 100644 --- a/mlir/test/Dialect/GPU/module-to-binary-nvvm.mlir +++ b/mlir/test/Dialect/GPU/module-to-binary-nvvm.mlir @@ -1,10 +1,12 @@ // 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 -check-prefix=CHECK-ISA +// RUN: mlir-opt %s --gpu-module-to-binary="section=__fatbin" | FileCheck %s -check-prefix=CHECK-SECTION module attributes {gpu.container_module} { // CHECK-LABEL:gpu.binary @kernel_module1 // CHECK:[#gpu.object<#nvvm.target, offload = "{{.*}}">] + // CHECK-SECTION: #gpu.object<#nvvm.target, properties = {section = "__fatbin"} gpu.module @kernel_module1 [#nvvm.target] { llvm.func @kernel(%arg0: i32, %arg1: !llvm.ptr, %arg2: !llvm.ptr, %arg3: i64, %arg4: i64, @@ -25,6 +27,7 @@ module attributes {gpu.container_module} { // CHECK-LABEL:gpu.binary @kernel_module3 <#gpu.select_object<1 : i64>> // CHECK:[#gpu.object<#nvvm.target, offload = "{{.*}}">, #gpu.object<#nvvm.target, offload = "{{.*}}">] + // CHECK-SECTION: [#gpu.object<#nvvm.target, properties = {section = "__fatbin"},{{.*}} #gpu.object<#nvvm.target, properties = {section = "__fatbin"} gpu.module @kernel_module3 <#gpu.select_object<1>> [ #nvvm.target, #nvvm.target] { diff --git a/mlir/unittests/Target/LLVM/SerializeNVVMTarget.cpp b/mlir/unittests/Target/LLVM/SerializeNVVMTarget.cpp index eee9bd5f23475..a92ad18c95682 100644 --- a/mlir/unittests/Target/LLVM/SerializeNVVMTarget.cpp +++ b/mlir/unittests/Target/LLVM/SerializeNVVMTarget.cpp @@ -81,7 +81,7 @@ TEST_F(MLIRTargetLLVMNVVM, SKIP_WITHOUT_NVPTX(SerializeNVVMMToLLVM)) { // Serialize the module. auto serializer = dyn_cast(target); ASSERT_TRUE(!!serializer); - gpu::TargetOptions options("", {}, "", gpu::CompilationTarget::Offload); + gpu::TargetOptions options("", {}, "", "", gpu::CompilationTarget::Offload); for (auto gpuModule : (*module).getBody()->getOps()) { std::optional> object = serializer.serializeToObject(gpuModule, options); @@ -117,7 +117,7 @@ TEST_F(MLIRTargetLLVMNVVM, SKIP_WITHOUT_NVPTX(SerializeNVVMToPTX)) { // Serialize the module. auto serializer = dyn_cast(target); ASSERT_TRUE(!!serializer); - gpu::TargetOptions options("", {}, "", gpu::CompilationTarget::Assembly); + gpu::TargetOptions options("", {}, "", "", gpu::CompilationTarget::Assembly); for (auto gpuModule : (*module).getBody()->getOps()) { std::optional> object = serializer.serializeToObject(gpuModule, options); @@ -147,7 +147,7 @@ TEST_F(MLIRTargetLLVMNVVM, SKIP_WITHOUT_NVPTX(SerializeNVVMToBinary)) { // Serialize the module. auto serializer = dyn_cast(target); ASSERT_TRUE(!!serializer); - gpu::TargetOptions options("", {}, "", gpu::CompilationTarget::Binary); + gpu::TargetOptions options("", {}, "", "", gpu::CompilationTarget::Binary); for (auto gpuModule : (*module).getBody()->getOps()) { std::optional> object = serializer.serializeToObject(gpuModule, options); @@ -194,9 +194,9 @@ TEST_F(MLIRTargetLLVMNVVM, isaResult = isa.str(); }; - gpu::TargetOptions options({}, {}, {}, gpu::CompilationTarget::Assembly, {}, - initialCallback, linkedCallback, optimizedCallback, - isaCallback); + gpu::TargetOptions options({}, {}, {}, {}, gpu::CompilationTarget::Assembly, + {}, initialCallback, linkedCallback, + optimizedCallback, isaCallback); for (auto gpuModule : (*module).getBody()->getOps()) { std::optional> object = diff --git a/mlir/unittests/Target/LLVM/SerializeROCDLTarget.cpp b/mlir/unittests/Target/LLVM/SerializeROCDLTarget.cpp index b02f34c812b3f..d5e72a1713174 100644 --- a/mlir/unittests/Target/LLVM/SerializeROCDLTarget.cpp +++ b/mlir/unittests/Target/LLVM/SerializeROCDLTarget.cpp @@ -83,7 +83,7 @@ TEST_F(MLIRTargetLLVMROCDL, SKIP_WITHOUT_AMDGPU(SerializeROCDLToLLVM)) { // Serialize the module. auto serializer = dyn_cast(target); ASSERT_TRUE(!!serializer); - gpu::TargetOptions options("", {}, "", gpu::CompilationTarget::Offload); + gpu::TargetOptions options("", {}, "", "", gpu::CompilationTarget::Offload); for (auto gpuModule : (*module).getBody()->getOps()) { std::optional> object = serializer.serializeToObject(gpuModule, options); @@ -119,7 +119,7 @@ TEST_F(MLIRTargetLLVMROCDL, // Serialize the module. auto serializer = dyn_cast(target); ASSERT_TRUE(!!serializer); - gpu::TargetOptions options("", {}, "", gpu::CompilationTarget::Assembly); + gpu::TargetOptions options("", {}, "", "", gpu::CompilationTarget::Assembly); for (auto gpuModule : (*module).getBody()->getOps()) { std::optional> object = serializer.serializeToObject(gpuModule, options); @@ -145,7 +145,7 @@ TEST_F(MLIRTargetLLVMROCDL, // Serialize the module. auto serializer = dyn_cast(target); ASSERT_TRUE(!!serializer); - gpu::TargetOptions options("", {}, "", gpu::CompilationTarget::Assembly); + gpu::TargetOptions options("", {}, "", "", gpu::CompilationTarget::Assembly); for (auto gpuModule : (*module).getBody()->getOps()) { std::optional> object = serializer.serializeToObject(gpuModule, options); @@ -169,7 +169,7 @@ TEST_F(MLIRTargetLLVMROCDL, SKIP_WITHOUT_AMDGPU(SerializeROCDLToPTX)) { // Serialize the module. auto serializer = dyn_cast(target); ASSERT_TRUE(!!serializer); - gpu::TargetOptions options("", {}, "", gpu::CompilationTarget::Assembly); + gpu::TargetOptions options("", {}, "", "", gpu::CompilationTarget::Assembly); for (auto gpuModule : (*module).getBody()->getOps()) { std::optional> object = serializer.serializeToObject(gpuModule, options); @@ -199,7 +199,7 @@ TEST_F(MLIRTargetLLVMROCDL, SKIP_WITHOUT_AMDGPU(SerializeROCDLToBinary)) { // Serialize the module. auto serializer = dyn_cast(target); ASSERT_TRUE(!!serializer); - gpu::TargetOptions options("", {}, "", gpu::CompilationTarget::Binary); + gpu::TargetOptions options("", {}, "", "", gpu::CompilationTarget::Binary); for (auto gpuModule : (*module).getBody()->getOps()) { std::optional> object = serializer.serializeToObject(gpuModule, options); @@ -243,7 +243,7 @@ TEST_F(MLIRTargetLLVMROCDL, SKIP_WITHOUT_AMDGPU(GetELFMetadata)) { // Serialize the module. auto serializer = dyn_cast(target); ASSERT_TRUE(!!serializer); - gpu::TargetOptions options("", {}, "", gpu::CompilationTarget::Binary); + gpu::TargetOptions options("", {}, "", "", gpu::CompilationTarget::Binary); for (auto gpuModule : (*module).getBody()->getOps()) { std::optional> object = serializer.serializeToObject(gpuModule, options); diff --git a/mlir/unittests/Target/LLVM/SerializeToLLVMBitcode.cpp b/mlir/unittests/Target/LLVM/SerializeToLLVMBitcode.cpp index 63d1dbd2519be..314f25833f38c 100644 --- a/mlir/unittests/Target/LLVM/SerializeToLLVMBitcode.cpp +++ b/mlir/unittests/Target/LLVM/SerializeToLLVMBitcode.cpp @@ -174,8 +174,8 @@ TEST_F(MLIRTargetLLVM, SKIP_WITHOUT_NATIVE(CallbackInvokedWithInitialLLVMIR)) { }; gpu::TargetOptions opts( - {}, {}, {}, mlir::gpu::TargetOptions::getDefaultCompilationTarget(), {}, - initialCallback); + {}, {}, {}, {}, mlir::gpu::TargetOptions::getDefaultCompilationTarget(), + {}, initialCallback); std::optional> serializedBinary = targetAttr.serializeToObject(*module, opts); @@ -202,8 +202,8 @@ TEST_F(MLIRTargetLLVM, SKIP_WITHOUT_NATIVE(CallbackInvokedWithLinkedLLVMIR)) { }; gpu::TargetOptions opts( - {}, {}, {}, mlir::gpu::TargetOptions::getDefaultCompilationTarget(), {}, - {}, linkedCallback); + {}, {}, {}, {}, mlir::gpu::TargetOptions::getDefaultCompilationTarget(), + {}, {}, linkedCallback); std::optional> serializedBinary = targetAttr.serializeToObject(*module, opts); @@ -231,8 +231,8 @@ TEST_F(MLIRTargetLLVM, }; gpu::TargetOptions opts( - {}, {}, {}, mlir::gpu::TargetOptions::getDefaultCompilationTarget(), {}, - {}, {}, optimizedCallback); + {}, {}, {}, {}, mlir::gpu::TargetOptions::getDefaultCompilationTarget(), + {}, {}, {}, optimizedCallback); std::optional> serializedBinary = targetAttr.serializeToObject(*module, opts);