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