Skip to content

Commit 08e7609

Browse files
MikaOvOMikaOvO
authored
[mlir][fix] Add callback functions for ModuleToObject (#116916)
Here is the [merged MR](#116007) which caused a failure and [was reverted](#116811). Thanks to @joker-eph for the help, I fix it (miss constructing `ModuleObject` with callback functions in `mlir/lib/Target/LLVM/NVVM/Target.cpp`) and split unit tests from origin test which don't need `ptxas` to make the test runs more widely.
1 parent 33fcd6a commit 08e7609

File tree

7 files changed

+275
-13
lines changed

7 files changed

+275
-13
lines changed

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

Lines changed: 42 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -14,6 +14,7 @@
1414
#define MLIR_DIALECT_GPU_IR_COMPILATIONINTERFACES_H
1515

1616
#include "mlir/IR/Attributes.h"
17+
#include "llvm/IR/Module.h"
1718

1819
namespace llvm {
1920
class IRBuilderBase;
@@ -52,7 +53,11 @@ class TargetOptions {
5253
StringRef toolkitPath = {}, ArrayRef<std::string> linkFiles = {},
5354
StringRef cmdOptions = {},
5455
CompilationTarget compilationTarget = getDefaultCompilationTarget(),
55-
function_ref<SymbolTable *()> getSymbolTableCallback = {});
56+
function_ref<SymbolTable *()> getSymbolTableCallback = {},
57+
function_ref<void(llvm::Module &)> initialLlvmIRCallback = {},
58+
function_ref<void(llvm::Module &)> linkedLlvmIRCallback = {},
59+
function_ref<void(llvm::Module &)> optimizedLlvmIRCallback = {},
60+
function_ref<void(StringRef)> isaCallback = {});
5661

5762
/// Returns the typeID.
5863
TypeID getTypeID() const;
@@ -80,6 +85,22 @@ class TargetOptions {
8085
/// table.
8186
SymbolTable *getSymbolTable() const;
8287

88+
/// Returns the callback invoked with the initial LLVM IR for the device
89+
/// module.
90+
function_ref<void(llvm::Module &)> getInitialLlvmIRCallback() const;
91+
92+
/// Returns the callback invoked with LLVM IR for the device module
93+
/// after linking the device libraries.
94+
function_ref<void(llvm::Module &)> getLinkedLlvmIRCallback() const;
95+
96+
/// Returns the callback invoked with LLVM IR for the device module after
97+
/// LLVM optimizations but before codegen.
98+
function_ref<void(llvm::Module &)> getOptimizedLlvmIRCallback() const;
99+
100+
/// Returns the callback invoked with the target ISA for the device,
101+
/// for example PTX assembly.
102+
function_ref<void(StringRef)> getISACallback() const;
103+
83104
/// Returns the default compilation target: `CompilationTarget::Fatbin`.
84105
static CompilationTarget getDefaultCompilationTarget();
85106

@@ -90,7 +111,11 @@ class TargetOptions {
90111
TypeID typeID, StringRef toolkitPath = {},
91112
ArrayRef<std::string> linkFiles = {}, StringRef cmdOptions = {},
92113
CompilationTarget compilationTarget = getDefaultCompilationTarget(),
93-
function_ref<SymbolTable *()> getSymbolTableCallback = {});
114+
function_ref<SymbolTable *()> getSymbolTableCallback = {},
115+
function_ref<void(llvm::Module &)> initialLlvmIRCallback = {},
116+
function_ref<void(llvm::Module &)> linkedLlvmIRCallback = {},
117+
function_ref<void(llvm::Module &)> optimizedLlvmIRCallback = {},
118+
function_ref<void(StringRef)> isaCallback = {});
94119

95120
/// Path to the target toolkit.
96121
std::string toolkitPath;
@@ -109,6 +134,21 @@ class TargetOptions {
109134
/// being serialized.
110135
function_ref<SymbolTable *()> getSymbolTableCallback;
111136

137+
/// Callback invoked with the initial LLVM IR for the device module.
138+
function_ref<void(llvm::Module &)> initialLlvmIRCallback;
139+
140+
/// Callback invoked with LLVM IR for the device module after
141+
/// linking the device libraries.
142+
function_ref<void(llvm::Module &)> linkedLlvmIRCallback;
143+
144+
/// Callback invoked with LLVM IR for the device module after
145+
/// LLVM optimizations but before codegen.
146+
function_ref<void(llvm::Module &)> optimizedLlvmIRCallback;
147+
148+
/// Callback invoked with the target ISA for the device,
149+
/// for example PTX assembly.
150+
function_ref<void(StringRef)> isaCallback;
151+
112152
private:
113153
TypeID typeID;
114154
};

‎mlir/include/mlir/Target/LLVM/ModuleToObject.h‎

Lines changed: 22 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -29,8 +29,13 @@ class ModuleTranslation;
2929
/// operations being transformed must be translatable into LLVM IR.
3030
class ModuleToObject {
3131
public:
32-
ModuleToObject(Operation &module, StringRef triple, StringRef chip,
33-
StringRef features = {}, int optLevel = 3);
32+
ModuleToObject(
33+
Operation &module, StringRef triple, StringRef chip,
34+
StringRef features = {}, int optLevel = 3,
35+
function_ref<void(llvm::Module &)> initialLlvmIRCallback = {},
36+
function_ref<void(llvm::Module &)> linkedLlvmIRCallback = {},
37+
function_ref<void(llvm::Module &)> optimizedLlvmIRCallback = {},
38+
function_ref<void(StringRef)> isaCallback = {});
3439
virtual ~ModuleToObject();
3540

3641
/// Returns the operation being serialized.
@@ -114,6 +119,21 @@ class ModuleToObject {
114119
/// Optimization level.
115120
int optLevel;
116121

122+
/// Callback invoked with the initial LLVM IR for the device module.
123+
function_ref<void(llvm::Module &)> initialLlvmIRCallback;
124+
125+
/// Callback invoked with LLVM IR for the device module after
126+
/// linking the device libraries.
127+
function_ref<void(llvm::Module &)> linkedLlvmIRCallback;
128+
129+
/// Callback invoked with LLVM IR for the device module after
130+
/// LLVM optimizations but before codegen.
131+
function_ref<void(llvm::Module &)> optimizedLlvmIRCallback;
132+
133+
/// Callback invoked with the target ISA for the device,
134+
/// for example PTX assembly.
135+
function_ref<void(StringRef)> isaCallback;
136+
117137
private:
118138
/// The TargetMachine created for the given Triple, if available.
119139
/// Accessible through `getOrCreateTargetMachine()`.

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

Lines changed: 37 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -2302,17 +2302,31 @@ KernelMetadataAttr KernelTableAttr::lookup(StringAttr key) const {
23022302
TargetOptions::TargetOptions(
23032303
StringRef toolkitPath, ArrayRef<std::string> linkFiles,
23042304
StringRef cmdOptions, CompilationTarget compilationTarget,
2305-
function_ref<SymbolTable *()> getSymbolTableCallback)
2305+
function_ref<SymbolTable *()> getSymbolTableCallback,
2306+
function_ref<void(llvm::Module &)> initialLlvmIRCallback,
2307+
function_ref<void(llvm::Module &)> linkedLlvmIRCallback,
2308+
function_ref<void(llvm::Module &)> optimizedLlvmIRCallback,
2309+
function_ref<void(StringRef)> isaCallback)
23062310
: TargetOptions(TypeID::get<TargetOptions>(), toolkitPath, linkFiles,
2307-
cmdOptions, compilationTarget, getSymbolTableCallback) {}
2311+
cmdOptions, compilationTarget, getSymbolTableCallback,
2312+
initialLlvmIRCallback, linkedLlvmIRCallback,
2313+
optimizedLlvmIRCallback, isaCallback) {}
23082314

23092315
TargetOptions::TargetOptions(
23102316
TypeID typeID, StringRef toolkitPath, ArrayRef<std::string> linkFiles,
23112317
StringRef cmdOptions, CompilationTarget compilationTarget,
2312-
function_ref<SymbolTable *()> getSymbolTableCallback)
2318+
function_ref<SymbolTable *()> getSymbolTableCallback,
2319+
function_ref<void(llvm::Module &)> initialLlvmIRCallback,
2320+
function_ref<void(llvm::Module &)> linkedLlvmIRCallback,
2321+
function_ref<void(llvm::Module &)> optimizedLlvmIRCallback,
2322+
function_ref<void(StringRef)> isaCallback)
23132323
: toolkitPath(toolkitPath.str()), linkFiles(linkFiles),
23142324
cmdOptions(cmdOptions.str()), compilationTarget(compilationTarget),
2315-
getSymbolTableCallback(getSymbolTableCallback), typeID(typeID) {}
2325+
getSymbolTableCallback(getSymbolTableCallback),
2326+
initialLlvmIRCallback(initialLlvmIRCallback),
2327+
linkedLlvmIRCallback(linkedLlvmIRCallback),
2328+
optimizedLlvmIRCallback(optimizedLlvmIRCallback),
2329+
isaCallback(isaCallback), typeID(typeID) {}
23162330

23172331
TypeID TargetOptions::getTypeID() const { return typeID; }
23182332

@@ -2326,6 +2340,25 @@ SymbolTable *TargetOptions::getSymbolTable() const {
23262340
return getSymbolTableCallback ? getSymbolTableCallback() : nullptr;
23272341
}
23282342

2343+
function_ref<void(llvm::Module &)>
2344+
TargetOptions::getInitialLlvmIRCallback() const {
2345+
return initialLlvmIRCallback;
2346+
}
2347+
2348+
function_ref<void(llvm::Module &)>
2349+
TargetOptions::getLinkedLlvmIRCallback() const {
2350+
return linkedLlvmIRCallback;
2351+
}
2352+
2353+
function_ref<void(llvm::Module &)>
2354+
TargetOptions::getOptimizedLlvmIRCallback() const {
2355+
return optimizedLlvmIRCallback;
2356+
}
2357+
2358+
function_ref<void(StringRef)> TargetOptions::getISACallback() const {
2359+
return isaCallback;
2360+
}
2361+
23292362
CompilationTarget TargetOptions::getCompilationTarget() const {
23302363
return compilationTarget;
23312364
}

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

Lines changed: 19 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -34,10 +34,17 @@
3434
using namespace mlir;
3535
using namespace mlir::LLVM;
3636

37-
ModuleToObject::ModuleToObject(Operation &module, StringRef triple,
38-
StringRef chip, StringRef features, int optLevel)
37+
ModuleToObject::ModuleToObject(
38+
Operation &module, StringRef triple, StringRef chip, StringRef features,
39+
int optLevel, function_ref<void(llvm::Module &)> initialLlvmIRCallback,
40+
function_ref<void(llvm::Module &)> linkedLlvmIRCallback,
41+
function_ref<void(llvm::Module &)> optimizedLlvmIRCallback,
42+
function_ref<void(StringRef)> isaCallback)
3943
: module(module), triple(triple), chip(chip), features(features),
40-
optLevel(optLevel) {}
44+
optLevel(optLevel), initialLlvmIRCallback(initialLlvmIRCallback),
45+
linkedLlvmIRCallback(linkedLlvmIRCallback),
46+
optimizedLlvmIRCallback(optimizedLlvmIRCallback),
47+
isaCallback(isaCallback) {}
4148

4249
ModuleToObject::~ModuleToObject() = default;
4350

@@ -215,6 +222,9 @@ std::optional<SmallVector<char, 0>> ModuleToObject::run() {
215222
}
216223
setDataLayoutAndTriple(*llvmModule);
217224

225+
if (initialLlvmIRCallback)
226+
initialLlvmIRCallback(*llvmModule);
227+
218228
// Link bitcode files.
219229
handleModulePreLink(*llvmModule);
220230
{
@@ -227,10 +237,16 @@ std::optional<SmallVector<char, 0>> ModuleToObject::run() {
227237
handleModulePostLink(*llvmModule);
228238
}
229239

240+
if (linkedLlvmIRCallback)
241+
linkedLlvmIRCallback(*llvmModule);
242+
230243
// Optimize the module.
231244
if (failed(optimizeModule(*llvmModule, optLevel)))
232245
return std::nullopt;
233246

247+
if (optimizedLlvmIRCallback)
248+
optimizedLlvmIRCallback(*llvmModule);
249+
234250
// Return the serialized object.
235251
return moduleToObject(*llvmModule);
236252
}

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

Lines changed: 8 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -86,7 +86,11 @@ SerializeGPUModuleBase::SerializeGPUModuleBase(
8686
Operation &module, NVVMTargetAttr target,
8787
const gpu::TargetOptions &targetOptions)
8888
: ModuleToObject(module, target.getTriple(), target.getChip(),
89-
target.getFeatures(), target.getO()),
89+
target.getFeatures(), target.getO(),
90+
targetOptions.getInitialLlvmIRCallback(),
91+
targetOptions.getLinkedLlvmIRCallback(),
92+
targetOptions.getOptimizedLlvmIRCallback(),
93+
targetOptions.getISACallback()),
9094
target(target), toolkitPath(targetOptions.getToolkitPath()),
9195
fileList(targetOptions.getLinkFiles()) {
9296

@@ -572,6 +576,9 @@ NVPTXSerializer::moduleToObject(llvm::Module &llvmModule) {
572576
getOperation().emitError() << "Failed translating the module to ISA.";
573577
return std::nullopt;
574578
}
579+
if (isaCallback)
580+
isaCallback(serializedISA.value());
581+
575582
#define DEBUG_TYPE "serialize-to-isa"
576583
LLVM_DEBUG({
577584
llvm::dbgs() << "PTX for module: " << getOperation().getNameAttr() << "\n";

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

Lines changed: 59 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -156,3 +156,62 @@ TEST_F(MLIRTargetLLVMNVVM, SKIP_WITHOUT_NVPTX(SerializeNVVMToBinary)) {
156156
ASSERT_TRUE(!object->empty());
157157
}
158158
}
159+
160+
// Test callback functions invoked with LLVM IR and ISA.
161+
TEST_F(MLIRTargetLLVMNVVM,
162+
SKIP_WITHOUT_NVPTX(CallbackInvokedWithLLVMIRAndISA)) {
163+
MLIRContext context(registry);
164+
165+
OwningOpRef<ModuleOp> module =
166+
parseSourceString<ModuleOp>(moduleStr, &context);
167+
ASSERT_TRUE(!!module);
168+
169+
NVVM::NVVMTargetAttr target = NVVM::NVVMTargetAttr::get(&context);
170+
171+
auto serializer = dyn_cast<gpu::TargetAttrInterface>(target);
172+
ASSERT_TRUE(!!serializer);
173+
174+
std::string initialLLVMIR;
175+
auto initialCallback = [&initialLLVMIR](llvm::Module &module) {
176+
llvm::raw_string_ostream ros(initialLLVMIR);
177+
module.print(ros, nullptr);
178+
};
179+
180+
std::string linkedLLVMIR;
181+
auto linkedCallback = [&linkedLLVMIR](llvm::Module &module) {
182+
llvm::raw_string_ostream ros(linkedLLVMIR);
183+
module.print(ros, nullptr);
184+
};
185+
186+
std::string optimizedLLVMIR;
187+
auto optimizedCallback = [&optimizedLLVMIR](llvm::Module &module) {
188+
llvm::raw_string_ostream ros(optimizedLLVMIR);
189+
module.print(ros, nullptr);
190+
};
191+
192+
std::string isaResult;
193+
auto isaCallback = [&isaResult](llvm::StringRef isa) {
194+
isaResult = isa.str();
195+
};
196+
197+
gpu::TargetOptions options({}, {}, {}, gpu::CompilationTarget::Assembly, {},
198+
initialCallback, linkedCallback, optimizedCallback,
199+
isaCallback);
200+
201+
for (auto gpuModule : (*module).getBody()->getOps<gpu::GPUModuleOp>()) {
202+
std::optional<SmallVector<char, 0>> object =
203+
serializer.serializeToObject(gpuModule, options);
204+
205+
ASSERT_TRUE(object != std::nullopt);
206+
ASSERT_TRUE(!object->empty());
207+
ASSERT_TRUE(!initialLLVMIR.empty());
208+
ASSERT_TRUE(!linkedLLVMIR.empty());
209+
ASSERT_TRUE(!optimizedLLVMIR.empty());
210+
ASSERT_TRUE(!isaResult.empty());
211+
212+
initialLLVMIR.clear();
213+
linkedLLVMIR.clear();
214+
optimizedLLVMIR.clear();
215+
isaResult.clear();
216+
}
217+
}

0 commit comments

Comments
 (0)