Home | History | Annotate | Download | only in Orc
      1 //===- CompileOnDemandLayer.h - Compile each function on demand -*- C++ -*-===//
      2 //
      3 //                     The LLVM Compiler Infrastructure
      4 //
      5 // This file is distributed under the University of Illinois Open Source
      6 // License. See LICENSE.TXT for details.
      7 //
      8 //===----------------------------------------------------------------------===//
      9 //
     10 // JIT layer for breaking up modules and inserting callbacks to allow
     11 // individual functions to be compiled on demand.
     12 //
     13 //===----------------------------------------------------------------------===//
     14 
     15 #ifndef LLVM_EXECUTIONENGINE_ORC_COMPILEONDEMANDLAYER_H
     16 #define LLVM_EXECUTIONENGINE_ORC_COMPILEONDEMANDLAYER_H
     17 
     18 #include "IndirectionUtils.h"
     19 #include "LambdaResolver.h"
     20 #include "LogicalDylib.h"
     21 #include "llvm/ADT/STLExtras.h"
     22 #include "llvm/ExecutionEngine/SectionMemoryManager.h"
     23 #include "llvm/Transforms/Utils/Cloning.h"
     24 #include <list>
     25 #include <memory>
     26 #include <set>
     27 
     28 #include "llvm/Support/Debug.h"
     29 
     30 namespace llvm {
     31 namespace orc {
     32 
     33 /// @brief Compile-on-demand layer.
     34 ///
     35 ///   When a module is added to this layer a stub is created for each of its
     36 /// function definitions. The stubs and other global values are immediately
     37 /// added to the layer below. When a stub is called it triggers the extraction
     38 /// of the function body from the original module. The extracted body is then
     39 /// compiled and executed.
     40 template <typename BaseLayerT,
     41           typename CompileCallbackMgrT = JITCompileCallbackManager,
     42           typename IndirectStubsMgrT = IndirectStubsManager>
     43 class CompileOnDemandLayer {
     44 private:
     45 
     46   template <typename MaterializerFtor>
     47   class LambdaMaterializer final : public ValueMaterializer {
     48   public:
     49     LambdaMaterializer(MaterializerFtor M) : M(std::move(M)) {}
     50     Value *materializeDeclFor(Value *V) final { return M(V); }
     51 
     52   private:
     53     MaterializerFtor M;
     54   };
     55 
     56   template <typename MaterializerFtor>
     57   LambdaMaterializer<MaterializerFtor>
     58   createLambdaMaterializer(MaterializerFtor M) {
     59     return LambdaMaterializer<MaterializerFtor>(std::move(M));
     60   }
     61 
     62   typedef typename BaseLayerT::ModuleSetHandleT BaseLayerModuleSetHandleT;
     63 
     64   class ModuleOwner {
     65   public:
     66     ModuleOwner() = default;
     67     ModuleOwner(const ModuleOwner&) = delete;
     68     ModuleOwner& operator=(const ModuleOwner&) = delete;
     69     virtual ~ModuleOwner() { }
     70     virtual Module& getModule() const = 0;
     71   };
     72 
     73   template <typename ModulePtrT>
     74   class ModuleOwnerImpl : public ModuleOwner {
     75   public:
     76     ModuleOwnerImpl(ModulePtrT ModulePtr) : ModulePtr(std::move(ModulePtr)) {}
     77     Module& getModule() const override { return *ModulePtr; }
     78   private:
     79     ModulePtrT ModulePtr;
     80   };
     81 
     82   template <typename ModulePtrT>
     83   std::unique_ptr<ModuleOwner> wrapOwnership(ModulePtrT ModulePtr) {
     84     return llvm::make_unique<ModuleOwnerImpl<ModulePtrT>>(std::move(ModulePtr));
     85   }
     86 
     87   struct LogicalModuleResources {
     88     std::unique_ptr<ModuleOwner> SourceModuleOwner;
     89     std::set<const Function*> StubsToClone;
     90     std::unique_ptr<IndirectStubsMgrT> StubsMgr;
     91 
     92     LogicalModuleResources() = default;
     93 
     94     // Explicit move constructor to make MSVC happy.
     95     LogicalModuleResources(LogicalModuleResources &&Other)
     96         : SourceModuleOwner(std::move(Other.SourceModuleOwner)),
     97           StubsToClone(std::move(Other.StubsToClone)),
     98           StubsMgr(std::move(Other.StubsMgr)) {}
     99 
    100     // Explicit move assignment to make MSVC happy.
    101     LogicalModuleResources& operator=(LogicalModuleResources &&Other) {
    102       SourceModuleOwner = std::move(Other.SourceModuleOwner);
    103       StubsToClone = std::move(Other.StubsToClone);
    104       StubsMgr = std::move(Other.StubsMgr);
    105     }
    106 
    107     JITSymbol findSymbol(StringRef Name, bool ExportedSymbolsOnly) {
    108       if (Name.endswith("$stub_ptr") && !ExportedSymbolsOnly) {
    109         assert(!ExportedSymbolsOnly && "Stubs are never exported");
    110         return StubsMgr->findPointer(Name.drop_back(9));
    111       }
    112       return StubsMgr->findStub(Name, ExportedSymbolsOnly);
    113     }
    114 
    115   };
    116 
    117 
    118 
    119   struct LogicalDylibResources {
    120     typedef std::function<RuntimeDyld::SymbolInfo(const std::string&)>
    121       SymbolResolverFtor;
    122     SymbolResolverFtor ExternalSymbolResolver;
    123   };
    124 
    125   typedef LogicalDylib<BaseLayerT, LogicalModuleResources,
    126                        LogicalDylibResources> CODLogicalDylib;
    127 
    128   typedef typename CODLogicalDylib::LogicalModuleHandle LogicalModuleHandle;
    129   typedef std::list<CODLogicalDylib> LogicalDylibList;
    130 
    131 public:
    132 
    133   /// @brief Handle to a set of loaded modules.
    134   typedef typename LogicalDylibList::iterator ModuleSetHandleT;
    135 
    136   /// @brief Module partitioning functor.
    137   typedef std::function<std::set<Function*>(Function&)> PartitioningFtor;
    138 
    139   /// @brief Builder for IndirectStubsManagers.
    140   typedef std::function<std::unique_ptr<IndirectStubsMgrT>()>
    141     IndirectStubsManagerBuilderT;
    142 
    143   /// @brief Construct a compile-on-demand layer instance.
    144   CompileOnDemandLayer(BaseLayerT &BaseLayer, PartitioningFtor Partition,
    145                        CompileCallbackMgrT &CallbackMgr,
    146                        IndirectStubsManagerBuilderT CreateIndirectStubsManager,
    147                        bool CloneStubsIntoPartitions = true)
    148       : BaseLayer(BaseLayer),  Partition(Partition),
    149         CompileCallbackMgr(CallbackMgr),
    150         CreateIndirectStubsManager(std::move(CreateIndirectStubsManager)),
    151         CloneStubsIntoPartitions(CloneStubsIntoPartitions) {}
    152 
    153   /// @brief Add a module to the compile-on-demand layer.
    154   template <typename ModuleSetT, typename MemoryManagerPtrT,
    155             typename SymbolResolverPtrT>
    156   ModuleSetHandleT addModuleSet(ModuleSetT Ms,
    157                                 MemoryManagerPtrT MemMgr,
    158                                 SymbolResolverPtrT Resolver) {
    159 
    160     assert(MemMgr == nullptr &&
    161            "User supplied memory managers not supported with COD yet.");
    162 
    163     LogicalDylibs.push_back(CODLogicalDylib(BaseLayer));
    164     auto &LDResources = LogicalDylibs.back().getDylibResources();
    165 
    166     LDResources.ExternalSymbolResolver =
    167       [Resolver](const std::string &Name) {
    168         return Resolver->findSymbol(Name);
    169       };
    170 
    171     // Process each of the modules in this module set.
    172     for (auto &M : Ms)
    173       addLogicalModule(LogicalDylibs.back(), std::move(M));
    174 
    175     return std::prev(LogicalDylibs.end());
    176   }
    177 
    178   /// @brief Remove the module represented by the given handle.
    179   ///
    180   ///   This will remove all modules in the layers below that were derived from
    181   /// the module represented by H.
    182   void removeModuleSet(ModuleSetHandleT H) {
    183     LogicalDylibs.erase(H);
    184   }
    185 
    186   /// @brief Search for the given named symbol.
    187   /// @param Name The name of the symbol to search for.
    188   /// @param ExportedSymbolsOnly If true, search only for exported symbols.
    189   /// @return A handle for the given named symbol, if it exists.
    190   JITSymbol findSymbol(StringRef Name, bool ExportedSymbolsOnly) {
    191     for (auto LDI = LogicalDylibs.begin(), LDE = LogicalDylibs.end();
    192          LDI != LDE; ++LDI)
    193       if (auto Symbol = findSymbolIn(LDI, Name, ExportedSymbolsOnly))
    194         return Symbol;
    195     return BaseLayer.findSymbol(Name, ExportedSymbolsOnly);
    196   }
    197 
    198   /// @brief Get the address of a symbol provided by this layer, or some layer
    199   ///        below this one.
    200   JITSymbol findSymbolIn(ModuleSetHandleT H, const std::string &Name,
    201                          bool ExportedSymbolsOnly) {
    202     return H->findSymbol(Name, ExportedSymbolsOnly);
    203   }
    204 
    205 private:
    206 
    207   template <typename ModulePtrT>
    208   void addLogicalModule(CODLogicalDylib &LD, ModulePtrT SrcMPtr) {
    209 
    210     // Bump the linkage and rename any anonymous/privote members in SrcM to
    211     // ensure that everything will resolve properly after we partition SrcM.
    212     makeAllSymbolsExternallyAccessible(*SrcMPtr);
    213 
    214     // Create a logical module handle for SrcM within the logical dylib.
    215     auto LMH = LD.createLogicalModule();
    216     auto &LMResources =  LD.getLogicalModuleResources(LMH);
    217 
    218     LMResources.SourceModuleOwner = wrapOwnership(std::move(SrcMPtr));
    219 
    220     Module &SrcM = LMResources.SourceModuleOwner->getModule();
    221 
    222     // Create the GlobalValues module.
    223     const DataLayout &DL = SrcM.getDataLayout();
    224     auto GVsM = llvm::make_unique<Module>((SrcM.getName() + ".globals").str(),
    225                                           SrcM.getContext());
    226     GVsM->setDataLayout(DL);
    227 
    228     // Create function stubs.
    229     ValueToValueMapTy VMap;
    230     {
    231       typename IndirectStubsMgrT::StubInitsMap StubInits;
    232       for (auto &F : SrcM) {
    233         // Skip declarations.
    234         if (F.isDeclaration())
    235           continue;
    236 
    237         // Record all functions defined by this module.
    238         if (CloneStubsIntoPartitions)
    239           LMResources.StubsToClone.insert(&F);
    240 
    241         // Create a callback, associate it with the stub for the function,
    242         // and set the compile action to compile the partition containing the
    243         // function.
    244         auto CCInfo = CompileCallbackMgr.getCompileCallback();
    245         StubInits[mangle(F.getName(), DL)] =
    246           std::make_pair(CCInfo.getAddress(),
    247                          JITSymbolBase::flagsFromGlobalValue(F));
    248         CCInfo.setCompileAction([this, &LD, LMH, &F]() {
    249           return this->extractAndCompile(LD, LMH, F);
    250         });
    251       }
    252 
    253       LMResources.StubsMgr = CreateIndirectStubsManager();
    254       auto EC = LMResources.StubsMgr->createStubs(StubInits);
    255       (void)EC;
    256       // FIXME: This should be propagated back to the user. Stub creation may
    257       //        fail for remote JITs.
    258       assert(!EC && "Error generating stubs");
    259     }
    260 
    261     // Clone global variable decls.
    262     for (auto &GV : SrcM.globals())
    263       if (!GV.isDeclaration() && !VMap.count(&GV))
    264         cloneGlobalVariableDecl(*GVsM, GV, &VMap);
    265 
    266     // And the aliases.
    267     for (auto &A : SrcM.aliases())
    268       if (!VMap.count(&A))
    269         cloneGlobalAliasDecl(*GVsM, A, VMap);
    270 
    271     // Now we need to clone the GV and alias initializers.
    272 
    273     // Initializers may refer to functions declared (but not defined) in this
    274     // module. Build a materializer to clone decls on demand.
    275     auto Materializer = createLambdaMaterializer(
    276       [this, &GVsM, &LMResources](Value *V) -> Value* {
    277         if (auto *F = dyn_cast<Function>(V)) {
    278           // Decls in the original module just get cloned.
    279           if (F->isDeclaration())
    280             return cloneFunctionDecl(*GVsM, *F);
    281 
    282           // Definitions in the original module (which we have emitted stubs
    283           // for at this point) get turned into a constant alias to the stub
    284           // instead.
    285           const DataLayout &DL = GVsM->getDataLayout();
    286           std::string FName = mangle(F->getName(), DL);
    287           auto StubSym = LMResources.StubsMgr->findStub(FName, false);
    288           unsigned PtrBitWidth = DL.getPointerTypeSizeInBits(F->getType());
    289           ConstantInt *StubAddr =
    290             ConstantInt::get(GVsM->getContext(),
    291                              APInt(PtrBitWidth, StubSym.getAddress()));
    292           Constant *Init = ConstantExpr::getCast(Instruction::IntToPtr,
    293                                                  StubAddr, F->getType());
    294           return GlobalAlias::create(F->getFunctionType(),
    295                                      F->getType()->getAddressSpace(),
    296                                      F->getLinkage(), F->getName(),
    297                                      Init, GVsM.get());
    298         }
    299         // else....
    300         return nullptr;
    301       });
    302 
    303     // Clone the global variable initializers.
    304     for (auto &GV : SrcM.globals())
    305       if (!GV.isDeclaration())
    306         moveGlobalVariableInitializer(GV, VMap, &Materializer);
    307 
    308     // Clone the global alias initializers.
    309     for (auto &A : SrcM.aliases()) {
    310       auto *NewA = cast<GlobalAlias>(VMap[&A]);
    311       assert(NewA && "Alias not cloned?");
    312       Value *Init = MapValue(A.getAliasee(), VMap, RF_None, nullptr,
    313                              &Materializer);
    314       NewA->setAliasee(cast<Constant>(Init));
    315     }
    316 
    317     // Build a resolver for the globals module and add it to the base layer.
    318     auto GVsResolver = createLambdaResolver(
    319         [&LD, LMH](const std::string &Name) {
    320           auto &LMResources = LD.getLogicalModuleResources(LMH);
    321           if (auto Sym = LMResources.StubsMgr->findStub(Name, false))
    322             return RuntimeDyld::SymbolInfo(Sym.getAddress(), Sym.getFlags());
    323           return LD.getDylibResources().ExternalSymbolResolver(Name);
    324         },
    325         [](const std::string &Name) {
    326           return RuntimeDyld::SymbolInfo(nullptr);
    327         });
    328 
    329     std::vector<std::unique_ptr<Module>> GVsMSet;
    330     GVsMSet.push_back(std::move(GVsM));
    331     auto GVsH =
    332       BaseLayer.addModuleSet(std::move(GVsMSet),
    333                              llvm::make_unique<SectionMemoryManager>(),
    334                              std::move(GVsResolver));
    335     LD.addToLogicalModule(LMH, GVsH);
    336   }
    337 
    338   static std::string mangle(StringRef Name, const DataLayout &DL) {
    339     std::string MangledName;
    340     {
    341       raw_string_ostream MangledNameStream(MangledName);
    342       Mangler::getNameWithPrefix(MangledNameStream, Name, DL);
    343     }
    344     return MangledName;
    345   }
    346 
    347   TargetAddress extractAndCompile(CODLogicalDylib &LD,
    348                                   LogicalModuleHandle LMH,
    349                                   Function &F) {
    350     auto &LMResources = LD.getLogicalModuleResources(LMH);
    351     Module &SrcM = LMResources.SourceModuleOwner->getModule();
    352 
    353     // If F is a declaration we must already have compiled it.
    354     if (F.isDeclaration())
    355       return 0;
    356 
    357     // Grab the name of the function being called here.
    358     std::string CalledFnName = mangle(F.getName(), SrcM.getDataLayout());
    359 
    360     auto Part = Partition(F);
    361     auto PartH = emitPartition(LD, LMH, Part);
    362 
    363     TargetAddress CalledAddr = 0;
    364     for (auto *SubF : Part) {
    365       std::string FnName = mangle(SubF->getName(), SrcM.getDataLayout());
    366       auto FnBodySym = BaseLayer.findSymbolIn(PartH, FnName, false);
    367       assert(FnBodySym && "Couldn't find function body.");
    368 
    369       TargetAddress FnBodyAddr = FnBodySym.getAddress();
    370 
    371       // If this is the function we're calling record the address so we can
    372       // return it from this function.
    373       if (SubF == &F)
    374         CalledAddr = FnBodyAddr;
    375 
    376       // Update the function body pointer for the stub.
    377       if (auto EC = LMResources.StubsMgr->updatePointer(FnName, FnBodyAddr))
    378         return 0;
    379     }
    380 
    381     return CalledAddr;
    382   }
    383 
    384   template <typename PartitionT>
    385   BaseLayerModuleSetHandleT emitPartition(CODLogicalDylib &LD,
    386                                           LogicalModuleHandle LMH,
    387                                           const PartitionT &Part) {
    388     auto &LMResources = LD.getLogicalModuleResources(LMH);
    389     Module &SrcM = LMResources.SourceModuleOwner->getModule();
    390 
    391     // Create the module.
    392     std::string NewName = SrcM.getName();
    393     for (auto *F : Part) {
    394       NewName += ".";
    395       NewName += F->getName();
    396     }
    397 
    398     auto M = llvm::make_unique<Module>(NewName, SrcM.getContext());
    399     M->setDataLayout(SrcM.getDataLayout());
    400     ValueToValueMapTy VMap;
    401 
    402     auto Materializer = createLambdaMaterializer([this, &LMResources, &M,
    403                                                   &VMap](Value *V) -> Value * {
    404       if (auto *GV = dyn_cast<GlobalVariable>(V))
    405         return cloneGlobalVariableDecl(*M, *GV);
    406 
    407       if (auto *F = dyn_cast<Function>(V)) {
    408         // Check whether we want to clone an available_externally definition.
    409         if (!LMResources.StubsToClone.count(F))
    410           return cloneFunctionDecl(*M, *F);
    411 
    412         // Ok - we want an inlinable stub. For that to work we need a decl
    413         // for the stub pointer.
    414         auto *StubPtr = createImplPointer(*F->getType(), *M,
    415                                           F->getName() + "$stub_ptr", nullptr);
    416         auto *ClonedF = cloneFunctionDecl(*M, *F);
    417         makeStub(*ClonedF, *StubPtr);
    418         ClonedF->setLinkage(GlobalValue::AvailableExternallyLinkage);
    419         ClonedF->addFnAttr(Attribute::AlwaysInline);
    420         return ClonedF;
    421       }
    422 
    423       if (auto *A = dyn_cast<GlobalAlias>(V)) {
    424         auto *Ty = A->getValueType();
    425         if (Ty->isFunctionTy())
    426           return Function::Create(cast<FunctionType>(Ty),
    427                                   GlobalValue::ExternalLinkage, A->getName(),
    428                                   M.get());
    429 
    430         return new GlobalVariable(*M, Ty, false, GlobalValue::ExternalLinkage,
    431                                   nullptr, A->getName(), nullptr,
    432                                   GlobalValue::NotThreadLocal,
    433                                   A->getType()->getAddressSpace());
    434       }
    435 
    436       return nullptr;
    437     });
    438 
    439     // Create decls in the new module.
    440     for (auto *F : Part)
    441       cloneFunctionDecl(*M, *F, &VMap);
    442 
    443     // Move the function bodies.
    444     for (auto *F : Part)
    445       moveFunctionBody(*F, VMap, &Materializer);
    446 
    447     // Create memory manager and symbol resolver.
    448     auto MemMgr = llvm::make_unique<SectionMemoryManager>();
    449     auto Resolver = createLambdaResolver(
    450         [this, &LD, LMH](const std::string &Name) {
    451           if (auto Symbol = LD.findSymbolInternally(LMH, Name))
    452             return RuntimeDyld::SymbolInfo(Symbol.getAddress(),
    453                                            Symbol.getFlags());
    454           return LD.getDylibResources().ExternalSymbolResolver(Name);
    455         },
    456         [this, &LD, LMH](const std::string &Name) {
    457           if (auto Symbol = LD.findSymbolInternally(LMH, Name))
    458             return RuntimeDyld::SymbolInfo(Symbol.getAddress(),
    459                                            Symbol.getFlags());
    460           return RuntimeDyld::SymbolInfo(nullptr);
    461         });
    462     std::vector<std::unique_ptr<Module>> PartMSet;
    463     PartMSet.push_back(std::move(M));
    464     return BaseLayer.addModuleSet(std::move(PartMSet), std::move(MemMgr),
    465                                   std::move(Resolver));
    466   }
    467 
    468   BaseLayerT &BaseLayer;
    469   PartitioningFtor Partition;
    470   CompileCallbackMgrT &CompileCallbackMgr;
    471   IndirectStubsManagerBuilderT CreateIndirectStubsManager;
    472 
    473   LogicalDylibList LogicalDylibs;
    474   bool CloneStubsIntoPartitions;
    475 };
    476 
    477 } // End namespace orc.
    478 } // End namespace llvm.
    479 
    480 #endif // LLVM_EXECUTIONENGINE_ORC_COMPILEONDEMANDLAYER_H
    481