[Orc] Refactor ObjectLinkingLayer::addObjectSet to defer loading objects until

they're needed.

Prior to this patch objects were loaded (via RuntimeDyld::loadObject) when they
were added to the ObjectLinkingLayer, but were not relocated and finalized until
a symbol address was requested. In the interim, another object could be loaded
and finalized with the same memory manager, causing relocation/finalization of
the first object to fail (as the first finalization call may have marked the
allocated memory for the first object read-only).

By deferring the loadObject call (and subsequent memory allocations) until an
object file is needed we can avoid prematurely finalizing memory.



git-svn-id: https://llvm.org/svn/llvm-project/llvm/trunk@258185 91177308-0d34-0410-b5e6-96231b3b80d8

include/llvm/ExecutionEngine/JITSymbolFlags.h

@@ -15,6 +15,7 @@
 #define LLVM_EXECUTIONENGINE_JITSYMBOLFLAGS_H
 
 #include "llvm/IR/GlobalValue.h"
+#include "llvm/Object/SymbolicFile.h"
 
 namespace llvm {
 
@@ -69,7 +70,16 @@ public:
     if (!GV.hasLocalLinkage() && !GV.hasHiddenVisibility())
       Flags |= JITSymbolFlags::Exported;
     return Flags;
+  }
 
+  static JITSymbolFlags
+  flagsFromObjectSymbol(const object::BasicSymbolRef &Symbol) {
+    JITSymbolFlags Flags = JITSymbolFlags::None;
+    if (Symbol.getFlags() & object::BasicSymbolRef::SF_Weak)
+      Flags |= JITSymbolFlags::Weak;
+    if (Symbol.getFlags() & object::BasicSymbolRef::SF_Exported)
+      Flags |= JITSymbolFlags::Exported;
+    return Flags;
   }
 
 private:

include/llvm/ExecutionEngine/Orc/IRCompileLayer.h

@@ -37,9 +37,6 @@ public:
 private:
   typedef typename BaseLayerT::ObjSetHandleT ObjSetHandleT;
 
-  typedef std::vector<std::unique_ptr<object::ObjectFile>> OwningObjectVec;
-  typedef std::vector<std::unique_ptr<MemoryBuffer>> OwningBufferVec;
-
 public:
   /// @brief Handle to a set of compiled modules.
   typedef ObjSetHandleT ModuleSetHandleT;
@@ -62,28 +59,29 @@ public:
   ModuleSetHandleT addModuleSet(ModuleSetT Ms,
                                 MemoryManagerPtrT MemMgr,
                                 SymbolResolverPtrT Resolver) {
-    OwningObjectVec Objects;
-    OwningBufferVec Buffers;
+    std::vector<std::unique_ptr<object::OwningBinary<object::ObjectFile>>>
+      Objects;
 
     for (const auto &M : Ms) {
-      std::unique_ptr<object::ObjectFile> Object;
-      std::unique_ptr<MemoryBuffer> Buffer;
+      auto Object =
+        llvm::make_unique<object::OwningBinary<object::ObjectFile>>();
 
       if (ObjCache)
-        std::tie(Object, Buffer) = tryToLoadFromObjectCache(*M).takeBinary();
+        *Object = tryToLoadFromObjectCache(*M);
 
-      if (!Object) {
-        std::tie(Object, Buffer) = Compile(*M).takeBinary();
+      if (!Object->getBinary()) {
+        *Object = Compile(*M);
         if (ObjCache)
-          ObjCache->notifyObjectCompiled(&*M, Buffer->getMemBufferRef());
+          ObjCache->notifyObjectCompiled(&*M,
+                                     Object->getBinary()->getMemoryBufferRef());
       }
 
       Objects.push_back(std::move(Object));
-      Buffers.push_back(std::move(Buffer));
     }
 
     ModuleSetHandleT H =
-      BaseLayer.addObjectSet(Objects, std::move(MemMgr), std::move(Resolver));
+      BaseLayer.addObjectSet(std::move(Objects), std::move(MemMgr),
+                             std::move(Resolver));
 
     return H;
   }

include/llvm/ExecutionEngine/Orc/ObjectLinkingLayer.h

@@ -26,7 +26,6 @@ namespace orc {
 
 class ObjectLinkingLayerBase {
 protected:
-
   /// @brief Holds a set of objects to be allocated/linked as a unit in the JIT.
   ///
   /// An instance of this class will be created for each set of objects added
@@ -38,38 +37,32 @@ protected:
     LinkedObjectSet(const LinkedObjectSet&) = delete;
     void operator=(const LinkedObjectSet&) = delete;
   public:
-    LinkedObjectSet(RuntimeDyld::MemoryManager &MemMgr,
-                    RuntimeDyld::SymbolResolver &Resolver,
-                    bool ProcessAllSections)
-        : RTDyld(llvm::make_unique<RuntimeDyld>(MemMgr, Resolver)),
-          State(Raw) {
-      RTDyld->setProcessAllSections(ProcessAllSections);
-    }
-
+    LinkedObjectSet() = default;
     virtual ~LinkedObjectSet() {}
 
-    std::unique_ptr<RuntimeDyld::LoadedObjectInfo>
-    addObject(const object::ObjectFile &Obj) {
-      return RTDyld->loadObject(Obj);
+    virtual void finalize() = 0;
+
+    virtual JITSymbol::GetAddressFtor
+    getSymbolMaterializer(std::string Name) = 0;
+
+    virtual void mapSectionAddress(const void *LocalAddress,
+                                   TargetAddress TargetAddr) const = 0;
+
+    JITSymbol getSymbol(StringRef Name, bool ExportedSymbolsOnly) {
+      auto SymEntry = SymbolTable.find(Name);
+      if (SymEntry == SymbolTable.end())
+        return nullptr;
+      if (!SymEntry->second.isExported() && ExportedSymbolsOnly)
+        return nullptr;
+      if (!Finalized)
+        return JITSymbol(getSymbolMaterializer(Name),
+                         SymEntry->second.getFlags());
+      return JITSymbol(SymEntry->second.getAddress(),
+                       SymEntry->second.getFlags());
     }
-
-    RuntimeDyld::SymbolInfo getSymbol(StringRef Name) const {
-      return RTDyld->getSymbol(Name);
-    }
-
-    bool NeedsFinalization() const { return (State == Raw); }
-
-    virtual void Finalize() = 0;
-
-    void mapSectionAddress(const void *LocalAddress, TargetAddress TargetAddr) {
-      assert((State != Finalized) &&
-             "Attempting to remap sections for finalized objects.");
-      RTDyld->mapSectionAddress(LocalAddress, TargetAddr);
-    }
-
   protected:
-    std::unique_ptr<RuntimeDyld> RTDyld;
-    enum { Raw, Finalizing, Finalized } State;
+    StringMap<RuntimeDyld::SymbolInfo> SymbolTable;
+    bool Finalized = false;
   };
 
   typedef std::list<std::unique_ptr<LinkedObjectSet>> LinkedObjectSetListT;
@@ -79,6 +72,7 @@ public:
   typedef LinkedObjectSetListT::iterator ObjSetHandleT;
 };
 
+
 /// @brief Default (no-op) action to perform when loading objects.
 class DoNothingOnNotifyLoaded {
 public:
@@ -95,34 +89,124 @@ public:
 /// symbols.
 template <typename NotifyLoadedFtor = DoNothingOnNotifyLoaded>
 class ObjectLinkingLayer : public ObjectLinkingLayerBase {
+public:
+
+  /// @brief Functor for receiving finalization notifications.
+  typedef std::function<void(ObjSetHandleT)> NotifyFinalizedFtor;
+
 private:
 
-  template <typename MemoryManagerPtrT, typename SymbolResolverPtrT>
+  template <typename ObjSetT, typename MemoryManagerPtrT,
+            typename SymbolResolverPtrT, typename FinalizerFtor>
   class ConcreteLinkedObjectSet : public LinkedObjectSet {
   public:
-    ConcreteLinkedObjectSet(MemoryManagerPtrT MemMgr,
+    ConcreteLinkedObjectSet(ObjSetT Objects, MemoryManagerPtrT MemMgr,
                             SymbolResolverPtrT Resolver,
+                            FinalizerFtor Finalizer,
                             bool ProcessAllSections)
-      : LinkedObjectSet(*MemMgr, *Resolver, ProcessAllSections),
-        MemMgr(std::move(MemMgr)), Resolver(std::move(Resolver)) { }
+      : MemMgr(std::move(MemMgr)),
+        PFC(make_unique<PreFinalizeContents>(std::move(Objects),
+                                             std::move(Resolver),
+                                             std::move(Finalizer),
+                                             ProcessAllSections)) {
+      buildInitialSymbolTable(PFC->Objects);
+    }
+
+    void setHandle(ObjSetHandleT H) {
+      PFC->Handle = H;
+    }
 
-    void Finalize() override {
-      State = Finalizing;
-      RTDyld->finalizeWithMemoryManagerLocking();
-      State = Finalized;
+    void finalize() override {
+      assert(PFC && "mapSectionAddress called on finalized LinkedObjectSet");
+
+      RuntimeDyld RTDyld(*MemMgr, *PFC->Resolver);
+      RTDyld.setProcessAllSections(PFC->ProcessAllSections);
+      PFC->RTDyld = &RTDyld;
+
+      PFC->Finalizer(PFC->Handle, RTDyld, std::move(PFC->Objects),
+                     [&]() {
+                       updateSymbolTable(RTDyld);
+                       Finalized = true;
+                     });
+
+      // Release resources.
+      PFC = nullptr;
+    }
+
+    JITSymbol::GetAddressFtor getSymbolMaterializer(std::string Name) override {
+      return
+        [this, Name]() {
+          // The symbol may be materialized between the creation of this lambda
+          // and its execution, so we need to double check.
+          if (!Finalized)
+            finalize();
+          return getSymbol(Name, false).getAddress();
+        };
+    }
+
+    void mapSectionAddress(const void *LocalAddress,
+                           TargetAddress TargetAddr) const override {
+      assert(PFC && "mapSectionAddress called on finalized LinkedObjectSet");
+      assert(PFC->RTDyld && "mapSectionAddress called on raw LinkedObjectSet");
+      PFC->RTDyld->mapSectionAddress(LocalAddress, TargetAddr);
     }
 
   private:
+
+    void buildInitialSymbolTable(const ObjSetT &Objects) {
+      for (const auto &Obj : Objects)
+        for (auto &Symbol : getObject(*Obj).symbols()) {
+          if (Symbol.getFlags() & object::SymbolRef::SF_Undefined)
+            continue;
+          ErrorOr<StringRef> SymbolName = Symbol.getName();
+          // FIXME: Raise an error for bad symbols.
+          if (!SymbolName)
+            continue;
+          auto Flags = JITSymbol::flagsFromObjectSymbol(Symbol);
+          SymbolTable.insert(
+            std::make_pair(*SymbolName, RuntimeDyld::SymbolInfo(0, Flags)));
+        }
+    }
+
+    void updateSymbolTable(const RuntimeDyld &RTDyld) {
+      for (auto &SymEntry : SymbolTable)
+        SymEntry.second = RTDyld.getSymbol(SymEntry.first());
+    }
+
+    // Contains the information needed prior to finalization: the object files,
+    // memory manager, resolver, and flags needed for RuntimeDyld.
+    struct PreFinalizeContents {
+      PreFinalizeContents(ObjSetT Objects, SymbolResolverPtrT Resolver,
+                          FinalizerFtor Finalizer, bool ProcessAllSections)
+        : Objects(std::move(Objects)), Resolver(std::move(Resolver)),
+          Finalizer(std::move(Finalizer)),
+          ProcessAllSections(ProcessAllSections) {}
+
+      ObjSetT Objects;
+      SymbolResolverPtrT Resolver;
+      FinalizerFtor Finalizer;
+      bool ProcessAllSections;
+      ObjSetHandleT Handle;
+      RuntimeDyld *RTDyld;
+    };
+
     MemoryManagerPtrT MemMgr;
-    SymbolResolverPtrT Resolver;
+    std::unique_ptr<PreFinalizeContents> PFC;
   };
 
-  template <typename MemoryManagerPtrT, typename SymbolResolverPtrT>
-  std::unique_ptr<LinkedObjectSet>
-  createLinkedObjectSet(MemoryManagerPtrT MemMgr, SymbolResolverPtrT Resolver,
+  template <typename ObjSetT, typename MemoryManagerPtrT,
+            typename SymbolResolverPtrT, typename FinalizerFtor>
+  std::unique_ptr<
+    ConcreteLinkedObjectSet<ObjSetT, MemoryManagerPtrT,
+                            SymbolResolverPtrT, FinalizerFtor>>
+  createLinkedObjectSet(ObjSetT Objects, MemoryManagerPtrT MemMgr,
+                        SymbolResolverPtrT Resolver,
+                        FinalizerFtor Finalizer,
                         bool ProcessAllSections) {
-    typedef ConcreteLinkedObjectSet<MemoryManagerPtrT, SymbolResolverPtrT> LOS;
-    return llvm::make_unique<LOS>(std::move(MemMgr), std::move(Resolver),
+    typedef ConcreteLinkedObjectSet<ObjSetT, MemoryManagerPtrT,
+                                    SymbolResolverPtrT, FinalizerFtor> LOS;
+    return llvm::make_unique<LOS>(std::move(Objects), std::move(MemMgr),
+                                  std::move(Resolver), std::move(Finalizer),
                                   ProcessAllSections);
   }
 
@@ -133,9 +217,6 @@ public:
   typedef std::vector<std::unique_ptr<RuntimeDyld::LoadedObjectInfo>>
       LoadedObjInfoList;
 
-  /// @brief Functor for receiving finalization notifications.
-  typedef std::function<void(ObjSetHandleT)> NotifyFinalizedFtor;
-
   /// @brief Construct an ObjectLinkingLayer with the given NotifyLoaded,
   ///        and NotifyFinalized functors.
   ObjectLinkingLayer(
@@ -169,22 +250,39 @@ public:
   template <typename ObjSetT,
             typename MemoryManagerPtrT,
             typename SymbolResolverPtrT>
-  ObjSetHandleT addObjectSet(const ObjSetT &Objects,
+  ObjSetHandleT addObjectSet(ObjSetT Objects,
                              MemoryManagerPtrT MemMgr,
                              SymbolResolverPtrT Resolver) {
-    ObjSetHandleT Handle =
-      LinkedObjSetList.insert(
-        LinkedObjSetList.end(),
-        createLinkedObjectSet(std::move(MemMgr), std::move(Resolver),
-                              ProcessAllSections));
 
-    LinkedObjectSet &LOS = **Handle;
-    LoadedObjInfoList LoadedObjInfos;
+    auto Finalizer = [&](ObjSetHandleT H, RuntimeDyld &RTDyld,
+                         const ObjSetT &Objs,
+                         std::function<void()> LOSHandleLoad) {
+      LoadedObjInfoList LoadedObjInfos;
+
+      for (auto &Obj : Objs)
+        LoadedObjInfos.push_back(RTDyld.loadObject(getObject(*Obj)));
 
-    for (auto &Obj : Objects)
-      LoadedObjInfos.push_back(LOS.addObject(*Obj));
+      LOSHandleLoad();
 
-    NotifyLoaded(Handle, Objects, LoadedObjInfos);
+      NotifyLoaded(H, Objs, LoadedObjInfos);
+
+      RTDyld.finalizeWithMemoryManagerLocking();
+
+      if (NotifyFinalized)
+        NotifyFinalized(H);
+    };
+
+    auto LOS =
+      createLinkedObjectSet(std::move(Objects), std::move(MemMgr),
+                            std::move(Resolver), std::move(Finalizer),
+                            ProcessAllSections);
+    // LOS is an owning-ptr. Keep a non-owning one so that we can set the handle
+    // below.
+    auto *LOSPtr = LOS.get();
+
+    ObjSetHandleT Handle = LinkedObjSetList.insert(LinkedObjSetList.end(),
+                                                   std::move(LOS));
+    LOSPtr->setHandle(Handle);
 
     return Handle;
   }
@@ -224,33 +322,7 @@ public:
   ///         given object set.
   JITSymbol findSymbolIn(ObjSetHandleT H, StringRef Name,
                          bool ExportedSymbolsOnly) {
-    if (auto Sym = (*H)->getSymbol(Name)) {
-      if (Sym.isExported() || !ExportedSymbolsOnly) {
-        auto Addr = Sym.getAddress();
-        auto Flags = Sym.getFlags();
-        if (!(*H)->NeedsFinalization()) {
-          // If this instance has already been finalized then we can just return
-          // the address.
-          return JITSymbol(Addr, Flags);
-        } else {
-          // If this instance needs finalization return a functor that will do
-          // it. The functor still needs to double-check whether finalization is
-          // required, in case someone else finalizes this set before the
-          // functor is called.
-          auto GetAddress =
-            [this, Addr, H]() {
-              if ((*H)->NeedsFinalization()) {
-                (*H)->Finalize();
-                if (NotifyFinalized)
-                  NotifyFinalized(H);
-              }
-              return Addr;
-            };
-          return JITSymbol(std::move(GetAddress), Flags);
-        }
-      }
-    }
-    return nullptr;
+    return (*H)->getSymbol(Name, ExportedSymbolsOnly);
   }
 
   /// @brief Map section addresses for the objects associated with the handle H.
@@ -263,12 +335,21 @@ public:
   ///        given handle.
   /// @param H Handle for object set to emit/finalize.
   void emitAndFinalize(ObjSetHandleT H) {
-    (*H)->Finalize();
-    if (NotifyFinalized)
-      NotifyFinalized(H);
+    (*H)->finalize();
   }
 
 private:
+
+  static const object::ObjectFile& getObject(const object::ObjectFile &Obj) {
+    return Obj;
+  }
+
+  template <typename ObjT>
+  static const object::ObjectFile&
+  getObject(const object::OwningBinary<ObjT> &Obj) {
+    return *Obj.getBinary();
+  }
+
   LinkedObjectSetListT LinkedObjSetList;
   NotifyLoadedFtor NotifyLoaded;
   NotifyFinalizedFtor NotifyFinalized;

lib/ExecutionEngine/Orc/OrcMCJITReplacement.h

@@ -178,11 +178,10 @@ public:
   }
 
   void addObjectFile(object::OwningBinary<object::ObjectFile> O) override {
-    std::unique_ptr<object::ObjectFile> Obj;
-    std::unique_ptr<MemoryBuffer> Buf;
-    std::tie(Obj, Buf) = O.takeBinary();
-    std::vector<std::unique_ptr<object::ObjectFile>> Objs;
-    Objs.push_back(std::move(Obj));
+    std::vector<std::unique_ptr<object::OwningBinary<object::ObjectFile>>> Objs;
+    Objs.push_back(
+      llvm::make_unique<object::OwningBinary<object::ObjectFile>>(
+        std::move(O)));
     ObjectLayer.addObjectSet(std::move(Objs), &MemMgr, &Resolver);
   }
 
@@ -284,12 +283,12 @@ private:
 
   class NotifyObjectLoadedT {
   public:
-    typedef std::vector<std::unique_ptr<object::ObjectFile>> ObjListT;
     typedef std::vector<std::unique_ptr<RuntimeDyld::LoadedObjectInfo>>
         LoadedObjInfoListT;
 
     NotifyObjectLoadedT(OrcMCJITReplacement &M) : M(M) {}
 
+    template <typename ObjListT>
     void operator()(ObjectLinkingLayerBase::ObjSetHandleT H,
                     const ObjListT &Objects,
                     const LoadedObjInfoListT &Infos) const {
@@ -298,10 +297,21 @@ private:
       assert(Objects.size() == Infos.size() &&
              "Incorrect number of Infos for Objects.");
       for (unsigned I = 0; I < Objects.size(); ++I)
-        M.MemMgr.notifyObjectLoaded(&M, *Objects[I]);
+        M.MemMgr.notifyObjectLoaded(&M, getObject(*Objects[I]));
     }
 
   private:
+
+    static const object::ObjectFile& getObject(const object::ObjectFile &Obj) {
+      return Obj;
+    }
+
+    template <typename ObjT>
+    static const object::ObjectFile&
+    getObject(const object::OwningBinary<ObjT> &Obj) {
+      return *Obj.getBinary();
+    }
+
     OrcMCJITReplacement &M;
   };
 

unittests/ExecutionEngine/Orc/ObjectLinkingLayerTest.cpp

@@ -12,6 +12,7 @@
 #include "llvm/ExecutionEngine/SectionMemoryManager.h"
 #include "llvm/ExecutionEngine/Orc/CompileUtils.h"
 #include "llvm/ExecutionEngine/Orc/LambdaResolver.h"
+#include "llvm/ExecutionEngine/Orc/NullResolver.h"
 #include "llvm/ExecutionEngine/Orc/ObjectLinkingLayer.h"
 #include "llvm/IR/Constants.h"
 #include "llvm/IR/LLVMContext.h"
@@ -29,6 +30,13 @@ class ObjectLinkingLayerExecutionTest : public testing::Test,
 class SectionMemoryManagerWrapper : public SectionMemoryManager {
 public:
   int FinalizationCount = 0;
+  int NeedsToReserveAllocationSpaceCount = 0;
+
+  bool needsToReserveAllocationSpace() override {
+    ++NeedsToReserveAllocationSpaceCount;
+    return SectionMemoryManager::needsToReserveAllocationSpace();
+  }
+
   bool finalizeMemory(std::string *ErrMsg = 0) override {
     ++FinalizationCount;
     return SectionMemoryManager::finalizeMemory(ErrMsg);
@@ -178,4 +186,68 @@ TEST_F(ObjectLinkingLayerExecutionTest, NoDuplicateFinalization) {
       << "Extra call to finalize";
 }
 
+TEST_F(ObjectLinkingLayerExecutionTest, NoPrematureAllocation) {
+
+  if (!TM)
+    return;
+
+  ObjectLinkingLayer<> ObjLayer;
+  SimpleCompiler Compile(*TM);
+
+  // Create a pair of unrelated modules:
+  //
+  // Module 1:
+  //   int foo() { return 42; }
+  // Module 2:
+  //   int bar() { return 7; }
+  //
+  // Both modules will share a memory manager. We want to verify that the
+  // second object is not loaded before the first one is finalized. To do this
+  // in a portable way, we abuse the
+  // RuntimeDyld::MemoryManager::needsToReserveAllocationSpace hook, which is
+  // called once per object before any sections are allocated.
+
+  ModuleBuilder MB1(getGlobalContext(), "", "dummy");
+  {
+    MB1.getModule()->setDataLayout(TM->createDataLayout());
+    Function *BarImpl = MB1.createFunctionDecl<int32_t(void)>("foo");
+    BasicBlock *BarEntry = BasicBlock::Create(getGlobalContext(), "entry",
+                                              BarImpl);
+    IRBuilder<> Builder(BarEntry);
+    IntegerType *Int32Ty = IntegerType::get(getGlobalContext(), 32);
+    Value *FourtyTwo = ConstantInt::getSigned(Int32Ty, 42);
+    Builder.CreateRet(FourtyTwo);
+  }
+
+  auto Obj1 = Compile(*MB1.getModule());
+  std::vector<object::ObjectFile*> Obj1Set;
+  Obj1Set.push_back(Obj1.getBinary());
+
+  ModuleBuilder MB2(getGlobalContext(), "", "dummy");
+  {
+    MB2.getModule()->setDataLayout(TM->createDataLayout());
+    Function *BarImpl = MB2.createFunctionDecl<int32_t(void)>("bar");
+    BasicBlock *BarEntry = BasicBlock::Create(getGlobalContext(), "entry",
+                                              BarImpl);
+    IRBuilder<> Builder(BarEntry);
+    IntegerType *Int32Ty = IntegerType::get(getGlobalContext(), 32);
+    Value *Seven = ConstantInt::getSigned(Int32Ty, 7);
+    Builder.CreateRet(Seven);
+  }
+  auto Obj2 = Compile(*MB2.getModule());
+  std::vector<object::ObjectFile*> Obj2Set;
+  Obj2Set.push_back(Obj2.getBinary());
+
+  SectionMemoryManagerWrapper SMMW;
+  NullResolver NR;
+  auto H = ObjLayer.addObjectSet(std::move(Obj1Set), &SMMW, &NR);
+  ObjLayer.addObjectSet(std::move(Obj2Set), &SMMW, &NR);
+  ObjLayer.emitAndFinalize(H);
+
+  // Only one call to needsToReserveAllocationSpace should have been made.
+  EXPECT_EQ(SMMW.NeedsToReserveAllocationSpaceCount, 1)
+      << "More than one call to needsToReserveAllocationSpace "
+         "(multiple unrelated objects loaded prior to finalization)";
+}
+
 }