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 "CloneSubModule.h"
  19 #include "IndirectionUtils.h"
  20 #include "LambdaResolver.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 <set>
  26
  27 #include "llvm/Support/Debug.h"
  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, typename CompileCallbackMgrT>
  40 class CompileOnDemandLayer {
  41 private:
  42
  43   // Utility class for MapValue. Only materializes declarations for global
  44   // variables.
  45   class GlobalDeclMaterializer : public ValueMaterializer {
  46   public:
  47     GlobalDeclMaterializer(Module &Dst) : Dst(Dst) {}
  48     Value* materializeValueFor(Value *V) final {
  49       if (auto *GV = dyn_cast<GlobalVariable>(V))
  50         return cloneGlobalVariableDecl(Dst, *GV);
  51       else if (auto *F = dyn_cast<Function>(V))
  52         return cloneFunctionDecl(Dst, *F);
  53       // Else.
  54       return nullptr;
  55     }
  56   private:
  57     Module &Dst;
  58   };
  59
  60   typedef typename BaseLayerT::ModuleSetHandleT BaseLayerModuleSetHandleT;
  61   class UncompiledPartition;
  62
  63   // Logical module.
  64   //
  65   //   This struct contains the handles for the global values and stubs (which
  66   // cover the external symbols of the original module), plus the handes for
  67   // each of the extracted partitions. These handleds are used for lookup (only
  68   // the globals/stubs module is searched) and memory management. The actual
  69   // searching and resource management are handled by the LogicalDylib that owns
  70   // the LogicalModule.
  71   struct LogicalModule {
  72     std::unique_ptr<Module> SrcM;
  73     BaseLayerModuleSetHandleT GVsAndStubsHandle;
  74     std::vector<BaseLayerModuleSetHandleT> ImplHandles;
  75   };
  76
  77   // Logical dylib.
  78   //
  79   //   This class handles symbol resolution and resource management for a set of
  80   // modules that were added together as a logical dylib.
  81   //
  82   //   A logical dylib contains one-or-more LogicalModules plus a set of
  83   // UncompiledPartitions. LogicalModules support symbol resolution and resource
  84   // management for for code that has already been emitted. UncompiledPartitions
  85   // represent code that has not yet been compiled.
  86   class LogicalDylib {
  87   private:
  88     friend class UncompiledPartition;
  89     typedef std::list<LogicalModule> LogicalModuleList;
  90   public:
  91
  92     typedef unsigned UncompiledPartitionID;
  93     typedef typename LogicalModuleList::iterator LMHandle;
  94
  95     // Construct a logical dylib.
  96     LogicalDylib(CompileOnDemandLayer &CODLayer) : CODLayer(CODLayer) { }
  97
  98     // Delete this logical dylib, release logical module resources.
  99     virtual ~LogicalDylib() {
 100       releaseLogicalModuleResources();
 101     }
 102
 103     // Get a reference to the containing layer.
 104     CompileOnDemandLayer& getCODLayer() { return CODLayer; }
 105
 106     // Get a reference to the base layer.
 107     BaseLayerT& getBaseLayer() { return CODLayer.BaseLayer; }
 108
 109     // Start a new context for a single logical module.
 110     LMHandle createLogicalModule() {
 111       LogicalModules.push_back(LogicalModule());
 112       return std::prev(LogicalModules.end());
 113     }
 114
 115     // Set the global-values-and-stubs module handle for this logical module.
 116     void setGVsAndStubsHandle(LMHandle LMH, BaseLayerModuleSetHandleT H) {
 117       LMH->GVsAndStubsHandle = H;
 118     }
 119
 120     // Return the global-values-and-stubs module handle for this logical module.
 121     BaseLayerModuleSetHandleT getGVsAndStubsHandle(LMHandle LMH) {
 122       return LMH->GVsAndStubsHandle;
 123     }
 124
 125     //   Add a handle to a module containing lazy function bodies to the given
 126     // logical module.
 127     void addToLogicalModule(LMHandle LMH, BaseLayerModuleSetHandleT H) {
 128       LMH->ImplHandles.push_back(H);
 129     }
 130
 131     // Create an UncompiledPartition attached to this LogicalDylib.
 132     UncompiledPartition& createUncompiledPartition(LMHandle LMH,
 133                                                    std::shared_ptr<Module> SrcM);
 134
 135     // Take ownership of the given UncompiledPartition from the logical dylib.
 136     std::unique_ptr<UncompiledPartition>
 137     takeUPOwnership(UncompiledPartitionID ID);
 138
 139     // Look up a symbol in this context.
 140     JITSymbol findSymbolInternally(LMHandle LMH, const std::string &Name) {
 141       if (auto Symbol = getBaseLayer().findSymbolIn(LMH->GVsAndStubsHandle,
 142                                                     Name, false))
 143         return Symbol;
 144
 145       for (auto I = LogicalModules.begin(), E = LogicalModules.end(); I != E;
 146            ++I)
 147         if (I != LMH)
 148           if (auto Symbol = getBaseLayer().findSymbolIn(I->GVsAndStubsHandle,
 149                                                         Name, false))
 150             return Symbol;
 151
 152       return nullptr;
 153     }
 154
 155     JITSymbol findSymbol(const std::string &Name, bool ExportedSymbolsOnly) {
 156       for (auto &LM : LogicalModules)
 157         if (auto Symbol = getBaseLayer().findSymbolIn(LM.GVsAndStubsHandle,
 158                                                       Name,
 159                                                       ExportedSymbolsOnly))
 160           return Symbol;
 161       return nullptr;
 162     }
 163
 164     // Find an external symbol (via the user supplied SymbolResolver).
 165     virtual RuntimeDyld::SymbolInfo
 166     findSymbolExternally(const std::string &Name) const = 0;
 167
 168   private:
 169
 170     void releaseLogicalModuleResources() {
 171       for (auto I = LogicalModules.begin(), E = LogicalModules.end(); I != E;
 172            ++I) {
 173         getBaseLayer().removeModuleSet(I->GVsAndStubsHandle);
 174         for (auto H : I->ImplHandles)
 175           getBaseLayer().removeModuleSet(H);
 176       }
 177     }
 178
 179     CompileOnDemandLayer &CODLayer;
 180     LogicalModuleList LogicalModules;
 181     std::vector<std::unique_ptr<UncompiledPartition>> UncompiledPartitions;
 182   };
 183
 184   template <typename ResolverPtrT>
 185   class LogicalDylibImpl : public LogicalDylib  {
 186   public:
 187     LogicalDylibImpl(CompileOnDemandLayer &CODLayer, ResolverPtrT Resolver)
 188       : LogicalDylib(CODLayer), Resolver(std::move(Resolver)) {}
 189
 190     RuntimeDyld::SymbolInfo
 191     findSymbolExternally(const std::string &Name) const override {
 192       return Resolver->findSymbol(Name);
 193     }
 194
 195   private:
 196     ResolverPtrT Resolver;
 197   };
 198
 199   template <typename ResolverPtrT>
 200   static std::unique_ptr<LogicalDylib>
 201   createLogicalDylib(CompileOnDemandLayer &CODLayer,
 202                      ResolverPtrT Resolver) {
 203     typedef LogicalDylibImpl<ResolverPtrT> Impl;
 204     return llvm::make_unique<Impl>(CODLayer, std::move(Resolver));
 205   }
 206
 207   // Uncompiled partition.
 208   //
 209   // Represents one as-yet uncompiled portion of a module.
 210   class UncompiledPartition {
 211   public:
 212
 213     struct PartitionEntry {
 214       PartitionEntry(Function *F, TargetAddress CallbackID)
 215           : F(F), CallbackID(CallbackID) {}
 216       Function *F;
 217       TargetAddress CallbackID;
 218     };
 219
 220     typedef std::vector<PartitionEntry> PartitionEntryList;
 221
 222     // Creates an uncompiled partition with the list of functions that make up
 223     // this partition.
 224     UncompiledPartition(LogicalDylib &LD, typename LogicalDylib::LMHandle LMH,
 225                         std::shared_ptr<Module> SrcM)
 226         : LD(LD), LMH(LMH), SrcM(std::move(SrcM)), ID(~0U) {}
 227
 228     ~UncompiledPartition() {
 229       // FIXME: When we want to support threaded lazy compilation we'll need to
 230       //        lock the callback manager here.
 231       auto &CCMgr = LD.getCODLayer().CompileCallbackMgr;
 232       for (auto PEntry : PartitionEntries)
 233         CCMgr.releaseCompileCallback(PEntry.CallbackID);
 234     }
 235
 236     // Set the ID for this partition.
 237     void setID(typename LogicalDylib::UncompiledPartitionID ID) {
 238       this->ID = ID;
 239     }
 240
 241     // Set the function set and callbacks for this partition.
 242     void setPartitionEntries(PartitionEntryList PartitionEntries) {
 243       this->PartitionEntries = std::move(PartitionEntries);
 244     }
 245
 246     // Handle a compile callback for the function at index FnIdx.
 247     TargetAddress compile(unsigned FnIdx) {
 248       // Take ownership of self. This will ensure we delete the partition and
 249       // free all its resources once we're done compiling.
 250       std::unique_ptr<UncompiledPartition> This = LD.takeUPOwnership(ID);
 251
 252       // Release all other compile callbacks for this partition.
 253       // We skip the callback for this function because that's the one that
 254       // called us, and the callback manager will already have removed it.
 255       auto &CCMgr = LD.getCODLayer().CompileCallbackMgr;
 256       for (unsigned I = 0; I < PartitionEntries.size(); ++I)
 257         if (I != FnIdx)
 258           CCMgr.releaseCompileCallback(PartitionEntries[I].CallbackID);
 259
 260       // Grab the name of the function being called here.
 261       Function *F = PartitionEntries[FnIdx].F;
 262       std::string CalledFnName = Mangle(F->getName(), SrcM->getDataLayout());
 263
 264       // Extract the function and add it to the base layer.
 265       auto PartitionImplH = emitPartition();
 266       LD.addToLogicalModule(LMH, PartitionImplH);
 267
 268       // Update body pointers.
 269       // FIXME: When we start supporting remote lazy jitting this will need to
 270       //        be replaced with a user-supplied callback for updating the
 271       //        remote pointers.
 272       TargetAddress CalledAddr = 0;
 273       for (unsigned I = 0; I < PartitionEntries.size(); ++I) {
 274         auto F = PartitionEntries[I].F;
 275         std::string FName(F->getName());
 276         auto FnBodySym =
 277           LD.getBaseLayer().findSymbolIn(PartitionImplH,
 278                                          Mangle(FName, SrcM->getDataLayout()),
 279                                          false);
 280         auto FnPtrSym =
 281           LD.getBaseLayer().findSymbolIn(LD.getGVsAndStubsHandle(LMH),
 282                                          Mangle(FName + "$orc_addr",
 283                                                 SrcM->getDataLayout()),
 284                                          false);
 285         assert(FnBodySym && "Couldn't find function body.");
 286         assert(FnPtrSym && "Couldn't find function body pointer.");
 287
 288         auto FnBodyAddr = FnBodySym.getAddress();
 289         void *FnPtrAddr = reinterpret_cast<void*>(
 290                             static_cast<uintptr_t>(FnPtrSym.getAddress()));
 291
 292         // If this is the function we're calling record the address so we can
 293         // return it from this function.
 294         if (I == FnIdx)
 295           CalledAddr = FnBodyAddr;
 296
 297         memcpy(FnPtrAddr, &FnBodyAddr, sizeof(uintptr_t));
 298       }
 299
 300       // Finally, clear the partition structure so we don't try to
 301       // double-release the callbacks in the UncompiledPartition destructor.
 302       PartitionEntries.clear();
 303
 304       return CalledAddr;
 305     }
 306
 307   private:
 308
 309     BaseLayerModuleSetHandleT emitPartition() {
 310       // Create the module.
 311       std::string NewName(SrcM->getName());
 312       for (auto &PEntry : PartitionEntries) {
 313         NewName += ".";
 314         NewName += PEntry.F->getName();
 315       }
 316       auto PM = llvm::make_unique<Module>(NewName, SrcM->getContext());
 317       PM->setDataLayout(SrcM->getDataLayout());
 318       ValueToValueMapTy VMap;
 319       GlobalDeclMaterializer GDM(*PM);
 320
 321       // Create decls in the new module.
 322       for (auto &PEntry : PartitionEntries)
 323         cloneFunctionDecl(*PM, *PEntry.F, &VMap);
 324
 325       // Move the function bodies.
 326       for (auto &PEntry : PartitionEntries)
 327         moveFunctionBody(*PEntry.F, VMap);
 328
 329       // Create memory manager and symbol resolver.
 330       auto MemMgr = llvm::make_unique<SectionMemoryManager>();
 331       auto Resolver = createLambdaResolver(
 332           [this](const std::string &Name) {
 333             if (auto Symbol = LD.findSymbolInternally(LMH, Name))
 334               return RuntimeDyld::SymbolInfo(Symbol.getAddress(),
 335                                              Symbol.getFlags());
 336             return LD.findSymbolExternally(Name);
 337           },
 338           [this](const std::string &Name) {
 339             if (auto Symbol = LD.findSymbolInternally(LMH, Name))
 340               return RuntimeDyld::SymbolInfo(Symbol.getAddress(),
 341                                              Symbol.getFlags());
 342             return RuntimeDyld::SymbolInfo(nullptr);
 343           });
 344       std::vector<std::unique_ptr<Module>> PartMSet;
 345       PartMSet.push_back(std::move(PM));
 346       return LD.getBaseLayer().addModuleSet(std::move(PartMSet),
 347                                             std::move(MemMgr),
 348                                             std::move(Resolver));
 349     }
 350
 351     LogicalDylib &LD;
 352     typename LogicalDylib::LMHandle LMH;
 353     std::shared_ptr<Module> SrcM;
 354     typename LogicalDylib::UncompiledPartitionID ID;
 355     PartitionEntryList PartitionEntries;
 356   };
 357
 358   typedef std::list<std::unique_ptr<LogicalDylib>> LogicalDylibList;
 359
 360 public:
 361   /// @brief Handle to a set of loaded modules.
 362   typedef typename LogicalDylibList::iterator ModuleSetHandleT;
 363
 364   /// @brief Construct a compile-on-demand layer instance.
 365   CompileOnDemandLayer(BaseLayerT &BaseLayer, CompileCallbackMgrT &CallbackMgr)
 366       : BaseLayer(BaseLayer), CompileCallbackMgr(CallbackMgr) {}
 367
 368   /// @brief Add a module to the compile-on-demand layer.
 369   template <typename ModuleSetT, typename MemoryManagerPtrT,
 370             typename SymbolResolverPtrT>
 371   ModuleSetHandleT addModuleSet(ModuleSetT Ms,
 372                                 MemoryManagerPtrT MemMgr,
 373                                 SymbolResolverPtrT Resolver) {
 374
 375     assert(MemMgr == nullptr &&
 376            "User supplied memory managers not supported with COD yet.");
 377
 378     LogicalDylibs.push_back(createLogicalDylib(*this, std::move(Resolver)));
 379
 380     // Process each of the modules in this module set.
 381     for (auto &M : Ms) {
 382       std::vector<std::vector<Function*>> Partitioning;
 383       for (auto &F : *M) {
 384         if (F.isDeclaration())
 385           continue;
 386         Partitioning.push_back(std::vector<Function*>());
 387         Partitioning.back().push_back(&F);
 388       }
 389       addLogicalModule(*LogicalDylibs.back(),
 390                        std::shared_ptr<Module>(std::move(M)),
 391                        std::move(Partitioning));
 392     }
 393
 394     return std::prev(LogicalDylibs.end());
 395   }
 396
 397   /// @brief Remove the module represented by the given handle.
 398   ///
 399   ///   This will remove all modules in the layers below that were derived from
 400   /// the module represented by H.
 401   void removeModuleSet(ModuleSetHandleT H) {
 402     LogicalDylibs.erase(H);
 403   }
 404
 405   /// @brief Search for the given named symbol.
 406   /// @param Name The name of the symbol to search for.
 407   /// @param ExportedSymbolsOnly If true, search only for exported symbols.
 408   /// @return A handle for the given named symbol, if it exists.
 409   JITSymbol findSymbol(StringRef Name, bool ExportedSymbolsOnly) {
 410     return BaseLayer.findSymbol(Name, ExportedSymbolsOnly);
 411   }
 412
 413   /// @brief Get the address of a symbol provided by this layer, or some layer
 414   ///        below this one.
 415   JITSymbol findSymbolIn(ModuleSetHandleT H, const std::string &Name,
 416                          bool ExportedSymbolsOnly) {
 417     return (*H)->findSymbol(Name, ExportedSymbolsOnly);
 418   }
 419
 420 private:
 421
 422   void addLogicalModule(LogicalDylib &LD, std::shared_ptr<Module> SrcM,
 423                         std::vector<std::vector<Function*>> Partitions) {
 424
 425     // Bump the linkage and rename any anonymous/privote members in SrcM to
 426     // ensure that everything will resolve properly after we partition SrcM.
 427     makeAllSymbolsExternallyAccessible(*SrcM);
 428
 429     // Create a logical module handle for SrcM within the logical dylib.
 430     auto LMH = LD.createLogicalModule();
 431
 432     // Create the GVs-and-stubs module.
 433     auto GVsAndStubsM = llvm::make_unique<Module>(
 434                           (SrcM->getName() + ".globals_and_stubs").str(),
 435                           SrcM->getContext());
 436     GVsAndStubsM->setDataLayout(SrcM->getDataLayout());
 437     ValueToValueMapTy VMap;
 438
 439     // Process partitions and create stubs.
 440     // We create the stubs before copying the global variables as we know the
 441     // stubs won't refer to any globals (they only refer to their implementation
 442     // pointer) so there's no ordering/value-mapping issues.
 443     for (auto& Partition : Partitions) {
 444       auto &UP = LD.createUncompiledPartition(LMH, SrcM);
 445       typename UncompiledPartition::PartitionEntryList PartitionEntries;
 446       for (auto &F : Partition) {
 447         assert(!F->isDeclaration() &&
 448                "Partition should only contain definitions");
 449         unsigned FnIdx = PartitionEntries.size();
 450         auto CCI = CompileCallbackMgr.getCompileCallback(SrcM->getContext());
 451         PartitionEntries.push_back(
 452           typename UncompiledPartition::PartitionEntry(F, CCI.getAddress()));
 453         Function *StubF = cloneFunctionDecl(*GVsAndStubsM, *F, &VMap);
 454         GlobalVariable *FnBodyPtr =
 455           createImplPointer(*StubF->getType(), *StubF->getParent(),
 456                             StubF->getName() + "$orc_addr",
 457                             createIRTypedAddress(*StubF->getFunctionType(),
 458                                                  CCI.getAddress()));
 459         makeStub(*StubF, *FnBodyPtr);
 460         CCI.setCompileAction([&UP, FnIdx]() { return UP.compile(FnIdx); });
 461       }
 462
 463       UP.setPartitionEntries(std::move(PartitionEntries));
 464     }
 465
 466     // Now clone the global variable declarations.
 467     GlobalDeclMaterializer GDMat(*GVsAndStubsM);
 468     for (auto &GV : SrcM->globals())
 469       if (!GV.isDeclaration())
 470         cloneGlobalVariableDecl(*GVsAndStubsM, GV, &VMap);
 471
 472     // Then clone the initializers.
 473     for (auto &GV : SrcM->globals())
 474       if (!GV.isDeclaration())
 475         moveGlobalVariableInitializer(GV, VMap, &GDMat);
 476
 477     // Build a resolver for the stubs module and add it to the base layer.
 478     auto GVsAndStubsResolver = createLambdaResolver(
 479         [&LD](const std::string &Name) {
 480           if (auto Symbol = LD.findSymbol(Name, false))
 481             return RuntimeDyld::SymbolInfo(Symbol.getAddress(),
 482                                            Symbol.getFlags());
 483           return LD.findSymbolExternally(Name);
 484         },
 485         [&LD](const std::string &Name) {
 486           return RuntimeDyld::SymbolInfo(nullptr);
 487         });
 488
 489     std::vector<std::unique_ptr<Module>> GVsAndStubsMSet;
 490     GVsAndStubsMSet.push_back(std::move(GVsAndStubsM));
 491     auto GVsAndStubsH =
 492       BaseLayer.addModuleSet(std::move(GVsAndStubsMSet),
 493                              llvm::make_unique<SectionMemoryManager>(),
 494                              std::move(GVsAndStubsResolver));
 495     LD.setGVsAndStubsHandle(LMH, GVsAndStubsH);
 496   }
 497
 498   static std::string Mangle(StringRef Name, const DataLayout &DL) {
 499     Mangler M(&DL);
 500     std::string MangledName;
 501     {
 502       raw_string_ostream MangledNameStream(MangledName);
 503       M.getNameWithPrefix(MangledNameStream, Name);
 504     }
 505     return MangledName;
 506   }
 507
 508   BaseLayerT &BaseLayer;
 509   CompileCallbackMgrT &CompileCallbackMgr;
 510   LogicalDylibList LogicalDylibs;
 511 };
 512
 513 template <typename BaseLayerT, typename CompileCallbackMgrT>
 514 typename CompileOnDemandLayer<BaseLayerT, CompileCallbackMgrT>::
 515            UncompiledPartition&
 516 CompileOnDemandLayer<BaseLayerT, CompileCallbackMgrT>::LogicalDylib::
 517   createUncompiledPartition(LMHandle LMH, std::shared_ptr<Module> SrcM) {
 518   UncompiledPartitions.push_back(
 519       llvm::make_unique<UncompiledPartition>(*this, LMH, std::move(SrcM)));
 520   UncompiledPartitions.back()->setID(UncompiledPartitions.size() - 1);
 521   return *UncompiledPartitions.back();
 522 }
 523
 524 template <typename BaseLayerT, typename CompileCallbackMgrT>
 525 std::unique_ptr<typename CompileOnDemandLayer<BaseLayerT, CompileCallbackMgrT>::
 526                   UncompiledPartition>
 527 CompileOnDemandLayer<BaseLayerT, CompileCallbackMgrT>::LogicalDylib::
 528   takeUPOwnership(UncompiledPartitionID ID) {
 529
 530   std::swap(UncompiledPartitions[ID], UncompiledPartitions.back());
 531   UncompiledPartitions[ID]->setID(ID);
 532   auto UP = std::move(UncompiledPartitions.back());
 533   UncompiledPartitions.pop_back();
 534   return UP;
 535 }
 536
 537 } // End namespace orc.
 538 } // End namespace llvm.
 539
 540 #endif // LLVM_EXECUTIONENGINE_ORC_COMPILEONDEMANDLAYER_H