Pass memory buffer to RuntimeDyld::MemoryManager factory

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<SectionMemoryManager>(); }),
+                    [](const MemoryBuffer &) {
+                      return std::make_unique<SectionMemoryManager>();
+                    }),
         CompileLayer(*this->ES, ObjectLayer,
                      std::make_unique<ConcurrentIRCompiler>(std::move(JTMB))),
         MainJD(this->ES->createBareJITDylib("<main>")) {

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<SectionMemoryManager>(); }),
+                    [](const MemoryBuffer &) {
+                      return std::make_unique<SectionMemoryManager>();
+                    }),
         CompileLayer(*this->ES, ObjectLayer,
                      std::make_unique<ConcurrentIRCompiler>(std::move(JTMB))),
         OptimizeLayer(*this->ES, CompileLayer, optimizeModule),

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<SectionMemoryManager>(); }),
+                    [](const MemoryBuffer &) {
+                      return std::make_unique<SectionMemoryManager>();
+                    }),
         CompileLayer(*this->ES, ObjectLayer,
                      std::make_unique<ConcurrentIRCompiler>(std::move(JTMB))),
         OptimizeLayer(*this->ES, CompileLayer, optimizeModule),

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<SectionMemoryManager>(); }),
+                    [](const MemoryBuffer &) {
+                      return std::make_unique<SectionMemoryManager>();
+                    }),
         CompileLayer(*this->ES, ObjectLayer,
                      std::make_unique<ConcurrentIRCompiler>(std::move(JTMB))),
         OptimizeLayer(*this->ES, CompileLayer, optimizeModule),

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<SectionMemoryManager>(); }),
+                    [](const MemoryBuffer &) {
+                      return std::make_unique<SectionMemoryManager>();
+                    }),
         CompileLayer(*this->ES, ObjectLayer,
                      std::make_unique<ConcurrentIRCompiler>(std::move(JTMB))),
         MainJD(this->ES->createBareJITDylib("<main>")) {

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.

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<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));
 

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<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;
 }

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

llvm/unittests/ExecutionEngine/Orc/RTDyldObjectLinkingLayerTest.cpp

@@ -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

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);
         });