Skip to content

Commit cd58586

Browse files
basioli-kbasioli-k
authored
Pass memory buffer to RuntimeDyld::MemoryManager factory (#142930)
`RTDyldObjectLinkingLayer` is currently creating a memory manager without any parameters. In this PR I am passing the MemoryBuffer that will be emitted to the MemoryManager so that the user can use it to configure the behaviour of the MemoryManager.
1 parent f961d6a commit cd58586

File tree

11 files changed

+106
-27
lines changed

11 files changed

+106
-27
lines changed

‎llvm/examples/Kaleidoscope/BuildingAJIT/Chapter1/KaleidoscopeJIT.h

Lines changed: 3 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -47,7 +47,9 @@ class KaleidoscopeJIT {
4747
JITTargetMachineBuilder JTMB, DataLayout DL)
4848
: ES(std::move(ES)), DL(std::move(DL)), Mangle(*this->ES, this->DL),
4949
ObjectLayer(*this->ES,
50-
[]() { return std::make_unique<SectionMemoryManager>(); }),
50+
[](const MemoryBuffer &) {
51+
return std::make_unique<SectionMemoryManager>();
52+
}),
5153
CompileLayer(*this->ES, ObjectLayer,
5254
std::make_unique<ConcurrentIRCompiler>(std::move(JTMB))),
5355
MainJD(this->ES->createBareJITDylib("<main>")) {

‎llvm/examples/Kaleidoscope/BuildingAJIT/Chapter2/KaleidoscopeJIT.h

Lines changed: 3 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -53,7 +53,9 @@ class KaleidoscopeJIT {
5353
JITTargetMachineBuilder JTMB, DataLayout DL)
5454
: ES(std::move(ES)), DL(std::move(DL)), Mangle(*this->ES, this->DL),
5555
ObjectLayer(*this->ES,
56-
[]() { return std::make_unique<SectionMemoryManager>(); }),
56+
[](const MemoryBuffer &) {
57+
return std::make_unique<SectionMemoryManager>();
58+
}),
5759
CompileLayer(*this->ES, ObjectLayer,
5860
std::make_unique<ConcurrentIRCompiler>(std::move(JTMB))),
5961
OptimizeLayer(*this->ES, CompileLayer, optimizeModule),

‎llvm/examples/Kaleidoscope/BuildingAJIT/Chapter3/KaleidoscopeJIT.h

Lines changed: 3 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -66,7 +66,9 @@ class KaleidoscopeJIT {
6666
: ES(std::move(ES)), EPCIU(std::move(EPCIU)), DL(std::move(DL)),
6767
Mangle(*this->ES, this->DL),
6868
ObjectLayer(*this->ES,
69-
[]() { return std::make_unique<SectionMemoryManager>(); }),
69+
[](const MemoryBuffer &) {
70+
return std::make_unique<SectionMemoryManager>();
71+
}),
7072
CompileLayer(*this->ES, ObjectLayer,
7173
std::make_unique<ConcurrentIRCompiler>(std::move(JTMB))),
7274
OptimizeLayer(*this->ES, CompileLayer, optimizeModule),

‎llvm/examples/Kaleidoscope/BuildingAJIT/Chapter4/KaleidoscopeJIT.h

Lines changed: 3 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -151,7 +151,9 @@ class KaleidoscopeJIT {
151151
: ES(std::move(ES)), EPCIU(std::move(EPCIU)), DL(std::move(DL)),
152152
Mangle(*this->ES, this->DL),
153153
ObjectLayer(*this->ES,
154-
[]() { return std::make_unique<SectionMemoryManager>(); }),
154+
[](const MemoryBuffer &) {
155+
return std::make_unique<SectionMemoryManager>();
156+
}),
155157
CompileLayer(*this->ES, ObjectLayer,
156158
std::make_unique<ConcurrentIRCompiler>(std::move(JTMB))),
157159
OptimizeLayer(*this->ES, CompileLayer, optimizeModule),

‎llvm/examples/Kaleidoscope/include/KaleidoscopeJIT.h

Lines changed: 3 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -48,7 +48,9 @@ class KaleidoscopeJIT {
4848
JITTargetMachineBuilder JTMB, DataLayout DL)
4949
: ES(std::move(ES)), DL(std::move(DL)), Mangle(*this->ES, this->DL),
5050
ObjectLayer(*this->ES,
51-
[]() { return std::make_unique<SectionMemoryManager>(); }),
51+
[](const MemoryBuffer &) {
52+
return std::make_unique<SectionMemoryManager>();
53+
}),
5254
CompileLayer(*this->ES, ObjectLayer,
5355
std::make_unique<ConcurrentIRCompiler>(std::move(JTMB))),
5456
MainJD(this->ES->createBareJITDylib("<main>")) {

‎llvm/include/llvm/ExecutionEngine/Orc/RTDyldObjectLinkingLayer.h

Lines changed: 2 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -50,7 +50,8 @@ class LLVM_ABI RTDyldObjectLinkingLayer
5050
MaterializationResponsibility &R, std::unique_ptr<MemoryBuffer>)>;
5151

5252
using GetMemoryManagerFunction =
53-
unique_function<std::unique_ptr<RuntimeDyld::MemoryManager>()>;
53+
unique_function<std::unique_ptr<RuntimeDyld::MemoryManager>(
54+
const MemoryBuffer &)>;
5455

5556
/// Construct an ObjectLinkingLayer with the given NotifyLoaded,
5657
/// and NotifyEmitted functors.

‎llvm/lib/ExecutionEngine/Orc/LLJIT.cpp

Lines changed: 7 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -269,9 +269,8 @@ class GenericLLVMIRPlatformSupport : public LLJIT::PlatformSupport {
269269
}
270270

271271
void registerInitFunc(JITDylib &JD, SymbolStringPtr InitName) {
272-
getExecutionSession().runSessionLocked([&]() {
273-
InitFunctions[&JD].add(InitName);
274-
});
272+
getExecutionSession().runSessionLocked(
273+
[&]() { InitFunctions[&JD].add(InitName); });
275274
}
276275

277276
void registerDeInitFunc(JITDylib &JD, SymbolStringPtr DeInitName) {
@@ -935,8 +934,8 @@ Error LLJIT::addObjectFile(JITDylib &JD, std::unique_ptr<MemoryBuffer> Obj) {
935934
Expected<ExecutorAddr> LLJIT::lookupLinkerMangled(JITDylib &JD,
936935
SymbolStringPtr Name) {
937936
if (auto Sym = ES->lookup(
938-
makeJITDylibSearchOrder(&JD, JITDylibLookupFlags::MatchAllSymbols),
939-
Name))
937+
makeJITDylibSearchOrder(&JD, JITDylibLookupFlags::MatchAllSymbols),
938+
Name))
940939
return Sym->getAddress();
941940
else
942941
return Sym.takeError();
@@ -951,7 +950,9 @@ LLJIT::createObjectLinkingLayer(LLJITBuilderState &S, ExecutionSession &ES) {
951950

952951
// Otherwise default to creating an RTDyldObjectLinkingLayer that constructs
953952
// a new SectionMemoryManager for each object.
954-
auto GetMemMgr = []() { return std::make_unique<SectionMemoryManager>(); };
953+
auto GetMemMgr = [](const MemoryBuffer &) {
954+
return std::make_unique<SectionMemoryManager>();
955+
};
955956
auto Layer =
956957
std::make_unique<RTDyldObjectLinkingLayer>(ES, std::move(GetMemMgr));
957958

‎llvm/lib/ExecutionEngine/Orc/OrcV2CBindings.cpp

Lines changed: 8 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -1021,8 +1021,10 @@ LLVMOrcObjectLayerRef
10211021
LLVMOrcCreateRTDyldObjectLinkingLayerWithSectionMemoryManager(
10221022
LLVMOrcExecutionSessionRef ES) {
10231023
assert(ES && "ES must not be null");
1024-
return wrap(new RTDyldObjectLinkingLayer(
1025-
*unwrap(ES), [] { return std::make_unique<SectionMemoryManager>(); }));
1024+
return wrap(
1025+
new RTDyldObjectLinkingLayer(*unwrap(ES), [](const MemoryBuffer &) {
1026+
return std::make_unique<SectionMemoryManager>();
1027+
}));
10261028
}
10271029

10281030
LLVMOrcObjectLayerRef
@@ -1128,9 +1130,10 @@ LLVMOrcCreateRTDyldObjectLinkingLayerWithMCJITMemoryManagerLikeCallbacks(
11281130
CreateContextCtx, CreateContext, NotifyTerminating, AllocateCodeSection,
11291131
AllocateDataSection, FinalizeMemory, Destroy);
11301132

1131-
return wrap(new RTDyldObjectLinkingLayer(*unwrap(ES), [CBs = std::move(CBs)] {
1132-
return std::make_unique<MCJITMemoryManagerLikeCallbacksMemMgr>(CBs);
1133-
}));
1133+
return wrap(new RTDyldObjectLinkingLayer(
1134+
*unwrap(ES), [CBs = std::move(CBs)](const MemoryBuffer &) {
1135+
return std::make_unique<MCJITMemoryManagerLikeCallbacksMemMgr>(CBs);
1136+
}));
11341137

11351138
return nullptr;
11361139
}

‎llvm/lib/ExecutionEngine/Orc/RTDyldObjectLinkingLayer.cpp

Lines changed: 3 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -22,7 +22,8 @@ class JITDylibSearchOrderResolver : public JITSymbolResolver {
2222
SymbolDependenceMap &Deps)
2323
: MR(MR), Deps(Deps) {}
2424

25-
void lookup(const LookupSet &Symbols, OnResolvedFunction OnResolved) override {
25+
void lookup(const LookupSet &Symbols,
26+
OnResolvedFunction OnResolved) override {
2627
auto &ES = MR.getTargetJITDylib().getExecutionSession();
2728
SymbolLookupSet InternedSymbols;
2829

@@ -181,7 +182,7 @@ void RTDyldObjectLinkingLayer::emit(
181182
}
182183
}
183184

184-
auto MemMgr = GetMemoryManager();
185+
auto MemMgr = GetMemoryManager(*O);
185186
auto &MemMgrRef = *MemMgr;
186187

187188
// Switch to shared ownership of MR so that it can be captured by both

‎llvm/unittests/ExecutionEngine/Orc/RTDyldObjectLinkingLayerTest.cpp

Lines changed: 69 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -50,9 +50,10 @@ static bool testSetProcessAllSections(std::unique_ptr<MemoryBuffer> Obj,
5050
auto &JD = ES.createBareJITDylib("main");
5151
auto Foo = ES.intern("foo");
5252

53-
RTDyldObjectLinkingLayer ObjLayer(ES, [&NonAllocSectionSeen]() {
54-
return std::make_unique<MemoryManagerWrapper>(NonAllocSectionSeen);
55-
});
53+
RTDyldObjectLinkingLayer ObjLayer(
54+
ES, [&NonAllocSectionSeen](const MemoryBuffer &) {
55+
return std::make_unique<MemoryManagerWrapper>(NonAllocSectionSeen);
56+
});
5657

5758
auto OnResolveDoNothing = [](Expected<SymbolMap> R) {
5859
cantFail(std::move(R));
@@ -156,8 +157,9 @@ TEST(RTDyldObjectLinkingLayerTest, TestOverrideObjectFlags) {
156157
ExecutionSession ES{std::make_unique<UnsupportedExecutorProcessControl>()};
157158
auto &JD = ES.createBareJITDylib("main");
158159
auto Foo = ES.intern("foo");
159-
RTDyldObjectLinkingLayer ObjLayer(
160-
ES, []() { return std::make_unique<SectionMemoryManager>(); });
160+
RTDyldObjectLinkingLayer ObjLayer(ES, [](const MemoryBuffer &) {
161+
return std::make_unique<SectionMemoryManager>();
162+
});
161163
IRCompileLayer CompileLayer(ES, ObjLayer,
162164
std::make_unique<FunkySimpleCompiler>(*TM));
163165

@@ -226,8 +228,9 @@ TEST(RTDyldObjectLinkingLayerTest, TestAutoClaimResponsibilityForSymbols) {
226228
ExecutionSession ES{std::make_unique<UnsupportedExecutorProcessControl>()};
227229
auto &JD = ES.createBareJITDylib("main");
228230
auto Foo = ES.intern("foo");
229-
RTDyldObjectLinkingLayer ObjLayer(
230-
ES, []() { return std::make_unique<SectionMemoryManager>(); });
231+
RTDyldObjectLinkingLayer ObjLayer(ES, [](const MemoryBuffer &) {
232+
return std::make_unique<SectionMemoryManager>();
233+
});
231234
IRCompileLayer CompileLayer(ES, ObjLayer,
232235
std::make_unique<FunkySimpleCompiler>(*TM));
233236

@@ -244,4 +247,63 @@ TEST(RTDyldObjectLinkingLayerTest, TestAutoClaimResponsibilityForSymbols) {
244247
ES.reportError(std::move(Err));
245248
}
246249

250+
TEST(RTDyldObjectLinkingLayerTest, TestMemoryBufferNamePropagation) {
251+
OrcNativeTarget::initialize();
252+
253+
std::unique_ptr<TargetMachine> TM(
254+
EngineBuilder().selectTarget(Triple("x86_64-unknown-linux-gnu"), "", "",
255+
SmallVector<std::string, 1>()));
256+
257+
if (!TM)
258+
GTEST_SKIP();
259+
260+
// Create a module with two void() functions: foo and bar.
261+
ThreadSafeContext TSCtx(std::make_unique<LLVMContext>());
262+
ThreadSafeModule M;
263+
{
264+
ModuleBuilder MB(*TSCtx.getContext(), TM->getTargetTriple().str(), "dummy");
265+
MB.getModule()->setDataLayout(TM->createDataLayout());
266+
267+
Function *FooImpl = MB.createFunctionDecl(
268+
FunctionType::get(Type::getVoidTy(*TSCtx.getContext()), {}, false),
269+
"foo");
270+
BasicBlock *FooEntry =
271+
BasicBlock::Create(*TSCtx.getContext(), "entry", FooImpl);
272+
IRBuilder<> B1(FooEntry);
273+
B1.CreateRetVoid();
274+
275+
M = ThreadSafeModule(MB.takeModule(), std::move(TSCtx));
276+
}
277+
278+
ExecutionSession ES{std::make_unique<UnsupportedExecutorProcessControl>()};
279+
auto &JD = ES.createBareJITDylib("main");
280+
auto Foo = ES.intern("foo");
281+
std::string ObjectIdentifer;
282+
283+
RTDyldObjectLinkingLayer ObjLayer(
284+
ES, [&ObjectIdentifer](const MemoryBuffer &Obj) {
285+
// Capture the name of the object so that we can confirm that it
286+
// contains the module name.
287+
ObjectIdentifer = Obj.getBufferIdentifier().str();
288+
return std::make_unique<SectionMemoryManager>();
289+
});
290+
IRCompileLayer CompileLayer(ES, ObjLayer,
291+
std::make_unique<SimpleCompiler>(*TM));
292+
293+
// Capture the module name before we move the module.
294+
std::string ModuleName = M.getModuleUnlocked()->getName().str();
295+
296+
cantFail(CompileLayer.add(JD, std::move(M)));
297+
ES.lookup(
298+
LookupKind::Static, makeJITDylibSearchOrder(&JD), SymbolLookupSet(Foo),
299+
SymbolState::Resolved,
300+
[](Expected<SymbolMap> R) { cantFail(std::move(R)); },
301+
NoDependenciesToRegister);
302+
303+
if (auto Err = ES.endSession())
304+
ES.reportError(std::move(Err));
305+
306+
EXPECT_TRUE(ObjectIdentifer.find(ModuleName) != std::string::npos);
307+
}
308+
247309
} // end anonymous namespace

‎mlir/lib/ExecutionEngine/ExecutionEngine.cpp

Lines changed: 2 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -315,7 +315,8 @@ ExecutionEngine::create(Operation *m, const ExecutionEngineOptions &options,
315315
// process and dynamically linked libraries.
316316
auto objectLinkingLayerCreator = [&](ExecutionSession &session) {
317317
auto objectLayer = std::make_unique<RTDyldObjectLinkingLayer>(
318-
session, [sectionMemoryMapper = options.sectionMemoryMapper]() {
318+
session, [sectionMemoryMapper =
319+
options.sectionMemoryMapper](const MemoryBuffer &) {
319320
return std::make_unique<SectionMemoryManager>(sectionMemoryMapper);
320321
});
321322

0 commit comments

Comments
 (0)