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