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