[llvm] [mlir] Pass memory buffer to RuntimeDyld::MemoryManager factory (PR #142930)

[Karlo Basioli via llvm-commits]

https://github.com/basioli-k created https://github.com/llvm/llvm-project/pull/142930

None

>From 069d089d8cefc3398c08188c2fea1da06d017072 Mon Sep 17 00:00:00 2001
From: basioli-k <k.basioli at gmail.com>
Date: Thu, 5 Jun 2025 08:38:28 +0000
Subject: [PATCH] Pass memory buffer to RuntimeDyld::MemoryManager factory

---
 .../ExecutionEngine/Orc/RTDyldObjectLinkingLayer.h |  3 ++-
 llvm/lib/ExecutionEngine/Orc/LLJIT.cpp             | 14 ++++++++------
 llvm/lib/ExecutionEngine/Orc/OrcV2CBindings.cpp    |  8 ++++++--
 .../Orc/RTDyldObjectLinkingLayer.cpp               |  4 ++--
 mlir/lib/ExecutionEngine/ExecutionEngine.cpp       |  3 ++-
 5 files changed, 20 insertions(+), 12 deletions(-)

diff --git a/llvm/include/llvm/ExecutionEngine/Orc/RTDyldObjectLinkingLayer.h b/llvm/include/llvm/ExecutionEngine/Orc/RTDyldObjectLinkingLayer.h
index 05c9b574aa0f0..e712d8553f592 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<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.
diff --git a/llvm/lib/ExecutionEngine/Orc/LLJIT.cpp b/llvm/lib/ExecutionEngine/Orc/LLJIT.cpp
index 21ebe82c8a71a..bb16fadc125fc 100644
--- a/llvm/lib/ExecutionEngine/Orc/LLJIT.cpp
+++ b/llvm/lib/ExecutionEngine/Orc/LLJIT.cpp
@@ -269,9 +269,9 @@ 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 +935,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 +951,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));
 
diff --git a/llvm/lib/ExecutionEngine/Orc/OrcV2CBindings.cpp b/llvm/lib/ExecutionEngine/Orc/OrcV2CBindings.cpp
index d44199f1698e9..234fa7bd29f5e 100644
--- a/llvm/lib/ExecutionEngine/Orc/OrcV2CBindings.cpp
+++ b/llvm/lib/ExecutionEngine/Orc/OrcV2CBindings.cpp
@@ -1022,7 +1022,10 @@ LLVMOrcCreateRTDyldObjectLinkingLayerWithSectionMemoryManager(
     LLVMOrcExecutionSessionRef ES) {
   assert(ES && "ES must not be null");
   return wrap(new RTDyldObjectLinkingLayer(
-      *unwrap(ES), [] { return std::make_unique<SectionMemoryManager>(); }));
+      *unwrap(ES), 
+      [] (const MemoryBuffer &) { 
+        return std::make_unique<SectionMemoryManager>(); 
+      }));
 }
 
 LLVMOrcObjectLayerRef
@@ -1128,7 +1131,8 @@ LLVMOrcCreateRTDyldObjectLinkingLayerWithMCJITMemoryManagerLikeCallbacks(
       CreateContextCtx, CreateContext, NotifyTerminating, AllocateCodeSection,
       AllocateDataSection, FinalizeMemory, Destroy);
 
-  return wrap(new RTDyldObjectLinkingLayer(*unwrap(ES), [CBs = std::move(CBs)] {
+  return wrap(new RTDyldObjectLinkingLayer(*unwrap(ES), 
+  [CBs = std::move(CBs)] (const MemoryBuffer &) {
     return std::make_unique<MCJITMemoryManagerLikeCallbacksMemMgr>(CBs);
   }));
 
diff --git a/llvm/lib/ExecutionEngine/Orc/RTDyldObjectLinkingLayer.cpp b/llvm/lib/ExecutionEngine/Orc/RTDyldObjectLinkingLayer.cpp
index 88cceacb71184..b8fad7a48dfa6 100644
--- a/llvm/lib/ExecutionEngine/Orc/RTDyldObjectLinkingLayer.cpp
+++ b/llvm/lib/ExecutionEngine/Orc/RTDyldObjectLinkingLayer.cpp
@@ -22,7 +22,7 @@ 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 +181,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/mlir/lib/ExecutionEngine/ExecutionEngine.cpp b/mlir/lib/ExecutionEngine/ExecutionEngine.cpp
index 39f5c003320da..c68a97d7f8343 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<RTDyldObjectLinkingLayer>(
-        session, [sectionMemoryMapper = options.sectionMemoryMapper]() {
+        session, [sectionMemoryMapper = options.sectionMemoryMapper]
+        (const MemoryBuffer&) {
           return std::make_unique<SectionMemoryManager>(sectionMemoryMapper);
         });