include/llvm/ExecutionEngine/Orc/CompileOnDemandLayer.h

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