From 8f5000668e7147cecdb68fb54d3637508c56f05f Mon Sep 17 00:00:00 2001 From: basioli-k Date: Thu, 5 Jun 2025 08:38:28 +0000 Subject: [PATCH] Pass memory buffer to RuntimeDyld::MemoryManager factory --- .../BuildingAJIT/Chapter1/KaleidoscopeJIT.h | 4 +- .../BuildingAJIT/Chapter2/KaleidoscopeJIT.h | 4 +- .../BuildingAJIT/Chapter3/KaleidoscopeJIT.h | 4 +- .../BuildingAJIT/Chapter4/KaleidoscopeJIT.h | 4 +- .../Kaleidoscope/include/KaleidoscopeJIT.h | 4 +- .../Orc/RTDyldObjectLinkingLayer.h | 3 +- llvm/lib/ExecutionEngine/Orc/LLJIT.cpp | 13 ++-- .../ExecutionEngine/Orc/OrcV2CBindings.cpp | 13 ++-- .../Orc/RTDyldObjectLinkingLayer.cpp | 5 +- .../Orc/RTDyldObjectLinkingLayerTest.cpp | 76 +++++++++++++++++-- mlir/lib/ExecutionEngine/ExecutionEngine.cpp | 3 +- 11 files changed, 106 insertions(+), 27 deletions(-) diff --git a/llvm/examples/Kaleidoscope/BuildingAJIT/Chapter1/KaleidoscopeJIT.h b/llvm/examples/Kaleidoscope/BuildingAJIT/Chapter1/KaleidoscopeJIT.h index bb9e6b8c99d28..73a7b5b38a385 100644 --- a/llvm/examples/Kaleidoscope/BuildingAJIT/Chapter1/KaleidoscopeJIT.h +++ b/llvm/examples/Kaleidoscope/BuildingAJIT/Chapter1/KaleidoscopeJIT.h @@ -47,7 +47,9 @@ class KaleidoscopeJIT { JITTargetMachineBuilder JTMB, DataLayout DL) : ES(std::move(ES)), DL(std::move(DL)), Mangle(*this->ES, this->DL), ObjectLayer(*this->ES, - []() { return std::make_unique(); }), + [](const MemoryBuffer &) { + return std::make_unique(); + }), CompileLayer(*this->ES, ObjectLayer, std::make_unique(std::move(JTMB))), MainJD(this->ES->createBareJITDylib("
")) { diff --git a/llvm/examples/Kaleidoscope/BuildingAJIT/Chapter2/KaleidoscopeJIT.h b/llvm/examples/Kaleidoscope/BuildingAJIT/Chapter2/KaleidoscopeJIT.h index 2f7846e7ee2cb..d2bedd2c7270f 100644 --- a/llvm/examples/Kaleidoscope/BuildingAJIT/Chapter2/KaleidoscopeJIT.h +++ b/llvm/examples/Kaleidoscope/BuildingAJIT/Chapter2/KaleidoscopeJIT.h @@ -53,7 +53,9 @@ class KaleidoscopeJIT { JITTargetMachineBuilder JTMB, DataLayout DL) : ES(std::move(ES)), DL(std::move(DL)), Mangle(*this->ES, this->DL), ObjectLayer(*this->ES, - []() { return std::make_unique(); }), + [](const MemoryBuffer &) { + return std::make_unique(); + }), CompileLayer(*this->ES, ObjectLayer, std::make_unique(std::move(JTMB))), OptimizeLayer(*this->ES, CompileLayer, optimizeModule), diff --git a/llvm/examples/Kaleidoscope/BuildingAJIT/Chapter3/KaleidoscopeJIT.h b/llvm/examples/Kaleidoscope/BuildingAJIT/Chapter3/KaleidoscopeJIT.h index fee2d26e5d925..09b87b0181f92 100644 --- a/llvm/examples/Kaleidoscope/BuildingAJIT/Chapter3/KaleidoscopeJIT.h +++ b/llvm/examples/Kaleidoscope/BuildingAJIT/Chapter3/KaleidoscopeJIT.h @@ -66,7 +66,9 @@ class KaleidoscopeJIT { : ES(std::move(ES)), EPCIU(std::move(EPCIU)), DL(std::move(DL)), Mangle(*this->ES, this->DL), ObjectLayer(*this->ES, - []() { return std::make_unique(); }), + [](const MemoryBuffer &) { + return std::make_unique(); + }), CompileLayer(*this->ES, ObjectLayer, std::make_unique(std::move(JTMB))), OptimizeLayer(*this->ES, CompileLayer, optimizeModule), diff --git a/llvm/examples/Kaleidoscope/BuildingAJIT/Chapter4/KaleidoscopeJIT.h b/llvm/examples/Kaleidoscope/BuildingAJIT/Chapter4/KaleidoscopeJIT.h index 136e88505a1d3..30c0e239a6cb4 100644 --- a/llvm/examples/Kaleidoscope/BuildingAJIT/Chapter4/KaleidoscopeJIT.h +++ b/llvm/examples/Kaleidoscope/BuildingAJIT/Chapter4/KaleidoscopeJIT.h @@ -151,7 +151,9 @@ class KaleidoscopeJIT { : ES(std::move(ES)), EPCIU(std::move(EPCIU)), DL(std::move(DL)), Mangle(*this->ES, this->DL), ObjectLayer(*this->ES, - []() { return std::make_unique(); }), + [](const MemoryBuffer &) { + return std::make_unique(); + }), CompileLayer(*this->ES, ObjectLayer, std::make_unique(std::move(JTMB))), OptimizeLayer(*this->ES, CompileLayer, optimizeModule), diff --git a/llvm/examples/Kaleidoscope/include/KaleidoscopeJIT.h b/llvm/examples/Kaleidoscope/include/KaleidoscopeJIT.h index 778437cd8a3d6..65f94ed7a80d6 100644 --- a/llvm/examples/Kaleidoscope/include/KaleidoscopeJIT.h +++ b/llvm/examples/Kaleidoscope/include/KaleidoscopeJIT.h @@ -48,7 +48,9 @@ class KaleidoscopeJIT { JITTargetMachineBuilder JTMB, DataLayout DL) : ES(std::move(ES)), DL(std::move(DL)), Mangle(*this->ES, this->DL), ObjectLayer(*this->ES, - []() { return std::make_unique(); }), + [](const MemoryBuffer &) { + return std::make_unique(); + }), CompileLayer(*this->ES, ObjectLayer, std::make_unique(std::move(JTMB))), MainJD(this->ES->createBareJITDylib("
")) { diff --git a/llvm/include/llvm/ExecutionEngine/Orc/RTDyldObjectLinkingLayer.h b/llvm/include/llvm/ExecutionEngine/Orc/RTDyldObjectLinkingLayer.h index 05c9b574aa0f0..1fb472a177d6d 100644 --- a/llvm/include/llvm/ExecutionEngine/Orc/RTDyldObjectLinkingLayer.h +++ b/llvm/include/llvm/ExecutionEngine/Orc/RTDyldObjectLinkingLayer.h @@ -50,7 +50,8 @@ class LLVM_ABI RTDyldObjectLinkingLayer MaterializationResponsibility &R, std::unique_ptr)>; using GetMemoryManagerFunction = - unique_function()>; + unique_function( + const MemoryBuffer &)>; /// Construct an ObjectLinkingLayer with the given NotifyLoaded, /// and NotifyEmitted functors. diff --git a/llvm/lib/ExecutionEngine/Orc/LLJIT.cpp b/llvm/lib/ExecutionEngine/Orc/LLJIT.cpp index 21ebe82c8a71a..f1a71e4acb46d 100644 --- a/llvm/lib/ExecutionEngine/Orc/LLJIT.cpp +++ b/llvm/lib/ExecutionEngine/Orc/LLJIT.cpp @@ -269,9 +269,8 @@ class GenericLLVMIRPlatformSupport : public LLJIT::PlatformSupport { } 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 Obj) { Expected 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(); }; + auto GetMemMgr = [](const MemoryBuffer &) { + return std::make_unique(); + }; auto Layer = std::make_unique(ES, std::move(GetMemMgr)); diff --git a/llvm/lib/ExecutionEngine/Orc/OrcV2CBindings.cpp b/llvm/lib/ExecutionEngine/Orc/OrcV2CBindings.cpp index d44199f1698e9..9999e1ff3bf00 100644 --- a/llvm/lib/ExecutionEngine/Orc/OrcV2CBindings.cpp +++ b/llvm/lib/ExecutionEngine/Orc/OrcV2CBindings.cpp @@ -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(); })); + return wrap( + new RTDyldObjectLinkingLayer(*unwrap(ES), [](const MemoryBuffer &) { + return std::make_unique(); + })); } 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(CBs); - })); + return wrap(new RTDyldObjectLinkingLayer( + *unwrap(ES), [CBs = std::move(CBs)](const MemoryBuffer &) { + return std::make_unique(CBs); + })); return nullptr; } diff --git a/llvm/lib/ExecutionEngine/Orc/RTDyldObjectLinkingLayer.cpp b/llvm/lib/ExecutionEngine/Orc/RTDyldObjectLinkingLayer.cpp index 88cceacb71184..4d4a705ced64e 100644 --- a/llvm/lib/ExecutionEngine/Orc/RTDyldObjectLinkingLayer.cpp +++ b/llvm/lib/ExecutionEngine/Orc/RTDyldObjectLinkingLayer.cpp @@ -22,7 +22,8 @@ class JITDylibSearchOrderResolver : public JITSymbolResolver { 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 diff --git a/llvm/unittests/ExecutionEngine/Orc/RTDyldObjectLinkingLayerTest.cpp b/llvm/unittests/ExecutionEngine/Orc/RTDyldObjectLinkingLayerTest.cpp index ed1feed6a5a9e..a0ae9308d33ab 100644 --- a/llvm/unittests/ExecutionEngine/Orc/RTDyldObjectLinkingLayerTest.cpp +++ b/llvm/unittests/ExecutionEngine/Orc/RTDyldObjectLinkingLayerTest.cpp @@ -50,9 +50,10 @@ static bool testSetProcessAllSections(std::unique_ptr Obj, auto &JD = ES.createBareJITDylib("main"); auto Foo = ES.intern("foo"); - RTDyldObjectLinkingLayer ObjLayer(ES, [&NonAllocSectionSeen]() { - return std::make_unique(NonAllocSectionSeen); - }); + RTDyldObjectLinkingLayer ObjLayer( + ES, [&NonAllocSectionSeen](const MemoryBuffer &) { + return std::make_unique(NonAllocSectionSeen); + }); auto OnResolveDoNothing = [](Expected R) { cantFail(std::move(R)); @@ -156,8 +157,9 @@ TEST(RTDyldObjectLinkingLayerTest, TestOverrideObjectFlags) { ExecutionSession ES{std::make_unique()}; auto &JD = ES.createBareJITDylib("main"); auto Foo = ES.intern("foo"); - RTDyldObjectLinkingLayer ObjLayer( - ES, []() { return std::make_unique(); }); + RTDyldObjectLinkingLayer ObjLayer(ES, [](const MemoryBuffer &) { + return std::make_unique(); + }); IRCompileLayer CompileLayer(ES, ObjLayer, std::make_unique(*TM)); @@ -226,8 +228,9 @@ TEST(RTDyldObjectLinkingLayerTest, TestAutoClaimResponsibilityForSymbols) { ExecutionSession ES{std::make_unique()}; auto &JD = ES.createBareJITDylib("main"); auto Foo = ES.intern("foo"); - RTDyldObjectLinkingLayer ObjLayer( - ES, []() { return std::make_unique(); }); + RTDyldObjectLinkingLayer ObjLayer(ES, [](const MemoryBuffer &) { + return std::make_unique(); + }); IRCompileLayer CompileLayer(ES, ObjLayer, std::make_unique(*TM)); @@ -244,4 +247,63 @@ TEST(RTDyldObjectLinkingLayerTest, TestAutoClaimResponsibilityForSymbols) { ES.reportError(std::move(Err)); } +TEST(RTDyldObjectLinkingLayerTest, TestMemoryBufferNamePropagation) { + OrcNativeTarget::initialize(); + + std::unique_ptr TM( + EngineBuilder().selectTarget(Triple("x86_64-unknown-linux-gnu"), "", "", + SmallVector())); + + if (!TM) + GTEST_SKIP(); + + // Create a module with two void() functions: foo and bar. + ThreadSafeContext TSCtx(std::make_unique()); + 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()}; + 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(); + }); + IRCompileLayer CompileLayer(ES, ObjLayer, + std::make_unique(*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 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 diff --git a/mlir/lib/ExecutionEngine/ExecutionEngine.cpp b/mlir/lib/ExecutionEngine/ExecutionEngine.cpp index 39f5c003320da..f704fbfbe8fff 100644 --- a/mlir/lib/ExecutionEngine/ExecutionEngine.cpp +++ b/mlir/lib/ExecutionEngine/ExecutionEngine.cpp @@ -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( - session, [sectionMemoryMapper = options.sectionMemoryMapper]() { + session, [sectionMemoryMapper = + options.sectionMemoryMapper](const MemoryBuffer &) { return std::make_unique(sectionMemoryMapper); });