Skip to content

Commit 9919295

Browse files
Renaud-KRenaud-K
authored
[mlir][gpu] Adding ELF section option to the gpu-module-to-binary pass (#119440)
This is a follow-up of #117246. I thought then it would be easy to edit a DictionaryAttr but it turns out that these attributes are immutable and need to be passed during the construction of the gpu.binary Op. The first commit was using the NVVMTargetAttr to pass the information. After feedback from @fabianmcg, this PR now passes the information through a new option of the gpu-module-to-binary pass. Please add reviewers, as you see fit.
1 parent f75c846 commit 9919295

File tree

10 files changed

+61
-31
lines changed

10 files changed

+61
-31
lines changed

‎mlir/include/mlir/Dialect/GPU/IR/CompilationInterfaces.h

Lines changed: 9 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -27,6 +27,7 @@ class ModuleTranslation;
2727
}
2828
namespace gpu {
2929
enum class CompilationTarget : uint32_t;
30+
constexpr StringLiteral elfSectionName = "section";
3031

3132
/// This class indicates that the attribute associated with this trait is a GPU
3233
/// offloading translation attribute. These kinds of attributes must implement
@@ -51,7 +52,7 @@ class TargetOptions {
5152
/// `Fatbin`.
5253
TargetOptions(
5354
StringRef toolkitPath = {}, ArrayRef<std::string> linkFiles = {},
54-
StringRef cmdOptions = {},
55+
StringRef cmdOptions = {}, StringRef elfSection = {},
5556
CompilationTarget compilationTarget = getDefaultCompilationTarget(),
5657
function_ref<SymbolTable *()> getSymbolTableCallback = {},
5758
function_ref<void(llvm::Module &)> initialLlvmIRCallback = {},
@@ -71,6 +72,9 @@ class TargetOptions {
7172
/// Returns the command line options.
7273
StringRef getCmdOptions() const;
7374

75+
/// Returns the ELF section.
76+
StringRef getELFSection() const;
77+
7478
/// Returns a tokenization of the command line options.
7579
std::pair<llvm::BumpPtrAllocator, SmallVector<const char *>>
7680
tokenizeCmdOptions() const;
@@ -110,6 +114,7 @@ class TargetOptions {
110114
TargetOptions(
111115
TypeID typeID, StringRef toolkitPath = {},
112116
ArrayRef<std::string> linkFiles = {}, StringRef cmdOptions = {},
117+
StringRef elfSection = {},
113118
CompilationTarget compilationTarget = getDefaultCompilationTarget(),
114119
function_ref<SymbolTable *()> getSymbolTableCallback = {},
115120
function_ref<void(llvm::Module &)> initialLlvmIRCallback = {},
@@ -127,6 +132,9 @@ class TargetOptions {
127132
/// process.
128133
std::string cmdOptions;
129134

135+
/// ELF Section where the binary needs to be located
136+
std::string elfSection;
137+
130138
/// Compilation process target format.
131139
CompilationTarget compilationTarget;
132140

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

Lines changed: 3 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -95,7 +95,9 @@ def GpuModuleToBinaryPass
9595
Option<"cmdOptions", "opts", "std::string", [{""}],
9696
"Command line options to pass to the tools.">,
9797
Option<"compilationTarget", "format", "std::string", [{"fatbin"}],
98-
"The target representation of the compilation process.">
98+
"The target representation of the compilation process.">,
99+
Option<"elfSection", "section", "std::string", [{""}],
100+
"ELF section where binary is to be located.">
99101
];
100102
}
101103

‎mlir/lib/Dialect/GPU/IR/GPUDialect.cpp

Lines changed: 12 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -2484,27 +2484,31 @@ KernelMetadataAttr KernelTableAttr::lookup(StringAttr key) const {
24842484

24852485
TargetOptions::TargetOptions(
24862486
StringRef toolkitPath, ArrayRef<std::string> linkFiles,
2487-
StringRef cmdOptions, CompilationTarget compilationTarget,
2487+
StringRef cmdOptions, StringRef elfSection,
2488+
CompilationTarget compilationTarget,
24882489
function_ref<SymbolTable *()> getSymbolTableCallback,
24892490
function_ref<void(llvm::Module &)> initialLlvmIRCallback,
24902491
function_ref<void(llvm::Module &)> linkedLlvmIRCallback,
24912492
function_ref<void(llvm::Module &)> optimizedLlvmIRCallback,
24922493
function_ref<void(StringRef)> isaCallback)
24932494
: TargetOptions(TypeID::get<TargetOptions>(), toolkitPath, linkFiles,
2494-
cmdOptions, compilationTarget, getSymbolTableCallback,
2495-
initialLlvmIRCallback, linkedLlvmIRCallback,
2496-
optimizedLlvmIRCallback, isaCallback) {}
2495+
cmdOptions, elfSection, compilationTarget,
2496+
getSymbolTableCallback, initialLlvmIRCallback,
2497+
linkedLlvmIRCallback, optimizedLlvmIRCallback,
2498+
isaCallback) {}
24972499

24982500
TargetOptions::TargetOptions(
24992501
TypeID typeID, StringRef toolkitPath, ArrayRef<std::string> linkFiles,
2500-
StringRef cmdOptions, CompilationTarget compilationTarget,
2502+
StringRef cmdOptions, StringRef elfSection,
2503+
CompilationTarget compilationTarget,
25012504
function_ref<SymbolTable *()> getSymbolTableCallback,
25022505
function_ref<void(llvm::Module &)> initialLlvmIRCallback,
25032506
function_ref<void(llvm::Module &)> linkedLlvmIRCallback,
25042507
function_ref<void(llvm::Module &)> optimizedLlvmIRCallback,
25052508
function_ref<void(StringRef)> isaCallback)
25062509
: toolkitPath(toolkitPath.str()), linkFiles(linkFiles),
2507-
cmdOptions(cmdOptions.str()), compilationTarget(compilationTarget),
2510+
cmdOptions(cmdOptions.str()), elfSection(elfSection.str()),
2511+
compilationTarget(compilationTarget),
25082512
getSymbolTableCallback(getSymbolTableCallback),
25092513
initialLlvmIRCallback(initialLlvmIRCallback),
25102514
linkedLlvmIRCallback(linkedLlvmIRCallback),
@@ -2519,6 +2523,8 @@ ArrayRef<std::string> TargetOptions::getLinkFiles() const { return linkFiles; }
25192523

25202524
StringRef TargetOptions::getCmdOptions() const { return cmdOptions; }
25212525

2526+
StringRef TargetOptions::getELFSection() const { return elfSection; }
2527+
25222528
SymbolTable *TargetOptions::getSymbolTable() const {
25232529
return getSymbolTableCallback ? getSymbolTableCallback() : nullptr;
25242530
}

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

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -69,8 +69,8 @@ void GpuModuleToBinaryPass::runOnOperation() {
6969
return &parentTable.value();
7070
};
7171

72-
TargetOptions targetOptions(toolkitPath, linkFiles, cmdOptions, *targetFormat,
73-
lazyTableBuilder);
72+
TargetOptions targetOptions(toolkitPath, linkFiles, cmdOptions, elfSection,
73+
*targetFormat, lazyTableBuilder);
7474
if (failed(transformGpuModulesToBinaries(
7575
getOperation(), OffloadingLLVMTranslationAttrInterface(nullptr),
7676
targetOptions)))

‎mlir/lib/Target/LLVM/NVVM/Target.cpp

Lines changed: 12 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -13,6 +13,7 @@
1313

1414
#include "mlir/Target/LLVM/NVVM/Target.h"
1515

16+
#include "mlir/Dialect/GPU/IR/CompilationInterfaces.h"
1617
#include "mlir/Dialect/GPU/IR/GPUDialect.h"
1718
#include "mlir/Dialect/LLVMIR/NVVMDialect.h"
1819
#include "mlir/Target/LLVM/NVVM/Utils.h"
@@ -664,9 +665,18 @@ NVVMTargetAttrImpl::createObject(Attribute attribute, Operation *module,
664665
gpu::CompilationTarget format = options.getCompilationTarget();
665666
DictionaryAttr objectProps;
666667
Builder builder(attribute.getContext());
668+
SmallVector<NamedAttribute, 2> properties;
667669
if (format == gpu::CompilationTarget::Assembly)
668-
objectProps = builder.getDictionaryAttr(
669-
{builder.getNamedAttr("O", builder.getI32IntegerAttr(target.getO()))});
670+
properties.push_back(
671+
builder.getNamedAttr("O", builder.getI32IntegerAttr(target.getO())));
672+
673+
if (StringRef section = options.getELFSection(); !section.empty())
674+
properties.push_back(builder.getNamedAttr(gpu::elfSectionName,
675+
builder.getStringAttr(section)));
676+
677+
if (!properties.empty())
678+
objectProps = builder.getDictionaryAttr(properties);
679+
670680
return builder.getAttr<gpu::ObjectAttr>(
671681
attribute, format,
672682
builder.getStringAttr(StringRef(object.data(), object.size())),

‎mlir/lib/Target/LLVMIR/Dialect/GPU/SelectObjectAttr.cpp

Lines changed: 2 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -11,6 +11,7 @@
1111
//
1212
//===----------------------------------------------------------------------===//
1313

14+
#include "mlir/Dialect/GPU/IR/CompilationInterfaces.h"
1415
#include "mlir/Dialect/GPU/IR/GPUDialect.h"
1516

1617
#include "mlir/Target/LLVMIR/Dialect/GPU/GPUToLLVMIRTranslation.h"
@@ -124,7 +125,7 @@ LogicalResult SelectObjectAttrImpl::embedBinary(
124125

125126
if (object.getProperties()) {
126127
if (auto section = mlir::dyn_cast_or_null<mlir::StringAttr>(
127-
object.getProperties().get("section"))) {
128+
object.getProperties().get(gpu::elfSectionName))) {
128129
serializedObj->setSection(section.getValue());
129130
}
130131
}

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

Lines changed: 3 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1,10 +1,12 @@
11
// REQUIRES: host-supports-nvptx
22
// RUN: mlir-opt %s --gpu-module-to-binary="format=llvm" | FileCheck %s
33
// RUN: mlir-opt %s --gpu-module-to-binary="format=isa" | FileCheck %s -check-prefix=CHECK-ISA
4+
// RUN: mlir-opt %s --gpu-module-to-binary="format=llvm section=__fatbin" | FileCheck %s -check-prefix=CHECK-SECTION
45

56
module attributes {gpu.container_module} {
67
// CHECK-LABEL:gpu.binary @kernel_module1
78
// CHECK:[#gpu.object<#nvvm.target<chip = "sm_70">, offload = "{{.*}}">]
9+
// CHECK-SECTION: #gpu.object<#nvvm.target<chip = "sm_70">, properties = {section = "__fatbin"}
810
gpu.module @kernel_module1 [#nvvm.target<chip = "sm_70">] {
911
llvm.func @kernel(%arg0: i32, %arg1: !llvm.ptr,
1012
%arg2: !llvm.ptr, %arg3: i64, %arg4: i64,
@@ -25,6 +27,7 @@ module attributes {gpu.container_module} {
2527

2628
// CHECK-LABEL:gpu.binary @kernel_module3 <#gpu.select_object<1 : i64>>
2729
// CHECK:[#gpu.object<#nvvm.target<chip = "sm_70">, offload = "{{.*}}">, #gpu.object<#nvvm.target<chip = "sm_80">, offload = "{{.*}}">]
30+
// CHECK-SECTION: [#gpu.object<#nvvm.target<chip = "sm_70">, properties = {section = "__fatbin"},{{.*}} #gpu.object<#nvvm.target<chip = "sm_80">, properties = {section = "__fatbin"}
2831
gpu.module @kernel_module3 <#gpu.select_object<1>> [
2932
#nvvm.target<chip = "sm_70">,
3033
#nvvm.target<chip = "sm_80">] {

‎mlir/unittests/Target/LLVM/SerializeNVVMTarget.cpp

Lines changed: 6 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -81,7 +81,7 @@ TEST_F(MLIRTargetLLVMNVVM, SKIP_WITHOUT_NVPTX(SerializeNVVMMToLLVM)) {
8181
// Serialize the module.
8282
auto serializer = dyn_cast<gpu::TargetAttrInterface>(target);
8383
ASSERT_TRUE(!!serializer);
84-
gpu::TargetOptions options("", {}, "", gpu::CompilationTarget::Offload);
84+
gpu::TargetOptions options("", {}, "", "", gpu::CompilationTarget::Offload);
8585
for (auto gpuModule : (*module).getBody()->getOps<gpu::GPUModuleOp>()) {
8686
std::optional<SmallVector<char, 0>> object =
8787
serializer.serializeToObject(gpuModule, options);
@@ -117,7 +117,7 @@ TEST_F(MLIRTargetLLVMNVVM, SKIP_WITHOUT_NVPTX(SerializeNVVMToPTX)) {
117117
// Serialize the module.
118118
auto serializer = dyn_cast<gpu::TargetAttrInterface>(target);
119119
ASSERT_TRUE(!!serializer);
120-
gpu::TargetOptions options("", {}, "", gpu::CompilationTarget::Assembly);
120+
gpu::TargetOptions options("", {}, "", "", gpu::CompilationTarget::Assembly);
121121
for (auto gpuModule : (*module).getBody()->getOps<gpu::GPUModuleOp>()) {
122122
std::optional<SmallVector<char, 0>> object =
123123
serializer.serializeToObject(gpuModule, options);
@@ -147,7 +147,7 @@ TEST_F(MLIRTargetLLVMNVVM, SKIP_WITHOUT_NVPTX(SerializeNVVMToBinary)) {
147147
// Serialize the module.
148148
auto serializer = dyn_cast<gpu::TargetAttrInterface>(target);
149149
ASSERT_TRUE(!!serializer);
150-
gpu::TargetOptions options("", {}, "", gpu::CompilationTarget::Binary);
150+
gpu::TargetOptions options("", {}, "", "", gpu::CompilationTarget::Binary);
151151
for (auto gpuModule : (*module).getBody()->getOps<gpu::GPUModuleOp>()) {
152152
std::optional<SmallVector<char, 0>> object =
153153
serializer.serializeToObject(gpuModule, options);
@@ -194,9 +194,9 @@ TEST_F(MLIRTargetLLVMNVVM,
194194
isaResult = isa.str();
195195
};
196196

197-
gpu::TargetOptions options({}, {}, {}, gpu::CompilationTarget::Assembly, {},
198-
initialCallback, linkedCallback, optimizedCallback,
199-
isaCallback);
197+
gpu::TargetOptions options({}, {}, {}, {}, gpu::CompilationTarget::Assembly,
198+
{}, initialCallback, linkedCallback,
199+
optimizedCallback, isaCallback);
200200

201201
for (auto gpuModule : (*module).getBody()->getOps<gpu::GPUModuleOp>()) {
202202
std::optional<SmallVector<char, 0>> object =

‎mlir/unittests/Target/LLVM/SerializeROCDLTarget.cpp

Lines changed: 6 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -83,7 +83,7 @@ TEST_F(MLIRTargetLLVMROCDL, SKIP_WITHOUT_AMDGPU(SerializeROCDLToLLVM)) {
8383
// Serialize the module.
8484
auto serializer = dyn_cast<gpu::TargetAttrInterface>(target);
8585
ASSERT_TRUE(!!serializer);
86-
gpu::TargetOptions options("", {}, "", gpu::CompilationTarget::Offload);
86+
gpu::TargetOptions options("", {}, "", "", gpu::CompilationTarget::Offload);
8787
for (auto gpuModule : (*module).getBody()->getOps<gpu::GPUModuleOp>()) {
8888
std::optional<SmallVector<char, 0>> object =
8989
serializer.serializeToObject(gpuModule, options);
@@ -119,7 +119,7 @@ TEST_F(MLIRTargetLLVMROCDL,
119119
// Serialize the module.
120120
auto serializer = dyn_cast<gpu::TargetAttrInterface>(target);
121121
ASSERT_TRUE(!!serializer);
122-
gpu::TargetOptions options("", {}, "", gpu::CompilationTarget::Assembly);
122+
gpu::TargetOptions options("", {}, "", "", gpu::CompilationTarget::Assembly);
123123
for (auto gpuModule : (*module).getBody()->getOps<gpu::GPUModuleOp>()) {
124124
std::optional<SmallVector<char, 0>> object =
125125
serializer.serializeToObject(gpuModule, options);
@@ -145,7 +145,7 @@ TEST_F(MLIRTargetLLVMROCDL,
145145
// Serialize the module.
146146
auto serializer = dyn_cast<gpu::TargetAttrInterface>(target);
147147
ASSERT_TRUE(!!serializer);
148-
gpu::TargetOptions options("", {}, "", gpu::CompilationTarget::Assembly);
148+
gpu::TargetOptions options("", {}, "", "", gpu::CompilationTarget::Assembly);
149149
for (auto gpuModule : (*module).getBody()->getOps<gpu::GPUModuleOp>()) {
150150
std::optional<SmallVector<char, 0>> object =
151151
serializer.serializeToObject(gpuModule, options);
@@ -169,7 +169,7 @@ TEST_F(MLIRTargetLLVMROCDL, SKIP_WITHOUT_AMDGPU(SerializeROCDLToPTX)) {
169169
// Serialize the module.
170170
auto serializer = dyn_cast<gpu::TargetAttrInterface>(target);
171171
ASSERT_TRUE(!!serializer);
172-
gpu::TargetOptions options("", {}, "", gpu::CompilationTarget::Assembly);
172+
gpu::TargetOptions options("", {}, "", "", gpu::CompilationTarget::Assembly);
173173
for (auto gpuModule : (*module).getBody()->getOps<gpu::GPUModuleOp>()) {
174174
std::optional<SmallVector<char, 0>> object =
175175
serializer.serializeToObject(gpuModule, options);
@@ -199,7 +199,7 @@ TEST_F(MLIRTargetLLVMROCDL, SKIP_WITHOUT_AMDGPU(SerializeROCDLToBinary)) {
199199
// Serialize the module.
200200
auto serializer = dyn_cast<gpu::TargetAttrInterface>(target);
201201
ASSERT_TRUE(!!serializer);
202-
gpu::TargetOptions options("", {}, "", gpu::CompilationTarget::Binary);
202+
gpu::TargetOptions options("", {}, "", "", gpu::CompilationTarget::Binary);
203203
for (auto gpuModule : (*module).getBody()->getOps<gpu::GPUModuleOp>()) {
204204
std::optional<SmallVector<char, 0>> object =
205205
serializer.serializeToObject(gpuModule, options);
@@ -243,7 +243,7 @@ TEST_F(MLIRTargetLLVMROCDL, SKIP_WITHOUT_AMDGPU(GetELFMetadata)) {
243243
// Serialize the module.
244244
auto serializer = dyn_cast<gpu::TargetAttrInterface>(target);
245245
ASSERT_TRUE(!!serializer);
246-
gpu::TargetOptions options("", {}, "", gpu::CompilationTarget::Binary);
246+
gpu::TargetOptions options("", {}, "", "", gpu::CompilationTarget::Binary);
247247
for (auto gpuModule : (*module).getBody()->getOps<gpu::GPUModuleOp>()) {
248248
std::optional<SmallVector<char, 0>> object =
249249
serializer.serializeToObject(gpuModule, options);

‎mlir/unittests/Target/LLVM/SerializeToLLVMBitcode.cpp

Lines changed: 6 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -174,8 +174,8 @@ TEST_F(MLIRTargetLLVM, SKIP_WITHOUT_NATIVE(CallbackInvokedWithInitialLLVMIR)) {
174174
};
175175

176176
gpu::TargetOptions opts(
177-
{}, {}, {}, mlir::gpu::TargetOptions::getDefaultCompilationTarget(), {},
178-
initialCallback);
177+
{}, {}, {}, {}, mlir::gpu::TargetOptions::getDefaultCompilationTarget(),
178+
{}, initialCallback);
179179
std::optional<SmallVector<char, 0>> serializedBinary =
180180
targetAttr.serializeToObject(*module, opts);
181181

@@ -202,8 +202,8 @@ TEST_F(MLIRTargetLLVM, SKIP_WITHOUT_NATIVE(CallbackInvokedWithLinkedLLVMIR)) {
202202
};
203203

204204
gpu::TargetOptions opts(
205-
{}, {}, {}, mlir::gpu::TargetOptions::getDefaultCompilationTarget(), {},
206-
{}, linkedCallback);
205+
{}, {}, {}, {}, mlir::gpu::TargetOptions::getDefaultCompilationTarget(),
206+
{}, {}, linkedCallback);
207207
std::optional<SmallVector<char, 0>> serializedBinary =
208208
targetAttr.serializeToObject(*module, opts);
209209

@@ -231,8 +231,8 @@ TEST_F(MLIRTargetLLVM,
231231
};
232232

233233
gpu::TargetOptions opts(
234-
{}, {}, {}, mlir::gpu::TargetOptions::getDefaultCompilationTarget(), {},
235-
{}, {}, optimizedCallback);
234+
{}, {}, {}, {}, mlir::gpu::TargetOptions::getDefaultCompilationTarget(),
235+
{}, {}, {}, optimizedCallback);
236236
std::optional<SmallVector<char, 0>> serializedBinary =
237237
targetAttr.serializeToObject(*module, opts);
238238

0 commit comments

Comments
 (0)