[Orc] Teach the CompileOnDemand layer to clone aliases.

[oota-llvm.git] / include / llvm / ExecutionEngine / Orc / CompileOnDemandLayer.h

diff --git a/include/llvm/ExecutionEngine/Orc/CompileOnDemandLayer.h b/include/llvm/ExecutionEngine/Orc/CompileOnDemandLayer.h
index cc93a6fa431982468f7c7cb6c7845d48bd6273e5..714ca2374dc359a3343307dacb8edf2e9994f4e0 100644 (file)
--- a/include/llvm/ExecutionEngine/Orc/CompileOnDemandLayer.h
+++ b/include/llvm/ExecutionEngine/Orc/CompileOnDemandLayer.h
@@ -46,23 +46,39 @@ private:
   // variables.
   class GlobalDeclMaterializer : public ValueMaterializer {
   public:
-    GlobalDeclMaterializer(Module &Dst) : Dst(Dst) {}
+    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))
-        return cloneFunctionDecl(Dst, *F);
+      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:
     Module &Dst;
+    const StubSet *StubsToClone;
   };
 
   typedef typename BaseLayerT::ModuleSetHandleT BaseLayerModuleSetHandleT;
 
   struct LogicalModuleResources {
     std::shared_ptr<Module> SourceModule;
+    std::set<const Function*> StubsToClone;
   };
 
   struct LogicalDylibResources {
@@ -83,8 +99,10 @@ public:
   typedef typename LogicalDylibList::iterator ModuleSetHandleT;
 
   /// @brief Construct a compile-on-demand layer instance.
-  CompileOnDemandLayer(BaseLayerT &BaseLayer, CompileCallbackMgrT &CallbackMgr)
-      : BaseLayer(BaseLayer), CompileCallbackMgr(CallbackMgr) {}
+  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, typename MemoryManagerPtrT,
@@ -97,14 +115,14 @@ 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 =
+    LDResources.Partitioner =
       [](Function &F) {
         std::set<Function*> Partition;
         Partition.insert(&F);
@@ -152,7 +170,8 @@ private:
 
     // 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);
+    LMResources.SourceModule = SrcM;
 
     // Create the GVs-and-stubs module.
     auto GVsAndStubsM = llvm::make_unique<Module>(
@@ -171,6 +190,10 @@ private:
       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.
@@ -184,7 +207,7 @@ private:
       makeStub(*StubF, *FnBodyPtr);
       CCInfo.setCompileAction(
         [this, &LD, LMH, &F]() {
-          return extractAndCompile(LD, LMH, F);
+          return this->extractAndCompile(LD, LMH, F);
         });
     }
 
@@ -194,6 +217,10 @@ private:
       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())
@@ -218,11 +245,10 @@ private:
   }
 
   static std::string Mangle(StringRef Name, const DataLayout &DL) {
-    Mangler M(&DL);
     std::string MangledName;
     {
       raw_string_ostream MangledNameStream(MangledName);
-      M.getNameWithPrefix(MangledNameStream, Name);
+      Mangler::getNameWithPrefix(MangledNameStream, Name, DL);
     }
     return MangledName;
   }
@@ -239,12 +265,12 @@ private:
     // Grab the name of the function being called here.
     std::string CalledFnName = Mangle(F.getName(), SrcM.getDataLayout());
 
-    const auto &Partition = LD.getDylibResources().Partitioner(F);
+    auto Partition = LD.getDylibResources().Partitioner(F);
     auto PartitionH = emitPartition(LD, LMH, Partition);
 
     TargetAddress CalledAddr = 0;
     for (auto *SubF : Partition) {
-      std::string FName(SubF->getName());
+      std::string FName = SubF->getName();
       auto FnBodySym =
         BaseLayer.findSymbolIn(PartitionH, Mangle(FName, SrcM.getDataLayout()),
                                false);
@@ -256,7 +282,7 @@ private:
       assert(FnBodySym && "Couldn't find function body.");
       assert(FnPtrSym && "Couldn't find function body pointer.");
 
-      auto FnBodyAddr = FnBodySym.getAddress();
+      TargetAddress FnBodyAddr = FnBodySym.getAddress();
       void *FnPtrAddr = reinterpret_cast<void*>(
           static_cast<uintptr_t>(FnPtrSym.getAddress()));
 
@@ -271,13 +297,15 @@ private:
     return CalledAddr;
   }
 
+  template <typename PartitionT>
   BaseLayerModuleSetHandleT emitPartition(CODLogicalDylib &LD,
                                           LogicalModuleHandle LMH,
-                                          const std::set<Function*> &Partition) {
-    Module &SrcM = *LD.getLogicalModuleResources(LMH).SourceModule;
+                                          const PartitionT &Partition) {
+    auto &LMResources = LD.getLogicalModuleResources(LMH);
+    Module &SrcM = *LMResources.SourceModule;
 
     // Create the module.
-    std::string NewName(SrcM.getName());
+    std::string NewName = SrcM.getName();
     for (auto *F : Partition) {
       NewName += ".";
       NewName += F->getName();
@@ -286,7 +314,7 @@ private:
     auto M = llvm::make_unique<Module>(NewName, SrcM.getContext());
     M->setDataLayout(SrcM.getDataLayout());
     ValueToValueMapTy VMap;
-    GlobalDeclMaterializer GDM(*M);
+    GlobalDeclMaterializer GDM(*M, &LMResources.StubsToClone);
 
     // Create decls in the new module.
     for (auto *F : Partition)
@@ -294,7 +322,7 @@ private:
 
     // Move the function bodies.
     for (auto *F : Partition)
-      moveFunctionBody(*F, VMap);
+      moveFunctionBody(*F, VMap, &GDM);
 
     // Create memory manager and symbol resolver.
     auto MemMgr = llvm::make_unique<SectionMemoryManager>();
@@ -320,6 +348,7 @@ private:
   BaseLayerT &BaseLayer;
   CompileCallbackMgrT &CompileCallbackMgr;
   LogicalDylibList LogicalDylibs;
+  bool CloneStubsIntoPartitions;
 };
 
 } // End namespace orc.