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