[Orc] Teach the CompileOnDemand layer to clone aliases.
[oota-llvm.git] / include / llvm / ExecutionEngine / Orc / CompileOnDemandLayer.h
index 25381f1bdd8d571e95bbda43f3db0cd538bbd67f..714ca2374dc359a3343307dacb8edf2e9994f4e0 100644 (file)
 #define LLVM_EXECUTIONENGINE_ORC_COMPILEONDEMANDLAYER_H
 
 #include "IndirectionUtils.h"
-#include "LookasideRTDyldMM.h"
+#include "LambdaResolver.h"
+#include "LogicalDylib.h"
 #include "llvm/ADT/STLExtras.h"
 #include "llvm/ExecutionEngine/SectionMemoryManager.h"
+#include "llvm/Transforms/Utils/Cloning.h"
 #include <list>
+#include <set>
+
+#include "llvm/Support/Debug.h"
 
 namespace llvm {
+namespace orc {
 
 /// @brief Compile-on-demand layer.
 ///
-///   Modules added to this layer have their calls indirected, and are then
-/// broken up into a set of single-function modules, each of which is added
-/// to the layer below in a singleton set. The lower layer can be any layer that
-/// accepts IR module sets.
-///
-/// It is expected that this layer will frequently be used on top of a
-/// LazyEmittingLayer. The combination of the two ensures that each function is
-/// compiled only when it is first called.
-template <typename BaseLayerT> class CompileOnDemandLayer {
-public:
-  /// @brief Lookup helper that provides compatibility with the classic
-  ///        static-compilation symbol resolution process.
-  ///
-  ///   The CompileOnDemand (COD) layer splits modules up into multiple
-  /// sub-modules, each held in its own llvm::Module instance, in order to
-  /// support lazy compilation. When a module that contains private symbols is
-  /// broken up symbol linkage changes may be required to enable access to
-  /// "private" data that now resides in a different llvm::Module instance. To
-  /// retain expected symbol resolution behavior for clients of the COD layer,
-  /// the CODScopedLookup class uses a two-tiered lookup system to resolve
-  /// symbols. Lookup first scans sibling modules that were split from the same
-  /// original module (logical-module scoped lookup), then scans all other
-  /// modules that have been added to the lookup scope (logical-dylib scoped
-  /// lookup).
-  class CODScopedLookup {
-  private:
-    typedef typename BaseLayerT::ModuleSetHandleT BaseLayerModuleSetHandleT;
-    typedef std::vector<BaseLayerModuleSetHandleT> SiblingHandlesList;
-    typedef std::list<SiblingHandlesList> PseudoDylibModuleSetHandlesList;
+///   When a module is added to this layer a stub is created for each of its
+/// function definitions. The stubs and other global values are immediately
+/// 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&)>>
+class CompileOnDemandLayer {
+private:
 
+  // Utility class for MapValue. Only materializes declarations for global
+  // variables.
+  class GlobalDeclMaterializer : public ValueMaterializer {
   public:
-    /// @brief Handle for a logical module.
-    typedef typename PseudoDylibModuleSetHandlesList::iterator LMHandle;
-
-    /// @brief Construct a scoped lookup.
-    CODScopedLookup(BaseLayerT &BaseLayer) : BaseLayer(BaseLayer) {}
-
-    /// @brief Start a new context for a single logical module.
-    LMHandle createLogicalModule() {
-      Handles.push_back(SiblingHandlesList());
-      return std::prev(Handles.end());
-    }
-
-    /// @brief Add a concrete Module's handle to the given logical Module's
-    ///        lookup scope.
-    void addToLogicalModule(LMHandle LMH, BaseLayerModuleSetHandleT H) {
-      LMH->push_back(H);
-    }
-
-    /// @brief Remove a logical Module from the CODScopedLookup entirely.
-    void removeLogicalModule(LMHandle LMH) { Handles.erase(LMH); }
-
-    /// @brief Look up a symbol in this context.
-    JITSymbol findSymbol(LMHandle LMH, const std::string &Name) {
-      if (auto Symbol = findSymbolIn(LMH, Name))
-        return Symbol;
-
-      for (auto I = Handles.begin(), E = Handles.end(); I != E; ++I)
-        if (I != LMH)
-          if (auto Symbol = findSymbolIn(I, Name))
-            return Symbol;
-
+    typedef std::set<const Function*> StubSet;
+
+    GlobalDeclMaterializer(Module &Dst, const StubSet *StubsToClone = nullptr)
+        : Dst(Dst), StubsToClone(StubsToClone) {}
+
+    Value* materializeValueFor(Value *V) final {
+      if (auto *GV = dyn_cast<GlobalVariable>(V))
+        return cloneGlobalVariableDecl(Dst, *GV);
+      else if (auto *F = dyn_cast<Function>(V)) {
+        auto *ClonedF = cloneFunctionDecl(Dst, *F);
+        if (StubsToClone && StubsToClone->count(F)) {
+          GlobalVariable *FnBodyPtr =
+            createImplPointer(*ClonedF->getType(), *ClonedF->getParent(),
+                              ClonedF->getName() + "$orc_addr", nullptr);
+          makeStub(*ClonedF, *FnBodyPtr);
+          ClonedF->setLinkage(GlobalValue::AvailableExternallyLinkage);
+          ClonedF->addFnAttr(Attribute::AlwaysInline);
+        }
+        return ClonedF;
+      }
+      // Else.
       return nullptr;
     }
-
   private:
-
-    JITSymbol findSymbolIn(LMHandle LMH, const std::string &Name) {
-      for (auto H : *LMH)
-        if (auto Symbol = BaseLayer.findSymbolIn(H, Name, false))
-          return Symbol;
-      return nullptr;
-    }
-
-    BaseLayerT &BaseLayer;
-    PseudoDylibModuleSetHandlesList Handles;
+    Module &Dst;
+    const StubSet *StubsToClone;
   };
 
-private:
   typedef typename BaseLayerT::ModuleSetHandleT BaseLayerModuleSetHandleT;
-  typedef std::vector<BaseLayerModuleSetHandleT> BaseLayerModuleSetHandleListT;
-
-  struct ModuleSetInfo {
-    // Symbol lookup - just one for the whole module set.
-    std::shared_ptr<CODScopedLookup> Lookup;
-
-    // Logical module handles.
-    std::vector<typename CODScopedLookup::LMHandle> LMHandles;
-
-    // Persistent manglers - one per TU.
-    std::vector<PersistentMangler> PersistentManglers;
 
-    // Symbol resolution callback handlers - one per TU.
-    std::vector<std::unique_ptr<JITResolveCallbackHandler>>
-        JITResolveCallbackHandlers;
-
-    // List of vectors of module set handles:
-    // One vector per logical module - each vector holds the handles for the
-    // exploded modules for that logical module in the base layer.
-    BaseLayerModuleSetHandleListT BaseLayerModuleSetHandles;
-
-    ModuleSetInfo(std::shared_ptr<CODScopedLookup> Lookup)
-        : Lookup(std::move(Lookup)) {}
+  struct LogicalModuleResources {
+    std::shared_ptr<Module> SourceModule;
+    std::set<const Function*> StubsToClone;
+  };
 
-    void releaseResources(BaseLayerT &BaseLayer) {
-      for (auto LMH : LMHandles)
-        Lookup->removeLogicalModule(LMH);
-      for (auto H : BaseLayerModuleSetHandles)
-        BaseLayer.removeModuleSet(H);
-    }
+  struct LogicalDylibResources {
+    typedef std::function<RuntimeDyld::SymbolInfo(const std::string&)>
+      SymbolResolverFtor;
+    SymbolResolverFtor ExternalSymbolResolver;
+    PartitioningFtor Partitioner;
   };
 
-  typedef std::list<ModuleSetInfo> ModuleSetInfoListT;
+  typedef LogicalDylib<BaseLayerT, LogicalModuleResources,
+                       LogicalDylibResources> CODLogicalDylib;
+
+  typedef typename CODLogicalDylib::LogicalModuleHandle LogicalModuleHandle;
+  typedef std::list<CODLogicalDylib> LogicalDylibList;
 
 public:
   /// @brief Handle to a set of loaded modules.
-  typedef typename ModuleSetInfoListT::iterator ModuleSetHandleT;
-
-  /// @brief Convenience typedef for callback inserter.
-  typedef std::function<void(Module&, JITResolveCallbackHandler&)>
-    InsertCallbackAsmFtor;
+  typedef typename LogicalDylibList::iterator ModuleSetHandleT;
 
   /// @brief Construct a compile-on-demand layer instance.
-  CompileOnDemandLayer(BaseLayerT &BaseLayer,
-                       InsertCallbackAsmFtor InsertCallbackAsm)
-    : BaseLayer(BaseLayer), InsertCallbackAsm(InsertCallbackAsm) {}
+  CompileOnDemandLayer(BaseLayerT &BaseLayer, CompileCallbackMgrT &CallbackMgr,
+                       bool CloneStubsIntoPartitions)
+      : BaseLayer(BaseLayer), CompileCallbackMgr(CallbackMgr),
+        CloneStubsIntoPartitions(CloneStubsIntoPartitions) {}
 
   /// @brief Add a module to the compile-on-demand layer.
-  template <typename ModuleSetT>
+  template <typename ModuleSetT, typename MemoryManagerPtrT,
+            typename SymbolResolverPtrT>
   ModuleSetHandleT addModuleSet(ModuleSetT Ms,
-                                std::unique_ptr<RTDyldMemoryManager> MM) {
-
-    const char *JITAddrSuffix = "$orc_addr";
-    const char *JITImplSuffix = "$orc_impl";
-
-    // Create a symbol lookup context and ModuleSetInfo for this module set.
-    auto DylibLookup = std::make_shared<CODScopedLookup>(BaseLayer);
-    ModuleSetHandleT H =
-        ModuleSetInfos.insert(ModuleSetInfos.end(), ModuleSetInfo(DylibLookup));
-    ModuleSetInfo &MSI = ModuleSetInfos.back();
-
-    // Process each of the modules in this module set. All modules share the
-    // same lookup context, but each will get its own TU lookup context.
-    for (auto &M : Ms) {
-
-      // Create a TU lookup context for this module.
-      auto LMH = DylibLookup->createLogicalModule();
-      MSI.LMHandles.push_back(LMH);
-
-      // Create a persistent mangler for this module.
-      MSI.PersistentManglers.emplace_back(*M->getDataLayout());
-
-      // Make all calls to functions defined in this module indirect.
-      JITIndirections Indirections =
-          makeCallsDoubleIndirect(*M, [](const Function &) { return true; },
-                                  JITImplSuffix, JITAddrSuffix);
-
-      // Then carve up the module into a bunch of single-function modules.
-      std::vector<std::unique_ptr<Module>> ExplodedModules =
-          explode(*M, Indirections);
-
-      // Add a resolve-callback handler for this module to look up symbol
-      // addresses when requested via a callback.
-      MSI.JITResolveCallbackHandlers.push_back(
-          createCallbackHandlerFromJITIndirections(
-              Indirections, MSI.PersistentManglers.back(),
-              [=](StringRef S) {
-                return DylibLookup->findSymbol(LMH, S).getAddress();
-              }));
-
-      // Insert callback asm code into the first module.
-      InsertCallbackAsm(*ExplodedModules[0],
-                        *MSI.JITResolveCallbackHandlers.back());
-
-      // Now we need to take each of the extracted Modules and add them to
-      // base layer. Each Module will be added individually to make sure they
-      // can be compiled separately, and each will get its own lookaside
-      // memory manager with lookup functors that resolve symbols in sibling
-      // modules first.OA
-      for (auto &M : ExplodedModules) {
-        std::vector<std::unique_ptr<Module>> MSet;
-        MSet.push_back(std::move(M));
-
-        BaseLayerModuleSetHandleT H = BaseLayer.addModuleSet(
-            std::move(MSet),
-            createLookasideRTDyldMM<SectionMemoryManager>(
-                [=](const std::string &Name) {
-                  if (auto Symbol = DylibLookup->findSymbol(LMH, Name))
-                    return Symbol.getAddress();
-                  return findSymbol(Name, true).getAddress();
-                },
-                [=](const std::string &Name) {
-                  return DylibLookup->findSymbol(LMH, Name).getAddress();
-                }));
-        DylibLookup->addToLogicalModule(LMH, H);
-        MSI.BaseLayerModuleSetHandles.push_back(H);
-      }
+                                MemoryManagerPtrT MemMgr,
+                                SymbolResolverPtrT Resolver) {
 
-      initializeFuncAddrs(*MSI.JITResolveCallbackHandlers.back(), Indirections,
-                          MSI.PersistentManglers.back(), [=](StringRef S) {
-                            return DylibLookup->findSymbol(LMH, S).getAddress();
-                          });
-    }
+    assert(MemMgr == nullptr &&
+           "User supplied memory managers not supported with COD yet.");
+
+    LogicalDylibs.push_back(CODLogicalDylib(BaseLayer));
+    auto &LDResources = LogicalDylibs.back().getDylibResources();
+
+    LDResources.ExternalSymbolResolver =
+      [Resolver](const std::string &Name) {
+        return Resolver->findSymbol(Name);
+      };
 
-    return H;
+    LDResources.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)));
+
+    return std::prev(LogicalDylibs.end());
   }
 
   /// @brief Remove the module represented by the given handle.
@@ -237,8 +142,7 @@ public:
   ///   This will remove all modules in the layers below that were derived from
   /// the module represented by H.
   void removeModuleSet(ModuleSetHandleT H) {
-    H->releaseResources(BaseLayer);
-    ModuleSetInfos.erase(H);
+    LogicalDylibs.erase(H);
   }
 
   /// @brief Search for the given named symbol.
@@ -253,19 +157,201 @@ public:
   ///        below this one.
   JITSymbol findSymbolIn(ModuleSetHandleT H, const std::string &Name,
                          bool ExportedSymbolsOnly) {
-    BaseLayerModuleSetHandleListT &BaseLayerHandles = H->second;
-    for (auto &BH : BaseLayerHandles) {
-      if (auto Symbol = BaseLayer.findSymbolIn(BH, Name, ExportedSymbolsOnly))
-        return Symbol;
-    }
-    return nullptr;
+    return H->findSymbol(Name, ExportedSymbolsOnly);
   }
 
 private:
+
+  void addLogicalModule(CODLogicalDylib &LD, std::shared_ptr<Module> SrcM) {
+
+    // Bump the linkage and rename any anonymous/privote members in SrcM to
+    // ensure that everything will resolve properly after we partition SrcM.
+    makeAllSymbolsExternallyAccessible(*SrcM);
+
+    // Create a logical module handle for SrcM within the logical dylib.
+    auto LMH = LD.createLogicalModule();
+    auto &LMResources =  LD.getLogicalModuleResources(LMH);
+    LMResources.SourceModule = SrcM;
+
+    // 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;
+
+    // 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;
+
+      // Record all functions defined by this module.
+      if (CloneStubsIntoPartitions)
+        LMResources.StubsToClone.insert(&F);
+
+      // 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]() {
+          return this->extractAndCompile(LD, LMH, F);
+        });
+    }
+
+    // Now clone the global variable declarations.
+    GlobalDeclMaterializer GDMat(*GVsAndStubsM);
+    for (auto &GV : SrcM->globals())
+      if (!GV.isDeclaration())
+        cloneGlobalVariableDecl(*GVsAndStubsM, GV, &VMap);
+
+    // And the aliases.
+    for (auto &Alias : SrcM->aliases())
+      cloneGlobalAlias(*GVsAndStubsM, Alias, VMap, &GDMat);
+
+    // Then clone the initializers.
+    for (auto &GV : SrcM->globals())
+      if (!GV.isDeclaration())
+        moveGlobalVariableInitializer(GV, VMap, &GDMat);
+
+    // Build a resolver for the stubs module and add it to the base layer.
+    auto GVsAndStubsResolver = createLambdaResolver(
+        [&LD](const std::string &Name) {
+          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),
+                             llvm::make_unique<SectionMemoryManager>(),
+                             std::move(GVsAndStubsResolver));
+    LD.addToLogicalModule(LMH, GVsAndStubsH);
+  }
+
+  static std::string Mangle(StringRef Name, const DataLayout &DL) {
+    std::string MangledName;
+    {
+      raw_string_ostream MangledNameStream(MangledName);
+      Mangler::getNameWithPrefix(MangledNameStream, Name, DL);
+    }
+    return MangledName;
+  }
+
+  TargetAddress extractAndCompile(CODLogicalDylib &LD,
+                                  LogicalModuleHandle LMH,
+                                  Function &F) {
+    Module &SrcM = *LD.getLogicalModuleResources(LMH).SourceModule;
+
+    // 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());
+
+    auto Partition = LD.getDylibResources().Partitioner(F);
+    auto PartitionH = emitPartition(LD, LMH, Partition);
+
+    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);
+      assert(FnBodySym && "Couldn't find function body.");
+      assert(FnPtrSym && "Couldn't find function body pointer.");
+
+      TargetAddress FnBodyAddr = FnBodySym.getAddress();
+      void *FnPtrAddr = reinterpret_cast<void*>(
+          static_cast<uintptr_t>(FnPtrSym.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));
+    }
+
+    return CalledAddr;
+  }
+
+  template <typename PartitionT>
+  BaseLayerModuleSetHandleT emitPartition(CODLogicalDylib &LD,
+                                          LogicalModuleHandle LMH,
+                                          const PartitionT &Partition) {
+    auto &LMResources = LD.getLogicalModuleResources(LMH);
+    Module &SrcM = *LMResources.SourceModule;
+
+    // Create the module.
+    std::string NewName = SrcM.getName();
+    for (auto *F : Partition) {
+      NewName += ".";
+      NewName += F->getName();
+    }
+
+    auto M = llvm::make_unique<Module>(NewName, SrcM.getContext());
+    M->setDataLayout(SrcM.getDataLayout());
+    ValueToValueMapTy VMap;
+    GlobalDeclMaterializer GDM(*M, &LMResources.StubsToClone);
+
+    // Create decls in the new module.
+    for (auto *F : Partition)
+      cloneFunctionDecl(*M, *F, &VMap);
+
+    // Move the function bodies.
+    for (auto *F : Partition)
+      moveFunctionBody(*F, VMap, &GDM);
+
+    // Create memory manager and symbol resolver.
+    auto MemMgr = llvm::make_unique<SectionMemoryManager>();
+    auto Resolver = createLambdaResolver(
+        [this, &LD, LMH](const std::string &Name) {
+          if (auto Symbol = LD.findSymbolInternally(LMH, Name))
+            return RuntimeDyld::SymbolInfo(Symbol.getAddress(),
+                                           Symbol.getFlags());
+          return LD.getDylibResources().ExternalSymbolResolver(Name);
+        },
+        [this, &LD, LMH](const std::string &Name) {
+          if (auto Symbol = LD.findSymbolInternally(LMH, Name))
+            return RuntimeDyld::SymbolInfo(Symbol.getAddress(),
+                                           Symbol.getFlags());
+          return RuntimeDyld::SymbolInfo(nullptr);
+        });
+    std::vector<std::unique_ptr<Module>> PartMSet;
+    PartMSet.push_back(std::move(M));
+    return BaseLayer.addModuleSet(std::move(PartMSet), std::move(MemMgr),
+                                  std::move(Resolver));
+  }
+
   BaseLayerT &BaseLayer;
-  InsertCallbackAsmFtor InsertCallbackAsm;
-  ModuleSetInfoListT ModuleSetInfos;
+  CompileCallbackMgrT &CompileCallbackMgr;
+  LogicalDylibList LogicalDylibs;
+  bool CloneStubsIntoPartitions;
 };
-}
+
+} // End namespace orc.
+} // End namespace llvm.
 
 #endif // LLVM_EXECUTIONENGINE_ORC_COMPILEONDEMANDLAYER_H