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 "LogicalDylib.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           typename PartitioningFtor =
  41             std::function<std::set<Function*>(Function&)>>
  42 class CompileOnDemandLayer {
  43 private:
  44
  45   // Utility class for MapValue. Only materializes declarations for global
  46   // variables.
  47   class GlobalDeclMaterializer : public ValueMaterializer {
  48   public:
  49     typedef std::set<const Function*> StubSet;
  50
  51     GlobalDeclMaterializer(Module &Dst, const StubSet *StubsToClone = nullptr)
  52         : Dst(Dst), StubsToClone(StubsToClone) {}
  53
  54     Value* materializeValueFor(Value *V) final {
  55       if (auto *GV = dyn_cast<GlobalVariable>(V))
  56         return cloneGlobalVariableDecl(Dst, *GV);
  57       else if (auto *F = dyn_cast<Function>(V)) {
  58         auto *ClonedF = cloneFunctionDecl(Dst, *F);
  59         if (StubsToClone && StubsToClone->count(F)) {
  60           GlobalVariable *FnBodyPtr =
  61             createImplPointer(*ClonedF->getType(), *ClonedF->getParent(),
  62                               ClonedF->getName() + "$orc_addr", nullptr);
  63           makeStub(*ClonedF, *FnBodyPtr);
  64           ClonedF->setLinkage(GlobalValue::AvailableExternallyLinkage);
  65           ClonedF->addFnAttr(Attribute::AlwaysInline);
  66         }
  67         return ClonedF;
  68       }
  69       // Else.
  70       return nullptr;
  71     }
  72   private:
  73     Module &Dst;
  74     const StubSet *StubsToClone;
  75   };
  76
  77   typedef typename BaseLayerT::ModuleSetHandleT BaseLayerModuleSetHandleT;
  78
  79   struct LogicalModuleResources {
  80     std::shared_ptr<Module> SourceModule;
  81     std::set<const Function*> StubsToClone;
  82   };
  83
  84   struct LogicalDylibResources {
  85     typedef std::function<RuntimeDyld::SymbolInfo(const std::string&)>
  86       SymbolResolverFtor;
  87     SymbolResolverFtor ExternalSymbolResolver;
  88   };
  89
  90   typedef LogicalDylib<BaseLayerT, LogicalModuleResources,
  91                        LogicalDylibResources> CODLogicalDylib;
  92
  93   typedef typename CODLogicalDylib::LogicalModuleHandle LogicalModuleHandle;
  94   typedef std::list<CODLogicalDylib> LogicalDylibList;
  95
  96 public:
  97   /// @brief Handle to a set of loaded modules.
  98   typedef typename LogicalDylibList::iterator ModuleSetHandleT;
  99
 100   /// @brief Construct a compile-on-demand layer instance.
 101   CompileOnDemandLayer(BaseLayerT &BaseLayer, CompileCallbackMgrT &CallbackMgr,
 102                        PartitioningFtor Partition,
 103                        bool CloneStubsIntoPartitions = true)
 104       : BaseLayer(BaseLayer), CompileCallbackMgr(CallbackMgr),
 105         Partition(Partition),
 106         CloneStubsIntoPartitions(CloneStubsIntoPartitions) {}
 107
 108   /// @brief Add a module to the compile-on-demand layer.
 109   template <typename ModuleSetT, typename MemoryManagerPtrT,
 110             typename SymbolResolverPtrT>
 111   ModuleSetHandleT addModuleSet(ModuleSetT Ms,
 112                                 MemoryManagerPtrT MemMgr,
 113                                 SymbolResolverPtrT Resolver) {
 114
 115     assert(MemMgr == nullptr &&
 116            "User supplied memory managers not supported with COD yet.");
 117
 118     LogicalDylibs.push_back(CODLogicalDylib(BaseLayer));
 119     auto &LDResources = LogicalDylibs.back().getDylibResources();
 120
 121     LDResources.ExternalSymbolResolver =
 122       [Resolver](const std::string &Name) {
 123         return Resolver->findSymbol(Name);
 124       };
 125
 126     // Process each of the modules in this module set.
 127     for (auto &M : Ms)
 128       addLogicalModule(LogicalDylibs.back(),
 129                        std::shared_ptr<Module>(std::move(M)));
 130
 131     return std::prev(LogicalDylibs.end());
 132   }
 133
 134   /// @brief Remove the module represented by the given handle.
 135   ///
 136   ///   This will remove all modules in the layers below that were derived from
 137   /// the module represented by H.
 138   void removeModuleSet(ModuleSetHandleT H) {
 139     LogicalDylibs.erase(H);
 140   }
 141
 142   /// @brief Search for the given named symbol.
 143   /// @param Name The name of the symbol to search for.
 144   /// @param ExportedSymbolsOnly If true, search only for exported symbols.
 145   /// @return A handle for the given named symbol, if it exists.
 146   JITSymbol findSymbol(StringRef Name, bool ExportedSymbolsOnly) {
 147     return BaseLayer.findSymbol(Name, ExportedSymbolsOnly);
 148   }
 149
 150   /// @brief Get the address of a symbol provided by this layer, or some layer
 151   ///        below this one.
 152   JITSymbol findSymbolIn(ModuleSetHandleT H, const std::string &Name,
 153                          bool ExportedSymbolsOnly) {
 154     return H->findSymbol(Name, ExportedSymbolsOnly);
 155   }
 156
 157 private:
 158
 159   void addLogicalModule(CODLogicalDylib &LD, std::shared_ptr<Module> SrcM) {
 160
 161     // Bump the linkage and rename any anonymous/privote members in SrcM to
 162     // ensure that everything will resolve properly after we partition SrcM.
 163     makeAllSymbolsExternallyAccessible(*SrcM);
 164
 165     // Create a logical module handle for SrcM within the logical dylib.
 166     auto LMH = LD.createLogicalModule();
 167     auto &LMResources =  LD.getLogicalModuleResources(LMH);
 168     LMResources.SourceModule = SrcM;
 169
 170     // Create the GVs-and-stubs module.
 171     auto GVsAndStubsM = llvm::make_unique<Module>(
 172                           (SrcM->getName() + ".globals_and_stubs").str(),
 173                           SrcM->getContext());
 174     GVsAndStubsM->setDataLayout(SrcM->getDataLayout());
 175     ValueToValueMapTy VMap;
 176
 177     // Process module and create stubs.
 178     // We create the stubs before copying the global variables as we know the
 179     // stubs won't refer to any globals (they only refer to their implementation
 180     // pointer) so there's no ordering/value-mapping issues.
 181     for (auto &F : *SrcM) {
 182
 183       // Skip declarations.
 184       if (F.isDeclaration())
 185         continue;
 186
 187       // Record all functions defined by this module.
 188       if (CloneStubsIntoPartitions)
 189         LMResources.StubsToClone.insert(&F);
 190
 191       // For each definition: create a callback, a stub, and a function body
 192       // pointer. Initialize the function body pointer to point at the callback,
 193       // and set the callback to compile the function body.
 194       auto CCInfo = CompileCallbackMgr.getCompileCallback(SrcM->getContext());
 195       Function *StubF = cloneFunctionDecl(*GVsAndStubsM, F, &VMap);
 196       GlobalVariable *FnBodyPtr =
 197         createImplPointer(*StubF->getType(), *StubF->getParent(),
 198                           StubF->getName() + "$orc_addr",
 199                           createIRTypedAddress(*StubF->getFunctionType(),
 200                                                CCInfo.getAddress()));
 201       makeStub(*StubF, *FnBodyPtr);
 202       CCInfo.setCompileAction(
 203         [this, &LD, LMH, &F]() {
 204           return this->extractAndCompile(LD, LMH, F);
 205         });
 206     }
 207
 208     // Now clone the global variable declarations.
 209     GlobalDeclMaterializer GDMat(*GVsAndStubsM);
 210     for (auto &GV : SrcM->globals())
 211       if (!GV.isDeclaration())
 212         cloneGlobalVariableDecl(*GVsAndStubsM, GV, &VMap);
 213
 214     // And the aliases.
 215     for (auto &Alias : SrcM->aliases())
 216       cloneGlobalAlias(*GVsAndStubsM, Alias, VMap, &GDMat);
 217
 218     // Then clone the initializers.
 219     for (auto &GV : SrcM->globals())
 220       if (!GV.isDeclaration())
 221         moveGlobalVariableInitializer(GV, VMap, &GDMat);
 222
 223     // Build a resolver for the stubs module and add it to the base layer.
 224     auto GVsAndStubsResolver = createLambdaResolver(
 225         [&LD](const std::string &Name) {
 226           return LD.getDylibResources().ExternalSymbolResolver(Name);
 227         },
 228         [](const std::string &Name) {
 229           return RuntimeDyld::SymbolInfo(nullptr);
 230         });
 231
 232     std::vector<std::unique_ptr<Module>> GVsAndStubsMSet;
 233     GVsAndStubsMSet.push_back(std::move(GVsAndStubsM));
 234     auto GVsAndStubsH =
 235       BaseLayer.addModuleSet(std::move(GVsAndStubsMSet),
 236                              llvm::make_unique<SectionMemoryManager>(),
 237                              std::move(GVsAndStubsResolver));
 238     LD.addToLogicalModule(LMH, GVsAndStubsH);
 239   }
 240
 241   static std::string Mangle(StringRef Name, const DataLayout &DL) {
 242     std::string MangledName;
 243     {
 244       raw_string_ostream MangledNameStream(MangledName);
 245       Mangler::getNameWithPrefix(MangledNameStream, Name, DL);
 246     }
 247     return MangledName;
 248   }
 249
 250   TargetAddress extractAndCompile(CODLogicalDylib &LD,
 251                                   LogicalModuleHandle LMH,
 252                                   Function &F) {
 253     Module &SrcM = *LD.getLogicalModuleResources(LMH).SourceModule;
 254
 255     // If F is a declaration we must already have compiled it.
 256     if (F.isDeclaration())
 257       return 0;
 258
 259     // Grab the name of the function being called here.
 260     std::string CalledFnName = Mangle(F.getName(), SrcM.getDataLayout());
 261
 262     auto Part = Partition(F);
 263     auto PartH = emitPartition(LD, LMH, Part);
 264
 265     TargetAddress CalledAddr = 0;
 266     for (auto *SubF : Part) {
 267       std::string FName = SubF->getName();
 268       auto FnBodySym =
 269         BaseLayer.findSymbolIn(PartH, Mangle(FName, SrcM.getDataLayout()),
 270                                false);
 271       auto FnPtrSym =
 272         BaseLayer.findSymbolIn(*LD.moduleHandlesBegin(LMH),
 273                                Mangle(FName + "$orc_addr",
 274                                       SrcM.getDataLayout()),
 275                                false);
 276       assert(FnBodySym && "Couldn't find function body.");
 277       assert(FnPtrSym && "Couldn't find function body pointer.");
 278
 279       TargetAddress FnBodyAddr = FnBodySym.getAddress();
 280       void *FnPtrAddr = reinterpret_cast<void*>(
 281           static_cast<uintptr_t>(FnPtrSym.getAddress()));
 282
 283       // If this is the function we're calling record the address so we can
 284       // return it from this function.
 285       if (SubF == &F)
 286         CalledAddr = FnBodyAddr;
 287
 288       memcpy(FnPtrAddr, &FnBodyAddr, sizeof(uintptr_t));
 289     }
 290
 291     return CalledAddr;
 292   }
 293
 294   template <typename PartitionT>
 295   BaseLayerModuleSetHandleT emitPartition(CODLogicalDylib &LD,
 296                                           LogicalModuleHandle LMH,
 297                                           const PartitionT &Part) {
 298     auto &LMResources = LD.getLogicalModuleResources(LMH);
 299     Module &SrcM = *LMResources.SourceModule;
 300
 301     // Create the module.
 302     std::string NewName = SrcM.getName();
 303     for (auto *F : Part) {
 304       NewName += ".";
 305       NewName += F->getName();
 306     }
 307
 308     auto M = llvm::make_unique<Module>(NewName, SrcM.getContext());
 309     M->setDataLayout(SrcM.getDataLayout());
 310     ValueToValueMapTy VMap;
 311     GlobalDeclMaterializer GDM(*M, &LMResources.StubsToClone);
 312
 313     // Create decls in the new module.
 314     for (auto *F : Part)
 315       cloneFunctionDecl(*M, *F, &VMap);
 316
 317     // Move the function bodies.
 318     for (auto *F : Part)
 319       moveFunctionBody(*F, VMap, &GDM);
 320
 321     // Create memory manager and symbol resolver.
 322     auto MemMgr = llvm::make_unique<SectionMemoryManager>();
 323     auto Resolver = createLambdaResolver(
 324         [this, &LD, LMH](const std::string &Name) {
 325           if (auto Symbol = LD.findSymbolInternally(LMH, Name))
 326             return RuntimeDyld::SymbolInfo(Symbol.getAddress(),
 327                                            Symbol.getFlags());
 328           return LD.getDylibResources().ExternalSymbolResolver(Name);
 329         },
 330         [this, &LD, LMH](const std::string &Name) {
 331           if (auto Symbol = LD.findSymbolInternally(LMH, Name))
 332             return RuntimeDyld::SymbolInfo(Symbol.getAddress(),
 333                                            Symbol.getFlags());
 334           return RuntimeDyld::SymbolInfo(nullptr);
 335         });
 336     std::vector<std::unique_ptr<Module>> PartMSet;
 337     PartMSet.push_back(std::move(M));
 338     return BaseLayer.addModuleSet(std::move(PartMSet), std::move(MemMgr),
 339                                   std::move(Resolver));
 340   }
 341
 342   BaseLayerT &BaseLayer;
 343   CompileCallbackMgrT &CompileCallbackMgr;
 344   LogicalDylibList LogicalDylibs;
 345   PartitioningFtor Partition;
 346   bool CloneStubsIntoPartitions;
 347 };
 348
 349 } // End namespace orc.
 350 } // End namespace llvm.
 351
 352 #endif // LLVM_EXECUTIONENGINE_ORC_COMPILEONDEMANDLAYER_H