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