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