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