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