1 //===- CompileOnDemandLayer.h - Compile each function on demand -*- C++ -*-===//
3 // The LLVM Compiler Infrastructure
5 // This file is distributed under the University of Illinois Open Source
6 // License. See LICENSE.TXT for details.
8 //===----------------------------------------------------------------------===//
10 // JIT layer for breaking up modules and inserting callbacks to allow
11 // individual functions to be compiled on demand.
13 //===----------------------------------------------------------------------===//
15 #ifndef LLVM_EXECUTIONENGINE_ORC_COMPILEONDEMANDLAYER_H
16 #define LLVM_EXECUTIONENGINE_ORC_COMPILEONDEMANDLAYER_H
18 #include "llvm/ADT/APInt.h"
19 #include "llvm/ADT/STLExtras.h"
20 #include "llvm/ADT/StringRef.h"
21 #include "llvm/ADT/Twine.h"
22 #include "llvm/ExecutionEngine/JITSymbol.h"
23 #include "llvm/ExecutionEngine/Orc/IndirectionUtils.h"
24 #include "llvm/ExecutionEngine/Orc/LambdaResolver.h"
25 #include "llvm/ExecutionEngine/RuntimeDyld.h"
26 #include "llvm/IR/Attributes.h"
27 #include "llvm/IR/Constants.h"
28 #include "llvm/IR/DataLayout.h"
29 #include "llvm/IR/Function.h"
30 #include "llvm/IR/GlobalAlias.h"
31 #include "llvm/IR/GlobalValue.h"
32 #include "llvm/IR/GlobalVariable.h"
33 #include "llvm/IR/Mangler.h"
34 #include "llvm/IR/Module.h"
35 #include "llvm/Support/Casting.h"
36 #include "llvm/Support/raw_ostream.h"
51 /// @brief Compile-on-demand layer.
53 /// When a module is added to this layer a stub is created for each of its
54 /// function definitions. The stubs and other global values are immediately
55 /// added to the layer below. When a stub is called it triggers the extraction
56 /// of the function body from the original module. The extracted body is then
57 /// compiled and executed.
58 template <typename BaseLayerT,
59 typename CompileCallbackMgrT = JITCompileCallbackManager,
60 typename IndirectStubsMgrT = IndirectStubsManager>
61 class CompileOnDemandLayer {
63 template <typename MaterializerFtor>
64 class LambdaMaterializer final : public ValueMaterializer {
66 LambdaMaterializer(MaterializerFtor M) : M(std::move(M)) {}
68 Value *materialize(Value *V) final { return M(V); }
74 template <typename MaterializerFtor>
75 LambdaMaterializer<MaterializerFtor>
76 createLambdaMaterializer(MaterializerFtor M) {
77 return LambdaMaterializer<MaterializerFtor>(std::move(M));
80 typedef typename BaseLayerT::ModuleSetHandleT BaseLayerModuleSetHandleT;
82 // Provide type-erasure for the Modules and MemoryManagers.
83 template <typename ResourceT>
86 ResourceOwner() = default;
87 ResourceOwner(const ResourceOwner&) = delete;
88 ResourceOwner& operator=(const ResourceOwner&) = delete;
89 virtual ~ResourceOwner() = default;
91 virtual ResourceT& getResource() const = 0;
94 template <typename ResourceT, typename ResourcePtrT>
95 class ResourceOwnerImpl : public ResourceOwner<ResourceT> {
97 ResourceOwnerImpl(ResourcePtrT ResourcePtr)
98 : ResourcePtr(std::move(ResourcePtr)) {}
100 ResourceT& getResource() const override { return *ResourcePtr; }
103 ResourcePtrT ResourcePtr;
106 template <typename ResourceT, typename ResourcePtrT>
107 std::unique_ptr<ResourceOwner<ResourceT>>
108 wrapOwnership(ResourcePtrT ResourcePtr) {
109 typedef ResourceOwnerImpl<ResourceT, ResourcePtrT> RO;
110 return llvm::make_unique<RO>(std::move(ResourcePtr));
113 class StaticGlobalRenamer {
115 StaticGlobalRenamer() = default;
116 StaticGlobalRenamer(StaticGlobalRenamer &&) = default;
117 StaticGlobalRenamer &operator=(StaticGlobalRenamer &&) = default;
119 void rename(Module &M) {
121 if (F.hasLocalLinkage())
122 F.setName("$static." + Twine(NextId++));
123 for (auto &G : M.globals())
124 if (G.hasLocalLinkage())
125 G.setName("$static." + Twine(NextId++));
132 struct LogicalDylib {
133 typedef std::function<JITSymbol(const std::string&)> SymbolResolverFtor;
135 typedef std::function<typename BaseLayerT::ModuleSetHandleT(
137 std::unique_ptr<Module>,
138 std::unique_ptr<JITSymbolResolver>)>
141 struct SourceModuleEntry {
142 std::unique_ptr<ResourceOwner<Module>> SourceMod;
143 std::set<Function*> StubsToClone;
146 typedef std::vector<SourceModuleEntry> SourceModulesList;
147 typedef typename SourceModulesList::size_type SourceModuleHandle;
150 addSourceModule(std::unique_ptr<ResourceOwner<Module>> M) {
151 SourceModuleHandle H = SourceModules.size();
152 SourceModules.push_back(SourceModuleEntry());
153 SourceModules.back().SourceMod = std::move(M);
157 Module& getSourceModule(SourceModuleHandle H) {
158 return SourceModules[H].SourceMod->getResource();
161 std::set<Function*>& getStubsToClone(SourceModuleHandle H) {
162 return SourceModules[H].StubsToClone;
165 JITSymbol findSymbol(BaseLayerT &BaseLayer, const std::string &Name,
166 bool ExportedSymbolsOnly) {
167 if (auto Sym = StubsMgr->findStub(Name, ExportedSymbolsOnly))
169 for (auto BLH : BaseLayerHandles)
170 if (auto Sym = BaseLayer.findSymbolIn(BLH, Name, ExportedSymbolsOnly))
175 void removeModulesFromBaseLayer(BaseLayerT &BaseLayer) {
176 for (auto &BLH : BaseLayerHandles)
177 BaseLayer.removeModuleSet(BLH);
180 std::unique_ptr<JITSymbolResolver> ExternalSymbolResolver;
181 std::unique_ptr<ResourceOwner<RuntimeDyld::MemoryManager>> MemMgr;
182 std::unique_ptr<IndirectStubsMgrT> StubsMgr;
183 StaticGlobalRenamer StaticRenamer;
184 ModuleAdderFtor ModuleAdder;
185 SourceModulesList SourceModules;
186 std::vector<BaseLayerModuleSetHandleT> BaseLayerHandles;
189 typedef std::list<LogicalDylib> LogicalDylibList;
192 /// @brief Handle to a set of loaded modules.
193 typedef typename LogicalDylibList::iterator ModuleSetHandleT;
195 /// @brief Module partitioning functor.
196 typedef std::function<std::set<Function*>(Function&)> PartitioningFtor;
198 /// @brief Builder for IndirectStubsManagers.
199 typedef std::function<std::unique_ptr<IndirectStubsMgrT>()>
200 IndirectStubsManagerBuilderT;
202 /// @brief Construct a compile-on-demand layer instance.
203 CompileOnDemandLayer(BaseLayerT &BaseLayer, PartitioningFtor Partition,
204 CompileCallbackMgrT &CallbackMgr,
205 IndirectStubsManagerBuilderT CreateIndirectStubsManager,
206 bool CloneStubsIntoPartitions = true)
207 : BaseLayer(BaseLayer), Partition(std::move(Partition)),
208 CompileCallbackMgr(CallbackMgr),
209 CreateIndirectStubsManager(std::move(CreateIndirectStubsManager)),
210 CloneStubsIntoPartitions(CloneStubsIntoPartitions) {}
212 ~CompileOnDemandLayer() {
213 while (!LogicalDylibs.empty())
214 removeModuleSet(LogicalDylibs.begin());
217 /// @brief Add a module to the compile-on-demand layer.
218 template <typename ModuleSetT, typename MemoryManagerPtrT,
219 typename SymbolResolverPtrT>
220 ModuleSetHandleT addModuleSet(ModuleSetT Ms,
221 MemoryManagerPtrT MemMgr,
222 SymbolResolverPtrT Resolver) {
224 LogicalDylibs.push_back(LogicalDylib());
225 auto &LD = LogicalDylibs.back();
226 LD.ExternalSymbolResolver = std::move(Resolver);
227 LD.StubsMgr = CreateIndirectStubsManager();
229 auto &MemMgrRef = *MemMgr;
230 LD.MemMgr = wrapOwnership<RuntimeDyld::MemoryManager>(std::move(MemMgr));
233 [&MemMgrRef](BaseLayerT &B, std::unique_ptr<Module> M,
234 std::unique_ptr<JITSymbolResolver> R) {
235 std::vector<std::unique_ptr<Module>> Ms;
236 Ms.push_back(std::move(M));
237 return B.addModuleSet(std::move(Ms), &MemMgrRef, std::move(R));
240 // Process each of the modules in this module set.
242 addLogicalModule(LogicalDylibs.back(), std::move(M));
244 return std::prev(LogicalDylibs.end());
247 /// @brief Remove the module represented by the given handle.
249 /// This will remove all modules in the layers below that were derived from
250 /// the module represented by H.
251 void removeModuleSet(ModuleSetHandleT H) {
252 H->removeModulesFromBaseLayer(BaseLayer);
253 LogicalDylibs.erase(H);
256 /// @brief Search for the given named symbol.
257 /// @param Name The name of the symbol to search for.
258 /// @param ExportedSymbolsOnly If true, search only for exported symbols.
259 /// @return A handle for the given named symbol, if it exists.
260 JITSymbol findSymbol(StringRef Name, bool ExportedSymbolsOnly) {
261 for (auto LDI = LogicalDylibs.begin(), LDE = LogicalDylibs.end();
263 if (auto Sym = LDI->StubsMgr->findStub(Name, ExportedSymbolsOnly))
265 if (auto Sym = findSymbolIn(LDI, Name, ExportedSymbolsOnly))
268 return BaseLayer.findSymbol(Name, ExportedSymbolsOnly);
271 /// @brief Get the address of a symbol provided by this layer, or some layer
273 JITSymbol findSymbolIn(ModuleSetHandleT H, const std::string &Name,
274 bool ExportedSymbolsOnly) {
275 return H->findSymbol(BaseLayer, Name, ExportedSymbolsOnly);
278 /// @brief Update the stub for the given function to point at FnBodyAddr.
279 /// This can be used to support re-optimization.
280 /// @return true if the function exists and the stub is updated, false
283 // FIXME: We should track and free associated resources (unused compile
284 // callbacks, uncompiled IR, and no-longer-needed/reachable function
286 // FIXME: Return Error once the JIT APIs are Errorized.
287 bool updatePointer(std::string FuncName, JITTargetAddress FnBodyAddr) {
288 //Find out which logical dylib contains our symbol
289 auto LDI = LogicalDylibs.begin();
290 for (auto LDE = LogicalDylibs.end(); LDI != LDE; ++LDI) {
291 if (auto LMResources = LDI->getLogicalModuleResourcesForSymbol(FuncName, false)) {
292 Module &SrcM = LMResources->SourceModule->getResource();
293 std::string CalledFnName = mangle(FuncName, SrcM.getDataLayout());
294 if (auto EC = LMResources->StubsMgr->updatePointer(CalledFnName, FnBodyAddr))
304 template <typename ModulePtrT>
305 void addLogicalModule(LogicalDylib &LD, ModulePtrT SrcMPtr) {
307 // Rename all static functions / globals to $static.X :
308 // This will unique the names across all modules in the logical dylib,
309 // simplifying symbol lookup.
310 LD.StaticRenamer.rename(*SrcMPtr);
312 // Bump the linkage and rename any anonymous/privote members in SrcM to
313 // ensure that everything will resolve properly after we partition SrcM.
314 makeAllSymbolsExternallyAccessible(*SrcMPtr);
316 // Create a logical module handle for SrcM within the logical dylib.
317 Module &SrcM = *SrcMPtr;
318 auto LMId = LD.addSourceModule(wrapOwnership<Module>(std::move(SrcMPtr)));
320 // Create stub functions.
321 const DataLayout &DL = SrcM.getDataLayout();
323 typename IndirectStubsMgrT::StubInitsMap StubInits;
324 for (auto &F : SrcM) {
325 // Skip declarations.
326 if (F.isDeclaration())
329 // Skip weak functions for which we already have definitions.
330 auto MangledName = mangle(F.getName(), DL);
331 if (F.hasWeakLinkage() || F.hasLinkOnceLinkage())
332 if (auto Sym = LD.findSymbol(BaseLayer, MangledName, false))
335 // Record all functions defined by this module.
336 if (CloneStubsIntoPartitions)
337 LD.getStubsToClone(LMId).insert(&F);
339 // Create a callback, associate it with the stub for the function,
340 // and set the compile action to compile the partition containing the
342 auto CCInfo = CompileCallbackMgr.getCompileCallback();
343 StubInits[MangledName] =
344 std::make_pair(CCInfo.getAddress(),
345 JITSymbolFlags::fromGlobalValue(F));
346 CCInfo.setCompileAction([this, &LD, LMId, &F]() {
347 return this->extractAndCompile(LD, LMId, F);
351 auto EC = LD.StubsMgr->createStubs(StubInits);
353 // FIXME: This should be propagated back to the user. Stub creation may
354 // fail for remote JITs.
355 assert(!EC && "Error generating stubs");
358 // If this module doesn't contain any globals, aliases, or module flags then
359 // we can bail out early and avoid the overhead of creating and managing an
360 // empty globals module.
361 if (SrcM.global_empty() && SrcM.alias_empty() &&
362 !SrcM.getModuleFlagsMetadata())
365 // Create the GlobalValues module.
366 auto GVsM = llvm::make_unique<Module>((SrcM.getName() + ".globals").str(),
368 GVsM->setDataLayout(DL);
370 ValueToValueMapTy VMap;
372 // Clone global variable decls.
373 for (auto &GV : SrcM.globals())
374 if (!GV.isDeclaration() && !VMap.count(&GV))
375 cloneGlobalVariableDecl(*GVsM, GV, &VMap);
378 for (auto &A : SrcM.aliases())
380 cloneGlobalAliasDecl(*GVsM, A, VMap);
382 // Clone the module flags.
383 cloneModuleFlagsMetadata(*GVsM, SrcM, VMap);
385 // Now we need to clone the GV and alias initializers.
387 // Initializers may refer to functions declared (but not defined) in this
388 // module. Build a materializer to clone decls on demand.
389 auto Materializer = createLambdaMaterializer(
390 [&LD, &GVsM](Value *V) -> Value* {
391 if (auto *F = dyn_cast<Function>(V)) {
392 // Decls in the original module just get cloned.
393 if (F->isDeclaration())
394 return cloneFunctionDecl(*GVsM, *F);
396 // Definitions in the original module (which we have emitted stubs
397 // for at this point) get turned into a constant alias to the stub
399 const DataLayout &DL = GVsM->getDataLayout();
400 std::string FName = mangle(F->getName(), DL);
401 auto StubSym = LD.StubsMgr->findStub(FName, false);
402 unsigned PtrBitWidth = DL.getPointerTypeSizeInBits(F->getType());
403 ConstantInt *StubAddr =
404 ConstantInt::get(GVsM->getContext(),
405 APInt(PtrBitWidth, StubSym.getAddress()));
406 Constant *Init = ConstantExpr::getCast(Instruction::IntToPtr,
407 StubAddr, F->getType());
408 return GlobalAlias::create(F->getFunctionType(),
409 F->getType()->getAddressSpace(),
410 F->getLinkage(), F->getName(),
417 // Clone the global variable initializers.
418 for (auto &GV : SrcM.globals())
419 if (!GV.isDeclaration())
420 moveGlobalVariableInitializer(GV, VMap, &Materializer);
422 // Clone the global alias initializers.
423 for (auto &A : SrcM.aliases()) {
424 auto *NewA = cast<GlobalAlias>(VMap[&A]);
425 assert(NewA && "Alias not cloned?");
426 Value *Init = MapValue(A.getAliasee(), VMap, RF_None, nullptr,
428 NewA->setAliasee(cast<Constant>(Init));
431 // Build a resolver for the globals module and add it to the base layer.
432 auto GVsResolver = createLambdaResolver(
433 [this, &LD](const std::string &Name) {
434 if (auto Sym = LD.StubsMgr->findStub(Name, false))
436 if (auto Sym = LD.findSymbol(BaseLayer, Name, false))
438 return LD.ExternalSymbolResolver->findSymbolInLogicalDylib(Name);
440 [&LD](const std::string &Name) {
441 return LD.ExternalSymbolResolver->findSymbol(Name);
444 auto GVsH = LD.ModuleAdder(BaseLayer, std::move(GVsM),
445 std::move(GVsResolver));
446 LD.BaseLayerHandles.push_back(GVsH);
449 static std::string mangle(StringRef Name, const DataLayout &DL) {
450 std::string MangledName;
452 raw_string_ostream MangledNameStream(MangledName);
453 Mangler::getNameWithPrefix(MangledNameStream, Name, DL);
459 extractAndCompile(LogicalDylib &LD,
460 typename LogicalDylib::SourceModuleHandle LMId,
462 Module &SrcM = LD.getSourceModule(LMId);
464 // If F is a declaration we must already have compiled it.
465 if (F.isDeclaration())
468 // Grab the name of the function being called here.
469 std::string CalledFnName = mangle(F.getName(), SrcM.getDataLayout());
471 auto Part = Partition(F);
472 auto PartH = emitPartition(LD, LMId, Part);
474 JITTargetAddress CalledAddr = 0;
475 for (auto *SubF : Part) {
476 std::string FnName = mangle(SubF->getName(), SrcM.getDataLayout());
477 auto FnBodySym = BaseLayer.findSymbolIn(PartH, FnName, false);
478 assert(FnBodySym && "Couldn't find function body.");
480 JITTargetAddress FnBodyAddr = FnBodySym.getAddress();
482 // If this is the function we're calling record the address so we can
483 // return it from this function.
485 CalledAddr = FnBodyAddr;
487 // Update the function body pointer for the stub.
488 if (auto EC = LD.StubsMgr->updatePointer(FnName, FnBodyAddr))
492 LD.BaseLayerHandles.push_back(PartH);
497 template <typename PartitionT>
498 BaseLayerModuleSetHandleT
499 emitPartition(LogicalDylib &LD,
500 typename LogicalDylib::SourceModuleHandle LMId,
501 const PartitionT &Part) {
502 Module &SrcM = LD.getSourceModule(LMId);
504 // Create the module.
505 std::string NewName = SrcM.getName();
506 for (auto *F : Part) {
508 NewName += F->getName();
511 auto M = llvm::make_unique<Module>(NewName, SrcM.getContext());
512 M->setDataLayout(SrcM.getDataLayout());
513 ValueToValueMapTy VMap;
515 auto Materializer = createLambdaMaterializer([&LD, &LMId,
516 &M](Value *V) -> Value * {
517 if (auto *GV = dyn_cast<GlobalVariable>(V))
518 return cloneGlobalVariableDecl(*M, *GV);
520 if (auto *F = dyn_cast<Function>(V)) {
521 // Check whether we want to clone an available_externally definition.
522 if (!LD.getStubsToClone(LMId).count(F))
523 return cloneFunctionDecl(*M, *F);
525 // Ok - we want an inlinable stub. For that to work we need a decl
526 // for the stub pointer.
527 auto *StubPtr = createImplPointer(*F->getType(), *M,
528 F->getName() + "$stub_ptr", nullptr);
529 auto *ClonedF = cloneFunctionDecl(*M, *F);
530 makeStub(*ClonedF, *StubPtr);
531 ClonedF->setLinkage(GlobalValue::AvailableExternallyLinkage);
532 ClonedF->addFnAttr(Attribute::AlwaysInline);
536 if (auto *A = dyn_cast<GlobalAlias>(V)) {
537 auto *Ty = A->getValueType();
538 if (Ty->isFunctionTy())
539 return Function::Create(cast<FunctionType>(Ty),
540 GlobalValue::ExternalLinkage, A->getName(),
543 return new GlobalVariable(*M, Ty, false, GlobalValue::ExternalLinkage,
544 nullptr, A->getName(), nullptr,
545 GlobalValue::NotThreadLocal,
546 A->getType()->getAddressSpace());
552 // Create decls in the new module.
554 cloneFunctionDecl(*M, *F, &VMap);
556 // Move the function bodies.
558 moveFunctionBody(*F, VMap, &Materializer);
560 // Create memory manager and symbol resolver.
561 auto Resolver = createLambdaResolver(
562 [this, &LD](const std::string &Name) {
563 if (auto Sym = LD.findSymbol(BaseLayer, Name, false))
565 return LD.ExternalSymbolResolver->findSymbolInLogicalDylib(Name);
567 [&LD](const std::string &Name) {
568 return LD.ExternalSymbolResolver->findSymbol(Name);
571 return LD.ModuleAdder(BaseLayer, std::move(M), std::move(Resolver));
574 BaseLayerT &BaseLayer;
575 PartitioningFtor Partition;
576 CompileCallbackMgrT &CompileCallbackMgr;
577 IndirectStubsManagerBuilderT CreateIndirectStubsManager;
579 LogicalDylibList LogicalDylibs;
580 bool CloneStubsIntoPartitions;
583 } // end namespace orc
584 } // end namespace llvm
586 #endif // LLVM_EXECUTIONENGINE_ORC_COMPILEONDEMANDLAYER_H