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 "llvm/ADT/APInt.h" 19 #include "llvm/ADT/STLExtras.h" 20 #include "llvm/ADT/StringRef.h" 21 #include "llvm/ADT/Twine.h" 22 #include "llvm/ExecutionEngine/JITSymbol.h" 23 #include "llvm/ExecutionEngine/RuntimeDyld.h" 24 #include "llvm/ExecutionEngine/Orc/IndirectionUtils.h" 25 #include "llvm/ExecutionEngine/Orc/LambdaResolver.h" 26 #include "llvm/IR/Attributes.h" 27 #include "llvm/IR/Constants.h" 28 #include "llvm/IR/DataLayout.h" 29 #include "llvm/IR/Function.h" 30 #include "llvm/IR/GlobalAlias.h" 31 #include "llvm/IR/GlobalValue.h" 32 #include "llvm/IR/GlobalVariable.h" 33 #include "llvm/IR/Mangler.h" 34 #include "llvm/IR/Module.h" 35 #include "llvm/Support/Casting.h" 36 #include "llvm/Support/raw_ostream.h" 37 #include <algorithm> 38 #include <cassert> 39 #include <functional> 40 #include <iterator> 41 #include <list> 42 #include <memory> 43 #include <set> 44 #include <string> 45 #include <utility> 46 #include <vector> 47 48 namespace llvm { 49 namespace orc { 50 51 /// @brief Compile-on-demand layer. 52 /// 53 /// When a module is added to this layer a stub is created for each of its 54 /// function definitions. The stubs and other global values are immediately 55 /// added to the layer below. When a stub is called it triggers the extraction 56 /// of the function body from the original module. The extracted body is then 57 /// compiled and executed. 58 template <typename BaseLayerT, 59 typename CompileCallbackMgrT = JITCompileCallbackManager, 60 typename IndirectStubsMgrT = IndirectStubsManager> 61 class CompileOnDemandLayer { 62 private: 63 template <typename MaterializerFtor> 64 class LambdaMaterializer final : public ValueMaterializer { 65 public: 66 LambdaMaterializer(MaterializerFtor M) : M(std::move(M)) {} 67 68 Value *materialize(Value *V) final { return M(V); } 69 70 private: 71 MaterializerFtor M; 72 }; 73 74 template <typename MaterializerFtor> 75 LambdaMaterializer<MaterializerFtor> 76 createLambdaMaterializer(MaterializerFtor M) { 77 return LambdaMaterializer<MaterializerFtor>(std::move(M)); 78 } 79 80 typedef typename BaseLayerT::ModuleSetHandleT BaseLayerModuleSetHandleT; 81 82 // Provide type-erasure for the Modules and MemoryManagers. 83 template <typename ResourceT> 84 class ResourceOwner { 85 public: 86 ResourceOwner() = default; 87 ResourceOwner(const ResourceOwner&) = delete; 88 ResourceOwner& operator=(const ResourceOwner&) = delete; 89 virtual ~ResourceOwner() = default; 90 91 virtual ResourceT& getResource() const = 0; 92 }; 93 94 template <typename ResourceT, typename ResourcePtrT> 95 class ResourceOwnerImpl : public ResourceOwner<ResourceT> { 96 public: 97 ResourceOwnerImpl(ResourcePtrT ResourcePtr) 98 : ResourcePtr(std::move(ResourcePtr)) {} 99 100 ResourceT& getResource() const override { return *ResourcePtr; } 101 102 private: 103 ResourcePtrT ResourcePtr; 104 }; 105 106 template <typename ResourceT, typename ResourcePtrT> 107 std::unique_ptr<ResourceOwner<ResourceT>> 108 wrapOwnership(ResourcePtrT ResourcePtr) { 109 typedef ResourceOwnerImpl<ResourceT, ResourcePtrT> RO; 110 return llvm::make_unique<RO>(std::move(ResourcePtr)); 111 } 112 113 class StaticGlobalRenamer { 114 public: 115 StaticGlobalRenamer() = default; 116 StaticGlobalRenamer(StaticGlobalRenamer &&) = default; 117 StaticGlobalRenamer &operator=(StaticGlobalRenamer &&) = default; 118 119 void rename(Module &M) { 120 for (auto &F : M) 121 if (F.hasLocalLinkage()) 122 F.setName("$static." + Twine(NextId++)); 123 for (auto &G : M.globals()) 124 if (G.hasLocalLinkage()) 125 G.setName("$static." + Twine(NextId++)); 126 } 127 128 private: 129 unsigned NextId = 0; 130 }; 131 132 struct LogicalDylib { 133 typedef std::function<JITSymbol(const std::string&)> SymbolResolverFtor; 134 135 typedef std::function<typename BaseLayerT::ModuleSetHandleT( 136 BaseLayerT&, 137 std::unique_ptr<Module>, 138 std::unique_ptr<JITSymbolResolver>)> 139 ModuleAdderFtor; 140 141 struct SourceModuleEntry { 142 std::unique_ptr<ResourceOwner<Module>> SourceMod; 143 std::set<Function*> StubsToClone; 144 }; 145 146 typedef std::vector<SourceModuleEntry> SourceModulesList; 147 typedef typename SourceModulesList::size_type SourceModuleHandle; 148 149 SourceModuleHandle 150 addSourceModule(std::unique_ptr<ResourceOwner<Module>> M) { 151 SourceModuleHandle H = SourceModules.size(); 152 SourceModules.push_back(SourceModuleEntry()); 153 SourceModules.back().SourceMod = std::move(M); 154 return H; 155 } 156 157 Module& getSourceModule(SourceModuleHandle H) { 158 return SourceModules[H].SourceMod->getResource(); 159 } 160 161 std::set<Function*>& getStubsToClone(SourceModuleHandle H) { 162 return SourceModules[H].StubsToClone; 163 } 164 165 JITSymbol findSymbol(BaseLayerT &BaseLayer, const std::string &Name, 166 bool ExportedSymbolsOnly) { 167 if (auto Sym = StubsMgr->findStub(Name, ExportedSymbolsOnly)) 168 return Sym; 169 for (auto BLH : BaseLayerHandles) 170 if (auto Sym = BaseLayer.findSymbolIn(BLH, Name, ExportedSymbolsOnly)) 171 return Sym; 172 return nullptr; 173 } 174 175 std::unique_ptr<JITSymbolResolver> ExternalSymbolResolver; 176 std::unique_ptr<ResourceOwner<RuntimeDyld::MemoryManager>> MemMgr; 177 std::unique_ptr<IndirectStubsMgrT> StubsMgr; 178 StaticGlobalRenamer StaticRenamer; 179 ModuleAdderFtor ModuleAdder; 180 SourceModulesList SourceModules; 181 std::vector<BaseLayerModuleSetHandleT> BaseLayerHandles; 182 }; 183 184 typedef std::list<LogicalDylib> LogicalDylibList; 185 186 public: 187 /// @brief Handle to a set of loaded modules. 188 typedef typename LogicalDylibList::iterator ModuleSetHandleT; 189 190 /// @brief Module partitioning functor. 191 typedef std::function<std::set<Function*>(Function&)> PartitioningFtor; 192 193 /// @brief Builder for IndirectStubsManagers. 194 typedef std::function<std::unique_ptr<IndirectStubsMgrT>()> 195 IndirectStubsManagerBuilderT; 196 197 /// @brief Construct a compile-on-demand layer instance. 198 CompileOnDemandLayer(BaseLayerT &BaseLayer, PartitioningFtor Partition, 199 CompileCallbackMgrT &CallbackMgr, 200 IndirectStubsManagerBuilderT CreateIndirectStubsManager, 201 bool CloneStubsIntoPartitions = true) 202 : BaseLayer(BaseLayer), Partition(std::move(Partition)), 203 CompileCallbackMgr(CallbackMgr), 204 CreateIndirectStubsManager(std::move(CreateIndirectStubsManager)), 205 CloneStubsIntoPartitions(CloneStubsIntoPartitions) {} 206 207 /// @brief Add a module to the compile-on-demand layer. 208 template <typename ModuleSetT, typename MemoryManagerPtrT, 209 typename SymbolResolverPtrT> 210 ModuleSetHandleT addModuleSet(ModuleSetT Ms, 211 MemoryManagerPtrT MemMgr, 212 SymbolResolverPtrT Resolver) { 213 214 LogicalDylibs.push_back(LogicalDylib()); 215 auto &LD = LogicalDylibs.back(); 216 LD.ExternalSymbolResolver = std::move(Resolver); 217 LD.StubsMgr = CreateIndirectStubsManager(); 218 219 auto &MemMgrRef = *MemMgr; 220 LD.MemMgr = wrapOwnership<RuntimeDyld::MemoryManager>(std::move(MemMgr)); 221 222 LD.ModuleAdder = 223 [&MemMgrRef](BaseLayerT &B, std::unique_ptr<Module> M, 224 std::unique_ptr<JITSymbolResolver> R) { 225 std::vector<std::unique_ptr<Module>> Ms; 226 Ms.push_back(std::move(M)); 227 return B.addModuleSet(std::move(Ms), &MemMgrRef, std::move(R)); 228 }; 229 230 // Process each of the modules in this module set. 231 for (auto &M : Ms) 232 addLogicalModule(LogicalDylibs.back(), std::move(M)); 233 234 return std::prev(LogicalDylibs.end()); 235 } 236 237 /// @brief Remove the module represented by the given handle. 238 /// 239 /// This will remove all modules in the layers below that were derived from 240 /// the module represented by H. 241 void removeModuleSet(ModuleSetHandleT H) { 242 LogicalDylibs.erase(H); 243 } 244 245 /// @brief Search for the given named symbol. 246 /// @param Name The name of the symbol to search for. 247 /// @param ExportedSymbolsOnly If true, search only for exported symbols. 248 /// @return A handle for the given named symbol, if it exists. 249 JITSymbol findSymbol(StringRef Name, bool ExportedSymbolsOnly) { 250 for (auto LDI = LogicalDylibs.begin(), LDE = LogicalDylibs.end(); 251 LDI != LDE; ++LDI) { 252 if (auto Sym = LDI->StubsMgr->findStub(Name, ExportedSymbolsOnly)) 253 return Sym; 254 if (auto Sym = findSymbolIn(LDI, Name, ExportedSymbolsOnly)) 255 return Sym; 256 } 257 return BaseLayer.findSymbol(Name, ExportedSymbolsOnly); 258 } 259 260 /// @brief Get the address of a symbol provided by this layer, or some layer 261 /// below this one. 262 JITSymbol findSymbolIn(ModuleSetHandleT H, const std::string &Name, 263 bool ExportedSymbolsOnly) { 264 return H->findSymbol(BaseLayer, Name, ExportedSymbolsOnly); 265 } 266 267 /// @brief Update the stub for the given function to point at FnBodyAddr. 268 /// This can be used to support re-optimization. 269 /// @return true if the function exists and the stub is updated, false 270 /// otherwise. 271 // 272 // FIXME: We should track and free associated resources (unused compile 273 // callbacks, uncompiled IR, and no-longer-needed/reachable function 274 // implementations). 275 // FIXME: Return Error once the JIT APIs are Errorized. 276 bool updatePointer(std::string FuncName, JITTargetAddress FnBodyAddr) { 277 //Find out which logical dylib contains our symbol 278 auto LDI = LogicalDylibs.begin(); 279 for (auto LDE = LogicalDylibs.end(); LDI != LDE; ++LDI) { 280 if (auto LMResources = LDI->getLogicalModuleResourcesForSymbol(FuncName, false)) { 281 Module &SrcM = LMResources->SourceModule->getResource(); 282 std::string CalledFnName = mangle(FuncName, SrcM.getDataLayout()); 283 if (auto EC = LMResources->StubsMgr->updatePointer(CalledFnName, FnBodyAddr)) 284 return false; 285 else 286 return true; 287 } 288 } 289 return false; 290 } 291 292 private: 293 template <typename ModulePtrT> 294 void addLogicalModule(LogicalDylib &LD, ModulePtrT SrcMPtr) { 295 296 // Rename all static functions / globals to $static.X : 297 // This will unique the names across all modules in the logical dylib, 298 // simplifying symbol lookup. 299 LD.StaticRenamer.rename(*SrcMPtr); 300 301 // Bump the linkage and rename any anonymous/privote members in SrcM to 302 // ensure that everything will resolve properly after we partition SrcM. 303 makeAllSymbolsExternallyAccessible(*SrcMPtr); 304 305 // Create a logical module handle for SrcM within the logical dylib. 306 Module &SrcM = *SrcMPtr; 307 auto LMId = LD.addSourceModule(wrapOwnership<Module>(std::move(SrcMPtr))); 308 309 // Create stub functions. 310 const DataLayout &DL = SrcM.getDataLayout(); 311 { 312 typename IndirectStubsMgrT::StubInitsMap StubInits; 313 for (auto &F : SrcM) { 314 // Skip declarations. 315 if (F.isDeclaration()) 316 continue; 317 318 // Skip weak functions for which we already have definitions. 319 auto MangledName = mangle(F.getName(), DL); 320 if (F.hasWeakLinkage() || F.hasLinkOnceLinkage()) 321 if (auto Sym = LD.findSymbol(BaseLayer, MangledName, false)) 322 continue; 323 324 // Record all functions defined by this module. 325 if (CloneStubsIntoPartitions) 326 LD.getStubsToClone(LMId).insert(&F); 327 328 // Create a callback, associate it with the stub for the function, 329 // and set the compile action to compile the partition containing the 330 // function. 331 auto CCInfo = CompileCallbackMgr.getCompileCallback(); 332 StubInits[MangledName] = 333 std::make_pair(CCInfo.getAddress(), 334 JITSymbolFlags::fromGlobalValue(F)); 335 CCInfo.setCompileAction([this, &LD, LMId, &F]() { 336 return this->extractAndCompile(LD, LMId, F); 337 }); 338 } 339 340 auto EC = LD.StubsMgr->createStubs(StubInits); 341 (void)EC; 342 // FIXME: This should be propagated back to the user. Stub creation may 343 // fail for remote JITs. 344 assert(!EC && "Error generating stubs"); 345 } 346 347 // If this module doesn't contain any globals, aliases, or module flags then 348 // we can bail out early and avoid the overhead of creating and managing an 349 // empty globals module. 350 if (SrcM.global_empty() && SrcM.alias_empty() && 351 !SrcM.getModuleFlagsMetadata()) 352 return; 353 354 // Create the GlobalValues module. 355 auto GVsM = llvm::make_unique<Module>((SrcM.getName() + ".globals").str(), 356 SrcM.getContext()); 357 GVsM->setDataLayout(DL); 358 359 ValueToValueMapTy VMap; 360 361 // Clone global variable decls. 362 for (auto &GV : SrcM.globals()) 363 if (!GV.isDeclaration() && !VMap.count(&GV)) 364 cloneGlobalVariableDecl(*GVsM, GV, &VMap); 365 366 // And the aliases. 367 for (auto &A : SrcM.aliases()) 368 if (!VMap.count(&A)) 369 cloneGlobalAliasDecl(*GVsM, A, VMap); 370 371 // Clone the module flags. 372 cloneModuleFlagsMetadata(*GVsM, SrcM, VMap); 373 374 // Now we need to clone the GV and alias initializers. 375 376 // Initializers may refer to functions declared (but not defined) in this 377 // module. Build a materializer to clone decls on demand. 378 auto Materializer = createLambdaMaterializer( 379 [&LD, &GVsM](Value *V) -> Value* { 380 if (auto *F = dyn_cast<Function>(V)) { 381 // Decls in the original module just get cloned. 382 if (F->isDeclaration()) 383 return cloneFunctionDecl(*GVsM, *F); 384 385 // Definitions in the original module (which we have emitted stubs 386 // for at this point) get turned into a constant alias to the stub 387 // instead. 388 const DataLayout &DL = GVsM->getDataLayout(); 389 std::string FName = mangle(F->getName(), DL); 390 auto StubSym = LD.StubsMgr->findStub(FName, false); 391 unsigned PtrBitWidth = DL.getPointerTypeSizeInBits(F->getType()); 392 ConstantInt *StubAddr = 393 ConstantInt::get(GVsM->getContext(), 394 APInt(PtrBitWidth, StubSym.getAddress())); 395 Constant *Init = ConstantExpr::getCast(Instruction::IntToPtr, 396 StubAddr, F->getType()); 397 return GlobalAlias::create(F->getFunctionType(), 398 F->getType()->getAddressSpace(), 399 F->getLinkage(), F->getName(), 400 Init, GVsM.get()); 401 } 402 // else.... 403 return nullptr; 404 }); 405 406 // Clone the global variable initializers. 407 for (auto &GV : SrcM.globals()) 408 if (!GV.isDeclaration()) 409 moveGlobalVariableInitializer(GV, VMap, &Materializer); 410 411 // Clone the global alias initializers. 412 for (auto &A : SrcM.aliases()) { 413 auto *NewA = cast<GlobalAlias>(VMap[&A]); 414 assert(NewA && "Alias not cloned?"); 415 Value *Init = MapValue(A.getAliasee(), VMap, RF_None, nullptr, 416 &Materializer); 417 NewA->setAliasee(cast<Constant>(Init)); 418 } 419 420 // Build a resolver for the globals module and add it to the base layer. 421 auto GVsResolver = createLambdaResolver( 422 [this, &LD](const std::string &Name) { 423 if (auto Sym = LD.StubsMgr->findStub(Name, false)) 424 return Sym; 425 if (auto Sym = LD.findSymbol(BaseLayer, Name, false)) 426 return Sym; 427 return LD.ExternalSymbolResolver->findSymbolInLogicalDylib(Name); 428 }, 429 [&LD](const std::string &Name) { 430 return LD.ExternalSymbolResolver->findSymbol(Name); 431 }); 432 433 auto GVsH = LD.ModuleAdder(BaseLayer, std::move(GVsM), 434 std::move(GVsResolver)); 435 LD.BaseLayerHandles.push_back(GVsH); 436 } 437 438 static std::string mangle(StringRef Name, const DataLayout &DL) { 439 std::string MangledName; 440 { 441 raw_string_ostream MangledNameStream(MangledName); 442 Mangler::getNameWithPrefix(MangledNameStream, Name, DL); 443 } 444 return MangledName; 445 } 446 447 JITTargetAddress 448 extractAndCompile(LogicalDylib &LD, 449 typename LogicalDylib::SourceModuleHandle LMId, 450 Function &F) { 451 Module &SrcM = LD.getSourceModule(LMId); 452 453 // If F is a declaration we must already have compiled it. 454 if (F.isDeclaration()) 455 return 0; 456 457 // Grab the name of the function being called here. 458 std::string CalledFnName = mangle(F.getName(), SrcM.getDataLayout()); 459 460 auto Part = Partition(F); 461 auto PartH = emitPartition(LD, LMId, Part); 462 463 JITTargetAddress CalledAddr = 0; 464 for (auto *SubF : Part) { 465 std::string FnName = mangle(SubF->getName(), SrcM.getDataLayout()); 466 auto FnBodySym = BaseLayer.findSymbolIn(PartH, FnName, false); 467 assert(FnBodySym && "Couldn't find function body."); 468 469 JITTargetAddress FnBodyAddr = FnBodySym.getAddress(); 470 471 // If this is the function we're calling record the address so we can 472 // return it from this function. 473 if (SubF == &F) 474 CalledAddr = FnBodyAddr; 475 476 // Update the function body pointer for the stub. 477 if (auto EC = LD.StubsMgr->updatePointer(FnName, FnBodyAddr)) 478 return 0; 479 } 480 481 return CalledAddr; 482 } 483 484 template <typename PartitionT> 485 BaseLayerModuleSetHandleT 486 emitPartition(LogicalDylib &LD, 487 typename LogicalDylib::SourceModuleHandle LMId, 488 const PartitionT &Part) { 489 Module &SrcM = LD.getSourceModule(LMId); 490 491 // Create the module. 492 std::string NewName = SrcM.getName(); 493 for (auto *F : Part) { 494 NewName += "."; 495 NewName += F->getName(); 496 } 497 498 auto M = llvm::make_unique<Module>(NewName, SrcM.getContext()); 499 M->setDataLayout(SrcM.getDataLayout()); 500 ValueToValueMapTy VMap; 501 502 auto Materializer = createLambdaMaterializer([&LD, &LMId, 503 &M](Value *V) -> Value * { 504 if (auto *GV = dyn_cast<GlobalVariable>(V)) 505 return cloneGlobalVariableDecl(*M, *GV); 506 507 if (auto *F = dyn_cast<Function>(V)) { 508 // Check whether we want to clone an available_externally definition. 509 if (!LD.getStubsToClone(LMId).count(F)) 510 return cloneFunctionDecl(*M, *F); 511 512 // Ok - we want an inlinable stub. For that to work we need a decl 513 // for the stub pointer. 514 auto *StubPtr = createImplPointer(*F->getType(), *M, 515 F->getName() + "$stub_ptr", nullptr); 516 auto *ClonedF = cloneFunctionDecl(*M, *F); 517 makeStub(*ClonedF, *StubPtr); 518 ClonedF->setLinkage(GlobalValue::AvailableExternallyLinkage); 519 ClonedF->addFnAttr(Attribute::AlwaysInline); 520 return ClonedF; 521 } 522 523 if (auto *A = dyn_cast<GlobalAlias>(V)) { 524 auto *Ty = A->getValueType(); 525 if (Ty->isFunctionTy()) 526 return Function::Create(cast<FunctionType>(Ty), 527 GlobalValue::ExternalLinkage, A->getName(), 528 M.get()); 529 530 return new GlobalVariable(*M, Ty, false, GlobalValue::ExternalLinkage, 531 nullptr, A->getName(), nullptr, 532 GlobalValue::NotThreadLocal, 533 A->getType()->getAddressSpace()); 534 } 535 536 return nullptr; 537 }); 538 539 // Create decls in the new module. 540 for (auto *F : Part) 541 cloneFunctionDecl(*M, *F, &VMap); 542 543 // Move the function bodies. 544 for (auto *F : Part) 545 moveFunctionBody(*F, VMap, &Materializer); 546 547 // Create memory manager and symbol resolver. 548 auto Resolver = createLambdaResolver( 549 [this, &LD](const std::string &Name) { 550 if (auto Sym = LD.findSymbol(BaseLayer, Name, false)) 551 return Sym; 552 return LD.ExternalSymbolResolver->findSymbolInLogicalDylib(Name); 553 }, 554 [&LD](const std::string &Name) { 555 return LD.ExternalSymbolResolver->findSymbol(Name); 556 }); 557 558 return LD.ModuleAdder(BaseLayer, std::move(M), std::move(Resolver)); 559 } 560 561 BaseLayerT &BaseLayer; 562 PartitioningFtor Partition; 563 CompileCallbackMgrT &CompileCallbackMgr; 564 IndirectStubsManagerBuilderT CreateIndirectStubsManager; 565 566 LogicalDylibList LogicalDylibs; 567 bool CloneStubsIntoPartitions; 568 }; 569 570 } // end namespace orc 571 } // end namespace llvm 572 573 #endif // LLVM_EXECUTIONENGINE_ORC_COMPILEONDEMANDLAYER_H 574