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     // And the aliases.
 221     for (auto &Alias : SrcM->aliases())
 222       cloneGlobalAlias(*GVsAndStubsM, Alias, VMap, &GDMat);
 223
 224     // Then clone the initializers.
 225     for (auto &GV : SrcM->globals())
 226       if (!GV.isDeclaration())
 227         moveGlobalVariableInitializer(GV, VMap, &GDMat);
 228
 229     // Build a resolver for the stubs module and add it to the base layer.
 230     auto GVsAndStubsResolver = createLambdaResolver(
 231         [&LD](const std::string &Name) {
 232           return LD.getDylibResources().ExternalSymbolResolver(Name);
 233         },
 234         [](const std::string &Name) {
 235           return RuntimeDyld::SymbolInfo(nullptr);
 236         });
 237
 238     std::vector<std::unique_ptr<Module>> GVsAndStubsMSet;
 239     GVsAndStubsMSet.push_back(std::move(GVsAndStubsM));
 240     auto GVsAndStubsH =
 241       BaseLayer.addModuleSet(std::move(GVsAndStubsMSet),
 242                              llvm::make_unique<SectionMemoryManager>(),
 243                              std::move(GVsAndStubsResolver));
 244     LD.addToLogicalModule(LMH, GVsAndStubsH);
 245   }
 246
 247   static std::string Mangle(StringRef Name, const DataLayout &DL) {
 248     std::string MangledName;
 249     {
 250       raw_string_ostream MangledNameStream(MangledName);
 251       Mangler::getNameWithPrefix(MangledNameStream, Name, DL);
 252     }
 253     return MangledName;
 254   }
 255
 256   TargetAddress extractAndCompile(CODLogicalDylib &LD,
 257                                   LogicalModuleHandle LMH,
 258                                   Function &F) {
 259     Module &SrcM = *LD.getLogicalModuleResources(LMH).SourceModule;
 260
 261     // If F is a declaration we must already have compiled it.
 262     if (F.isDeclaration())
 263       return 0;
 264
 265     // Grab the name of the function being called here.
 266     std::string CalledFnName = Mangle(F.getName(), SrcM.getDataLayout());
 267
 268     auto Partition = LD.getDylibResources().Partitioner(F);
 269     auto PartitionH = emitPartition(LD, LMH, Partition);
 270
 271     TargetAddress CalledAddr = 0;
 272     for (auto *SubF : Partition) {
 273       std::string FName = SubF->getName();
 274       auto FnBodySym =
 275         BaseLayer.findSymbolIn(PartitionH, Mangle(FName, SrcM.getDataLayout()),
 276                                false);
 277       auto FnPtrSym =
 278         BaseLayer.findSymbolIn(*LD.moduleHandlesBegin(LMH),
 279                                Mangle(FName + "$orc_addr",
 280                                       SrcM.getDataLayout()),
 281                                false);
 282       assert(FnBodySym && "Couldn't find function body.");
 283       assert(FnPtrSym && "Couldn't find function body pointer.");
 284
 285       TargetAddress FnBodyAddr = FnBodySym.getAddress();
 286       void *FnPtrAddr = reinterpret_cast<void*>(
 287           static_cast<uintptr_t>(FnPtrSym.getAddress()));
 288
 289       // If this is the function we're calling record the address so we can
 290       // return it from this function.
 291       if (SubF == &F)
 292         CalledAddr = FnBodyAddr;
 293
 294       memcpy(FnPtrAddr, &FnBodyAddr, sizeof(uintptr_t));
 295     }
 296
 297     return CalledAddr;
 298   }
 299
 300   template <typename PartitionT>
 301   BaseLayerModuleSetHandleT emitPartition(CODLogicalDylib &LD,
 302                                           LogicalModuleHandle LMH,
 303                                           const PartitionT &Partition) {
 304     auto &LMResources = LD.getLogicalModuleResources(LMH);
 305     Module &SrcM = *LMResources.SourceModule;
 306
 307     // Create the module.
 308     std::string NewName = SrcM.getName();
 309     for (auto *F : Partition) {
 310       NewName += ".";
 311       NewName += F->getName();
 312     }
 313
 314     auto M = llvm::make_unique<Module>(NewName, SrcM.getContext());
 315     M->setDataLayout(SrcM.getDataLayout());
 316     ValueToValueMapTy VMap;
 317     GlobalDeclMaterializer GDM(*M, &LMResources.StubsToClone);
 318
 319     // Create decls in the new module.
 320     for (auto *F : Partition)
 321       cloneFunctionDecl(*M, *F, &VMap);
 322
 323     // Move the function bodies.
 324     for (auto *F : Partition)
 325       moveFunctionBody(*F, VMap, &GDM);
 326
 327     // Create memory manager and symbol resolver.
 328     auto MemMgr = llvm::make_unique<SectionMemoryManager>();
 329     auto Resolver = createLambdaResolver(
 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 LD.getDylibResources().ExternalSymbolResolver(Name);
 335         },
 336         [this, &LD, LMH](const std::string &Name) {
 337           if (auto Symbol = LD.findSymbolInternally(LMH, Name))
 338             return RuntimeDyld::SymbolInfo(Symbol.getAddress(),
 339                                            Symbol.getFlags());
 340           return RuntimeDyld::SymbolInfo(nullptr);
 341         });
 342     std::vector<std::unique_ptr<Module>> PartMSet;
 343     PartMSet.push_back(std::move(M));
 344     return BaseLayer.addModuleSet(std::move(PartMSet), std::move(MemMgr),
 345                                   std::move(Resolver));
 346   }
 347
 348   BaseLayerT &BaseLayer;
 349   CompileCallbackMgrT &CompileCallbackMgr;
 350   LogicalDylibList LogicalDylibs;
 351   bool CloneStubsIntoPartitions;
 352 };
 353
 354 } // End namespace orc.
 355 } // End namespace llvm.
 356
 357 #endif // LLVM_EXECUTIONENGINE_ORC_COMPILEONDEMANDLAYER_H