Orc: Drop some else-after-return, reflow a few spots, and avoid use of pointee types
[oota-llvm.git] / include / llvm / ExecutionEngine / Orc / CompileOnDemandLayer.h
index d631e1d5ce0b36cb5c0cd826422f3527b7fe2fe9..07914e386ad213f311231a59e8b7eecfb5ed3613 100644 (file)
@@ -22,6 +22,7 @@
 #include "llvm/ExecutionEngine/SectionMemoryManager.h"
 #include "llvm/Transforms/Utils/Cloning.h"
 #include <list>
+#include <memory>
 #include <set>
 
 #include "llvm/Support/Debug.h"
@@ -36,40 +37,90 @@ namespace orc {
 /// added to the layer below. When a stub is called it triggers the extraction
 /// of the function body from the original module. The extracted body is then
 /// compiled and executed.
-template <typename BaseLayerT, typename CompileCallbackMgrT,
-          typename PartitioningFtor =
-            std::function<std::set<Function*>(Function&)>>
+template <typename BaseLayerT,
+          typename CompileCallbackMgrT = JITCompileCallbackManagerBase,
+          typename IndirectStubsMgrT = IndirectStubsManagerBase>
 class CompileOnDemandLayer {
 private:
 
-  // Utility class for MapValue. Only materializes declarations for global
-  // variables.
-  class GlobalDeclMaterializer : public ValueMaterializer {
+  template <typename MaterializerFtor>
+  class LambdaMaterializer final : public ValueMaterializer {
   public:
-    GlobalDeclMaterializer(Module &Dst) : Dst(Dst) {}
+    LambdaMaterializer(MaterializerFtor M) : M(std::move(M)) {}
     Value* materializeValueFor(Value *V) final {
-      if (auto *GV = dyn_cast<GlobalVariable>(V))
-        return cloneGlobalVariableDecl(Dst, *GV);
-      else if (auto *F = dyn_cast<Function>(V))
-        return cloneFunctionDecl(Dst, *F);
-      // Else.
-      return nullptr;
+      return M(V);
     }
   private:
-    Module &Dst;
+    MaterializerFtor M;
   };
 
+  template <typename MaterializerFtor>
+  LambdaMaterializer<MaterializerFtor>
+  createLambdaMaterializer(MaterializerFtor M) {
+    return LambdaMaterializer<MaterializerFtor>(std::move(M));
+  }
+
   typedef typename BaseLayerT::ModuleSetHandleT BaseLayerModuleSetHandleT;
 
+  class ModuleOwner {
+  public:
+    ModuleOwner() = default;
+    ModuleOwner(const ModuleOwner&) = delete;
+    ModuleOwner& operator=(const ModuleOwner&) = delete;
+    virtual ~ModuleOwner() { }
+    virtual Module& getModule() const = 0;
+  };
+
+  template <typename ModulePtrT>
+  class ModuleOwnerImpl : public ModuleOwner {
+  public:
+    ModuleOwnerImpl(ModulePtrT ModulePtr) : ModulePtr(std::move(ModulePtr)) {}
+    Module& getModule() const override { return *ModulePtr; }
+  private:
+    ModulePtrT ModulePtr;
+  };
+
+  template <typename ModulePtrT>
+  std::unique_ptr<ModuleOwner> wrapOwnership(ModulePtrT ModulePtr) {
+    return llvm::make_unique<ModuleOwnerImpl<ModulePtrT>>(std::move(ModulePtr));
+  }
+
   struct LogicalModuleResources {
-    std::shared_ptr<Module> SourceModule;
+    std::unique_ptr<ModuleOwner> SourceModuleOwner;
+    std::set<const Function*> StubsToClone;
+    std::unique_ptr<IndirectStubsMgrT> StubsMgr;
+
+    LogicalModuleResources() = default;
+
+    // Explicit move constructor to make MSVC happy.
+    LogicalModuleResources(LogicalModuleResources &&Other)
+        : SourceModuleOwner(std::move(Other.SourceModuleOwner)),
+          StubsToClone(std::move(Other.StubsToClone)),
+          StubsMgr(std::move(Other.StubsMgr)) {}
+
+    // Explicit move assignment to make MSVC happy.
+    LogicalModuleResources& operator=(LogicalModuleResources &&Other) {
+      SourceModuleOwner = std::move(Other.SourceModuleOwner);
+      StubsToClone = std::move(Other.StubsToClone);
+      StubsMgr = std::move(Other.StubsMgr);
+    }
+
+    JITSymbol findSymbol(StringRef Name, bool ExportedSymbolsOnly) {
+      if (Name.endswith("$stub_ptr") && !ExportedSymbolsOnly) {
+        assert(!ExportedSymbolsOnly && "Stubs are never exported");
+        return StubsMgr->findPointer(Name.drop_back(9));
+      }
+      return StubsMgr->findStub(Name, ExportedSymbolsOnly);
+    }
+
   };
 
+
+
   struct LogicalDylibResources {
     typedef std::function<RuntimeDyld::SymbolInfo(const std::string&)>
       SymbolResolverFtor;
     SymbolResolverFtor ExternalSymbolResolver;
-    PartitioningFtor Partitioner;
   };
 
   typedef LogicalDylib<BaseLayerT, LogicalModuleResources,
@@ -79,12 +130,26 @@ private:
   typedef std::list<CODLogicalDylib> LogicalDylibList;
 
 public:
+
   /// @brief Handle to a set of loaded modules.
   typedef typename LogicalDylibList::iterator ModuleSetHandleT;
 
+  /// @brief Module partitioning functor.
+  typedef std::function<std::set<Function*>(Function&)> PartitioningFtor;
+
+  /// @brief Builder for IndirectStubsManagers.
+  typedef std::function<std::unique_ptr<IndirectStubsMgrT>()>
+    IndirectStubsManagerBuilderT;
+
   /// @brief Construct a compile-on-demand layer instance.
-  CompileOnDemandLayer(BaseLayerT &BaseLayer, CompileCallbackMgrT &CallbackMgr)
-      : BaseLayer(BaseLayer), CompileCallbackMgr(CallbackMgr) {}
+  CompileOnDemandLayer(BaseLayerT &BaseLayer, PartitioningFtor Partition,
+                       CompileCallbackMgrT &CallbackMgr,
+                       IndirectStubsManagerBuilderT CreateIndirectStubsManager,
+                       bool CloneStubsIntoPartitions = true)
+      : BaseLayer(BaseLayer),  Partition(Partition),
+        CompileCallbackMgr(CallbackMgr),
+        CreateIndirectStubsManager(std::move(CreateIndirectStubsManager)),
+        CloneStubsIntoPartitions(CloneStubsIntoPartitions) {}
 
   /// @brief Add a module to the compile-on-demand layer.
   template <typename ModuleSetT, typename MemoryManagerPtrT,
@@ -97,24 +162,16 @@ public:
            "User supplied memory managers not supported with COD yet.");
 
     LogicalDylibs.push_back(CODLogicalDylib(BaseLayer));
-    auto &LDLResources = LogicalDylibs.back().getDylibResources();
+    auto &LDResources = LogicalDylibs.back().getDylibResources();
 
-    LDLResources.ExternalSymbolResolver =
+    LDResources.ExternalSymbolResolver =
       [Resolver](const std::string &Name) {
         return Resolver->findSymbol(Name);
       };
 
-    LDLResources.Partitioner =
-      [](Function &F) {
-        std::set<Function*> Partition;
-        Partition.insert(&F);
-        return Partition;
-      };
-
     // Process each of the modules in this module set.
     for (auto &M : Ms)
-      addLogicalModule(LogicalDylibs.back(),
-                       std::shared_ptr<Module>(std::move(M)));
+      addLogicalModule(LogicalDylibs.back(), std::move(M));
 
     return std::prev(LogicalDylibs.end());
   }
@@ -132,6 +189,10 @@ public:
   /// @param ExportedSymbolsOnly If true, search only for exported symbols.
   /// @return A handle for the given named symbol, if it exists.
   JITSymbol findSymbol(StringRef Name, bool ExportedSymbolsOnly) {
+    for (auto LDI = LogicalDylibs.begin(), LDE = LogicalDylibs.end();
+         LDI != LDE; ++LDI)
+      if (auto Symbol = findSymbolIn(LDI, Name, ExportedSymbolsOnly))
+        return Symbol;
     return BaseLayer.findSymbol(Name, ExportedSymbolsOnly);
   }
 
@@ -144,141 +205,190 @@ public:
 
 private:
 
-  void addLogicalModule(CODLogicalDylib &LD, std::shared_ptr<Module> SrcM) {
+  template <typename ModulePtrT>
+  void addLogicalModule(CODLogicalDylib &LD, ModulePtrT SrcMPtr) {
 
     // Bump the linkage and rename any anonymous/privote members in SrcM to
     // ensure that everything will resolve properly after we partition SrcM.
-    makeAllSymbolsExternallyAccessible(*SrcM);
+    makeAllSymbolsExternallyAccessible(*SrcMPtr);
 
     // Create a logical module handle for SrcM within the logical dylib.
     auto LMH = LD.createLogicalModule();
-    LD.getLogicalModuleResources(LMH).SourceModule = SrcM;
+    auto &LMResources =  LD.getLogicalModuleResources(LMH);
 
-    // Create the GVs-and-stubs module.
-    auto GVsAndStubsM = llvm::make_unique<Module>(
-                          (SrcM->getName() + ".globals_and_stubs").str(),
-                          SrcM->getContext());
-    GVsAndStubsM->setDataLayout(SrcM->getDataLayout());
-    ValueToValueMapTy VMap;
+    LMResources.SourceModuleOwner = wrapOwnership(std::move(SrcMPtr));
 
-    // Process module and create stubs.
-    // We create the stubs before copying the global variables as we know the
-    // stubs won't refer to any globals (they only refer to their implementation
-    // pointer) so there's no ordering/value-mapping issues.
-    for (auto &F : *SrcM) {
-
-      // Skip declarations.
-      if (F.isDeclaration())
-        continue;
-
-      // For each definition: create a callback, a stub, and a function body
-      // pointer. Initialize the function body pointer to point at the callback,
-      // and set the callback to compile the function body.
-      auto CCInfo = CompileCallbackMgr.getCompileCallback(SrcM->getContext());
-      Function *StubF = cloneFunctionDecl(*GVsAndStubsM, F, &VMap);
-      GlobalVariable *FnBodyPtr =
-        createImplPointer(*StubF->getType(), *StubF->getParent(),
-                          StubF->getName() + "$orc_addr",
-                          createIRTypedAddress(*StubF->getFunctionType(),
-                                               CCInfo.getAddress()));
-      makeStub(*StubF, *FnBodyPtr);
-      CCInfo.setCompileAction(
-        [this, &LD, LMH, &F]() {
+    Module &SrcM = LMResources.SourceModuleOwner->getModule();
+
+    // Create the GlobalValues module.
+    const DataLayout &DL = SrcM.getDataLayout();
+    auto GVsM = llvm::make_unique<Module>((SrcM.getName() + ".globals").str(),
+                                          SrcM.getContext());
+    GVsM->setDataLayout(DL);
+
+    // Create function stubs.
+    ValueToValueMapTy VMap;
+    {
+      typename IndirectStubsMgrT::StubInitsMap StubInits;
+      for (auto &F : SrcM) {
+        // Skip declarations.
+        if (F.isDeclaration())
+          continue;
+
+        // Record all functions defined by this module.
+        if (CloneStubsIntoPartitions)
+          LMResources.StubsToClone.insert(&F);
+
+        // Create a callback, associate it with the stub for the function,
+        // and set the compile action to compile the partition containing the
+        // function.
+        auto CCInfo = CompileCallbackMgr.getCompileCallback(SrcM.getContext());
+        StubInits[mangle(F.getName(), DL)] =
+          std::make_pair(CCInfo.getAddress(),
+                         JITSymbolBase::flagsFromGlobalValue(F));
+        CCInfo.setCompileAction([this, &LD, LMH, &F]() {
           return this->extractAndCompile(LD, LMH, F);
         });
+      }
+
+      LMResources.StubsMgr = CreateIndirectStubsManager();
+      auto EC = LMResources.StubsMgr->createStubs(StubInits);
+      (void)EC;
+      // FIXME: This should be propagated back to the user. Stub creation may
+      //        fail for remote JITs.
+      assert(!EC && "Error generating stubs");
     }
 
-    // Now clone the global variable declarations.
-    GlobalDeclMaterializer GDMat(*GVsAndStubsM);
-    for (auto &GV : SrcM->globals())
+    // Clone global variable decls.
+    for (auto &GV : SrcM.globals())
+      if (!GV.isDeclaration() && !VMap.count(&GV))
+        cloneGlobalVariableDecl(*GVsM, GV, &VMap);
+
+    // And the aliases.
+    for (auto &A : SrcM.aliases())
+      if (!VMap.count(&A))
+        cloneGlobalAliasDecl(*GVsM, A, VMap);
+
+    // Now we need to clone the GV and alias initializers.
+
+    // Initializers may refer to functions declared (but not defined) in this
+    // module. Build a materializer to clone decls on demand.
+    auto Materializer = createLambdaMaterializer(
+      [this, &GVsM, &LMResources](Value *V) -> Value* {
+        if (auto *F = dyn_cast<Function>(V)) {
+          // Decls in the original module just get cloned.
+          if (F->isDeclaration())
+            return cloneFunctionDecl(*GVsM, *F);
+
+          // Definitions in the original module (which we have emitted stubs
+          // for at this point) get turned into a constant alias to the stub
+          // instead.
+          const DataLayout &DL = GVsM->getDataLayout();
+          std::string FName = mangle(F->getName(), DL);
+          auto StubSym = LMResources.StubsMgr->findStub(FName, false);
+          unsigned PtrBitWidth = DL.getPointerTypeSizeInBits(F->getType());
+          ConstantInt *StubAddr =
+            ConstantInt::get(GVsM->getContext(),
+                             APInt(PtrBitWidth, StubSym.getAddress()));
+          Constant *Init = ConstantExpr::getCast(Instruction::IntToPtr,
+                                                 StubAddr, F->getType());
+          return GlobalAlias::create(F->getFunctionType(),
+                                     F->getType()->getAddressSpace(),
+                                     F->getLinkage(), F->getName(),
+                                     Init, GVsM.get());
+        }
+        // else....
+        return nullptr;
+      });
+
+    // Clone the global variable initializers.
+    for (auto &GV : SrcM.globals())
       if (!GV.isDeclaration())
-        cloneGlobalVariableDecl(*GVsAndStubsM, GV, &VMap);
-
-    // Then clone the initializers.
-    for (auto &GV : SrcM->globals())
-      if (!GV.isDeclaration())
-        moveGlobalVariableInitializer(GV, VMap, &GDMat);
+        moveGlobalVariableInitializer(GV, VMap, &Materializer);
+
+    // Clone the global alias initializers.
+    for (auto &A : SrcM.aliases()) {
+      auto *NewA = cast<GlobalAlias>(VMap[&A]);
+      assert(NewA && "Alias not cloned?");
+      Value *Init = MapValue(A.getAliasee(), VMap, RF_None, nullptr,
+                             &Materializer);
+      NewA->setAliasee(cast<Constant>(Init));
+    }
 
-    // Build a resolver for the stubs module and add it to the base layer.
-    auto GVsAndStubsResolver = createLambdaResolver(
-        [&LD](const std::string &Name) {
+    // Build a resolver for the globals module and add it to the base layer.
+    auto GVsResolver = createLambdaResolver(
+        [&LD, LMH](const std::string &Name) {
+          auto &LMResources = LD.getLogicalModuleResources(LMH);
+          if (auto Sym = LMResources.StubsMgr->findStub(Name, false))
+            return RuntimeDyld::SymbolInfo(Sym.getAddress(), Sym.getFlags());
           return LD.getDylibResources().ExternalSymbolResolver(Name);
         },
         [](const std::string &Name) {
           return RuntimeDyld::SymbolInfo(nullptr);
         });
 
-    std::vector<std::unique_ptr<Module>> GVsAndStubsMSet;
-    GVsAndStubsMSet.push_back(std::move(GVsAndStubsM));
-    auto GVsAndStubsH =
-      BaseLayer.addModuleSet(std::move(GVsAndStubsMSet),
+    std::vector<std::unique_ptr<Module>> GVsMSet;
+    GVsMSet.push_back(std::move(GVsM));
+    auto GVsH =
+      BaseLayer.addModuleSet(std::move(GVsMSet),
                              llvm::make_unique<SectionMemoryManager>(),
-                             std::move(GVsAndStubsResolver));
-    LD.addToLogicalModule(LMH, GVsAndStubsH);
+                             std::move(GVsResolver));
+    LD.addToLogicalModule(LMH, GVsH);
   }
 
-  static std::string Mangle(StringRef Name, const DataLayout &DL) {
-    Mangler M(&DL);
+  static std::string mangle(StringRef Name, const DataLayout &DL) {
     std::string MangledName;
-    {
-      raw_string_ostream MangledNameStream(MangledName);
-      M.getNameWithPrefix(MangledNameStream, Name);
-    }
+    Mangler::getNameWithPrefix(raw_string_ostream(MangledName), Name, DL);
     return MangledName;
   }
 
   TargetAddress extractAndCompile(CODLogicalDylib &LD,
                                   LogicalModuleHandle LMH,
                                   Function &F) {
-    Module &SrcM = *LD.getLogicalModuleResources(LMH).SourceModule;
+    auto &LMResources = LD.getLogicalModuleResources(LMH);
+    Module &SrcM = LMResources.SourceModuleOwner->getModule();
 
     // If F is a declaration we must already have compiled it.
     if (F.isDeclaration())
       return 0;
 
     // Grab the name of the function being called here.
-    std::string CalledFnName = Mangle(F.getName(), SrcM.getDataLayout());
+    std::string CalledFnName = mangle(F.getName(), SrcM.getDataLayout());
 
-    const auto &Partition = LD.getDylibResources().Partitioner(F);
-    auto PartitionH = emitPartition(LD, LMH, Partition);
+    auto Part = Partition(F);
+    auto PartH = emitPartition(LD, LMH, Part);
 
     TargetAddress CalledAddr = 0;
-    for (auto *SubF : Partition) {
-      std::string FName(SubF->getName());
-      auto FnBodySym =
-        BaseLayer.findSymbolIn(PartitionH, Mangle(FName, SrcM.getDataLayout()),
-                               false);
-      auto FnPtrSym =
-        BaseLayer.findSymbolIn(*LD.moduleHandlesBegin(LMH),
-                               Mangle(FName + "$orc_addr",
-                                      SrcM.getDataLayout()),
-                               false);
+    for (auto *SubF : Part) {
+      std::string FnName = mangle(SubF->getName(), SrcM.getDataLayout());
+      auto FnBodySym = BaseLayer.findSymbolIn(PartH, FnName, false);
       assert(FnBodySym && "Couldn't find function body.");
-      assert(FnPtrSym && "Couldn't find function body pointer.");
 
-      auto FnBodyAddr = FnBodySym.getAddress();
-      void *FnPtrAddr = reinterpret_cast<void*>(
-          static_cast<uintptr_t>(FnPtrSym.getAddress()));
+      TargetAddress FnBodyAddr = FnBodySym.getAddress();
 
       // If this is the function we're calling record the address so we can
       // return it from this function.
       if (SubF == &F)
         CalledAddr = FnBodyAddr;
 
-      memcpy(FnPtrAddr, &FnBodyAddr, sizeof(uintptr_t));
+      // Update the function body pointer for the stub.
+      if (auto EC = LMResources.StubsMgr->updatePointer(FnName, FnBodyAddr))
+        return 0;
     }
 
     return CalledAddr;
   }
 
+  template <typename PartitionT>
   BaseLayerModuleSetHandleT emitPartition(CODLogicalDylib &LD,
                                           LogicalModuleHandle LMH,
-                                          const std::set<Function*> &Partition) {
-    Module &SrcM = *LD.getLogicalModuleResources(LMH).SourceModule;
+                                          const PartitionT &Part) {
+    auto &LMResources = LD.getLogicalModuleResources(LMH);
+    Module &SrcM = LMResources.SourceModuleOwner->getModule();
 
     // Create the module.
-    std::string NewName(SrcM.getName());
-    for (auto *F : Partition) {
+    std::string NewName = SrcM.getName();
+    for (auto *F : Part) {
       NewName += ".";
       NewName += F->getName();
     }
@@ -286,15 +396,51 @@ private:
     auto M = llvm::make_unique<Module>(NewName, SrcM.getContext());
     M->setDataLayout(SrcM.getDataLayout());
     ValueToValueMapTy VMap;
-    GlobalDeclMaterializer GDM(*M);
+
+    auto Materializer = createLambdaMaterializer([this, &LMResources, &M,
+                                                  &VMap](Value *V) -> Value * {
+      if (auto *GV = dyn_cast<GlobalVariable>(V))
+        return cloneGlobalVariableDecl(*M, *GV);
+
+      if (auto *F = dyn_cast<Function>(V)) {
+        // Check whether we want to clone an available_externally definition.
+        if (!LMResources.StubsToClone.count(F))
+          return cloneFunctionDecl(*M, *F);
+
+        // Ok - we want an inlinable stub. For that to work we need a decl
+        // for the stub pointer.
+        auto *StubPtr = createImplPointer(*F->getType(), *M,
+                                          F->getName() + "$stub_ptr", nullptr);
+        auto *ClonedF = cloneFunctionDecl(*M, *F);
+        makeStub(*ClonedF, *StubPtr);
+        ClonedF->setLinkage(GlobalValue::AvailableExternallyLinkage);
+        ClonedF->addFnAttr(Attribute::AlwaysInline);
+        return ClonedF;
+      }
+
+      if (auto *A = dyn_cast<GlobalAlias>(V)) {
+        auto *Ty = A->getValueType();
+        if (Ty->isFunctionTy())
+          return Function::Create(cast<FunctionType>(Ty),
+                                  GlobalValue::ExternalLinkage, A->getName(),
+                                  M.get());
+
+        return new GlobalVariable(*M, Ty, false, GlobalValue::ExternalLinkage,
+                                  nullptr, A->getName(), nullptr,
+                                  GlobalValue::NotThreadLocal,
+                                  A->getType()->getAddressSpace());
+      }
+
+      return nullptr;
+    });
 
     // Create decls in the new module.
-    for (auto *F : Partition)
+    for (auto *F : Part)
       cloneFunctionDecl(*M, *F, &VMap);
 
     // Move the function bodies.
-    for (auto *F : Partition)
-      moveFunctionBody(*F, VMap);
+    for (auto *F : Part)
+      moveFunctionBody(*F, VMap, &Materializer);
 
     // Create memory manager and symbol resolver.
     auto MemMgr = llvm::make_unique<SectionMemoryManager>();
@@ -318,8 +464,12 @@ private:
   }
 
   BaseLayerT &BaseLayer;
+  PartitioningFtor Partition;
   CompileCallbackMgrT &CompileCallbackMgr;
+  IndirectStubsManagerBuilderT CreateIndirectStubsManager;
+
   LogicalDylibList LogicalDylibs;
+  bool CloneStubsIntoPartitions;
 };
 
 } // End namespace orc.