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.
This commit is contained in:
Karlo Basioli
2025-06-06 00:44:39 +01:00
committed by GitHub
parent f961d6a89a
commit cd585864c0
11 changed files with 106 additions and 27 deletions

View File

@@ -47,7 +47,9 @@ public:
JITTargetMachineBuilder JTMB, DataLayout DL)
: ES(std::move(ES)), DL(std::move(DL)), Mangle(*this->ES, this->DL),
ObjectLayer(*this->ES,
[]() { return std::make_unique<SectionMemoryManager>(); }),
[](const MemoryBuffer &) {
return std::make_unique<SectionMemoryManager>();
}),
CompileLayer(*this->ES, ObjectLayer,
std::make_unique<ConcurrentIRCompiler>(std::move(JTMB))),
MainJD(this->ES->createBareJITDylib("<main>")) {

View File

@@ -53,7 +53,9 @@ public:
JITTargetMachineBuilder JTMB, DataLayout DL)
: ES(std::move(ES)), DL(std::move(DL)), Mangle(*this->ES, this->DL),
ObjectLayer(*this->ES,
[]() { return std::make_unique<SectionMemoryManager>(); }),
[](const MemoryBuffer &) {
return std::make_unique<SectionMemoryManager>();
}),
CompileLayer(*this->ES, ObjectLayer,
std::make_unique<ConcurrentIRCompiler>(std::move(JTMB))),
OptimizeLayer(*this->ES, CompileLayer, optimizeModule),

View File

@@ -66,7 +66,9 @@ public:
: ES(std::move(ES)), EPCIU(std::move(EPCIU)), DL(std::move(DL)),
Mangle(*this->ES, this->DL),
ObjectLayer(*this->ES,
[]() { return std::make_unique<SectionMemoryManager>(); }),
[](const MemoryBuffer &) {
return std::make_unique<SectionMemoryManager>();
}),
CompileLayer(*this->ES, ObjectLayer,
std::make_unique<ConcurrentIRCompiler>(std::move(JTMB))),
OptimizeLayer(*this->ES, CompileLayer, optimizeModule),

View File

@@ -151,7 +151,9 @@ public:
: ES(std::move(ES)), EPCIU(std::move(EPCIU)), DL(std::move(DL)),
Mangle(*this->ES, this->DL),
ObjectLayer(*this->ES,
[]() { return std::make_unique<SectionMemoryManager>(); }),
[](const MemoryBuffer &) {
return std::make_unique<SectionMemoryManager>();
}),
CompileLayer(*this->ES, ObjectLayer,
std::make_unique<ConcurrentIRCompiler>(std::move(JTMB))),
OptimizeLayer(*this->ES, CompileLayer, optimizeModule),

View File

@@ -48,7 +48,9 @@ public:
JITTargetMachineBuilder JTMB, DataLayout DL)
: ES(std::move(ES)), DL(std::move(DL)), Mangle(*this->ES, this->DL),
ObjectLayer(*this->ES,
[]() { return std::make_unique<SectionMemoryManager>(); }),
[](const MemoryBuffer &) {
return std::make_unique<SectionMemoryManager>();
}),
CompileLayer(*this->ES, ObjectLayer,
std::make_unique<ConcurrentIRCompiler>(std::move(JTMB))),
MainJD(this->ES->createBareJITDylib("<main>")) {

View File

@@ -50,7 +50,8 @@ public:
MaterializationResponsibility &R, std::unique_ptr<MemoryBuffer>)>;
using GetMemoryManagerFunction =
unique_function<std::unique_ptr<RuntimeDyld::MemoryManager>()>;
unique_function<std::unique_ptr<RuntimeDyld::MemoryManager>(
const MemoryBuffer &)>;
/// Construct an ObjectLinkingLayer with the given NotifyLoaded,
/// and NotifyEmitted functors.

View File

@@ -269,9 +269,8 @@ public:
}
void registerInitFunc(JITDylib &JD, SymbolStringPtr InitName) {
getExecutionSession().runSessionLocked([&]() {
InitFunctions[&JD].add(InitName);
});
getExecutionSession().runSessionLocked(
[&]() { InitFunctions[&JD].add(InitName); });
}
void registerDeInitFunc(JITDylib &JD, SymbolStringPtr DeInitName) {
@@ -935,8 +934,8 @@ Error LLJIT::addObjectFile(JITDylib &JD, std::unique_ptr<MemoryBuffer> Obj) {
Expected<ExecutorAddr> LLJIT::lookupLinkerMangled(JITDylib &JD,
SymbolStringPtr Name) {
if (auto Sym = ES->lookup(
makeJITDylibSearchOrder(&JD, JITDylibLookupFlags::MatchAllSymbols),
Name))
makeJITDylibSearchOrder(&JD, JITDylibLookupFlags::MatchAllSymbols),
Name))
return Sym->getAddress();
else
return Sym.takeError();
@@ -951,7 +950,9 @@ LLJIT::createObjectLinkingLayer(LLJITBuilderState &S, ExecutionSession &ES) {
// Otherwise default to creating an RTDyldObjectLinkingLayer that constructs
// a new SectionMemoryManager for each object.
auto GetMemMgr = []() { return std::make_unique<SectionMemoryManager>(); };
auto GetMemMgr = [](const MemoryBuffer &) {
return std::make_unique<SectionMemoryManager>();
};
auto Layer =
std::make_unique<RTDyldObjectLinkingLayer>(ES, std::move(GetMemMgr));

View File

@@ -1021,8 +1021,10 @@ LLVMOrcObjectLayerRef
LLVMOrcCreateRTDyldObjectLinkingLayerWithSectionMemoryManager(
LLVMOrcExecutionSessionRef ES) {
assert(ES && "ES must not be null");
return wrap(new RTDyldObjectLinkingLayer(
*unwrap(ES), [] { return std::make_unique<SectionMemoryManager>(); }));
return wrap(
new RTDyldObjectLinkingLayer(*unwrap(ES), [](const MemoryBuffer &) {
return std::make_unique<SectionMemoryManager>();
}));
}
LLVMOrcObjectLayerRef
@@ -1128,9 +1130,10 @@ LLVMOrcCreateRTDyldObjectLinkingLayerWithMCJITMemoryManagerLikeCallbacks(
CreateContextCtx, CreateContext, NotifyTerminating, AllocateCodeSection,
AllocateDataSection, FinalizeMemory, Destroy);
return wrap(new RTDyldObjectLinkingLayer(*unwrap(ES), [CBs = std::move(CBs)] {
return std::make_unique<MCJITMemoryManagerLikeCallbacksMemMgr>(CBs);
}));
return wrap(new RTDyldObjectLinkingLayer(
*unwrap(ES), [CBs = std::move(CBs)](const MemoryBuffer &) {
return std::make_unique<MCJITMemoryManagerLikeCallbacksMemMgr>(CBs);
}));
return nullptr;
}

View File

@@ -22,7 +22,8 @@ public:
SymbolDependenceMap &Deps)
: MR(MR), Deps(Deps) {}
void lookup(const LookupSet &Symbols, OnResolvedFunction OnResolved) override {
void lookup(const LookupSet &Symbols,
OnResolvedFunction OnResolved) override {
auto &ES = MR.getTargetJITDylib().getExecutionSession();
SymbolLookupSet InternedSymbols;
@@ -181,7 +182,7 @@ void RTDyldObjectLinkingLayer::emit(
}
}
auto MemMgr = GetMemoryManager();
auto MemMgr = GetMemoryManager(*O);
auto &MemMgrRef = *MemMgr;
// Switch to shared ownership of MR so that it can be captured by both

View File

@@ -50,9 +50,10 @@ static bool testSetProcessAllSections(std::unique_ptr<MemoryBuffer> Obj,
auto &JD = ES.createBareJITDylib("main");
auto Foo = ES.intern("foo");
RTDyldObjectLinkingLayer ObjLayer(ES, [&NonAllocSectionSeen]() {
return std::make_unique<MemoryManagerWrapper>(NonAllocSectionSeen);
});
RTDyldObjectLinkingLayer ObjLayer(
ES, [&NonAllocSectionSeen](const MemoryBuffer &) {
return std::make_unique<MemoryManagerWrapper>(NonAllocSectionSeen);
});
auto OnResolveDoNothing = [](Expected<SymbolMap> R) {
cantFail(std::move(R));
@@ -156,8 +157,9 @@ TEST(RTDyldObjectLinkingLayerTest, TestOverrideObjectFlags) {
ExecutionSession ES{std::make_unique<UnsupportedExecutorProcessControl>()};
auto &JD = ES.createBareJITDylib("main");
auto Foo = ES.intern("foo");
RTDyldObjectLinkingLayer ObjLayer(
ES, []() { return std::make_unique<SectionMemoryManager>(); });
RTDyldObjectLinkingLayer ObjLayer(ES, [](const MemoryBuffer &) {
return std::make_unique<SectionMemoryManager>();
});
IRCompileLayer CompileLayer(ES, ObjLayer,
std::make_unique<FunkySimpleCompiler>(*TM));
@@ -226,8 +228,9 @@ TEST(RTDyldObjectLinkingLayerTest, TestAutoClaimResponsibilityForSymbols) {
ExecutionSession ES{std::make_unique<UnsupportedExecutorProcessControl>()};
auto &JD = ES.createBareJITDylib("main");
auto Foo = ES.intern("foo");
RTDyldObjectLinkingLayer ObjLayer(
ES, []() { return std::make_unique<SectionMemoryManager>(); });
RTDyldObjectLinkingLayer ObjLayer(ES, [](const MemoryBuffer &) {
return std::make_unique<SectionMemoryManager>();
});
IRCompileLayer CompileLayer(ES, ObjLayer,
std::make_unique<FunkySimpleCompiler>(*TM));
@@ -244,4 +247,63 @@ TEST(RTDyldObjectLinkingLayerTest, TestAutoClaimResponsibilityForSymbols) {
ES.reportError(std::move(Err));
}
TEST(RTDyldObjectLinkingLayerTest, TestMemoryBufferNamePropagation) {
OrcNativeTarget::initialize();
std::unique_ptr<TargetMachine> TM(
EngineBuilder().selectTarget(Triple("x86_64-unknown-linux-gnu"), "", "",
SmallVector<std::string, 1>()));
if (!TM)
GTEST_SKIP();
// Create a module with two void() functions: foo and bar.
ThreadSafeContext TSCtx(std::make_unique<LLVMContext>());
ThreadSafeModule M;
{
ModuleBuilder MB(*TSCtx.getContext(), TM->getTargetTriple().str(), "dummy");
MB.getModule()->setDataLayout(TM->createDataLayout());
Function *FooImpl = MB.createFunctionDecl(
FunctionType::get(Type::getVoidTy(*TSCtx.getContext()), {}, false),
"foo");
BasicBlock *FooEntry =
BasicBlock::Create(*TSCtx.getContext(), "entry", FooImpl);
IRBuilder<> B1(FooEntry);
B1.CreateRetVoid();
M = ThreadSafeModule(MB.takeModule(), std::move(TSCtx));
}
ExecutionSession ES{std::make_unique<UnsupportedExecutorProcessControl>()};
auto &JD = ES.createBareJITDylib("main");
auto Foo = ES.intern("foo");
std::string ObjectIdentifer;
RTDyldObjectLinkingLayer ObjLayer(
ES, [&ObjectIdentifer](const MemoryBuffer &Obj) {
// Capture the name of the object so that we can confirm that it
// contains the module name.
ObjectIdentifer = Obj.getBufferIdentifier().str();
return std::make_unique<SectionMemoryManager>();
});
IRCompileLayer CompileLayer(ES, ObjLayer,
std::make_unique<SimpleCompiler>(*TM));
// Capture the module name before we move the module.
std::string ModuleName = M.getModuleUnlocked()->getName().str();
cantFail(CompileLayer.add(JD, std::move(M)));
ES.lookup(
LookupKind::Static, makeJITDylibSearchOrder(&JD), SymbolLookupSet(Foo),
SymbolState::Resolved,
[](Expected<SymbolMap> R) { cantFail(std::move(R)); },
NoDependenciesToRegister);
if (auto Err = ES.endSession())
ES.reportError(std::move(Err));
EXPECT_TRUE(ObjectIdentifer.find(ModuleName) != std::string::npos);
}
} // end anonymous namespace

View File

@@ -315,7 +315,8 @@ ExecutionEngine::create(Operation *m, const ExecutionEngineOptions &options,
// process and dynamically linked libraries.
auto objectLinkingLayerCreator = [&](ExecutionSession &session) {
auto objectLayer = std::make_unique<RTDyldObjectLinkingLayer>(
session, [sectionMemoryMapper = options.sectionMemoryMapper]() {
session, [sectionMemoryMapper =
options.sectionMemoryMapper](const MemoryBuffer &) {
return std::make_unique<SectionMemoryManager>(sectionMemoryMapper);
});