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