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 <memory>
26 #include <set>
27 
28 #include "llvm/Support/Debug.h"
29 
30 namespace llvm {
31 namespace orc {
32 
33 /// @brief Compile-on-demand layer.
34 ///
35 ///   When a module is added to this layer a stub is created for each of its
36 /// function definitions. The stubs and other global values are immediately
37 /// added to the layer below. When a stub is called it triggers the extraction
38 /// of the function body from the original module. The extracted body is then
39 /// compiled and executed.
40 template <typename BaseLayerT,
41           typename CompileCallbackMgrT = JITCompileCallbackManager,
42           typename IndirectStubsMgrT = IndirectStubsManager>
43 class CompileOnDemandLayer {
44 private:
45 
46   template <typename MaterializerFtor>
47   class LambdaMaterializer final : public ValueMaterializer {
48   public:
LambdaMaterializer(MaterializerFtor M)49     LambdaMaterializer(MaterializerFtor M) : M(std::move(M)) {}
materializeDeclFor(Value * V)50     Value *materializeDeclFor(Value *V) final { return M(V); }
51 
52   private:
53     MaterializerFtor M;
54   };
55 
56   template <typename MaterializerFtor>
57   LambdaMaterializer<MaterializerFtor>
createLambdaMaterializer(MaterializerFtor M)58   createLambdaMaterializer(MaterializerFtor M) {
59     return LambdaMaterializer<MaterializerFtor>(std::move(M));
60   }
61 
62   typedef typename BaseLayerT::ModuleSetHandleT BaseLayerModuleSetHandleT;
63 
64   class ModuleOwner {
65   public:
66     ModuleOwner() = default;
67     ModuleOwner(const ModuleOwner&) = delete;
68     ModuleOwner& operator=(const ModuleOwner&) = delete;
~ModuleOwner()69     virtual ~ModuleOwner() { }
70     virtual Module& getModule() const = 0;
71   };
72 
73   template <typename ModulePtrT>
74   class ModuleOwnerImpl : public ModuleOwner {
75   public:
ModuleOwnerImpl(ModulePtrT ModulePtr)76     ModuleOwnerImpl(ModulePtrT ModulePtr) : ModulePtr(std::move(ModulePtr)) {}
getModule()77     Module& getModule() const override { return *ModulePtr; }
78   private:
79     ModulePtrT ModulePtr;
80   };
81 
82   template <typename ModulePtrT>
wrapOwnership(ModulePtrT ModulePtr)83   std::unique_ptr<ModuleOwner> wrapOwnership(ModulePtrT ModulePtr) {
84     return llvm::make_unique<ModuleOwnerImpl<ModulePtrT>>(std::move(ModulePtr));
85   }
86 
87   struct LogicalModuleResources {
88     std::unique_ptr<ModuleOwner> SourceModuleOwner;
89     std::set<const Function*> StubsToClone;
90     std::unique_ptr<IndirectStubsMgrT> StubsMgr;
91 
92     LogicalModuleResources() = default;
93 
94     // Explicit move constructor to make MSVC happy.
LogicalModuleResourcesLogicalModuleResources95     LogicalModuleResources(LogicalModuleResources &&Other)
96         : SourceModuleOwner(std::move(Other.SourceModuleOwner)),
97           StubsToClone(std::move(Other.StubsToClone)),
98           StubsMgr(std::move(Other.StubsMgr)) {}
99 
100     // Explicit move assignment to make MSVC happy.
101     LogicalModuleResources& operator=(LogicalModuleResources &&Other) {
102       SourceModuleOwner = std::move(Other.SourceModuleOwner);
103       StubsToClone = std::move(Other.StubsToClone);
104       StubsMgr = std::move(Other.StubsMgr);
105     }
106 
findSymbolLogicalModuleResources107     JITSymbol findSymbol(StringRef Name, bool ExportedSymbolsOnly) {
108       if (Name.endswith("$stub_ptr") && !ExportedSymbolsOnly) {
109         assert(!ExportedSymbolsOnly && "Stubs are never exported");
110         return StubsMgr->findPointer(Name.drop_back(9));
111       }
112       return StubsMgr->findStub(Name, ExportedSymbolsOnly);
113     }
114 
115   };
116 
117 
118 
119   struct LogicalDylibResources {
120     typedef std::function<RuntimeDyld::SymbolInfo(const std::string&)>
121       SymbolResolverFtor;
122     SymbolResolverFtor ExternalSymbolResolver;
123   };
124 
125   typedef LogicalDylib<BaseLayerT, LogicalModuleResources,
126                        LogicalDylibResources> CODLogicalDylib;
127 
128   typedef typename CODLogicalDylib::LogicalModuleHandle LogicalModuleHandle;
129   typedef std::list<CODLogicalDylib> LogicalDylibList;
130 
131 public:
132 
133   /// @brief Handle to a set of loaded modules.
134   typedef typename LogicalDylibList::iterator ModuleSetHandleT;
135 
136   /// @brief Module partitioning functor.
137   typedef std::function<std::set<Function*>(Function&)> PartitioningFtor;
138 
139   /// @brief Builder for IndirectStubsManagers.
140   typedef std::function<std::unique_ptr<IndirectStubsMgrT>()>
141     IndirectStubsManagerBuilderT;
142 
143   /// @brief Construct a compile-on-demand layer instance.
144   CompileOnDemandLayer(BaseLayerT &BaseLayer, PartitioningFtor Partition,
145                        CompileCallbackMgrT &CallbackMgr,
146                        IndirectStubsManagerBuilderT CreateIndirectStubsManager,
147                        bool CloneStubsIntoPartitions = true)
BaseLayer(BaseLayer)148       : BaseLayer(BaseLayer),  Partition(Partition),
149         CompileCallbackMgr(CallbackMgr),
150         CreateIndirectStubsManager(std::move(CreateIndirectStubsManager)),
151         CloneStubsIntoPartitions(CloneStubsIntoPartitions) {}
152 
153   /// @brief Add a module to the compile-on-demand layer.
154   template <typename ModuleSetT, typename MemoryManagerPtrT,
155             typename SymbolResolverPtrT>
addModuleSet(ModuleSetT Ms,MemoryManagerPtrT MemMgr,SymbolResolverPtrT Resolver)156   ModuleSetHandleT addModuleSet(ModuleSetT Ms,
157                                 MemoryManagerPtrT MemMgr,
158                                 SymbolResolverPtrT Resolver) {
159 
160     assert(MemMgr == nullptr &&
161            "User supplied memory managers not supported with COD yet.");
162 
163     LogicalDylibs.push_back(CODLogicalDylib(BaseLayer));
164     auto &LDResources = LogicalDylibs.back().getDylibResources();
165 
166     LDResources.ExternalSymbolResolver =
167       [Resolver](const std::string &Name) {
168         return Resolver->findSymbol(Name);
169       };
170 
171     // Process each of the modules in this module set.
172     for (auto &M : Ms)
173       addLogicalModule(LogicalDylibs.back(), std::move(M));
174 
175     return std::prev(LogicalDylibs.end());
176   }
177 
178   /// @brief Remove the module represented by the given handle.
179   ///
180   ///   This will remove all modules in the layers below that were derived from
181   /// the module represented by H.
removeModuleSet(ModuleSetHandleT H)182   void removeModuleSet(ModuleSetHandleT H) {
183     LogicalDylibs.erase(H);
184   }
185 
186   /// @brief Search for the given named symbol.
187   /// @param Name The name of the symbol to search for.
188   /// @param ExportedSymbolsOnly If true, search only for exported symbols.
189   /// @return A handle for the given named symbol, if it exists.
findSymbol(StringRef Name,bool ExportedSymbolsOnly)190   JITSymbol findSymbol(StringRef Name, bool ExportedSymbolsOnly) {
191     for (auto LDI = LogicalDylibs.begin(), LDE = LogicalDylibs.end();
192          LDI != LDE; ++LDI)
193       if (auto Symbol = findSymbolIn(LDI, Name, ExportedSymbolsOnly))
194         return Symbol;
195     return BaseLayer.findSymbol(Name, ExportedSymbolsOnly);
196   }
197 
198   /// @brief Get the address of a symbol provided by this layer, or some layer
199   ///        below this one.
findSymbolIn(ModuleSetHandleT H,const std::string & Name,bool ExportedSymbolsOnly)200   JITSymbol findSymbolIn(ModuleSetHandleT H, const std::string &Name,
201                          bool ExportedSymbolsOnly) {
202     return H->findSymbol(Name, ExportedSymbolsOnly);
203   }
204 
205 private:
206 
207   template <typename ModulePtrT>
addLogicalModule(CODLogicalDylib & LD,ModulePtrT SrcMPtr)208   void addLogicalModule(CODLogicalDylib &LD, ModulePtrT SrcMPtr) {
209 
210     // Bump the linkage and rename any anonymous/privote members in SrcM to
211     // ensure that everything will resolve properly after we partition SrcM.
212     makeAllSymbolsExternallyAccessible(*SrcMPtr);
213 
214     // Create a logical module handle for SrcM within the logical dylib.
215     auto LMH = LD.createLogicalModule();
216     auto &LMResources =  LD.getLogicalModuleResources(LMH);
217 
218     LMResources.SourceModuleOwner = wrapOwnership(std::move(SrcMPtr));
219 
220     Module &SrcM = LMResources.SourceModuleOwner->getModule();
221 
222     // Create the GlobalValues module.
223     const DataLayout &DL = SrcM.getDataLayout();
224     auto GVsM = llvm::make_unique<Module>((SrcM.getName() + ".globals").str(),
225                                           SrcM.getContext());
226     GVsM->setDataLayout(DL);
227 
228     // Create function stubs.
229     ValueToValueMapTy VMap;
230     {
231       typename IndirectStubsMgrT::StubInitsMap StubInits;
232       for (auto &F : SrcM) {
233         // Skip declarations.
234         if (F.isDeclaration())
235           continue;
236 
237         // Record all functions defined by this module.
238         if (CloneStubsIntoPartitions)
239           LMResources.StubsToClone.insert(&F);
240 
241         // Create a callback, associate it with the stub for the function,
242         // and set the compile action to compile the partition containing the
243         // function.
244         auto CCInfo = CompileCallbackMgr.getCompileCallback();
245         StubInits[mangle(F.getName(), DL)] =
246           std::make_pair(CCInfo.getAddress(),
247                          JITSymbolBase::flagsFromGlobalValue(F));
248         CCInfo.setCompileAction([this, &LD, LMH, &F]() {
249           return this->extractAndCompile(LD, LMH, F);
250         });
251       }
252 
253       LMResources.StubsMgr = CreateIndirectStubsManager();
254       auto EC = LMResources.StubsMgr->createStubs(StubInits);
255       (void)EC;
256       // FIXME: This should be propagated back to the user. Stub creation may
257       //        fail for remote JITs.
258       assert(!EC && "Error generating stubs");
259     }
260 
261     // Clone global variable decls.
262     for (auto &GV : SrcM.globals())
263       if (!GV.isDeclaration() && !VMap.count(&GV))
264         cloneGlobalVariableDecl(*GVsM, GV, &VMap);
265 
266     // And the aliases.
267     for (auto &A : SrcM.aliases())
268       if (!VMap.count(&A))
269         cloneGlobalAliasDecl(*GVsM, A, VMap);
270 
271     // Now we need to clone the GV and alias initializers.
272 
273     // Initializers may refer to functions declared (but not defined) in this
274     // module. Build a materializer to clone decls on demand.
275     auto Materializer = createLambdaMaterializer(
276       [this, &GVsM, &LMResources](Value *V) -> Value* {
277         if (auto *F = dyn_cast<Function>(V)) {
278           // Decls in the original module just get cloned.
279           if (F->isDeclaration())
280             return cloneFunctionDecl(*GVsM, *F);
281 
282           // Definitions in the original module (which we have emitted stubs
283           // for at this point) get turned into a constant alias to the stub
284           // instead.
285           const DataLayout &DL = GVsM->getDataLayout();
286           std::string FName = mangle(F->getName(), DL);
287           auto StubSym = LMResources.StubsMgr->findStub(FName, false);
288           unsigned PtrBitWidth = DL.getPointerTypeSizeInBits(F->getType());
289           ConstantInt *StubAddr =
290             ConstantInt::get(GVsM->getContext(),
291                              APInt(PtrBitWidth, StubSym.getAddress()));
292           Constant *Init = ConstantExpr::getCast(Instruction::IntToPtr,
293                                                  StubAddr, F->getType());
294           return GlobalAlias::create(F->getFunctionType(),
295                                      F->getType()->getAddressSpace(),
296                                      F->getLinkage(), F->getName(),
297                                      Init, GVsM.get());
298         }
299         // else....
300         return nullptr;
301       });
302 
303     // Clone the global variable initializers.
304     for (auto &GV : SrcM.globals())
305       if (!GV.isDeclaration())
306         moveGlobalVariableInitializer(GV, VMap, &Materializer);
307 
308     // Clone the global alias initializers.
309     for (auto &A : SrcM.aliases()) {
310       auto *NewA = cast<GlobalAlias>(VMap[&A]);
311       assert(NewA && "Alias not cloned?");
312       Value *Init = MapValue(A.getAliasee(), VMap, RF_None, nullptr,
313                              &Materializer);
314       NewA->setAliasee(cast<Constant>(Init));
315     }
316 
317     // Build a resolver for the globals module and add it to the base layer.
318     auto GVsResolver = createLambdaResolver(
319         [&LD, LMH](const std::string &Name) {
320           auto &LMResources = LD.getLogicalModuleResources(LMH);
321           if (auto Sym = LMResources.StubsMgr->findStub(Name, false))
322             return RuntimeDyld::SymbolInfo(Sym.getAddress(), Sym.getFlags());
323           return LD.getDylibResources().ExternalSymbolResolver(Name);
324         },
325         [](const std::string &Name) {
326           return RuntimeDyld::SymbolInfo(nullptr);
327         });
328 
329     std::vector<std::unique_ptr<Module>> GVsMSet;
330     GVsMSet.push_back(std::move(GVsM));
331     auto GVsH =
332       BaseLayer.addModuleSet(std::move(GVsMSet),
333                              llvm::make_unique<SectionMemoryManager>(),
334                              std::move(GVsResolver));
335     LD.addToLogicalModule(LMH, GVsH);
336   }
337 
mangle(StringRef Name,const DataLayout & DL)338   static std::string mangle(StringRef Name, const DataLayout &DL) {
339     std::string MangledName;
340     {
341       raw_string_ostream MangledNameStream(MangledName);
342       Mangler::getNameWithPrefix(MangledNameStream, Name, DL);
343     }
344     return MangledName;
345   }
346 
extractAndCompile(CODLogicalDylib & LD,LogicalModuleHandle LMH,Function & F)347   TargetAddress extractAndCompile(CODLogicalDylib &LD,
348                                   LogicalModuleHandle LMH,
349                                   Function &F) {
350     auto &LMResources = LD.getLogicalModuleResources(LMH);
351     Module &SrcM = LMResources.SourceModuleOwner->getModule();
352 
353     // If F is a declaration we must already have compiled it.
354     if (F.isDeclaration())
355       return 0;
356 
357     // Grab the name of the function being called here.
358     std::string CalledFnName = mangle(F.getName(), SrcM.getDataLayout());
359 
360     auto Part = Partition(F);
361     auto PartH = emitPartition(LD, LMH, Part);
362 
363     TargetAddress CalledAddr = 0;
364     for (auto *SubF : Part) {
365       std::string FnName = mangle(SubF->getName(), SrcM.getDataLayout());
366       auto FnBodySym = BaseLayer.findSymbolIn(PartH, FnName, false);
367       assert(FnBodySym && "Couldn't find function body.");
368 
369       TargetAddress FnBodyAddr = FnBodySym.getAddress();
370 
371       // If this is the function we're calling record the address so we can
372       // return it from this function.
373       if (SubF == &F)
374         CalledAddr = FnBodyAddr;
375 
376       // Update the function body pointer for the stub.
377       if (auto EC = LMResources.StubsMgr->updatePointer(FnName, FnBodyAddr))
378         return 0;
379     }
380 
381     return CalledAddr;
382   }
383 
384   template <typename PartitionT>
emitPartition(CODLogicalDylib & LD,LogicalModuleHandle LMH,const PartitionT & Part)385   BaseLayerModuleSetHandleT emitPartition(CODLogicalDylib &LD,
386                                           LogicalModuleHandle LMH,
387                                           const PartitionT &Part) {
388     auto &LMResources = LD.getLogicalModuleResources(LMH);
389     Module &SrcM = LMResources.SourceModuleOwner->getModule();
390 
391     // Create the module.
392     std::string NewName = SrcM.getName();
393     for (auto *F : Part) {
394       NewName += ".";
395       NewName += F->getName();
396     }
397 
398     auto M = llvm::make_unique<Module>(NewName, SrcM.getContext());
399     M->setDataLayout(SrcM.getDataLayout());
400     ValueToValueMapTy VMap;
401 
402     auto Materializer = createLambdaMaterializer([this, &LMResources, &M,
403                                                   &VMap](Value *V) -> Value * {
404       if (auto *GV = dyn_cast<GlobalVariable>(V))
405         return cloneGlobalVariableDecl(*M, *GV);
406 
407       if (auto *F = dyn_cast<Function>(V)) {
408         // Check whether we want to clone an available_externally definition.
409         if (!LMResources.StubsToClone.count(F))
410           return cloneFunctionDecl(*M, *F);
411 
412         // Ok - we want an inlinable stub. For that to work we need a decl
413         // for the stub pointer.
414         auto *StubPtr = createImplPointer(*F->getType(), *M,
415                                           F->getName() + "$stub_ptr", nullptr);
416         auto *ClonedF = cloneFunctionDecl(*M, *F);
417         makeStub(*ClonedF, *StubPtr);
418         ClonedF->setLinkage(GlobalValue::AvailableExternallyLinkage);
419         ClonedF->addFnAttr(Attribute::AlwaysInline);
420         return ClonedF;
421       }
422 
423       if (auto *A = dyn_cast<GlobalAlias>(V)) {
424         auto *Ty = A->getValueType();
425         if (Ty->isFunctionTy())
426           return Function::Create(cast<FunctionType>(Ty),
427                                   GlobalValue::ExternalLinkage, A->getName(),
428                                   M.get());
429 
430         return new GlobalVariable(*M, Ty, false, GlobalValue::ExternalLinkage,
431                                   nullptr, A->getName(), nullptr,
432                                   GlobalValue::NotThreadLocal,
433                                   A->getType()->getAddressSpace());
434       }
435 
436       return nullptr;
437     });
438 
439     // Create decls in the new module.
440     for (auto *F : Part)
441       cloneFunctionDecl(*M, *F, &VMap);
442 
443     // Move the function bodies.
444     for (auto *F : Part)
445       moveFunctionBody(*F, VMap, &Materializer);
446 
447     // Create memory manager and symbol resolver.
448     auto MemMgr = llvm::make_unique<SectionMemoryManager>();
449     auto Resolver = createLambdaResolver(
450         [this, &LD, LMH](const std::string &Name) {
451           if (auto Symbol = LD.findSymbolInternally(LMH, Name))
452             return RuntimeDyld::SymbolInfo(Symbol.getAddress(),
453                                            Symbol.getFlags());
454           return LD.getDylibResources().ExternalSymbolResolver(Name);
455         },
456         [this, &LD, LMH](const std::string &Name) {
457           if (auto Symbol = LD.findSymbolInternally(LMH, Name))
458             return RuntimeDyld::SymbolInfo(Symbol.getAddress(),
459                                            Symbol.getFlags());
460           return RuntimeDyld::SymbolInfo(nullptr);
461         });
462     std::vector<std::unique_ptr<Module>> PartMSet;
463     PartMSet.push_back(std::move(M));
464     return BaseLayer.addModuleSet(std::move(PartMSet), std::move(MemMgr),
465                                   std::move(Resolver));
466   }
467 
468   BaseLayerT &BaseLayer;
469   PartitioningFtor Partition;
470   CompileCallbackMgrT &CompileCallbackMgr;
471   IndirectStubsManagerBuilderT CreateIndirectStubsManager;
472 
473   LogicalDylibList LogicalDylibs;
474   bool CloneStubsIntoPartitions;
475 };
476 
477 } // End namespace orc.
478 } // End namespace llvm.
479 
480 #endif // LLVM_EXECUTIONENGINE_ORC_COMPILEONDEMANDLAYER_H
481