Skip to content

Commit 8b1c249

Browse files
Renaud-KRenaud-K
committed
Adding ELF section option to the gpu-module-to-binary pass
1 parent 6d3cb91 commit 8b1c249

File tree

9 files changed

+55
-39
lines changed

9 files changed

+55
-39
lines changed

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

Lines changed: 8 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -51,7 +51,7 @@ class TargetOptions {
5151
/// `Fatbin`.
5252
TargetOptions(
5353
StringRef toolkitPath = {}, ArrayRef<std::string> linkFiles = {},
54-
StringRef cmdOptions = {},
54+
StringRef cmdOptions = {}, StringRef elfSection = {},
5555
CompilationTarget compilationTarget = getDefaultCompilationTarget(),
5656
function_ref<SymbolTable *()> getSymbolTableCallback = {},
5757
function_ref<void(llvm::Module &)> initialLlvmIRCallback = {},
@@ -71,6 +71,9 @@ class TargetOptions {
7171
/// Returns the command line options.
7272
StringRef getCmdOptions() const;
7373

74+
/// Returns the ELF section.
75+
StringRef getELFSection() const;
76+
7477
/// Returns a tokenization of the command line options.
7578
std::pair<llvm::BumpPtrAllocator, SmallVector<const char *>>
7679
tokenizeCmdOptions() const;
@@ -110,6 +113,7 @@ class TargetOptions {
110113
TargetOptions(
111114
TypeID typeID, StringRef toolkitPath = {},
112115
ArrayRef<std::string> linkFiles = {}, StringRef cmdOptions = {},
116+
StringRef elfSection = {},
113117
CompilationTarget compilationTarget = getDefaultCompilationTarget(),
114118
function_ref<SymbolTable *()> getSymbolTableCallback = {},
115119
function_ref<void(llvm::Module &)> initialLlvmIRCallback = {},
@@ -127,6 +131,9 @@ class TargetOptions {
127131
/// process.
128132
std::string cmdOptions;
129133

134+
/// ELF Section where the binary needs to be located
135+
std::string elfSection;
136+
130137
/// Compilation process target format.
131138
CompilationTarget compilationTarget;
132139

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

Lines changed: 9 additions & 10 deletions
Original file line numberDiff line numberDiff line change
@@ -87,16 +87,15 @@ def GpuModuleToBinaryPass
8787
3. `binary`, `bin`: produces binaries.
8888
4. `fatbinary`, `fatbin`: produces fatbinaries.
8989
}];
90-
let options = [
91-
Option<"toolkitPath", "toolkit", "std::string", [{""}],
92-
"Toolkit path.">,
93-
ListOption<"linkFiles", "l", "std::string",
94-
"Extra files to link to.">,
95-
Option<"cmdOptions", "opts", "std::string", [{""}],
96-
"Command line options to pass to the tools.">,
97-
Option<"compilationTarget", "format", "std::string", [{"fatbin"}],
98-
"The target representation of the compilation process.">
99-
];
90+
let options =
91+
[Option<"toolkitPath", "toolkit", "std::string", [{""}], "Toolkit path.">,
92+
ListOption<"linkFiles", "l", "std::string", "Extra files to link to.">,
93+
Option<"cmdOptions", "opts", "std::string", [{""}],
94+
"Command line options to pass to the tools.">,
95+
Option<"compilationTarget", "format", "std::string", [{"fatbin"}],
96+
"The target representation of the compilation process.">,
97+
Option<"elfSection", "section", "std::string", [{""}],
98+
"ELF section where binary is to be located.">];
10099
}
101100

102101
def GpuNVVMAttachTarget: Pass<"nvvm-attach-target", ""> {

‎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: 3 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -669,8 +669,9 @@ NVVMTargetAttrImpl::createObject(Attribute attribute, Operation *module,
669669
properties.push_back(
670670
builder.getNamedAttr("O", builder.getI32IntegerAttr(target.getO())));
671671

672-
if (StringAttr section = target.getSection())
673-
properties.push_back(builder.getNamedAttr("section", section));
672+
if (auto section = options.getELFSection(); !section.empty())
673+
properties.push_back(
674+
builder.getNamedAttr("section", builder.getStringAttr(section)));
674675

675676
if (!properties.empty())
676677
objectProps = builder.getDictionaryAttr(properties);

‎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="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)