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