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 "llvm/ADT/STLExtras.h"
  21 #include "llvm/ExecutionEngine/SectionMemoryManager.h"
  22 #include <list>
  23
  24 namespace llvm {
  25 namespace orc {
  26
  27 /// @brief Compile-on-demand layer.
  28 ///
  29 ///   Modules added to this layer have their calls indirected, and are then
  30 /// broken up into a set of single-function modules, each of which is added
  31 /// to the layer below in a singleton set. The lower layer can be any layer that
  32 /// accepts IR module sets.
  33 ///
  34 /// It is expected that this layer will frequently be used on top of a
  35 /// LazyEmittingLayer. The combination of the two ensures that each function is
  36 /// compiled only when it is first called.
  37 template <typename BaseLayerT, typename CompileCallbackMgrT>
  38 class CompileOnDemandLayer {
  39 private:
  40   /// @brief Lookup helper that provides compatibility with the classic
  41   ///        static-compilation symbol resolution process.
  42   ///
  43   ///   The CompileOnDemand (COD) layer splits modules up into multiple
  44   /// sub-modules, each held in its own llvm::Module instance, in order to
  45   /// support lazy compilation. When a module that contains private symbols is
  46   /// broken up symbol linkage changes may be required to enable access to
  47   /// "private" data that now resides in a different llvm::Module instance. To
  48   /// retain expected symbol resolution behavior for clients of the COD layer,
  49   /// the CODScopedLookup class uses a two-tiered lookup system to resolve
  50   /// symbols. Lookup first scans sibling modules that were split from the same
  51   /// original module (logical-module scoped lookup), then scans all other
  52   /// modules that have been added to the lookup scope (logical-dylib scoped
  53   /// lookup).
  54   class CODScopedLookup {
  55   private:
  56     typedef typename BaseLayerT::ModuleSetHandleT BaseLayerModuleSetHandleT;
  57     typedef std::vector<BaseLayerModuleSetHandleT> SiblingHandlesList;
  58     typedef std::list<SiblingHandlesList> PseudoDylibModuleSetHandlesList;
  59
  60   public:
  61     /// @brief Handle for a logical module.
  62     typedef typename PseudoDylibModuleSetHandlesList::iterator LMHandle;
  63
  64     /// @brief Construct a scoped lookup.
  65     CODScopedLookup(BaseLayerT &BaseLayer) : BaseLayer(BaseLayer) {}
  66
  67     virtual ~CODScopedLookup() {}
  68
  69     /// @brief Start a new context for a single logical module.
  70     LMHandle createLogicalModule() {
  71       Handles.push_back(SiblingHandlesList());
  72       return std::prev(Handles.end());
  73     }
  74
  75     /// @brief Add a concrete Module's handle to the given logical Module's
  76     ///        lookup scope.
  77     void addToLogicalModule(LMHandle LMH, BaseLayerModuleSetHandleT H) {
  78       LMH->push_back(H);
  79     }
  80
  81     /// @brief Remove a logical Module from the CODScopedLookup entirely.
  82     void removeLogicalModule(LMHandle LMH) { Handles.erase(LMH); }
  83
  84     /// @brief Look up a symbol in this context.
  85     JITSymbol findSymbol(LMHandle LMH, const std::string &Name) {
  86       if (auto Symbol = findSymbolIn(LMH, Name))
  87         return Symbol;
  88
  89       for (auto I = Handles.begin(), E = Handles.end(); I != E; ++I)
  90         if (I != LMH)
  91           if (auto Symbol = findSymbolIn(I, Name))
  92             return Symbol;
  93
  94       return nullptr;
  95     }
  96
  97     /// @brief Find an external symbol (via the user supplied SymbolResolver).
  98     virtual RuntimeDyld::SymbolInfo
  99     externalLookup(const std::string &Name) const = 0;
 100
 101   private:
 102
 103     JITSymbol findSymbolIn(LMHandle LMH, const std::string &Name) {
 104       for (auto H : *LMH)
 105         if (auto Symbol = BaseLayer.findSymbolIn(H, Name, false))
 106           return Symbol;
 107       return nullptr;
 108     }
 109
 110     BaseLayerT &BaseLayer;
 111     PseudoDylibModuleSetHandlesList Handles;
 112   };
 113
 114   template <typename ResolverPtrT>
 115   class CODScopedLookupImpl : public CODScopedLookup {
 116   public:
 117     CODScopedLookupImpl(BaseLayerT &BaseLayer, ResolverPtrT Resolver)
 118       : CODScopedLookup(BaseLayer), Resolver(std::move(Resolver)) {}
 119
 120     RuntimeDyld::SymbolInfo
 121     externalLookup(const std::string &Name) const override {
 122       return Resolver->findSymbol(Name);
 123     }
 124
 125   private:
 126     ResolverPtrT Resolver;
 127   };
 128
 129   template <typename ResolverPtrT>
 130   static std::shared_ptr<CODScopedLookup>
 131   createCODScopedLookup(BaseLayerT &BaseLayer,
 132                         ResolverPtrT Resolver) {
 133     typedef CODScopedLookupImpl<ResolverPtrT> Impl;
 134     return std::make_shared<Impl>(BaseLayer, std::move(Resolver));
 135   }
 136
 137   typedef typename BaseLayerT::ModuleSetHandleT BaseLayerModuleSetHandleT;
 138   typedef std::vector<BaseLayerModuleSetHandleT> BaseLayerModuleSetHandleListT;
 139
 140   struct ModuleSetInfo {
 141     // Symbol lookup - just one for the whole module set.
 142     std::shared_ptr<CODScopedLookup> Lookup;
 143
 144     // Logical module handles.
 145     std::vector<typename CODScopedLookup::LMHandle> LMHandles;
 146
 147     // List of vectors of module set handles:
 148     // One vector per logical module - each vector holds the handles for the
 149     // exploded modules for that logical module in the base layer.
 150     BaseLayerModuleSetHandleListT BaseLayerModuleSetHandles;
 151
 152     ModuleSetInfo(std::shared_ptr<CODScopedLookup> Lookup)
 153         : Lookup(std::move(Lookup)) {}
 154
 155     void releaseResources(BaseLayerT &BaseLayer) {
 156       for (auto LMH : LMHandles)
 157         Lookup->removeLogicalModule(LMH);
 158       for (auto H : BaseLayerModuleSetHandles)
 159         BaseLayer.removeModuleSet(H);
 160     }
 161   };
 162
 163   typedef std::list<ModuleSetInfo> ModuleSetInfoListT;
 164
 165 public:
 166   /// @brief Handle to a set of loaded modules.
 167   typedef typename ModuleSetInfoListT::iterator ModuleSetHandleT;
 168
 169   /// @brief Construct a compile-on-demand layer instance.
 170   CompileOnDemandLayer(BaseLayerT &BaseLayer, CompileCallbackMgrT &CallbackMgr)
 171       : BaseLayer(BaseLayer), CompileCallbackMgr(CallbackMgr) {}
 172
 173   /// @brief Add a module to the compile-on-demand layer.
 174   template <typename ModuleSetT, typename MemoryManagerPtrT,
 175             typename SymbolResolverPtrT>
 176   ModuleSetHandleT addModuleSet(ModuleSetT Ms,
 177                                 MemoryManagerPtrT MemMgr,
 178                                 SymbolResolverPtrT Resolver) {
 179
 180     assert(MemMgr == nullptr &&
 181            "User supplied memory managers not supported with COD yet.");
 182
 183     // Create a lookup context and ModuleSetInfo for this module set.
 184     // For the purposes of symbol resolution the set Ms will be treated as if
 185     // the modules it contained had been linked together as a dylib.
 186     auto DylibLookup = createCODScopedLookup(BaseLayer, std::move(Resolver));
 187     ModuleSetHandleT H =
 188         ModuleSetInfos.insert(ModuleSetInfos.end(), ModuleSetInfo(DylibLookup));
 189     ModuleSetInfo &MSI = ModuleSetInfos.back();
 190
 191     // Process each of the modules in this module set.
 192     for (auto &M : Ms)
 193       partitionAndAdd(*M, MSI);
 194
 195     return H;
 196   }
 197
 198   /// @brief Remove the module represented by the given handle.
 199   ///
 200   ///   This will remove all modules in the layers below that were derived from
 201   /// the module represented by H.
 202   void removeModuleSet(ModuleSetHandleT H) {
 203     H->releaseResources(BaseLayer);
 204     ModuleSetInfos.erase(H);
 205   }
 206
 207   /// @brief Search for the given named symbol.
 208   /// @param Name The name of the symbol to search for.
 209   /// @param ExportedSymbolsOnly If true, search only for exported symbols.
 210   /// @return A handle for the given named symbol, if it exists.
 211   JITSymbol findSymbol(StringRef Name, bool ExportedSymbolsOnly) {
 212     return BaseLayer.findSymbol(Name, ExportedSymbolsOnly);
 213   }
 214
 215   /// @brief Get the address of a symbol provided by this layer, or some layer
 216   ///        below this one.
 217   JITSymbol findSymbolIn(ModuleSetHandleT H, const std::string &Name,
 218                          bool ExportedSymbolsOnly) {
 219
 220     for (auto &BH : H->BaseLayerModuleSetHandles) {
 221       if (auto Symbol = BaseLayer.findSymbolIn(BH, Name, ExportedSymbolsOnly))
 222         return Symbol;
 223     }
 224     return nullptr;
 225   }
 226
 227 private:
 228
 229   void partitionAndAdd(Module &M, ModuleSetInfo &MSI) {
 230     const char *AddrSuffix = "$orc_addr";
 231     const char *BodySuffix = "$orc_body";
 232
 233     // We're going to break M up into a bunch of sub-modules, but we want
 234     // internal linkage symbols to still resolve sensibly. CODScopedLookup
 235     // provides the "logical module" concept to make this work, so create a
 236     // new logical module for M.
 237     auto DylibLookup = MSI.Lookup;
 238     auto LogicalModule = DylibLookup->createLogicalModule();
 239     MSI.LMHandles.push_back(LogicalModule);
 240
 241     // Partition M into a "globals and stubs" module, a "common symbols" module,
 242     // and a list of single-function modules.
 243     auto PartitionedModule = fullyPartition(M);
 244     auto StubsModule = std::move(PartitionedModule.GlobalVars);
 245     auto CommonsModule = std::move(PartitionedModule.Commons);
 246     auto FunctionModules = std::move(PartitionedModule.Functions);
 247
 248     // Emit the commons stright away.
 249     auto CommonHandle = addModule(std::move(CommonsModule), MSI, LogicalModule);
 250     BaseLayer.emitAndFinalize(CommonHandle);
 251
 252     // Map of definition names to callback-info data structures. We'll use
 253     // this to build the compile actions for the stubs below.
 254     typedef std::map<std::string,
 255                      typename CompileCallbackMgrT::CompileCallbackInfo>
 256       StubInfoMap;
 257     StubInfoMap StubInfos;
 258
 259     // Now we need to take each of the extracted Modules and add them to
 260     // base layer. Each Module will be added individually to make sure they
 261     // can be compiled separately, and each will get its own lookaside
 262     // memory manager that will resolve within this logical module first.
 263     for (auto &SubM : FunctionModules) {
 264
 265       // Keep track of the stubs we create for this module so that we can set
 266       // their compile actions.
 267       std::vector<typename StubInfoMap::iterator> NewStubInfos;
 268
 269       // Search for function definitions and insert stubs into the stubs
 270       // module.
 271       for (auto &F : *SubM) {
 272         if (F.isDeclaration())
 273           continue;
 274
 275         std::string Name = F.getName();
 276         Function *Proto = StubsModule->getFunction(Name);
 277         assert(Proto && "Failed to clone function decl into stubs module.");
 278         auto CallbackInfo =
 279           CompileCallbackMgr.getCompileCallback(Proto->getContext());
 280         GlobalVariable *FunctionBodyPointer =
 281           createImplPointer(*Proto->getType(), *Proto->getParent(),
 282                             Name + AddrSuffix,
 283                             createIRTypedAddress(*Proto->getFunctionType(),
 284                                                  CallbackInfo.getAddress()));
 285         makeStub(*Proto, *FunctionBodyPointer);
 286
 287         F.setName(Name + BodySuffix);
 288         F.setVisibility(GlobalValue::HiddenVisibility);
 289
 290         auto KV = std::make_pair(std::move(Name), std::move(CallbackInfo));
 291         NewStubInfos.push_back(StubInfos.insert(StubInfos.begin(), KV));
 292       }
 293
 294       auto H = addModule(std::move(SubM), MSI, LogicalModule);
 295
 296       // Set the compile actions for this module:
 297       for (auto &KVPair : NewStubInfos) {
 298         std::string BodyName = Mangle(KVPair->first + BodySuffix,
 299                                       M.getDataLayout());
 300         auto &CCInfo = KVPair->second;
 301         CCInfo.setCompileAction(
 302           [=](){
 303             return BaseLayer.findSymbolIn(H, BodyName, false).getAddress();
 304           });
 305       }
 306
 307     }
 308
 309     // Ok - we've processed all the partitioned modules. Now add the
 310     // stubs/globals module and set the update actions.
 311     auto StubsH =
 312       addModule(std::move(StubsModule), MSI, LogicalModule);
 313
 314     for (auto &KVPair : StubInfos) {
 315       std::string AddrName = Mangle(KVPair.first + AddrSuffix,
 316                                     M.getDataLayout());
 317       auto &CCInfo = KVPair.second;
 318       CCInfo.setUpdateAction(
 319         getLocalFPUpdater(BaseLayer, StubsH, AddrName));
 320     }
 321   }
 322
 323   // Add the given Module to the base layer using a memory manager that will
 324   // perform the appropriate scoped lookup (i.e. will look first with in the
 325   // module from which it was extracted, then into the set to which that module
 326   // belonged, and finally externally).
 327   BaseLayerModuleSetHandleT addModule(
 328                                std::unique_ptr<Module> M,
 329                                ModuleSetInfo &MSI,
 330                                typename CODScopedLookup::LMHandle LogicalModule) {
 331
 332     // Add this module to the JIT with a memory manager that uses the
 333     // DylibLookup to resolve symbols.
 334     std::vector<std::unique_ptr<Module>> MSet;
 335     MSet.push_back(std::move(M));
 336
 337     auto DylibLookup = MSI.Lookup;
 338     auto Resolver =
 339       createLambdaResolver(
 340         [=](const std::string &Name) {
 341           if (auto Symbol = DylibLookup->findSymbol(LogicalModule, Name))
 342             return RuntimeDyld::SymbolInfo(Symbol.getAddress(),
 343                                            Symbol.getFlags());
 344           return DylibLookup->externalLookup(Name);
 345         },
 346         [=](const std::string &Name) -> RuntimeDyld::SymbolInfo {
 347           if (auto Symbol = DylibLookup->findSymbol(LogicalModule, Name))
 348             return RuntimeDyld::SymbolInfo(Symbol.getAddress(),
 349                                            Symbol.getFlags());
 350           return nullptr;
 351         });
 352
 353     BaseLayerModuleSetHandleT H =
 354       BaseLayer.addModuleSet(std::move(MSet),
 355                              make_unique<SectionMemoryManager>(),
 356                              std::move(Resolver));
 357     // Add this module to the logical module lookup.
 358     DylibLookup->addToLogicalModule(LogicalModule, H);
 359     MSI.BaseLayerModuleSetHandles.push_back(H);
 360
 361     return H;
 362   }
 363
 364   static std::string Mangle(StringRef Name, const DataLayout &DL) {
 365     Mangler M(&DL);
 366     std::string MangledName;
 367     {
 368       raw_string_ostream MangledNameStream(MangledName);
 369       M.getNameWithPrefix(MangledNameStream, Name);
 370     }
 371     return MangledName;
 372   }
 373
 374   BaseLayerT &BaseLayer;
 375   CompileCallbackMgrT &CompileCallbackMgr;
 376   ModuleSetInfoListT ModuleSetInfos;
 377 };
 378
 379 } // End namespace orc.
 380 } // End namespace llvm.
 381
 382 #endif // LLVM_EXECUTIONENGINE_ORC_COMPILEONDEMANDLAYER_H