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