Skip to content

Commit 387aac1

Browse files
basioli-kbasioli-k
committed
Pass memory buffer to RuntimeDyld::MemoryManager factory
1 parent 8b167db commit 387aac1

File tree

11 files changed

+94
-24
lines changed

11 files changed

+94
-24
lines changed

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

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -47,7 +47,7 @@ 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&) { return std::make_unique<SectionMemoryManager>(); }),
5151
CompileLayer(*this->ES, ObjectLayer,
5252
std::make_unique<ConcurrentIRCompiler>(std::move(JTMB))),
5353
MainJD(this->ES->createBareJITDylib("<main>")) {

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

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -53,7 +53,7 @@ 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&) { return std::make_unique<SectionMemoryManager>(); }),
5757
CompileLayer(*this->ES, ObjectLayer,
5858
std::make_unique<ConcurrentIRCompiler>(std::move(JTMB))),
5959
OptimizeLayer(*this->ES, CompileLayer, optimizeModule),

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

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -66,7 +66,7 @@ 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&) { return std::make_unique<SectionMemoryManager>(); }),
7070
CompileLayer(*this->ES, ObjectLayer,
7171
std::make_unique<ConcurrentIRCompiler>(std::move(JTMB))),
7272
OptimizeLayer(*this->ES, CompileLayer, optimizeModule),

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

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -151,7 +151,7 @@ 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&) { return std::make_unique<SectionMemoryManager>(); }),
155155
CompileLayer(*this->ES, ObjectLayer,
156156
std::make_unique<ConcurrentIRCompiler>(std::move(JTMB))),
157157
OptimizeLayer(*this->ES, CompileLayer, optimizeModule),

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

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -48,7 +48,7 @@ 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&) { return std::make_unique<SectionMemoryManager>(); }),
5252
CompileLayer(*this->ES, ObjectLayer,
5353
std::make_unique<ConcurrentIRCompiler>(std::move(JTMB))),
5454
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: 8 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -269,9 +269,9 @@ 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);
274+
});
275275
}
276276

277277
void registerDeInitFunc(JITDylib &JD, SymbolStringPtr DeInitName) {
@@ -935,8 +935,8 @@ Error LLJIT::addObjectFile(JITDylib &JD, std::unique_ptr<MemoryBuffer> Obj) {
935935
Expected<ExecutorAddr> LLJIT::lookupLinkerMangled(JITDylib &JD,
936936
SymbolStringPtr Name) {
937937
if (auto Sym = ES->lookup(
938-
makeJITDylibSearchOrder(&JD, JITDylibLookupFlags::MatchAllSymbols),
939-
Name))
938+
makeJITDylibSearchOrder(&JD, JITDylibLookupFlags::MatchAllSymbols),
939+
Name))
940940
return Sym->getAddress();
941941
else
942942
return Sym.takeError();
@@ -951,7 +951,9 @@ LLJIT::createObjectLinkingLayer(LLJITBuilderState &S, ExecutionSession &ES) {
951951

952952
// Otherwise default to creating an RTDyldObjectLinkingLayer that constructs
953953
// a new SectionMemoryManager for each object.
954-
auto GetMemMgr = []() { return std::make_unique<SectionMemoryManager>(); };
954+
auto GetMemMgr = [](const MemoryBuffer&) {
955+
return std::make_unique<SectionMemoryManager>();
956+
};
955957
auto Layer =
956958
std::make_unique<RTDyldObjectLinkingLayer>(ES, std::move(GetMemMgr));
957959

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

Lines changed: 6 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -1022,7 +1022,10 @@ LLVMOrcCreateRTDyldObjectLinkingLayerWithSectionMemoryManager(
10221022
LLVMOrcExecutionSessionRef ES) {
10231023
assert(ES && "ES must not be null");
10241024
return wrap(new RTDyldObjectLinkingLayer(
1025-
*unwrap(ES), [] { return std::make_unique<SectionMemoryManager>(); }));
1025+
*unwrap(ES),
1026+
[] (const MemoryBuffer &) {
1027+
return std::make_unique<SectionMemoryManager>();
1028+
}));
10261029
}
10271030

10281031
LLVMOrcObjectLayerRef
@@ -1128,7 +1131,8 @@ LLVMOrcCreateRTDyldObjectLinkingLayerWithMCJITMemoryManagerLikeCallbacks(
11281131
CreateContextCtx, CreateContext, NotifyTerminating, AllocateCodeSection,
11291132
AllocateDataSection, FinalizeMemory, Destroy);
11301133

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

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

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -22,7 +22,7 @@ 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, OnResolvedFunction OnResolved) override {
2626
auto &ES = MR.getTargetJITDylib().getExecutionSession();
2727
SymbolLookupSet InternedSymbols;
2828

@@ -181,7 +181,7 @@ void RTDyldObjectLinkingLayer::emit(
181181
}
182182
}
183183

184-
auto MemMgr = GetMemoryManager();
184+
auto MemMgr = GetMemoryManager(*O);
185185
auto &MemMgrRef = *MemMgr;
186186

187187
// 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 = options.sectionMemoryMapper]
319+
(const MemoryBuffer&) {
319320
return std::make_unique<SectionMemoryManager>(sectionMemoryMapper);
320321
});
321322

0 commit comments

Comments
 (0)