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 "IndirectionUtils.h"
19 #include "LambdaResolver.h"
20 #include "llvm/ADT/STLExtras.h"
21 #include "llvm/Support/Debug.h"
22 #include "llvm/Transforms/Utils/Cloning.h"
31 /// @brief Compile-on-demand layer.
33 /// When a module is added to this layer a stub is created for each of its
34 /// function definitions. The stubs and other global values are immediately
35 /// added to the layer below. When a stub is called it triggers the extraction
36 /// of the function body from the original module. The extracted body is then
37 /// compiled and executed.
38 template <typename BaseLayerT,
39 typename CompileCallbackMgrT = JITCompileCallbackManager,
40 typename IndirectStubsMgrT = IndirectStubsManager>
41 class CompileOnDemandLayer {
44 template <typename MaterializerFtor>
45 class LambdaMaterializer final : public ValueMaterializer {
47 LambdaMaterializer(MaterializerFtor M) : M(std::move(M)) {}
48 Value *materialize(Value *V) final { return M(V); }
54 template <typename MaterializerFtor>
55 LambdaMaterializer<MaterializerFtor>
56 createLambdaMaterializer(MaterializerFtor M) {
57 return LambdaMaterializer<MaterializerFtor>(std::move(M));
60 typedef typename BaseLayerT::ModuleSetHandleT BaseLayerModuleSetHandleT;
62 // Provide type-erasure for the Modules and MemoryManagers.
63 template <typename ResourceT>
66 ResourceOwner() = default;
67 ResourceOwner(const ResourceOwner&) = delete;
68 ResourceOwner& operator=(const ResourceOwner&) = delete;
69 virtual ~ResourceOwner() { }
70 virtual ResourceT& getResource() const = 0;
73 template <typename ResourceT, typename ResourcePtrT>
74 class ResourceOwnerImpl : public ResourceOwner<ResourceT> {
76 ResourceOwnerImpl(ResourcePtrT ResourcePtr)
77 : ResourcePtr(std::move(ResourcePtr)) {}
78 ResourceT& getResource() const override { return *ResourcePtr; }
80 ResourcePtrT ResourcePtr;
83 template <typename ResourceT, typename ResourcePtrT>
84 std::unique_ptr<ResourceOwner<ResourceT>>
85 wrapOwnership(ResourcePtrT ResourcePtr) {
86 typedef ResourceOwnerImpl<ResourceT, ResourcePtrT> RO;
87 return llvm::make_unique<RO>(std::move(ResourcePtr));
90 class StaticGlobalRenamer {
92 StaticGlobalRenamer() {}
94 StaticGlobalRenamer(StaticGlobalRenamer &&Other)
95 : NextId(Other.NextId) {}
97 StaticGlobalRenamer& operator=(StaticGlobalRenamer &&Other) {
98 NextId = Other.NextId;
102 void rename(Module &M) {
104 if (F.hasLocalLinkage())
105 F.setName("$static." + Twine(NextId++));
106 for (auto &G : M.globals())
107 if (G.hasLocalLinkage())
108 G.setName("$static." + Twine(NextId++));
115 struct LogicalDylib {
116 typedef std::function<JITSymbol(const std::string&)> SymbolResolverFtor;
118 typedef std::function<typename BaseLayerT::ModuleSetHandleT(
120 std::unique_ptr<Module>,
121 std::unique_ptr<JITSymbolResolver>)>
124 struct SourceModuleEntry {
125 std::unique_ptr<ResourceOwner<Module>> SourceMod;
126 std::set<Function*> StubsToClone;
128 SourceModuleEntry() = default;
129 SourceModuleEntry(SourceModuleEntry &&Other)
130 : SourceMod(std::move(Other.SourceMod)),
131 StubsToClone(std::move(Other.StubsToClone)) {}
132 SourceModuleEntry& operator=(SourceModuleEntry &&Other) {
133 SourceMod = std::move(Other.SourceMod);
134 StubsToClone = std::move(Other.StubsToClone);
139 typedef std::vector<SourceModuleEntry> SourceModulesList;
140 typedef typename SourceModulesList::size_type SourceModuleHandle;
142 LogicalDylib() = default;
144 // Explicit move constructor to make MSVC happy.
145 LogicalDylib(LogicalDylib &&Other)
146 : ExternalSymbolResolver(std::move(Other.ExternalSymbolResolver)),
147 MemMgr(std::move(Other.MemMgr)),
148 StubsMgr(std::move(Other.StubsMgr)),
149 StaticRenamer(std::move(Other.StaticRenamer)),
150 ModuleAdder(std::move(Other.ModuleAdder)),
151 SourceModules(std::move(Other.SourceModules)),
152 BaseLayerHandles(std::move(Other.BaseLayerHandles)) {}
154 // Explicit move assignment operator to make MSVC happy.
155 LogicalDylib& operator=(LogicalDylib &&Other) {
156 ExternalSymbolResolver = std::move(Other.ExternalSymbolResolver);
157 MemMgr = std::move(Other.MemMgr);
158 StubsMgr = std::move(Other.StubsMgr);
159 StaticRenamer = std::move(Other.StaticRenamer);
160 ModuleAdder = std::move(Other.ModuleAdder);
161 SourceModules = std::move(Other.SourceModules);
162 BaseLayerHandles = std::move(Other.BaseLayerHandles);
167 addSourceModule(std::unique_ptr<ResourceOwner<Module>> M) {
168 SourceModuleHandle H = SourceModules.size();
169 SourceModules.push_back(SourceModuleEntry());
170 SourceModules.back().SourceMod = std::move(M);
174 Module& getSourceModule(SourceModuleHandle H) {
175 return SourceModules[H].SourceMod->getResource();
178 std::set<Function*>& getStubsToClone(SourceModuleHandle H) {
179 return SourceModules[H].StubsToClone;
182 JITSymbol findSymbol(BaseLayerT &BaseLayer, const std::string &Name,
183 bool ExportedSymbolsOnly) {
184 if (auto Sym = StubsMgr->findStub(Name, ExportedSymbolsOnly))
186 for (auto BLH : BaseLayerHandles)
187 if (auto Sym = BaseLayer.findSymbolIn(BLH, Name, ExportedSymbolsOnly))
192 std::unique_ptr<JITSymbolResolver> ExternalSymbolResolver;
193 std::unique_ptr<ResourceOwner<RuntimeDyld::MemoryManager>> MemMgr;
194 std::unique_ptr<IndirectStubsMgrT> StubsMgr;
195 StaticGlobalRenamer StaticRenamer;
196 ModuleAdderFtor ModuleAdder;
197 SourceModulesList SourceModules;
198 std::vector<BaseLayerModuleSetHandleT> BaseLayerHandles;
201 typedef std::list<LogicalDylib> LogicalDylibList;
205 /// @brief Handle to a set of loaded modules.
206 typedef typename LogicalDylibList::iterator ModuleSetHandleT;
208 /// @brief Module partitioning functor.
209 typedef std::function<std::set<Function*>(Function&)> PartitioningFtor;
211 /// @brief Builder for IndirectStubsManagers.
212 typedef std::function<std::unique_ptr<IndirectStubsMgrT>()>
213 IndirectStubsManagerBuilderT;
215 /// @brief Construct a compile-on-demand layer instance.
216 CompileOnDemandLayer(BaseLayerT &BaseLayer, PartitioningFtor Partition,
217 CompileCallbackMgrT &CallbackMgr,
218 IndirectStubsManagerBuilderT CreateIndirectStubsManager,
219 bool CloneStubsIntoPartitions = true)
220 : BaseLayer(BaseLayer), Partition(std::move(Partition)),
221 CompileCallbackMgr(CallbackMgr),
222 CreateIndirectStubsManager(std::move(CreateIndirectStubsManager)),
223 CloneStubsIntoPartitions(CloneStubsIntoPartitions) {}
225 /// @brief Add a module to the compile-on-demand layer.
226 template <typename ModuleSetT, typename MemoryManagerPtrT,
227 typename SymbolResolverPtrT>
228 ModuleSetHandleT addModuleSet(ModuleSetT Ms,
229 MemoryManagerPtrT MemMgr,
230 SymbolResolverPtrT Resolver) {
232 LogicalDylibs.push_back(LogicalDylib());
233 auto &LD = LogicalDylibs.back();
234 LD.ExternalSymbolResolver = std::move(Resolver);
235 LD.StubsMgr = CreateIndirectStubsManager();
237 auto &MemMgrRef = *MemMgr;
238 LD.MemMgr = wrapOwnership<RuntimeDyld::MemoryManager>(std::move(MemMgr));
241 [&MemMgrRef](BaseLayerT &B, std::unique_ptr<Module> M,
242 std::unique_ptr<JITSymbolResolver> R) {
243 std::vector<std::unique_ptr<Module>> Ms;
244 Ms.push_back(std::move(M));
245 return B.addModuleSet(std::move(Ms), &MemMgrRef, std::move(R));
248 // Process each of the modules in this module set.
250 addLogicalModule(LogicalDylibs.back(), std::move(M));
252 return std::prev(LogicalDylibs.end());
255 /// @brief Remove the module represented by the given handle.
257 /// This will remove all modules in the layers below that were derived from
258 /// the module represented by H.
259 void removeModuleSet(ModuleSetHandleT H) {
260 LogicalDylibs.erase(H);
263 /// @brief Search for the given named symbol.
264 /// @param Name The name of the symbol to search for.
265 /// @param ExportedSymbolsOnly If true, search only for exported symbols.
266 /// @return A handle for the given named symbol, if it exists.
267 JITSymbol findSymbol(StringRef Name, bool ExportedSymbolsOnly) {
268 for (auto LDI = LogicalDylibs.begin(), LDE = LogicalDylibs.end();
270 if (auto Sym = LDI->StubsMgr->findStub(Name, ExportedSymbolsOnly))
272 if (auto Sym = findSymbolIn(LDI, Name, ExportedSymbolsOnly))
275 return BaseLayer.findSymbol(Name, ExportedSymbolsOnly);
278 /// @brief Get the address of a symbol provided by this layer, or some layer
280 JITSymbol findSymbolIn(ModuleSetHandleT H, const std::string &Name,
281 bool ExportedSymbolsOnly) {
282 return H->findSymbol(BaseLayer, Name, ExportedSymbolsOnly);
285 /// @brief Update the stub for the given function to point at FnBodyAddr.
286 /// This can be used to support re-optimization.
287 /// @return true if the function exists and the stub is updated, false
290 // FIXME: We should track and free associated resources (unused compile
291 // callbacks, uncompiled IR, and no-longer-needed/reachable function
293 // FIXME: Return Error once the JIT APIs are Errorized.
294 bool updatePointer(std::string FuncName, JITTargetAddress FnBodyAddr) {
295 //Find out which logical dylib contains our symbol
296 auto LDI = LogicalDylibs.begin();
297 for (auto LDE = LogicalDylibs.end(); LDI != LDE; ++LDI) {
298 if (auto LMResources = LDI->getLogicalModuleResourcesForSymbol(FuncName, false)) {
299 Module &SrcM = LMResources->SourceModule->getResource();
300 std::string CalledFnName = mangle(FuncName, SrcM.getDataLayout());
301 if (auto EC = LMResources->StubsMgr->updatePointer(CalledFnName, FnBodyAddr)) {
313 template <typename ModulePtrT>
314 void addLogicalModule(LogicalDylib &LD, ModulePtrT SrcMPtr) {
316 // Rename all static functions / globals to $static.X :
317 // This will unique the names across all modules in the logical dylib,
318 // simplifying symbol lookup.
319 LD.StaticRenamer.rename(*SrcMPtr);
321 // Bump the linkage and rename any anonymous/privote members in SrcM to
322 // ensure that everything will resolve properly after we partition SrcM.
323 makeAllSymbolsExternallyAccessible(*SrcMPtr);
325 // Create a logical module handle for SrcM within the logical dylib.
326 Module &SrcM = *SrcMPtr;
327 auto LMId = LD.addSourceModule(wrapOwnership<Module>(std::move(SrcMPtr)));
329 // Create stub functions.
330 const DataLayout &DL = SrcM.getDataLayout();
332 typename IndirectStubsMgrT::StubInitsMap StubInits;
333 for (auto &F : SrcM) {
334 // Skip declarations.
335 if (F.isDeclaration())
338 // Skip weak functions for which we already have definitions.
339 auto MangledName = mangle(F.getName(), DL);
340 if (F.hasWeakLinkage() || F.hasLinkOnceLinkage())
341 if (auto Sym = LD.findSymbol(BaseLayer, MangledName, false))
344 // Record all functions defined by this module.
345 if (CloneStubsIntoPartitions)
346 LD.getStubsToClone(LMId).insert(&F);
348 // Create a callback, associate it with the stub for the function,
349 // and set the compile action to compile the partition containing the
351 auto CCInfo = CompileCallbackMgr.getCompileCallback();
352 StubInits[MangledName] =
353 std::make_pair(CCInfo.getAddress(),
354 JITSymbolFlags::fromGlobalValue(F));
355 CCInfo.setCompileAction([this, &LD, LMId, &F]() {
356 return this->extractAndCompile(LD, LMId, F);
360 auto EC = LD.StubsMgr->createStubs(StubInits);
362 // FIXME: This should be propagated back to the user. Stub creation may
363 // fail for remote JITs.
364 assert(!EC && "Error generating stubs");
367 // If this module doesn't contain any globals or aliases we can bail out
368 // early and avoid the overhead of creating and managing an empty globals
370 if (SrcM.global_empty() && SrcM.alias_empty())
373 // Create the GlobalValues module.
374 auto GVsM = llvm::make_unique<Module>((SrcM.getName() + ".globals").str(),
376 GVsM->setDataLayout(DL);
378 ValueToValueMapTy VMap;
380 // Clone global variable decls.
381 for (auto &GV : SrcM.globals())
382 if (!GV.isDeclaration() && !VMap.count(&GV))
383 cloneGlobalVariableDecl(*GVsM, GV, &VMap);
386 for (auto &A : SrcM.aliases())
388 cloneGlobalAliasDecl(*GVsM, A, VMap);
390 // Now we need to clone the GV and alias initializers.
392 // Initializers may refer to functions declared (but not defined) in this
393 // module. Build a materializer to clone decls on demand.
394 auto Materializer = createLambdaMaterializer(
395 [this, &LD, &GVsM](Value *V) -> Value* {
396 if (auto *F = dyn_cast<Function>(V)) {
397 // Decls in the original module just get cloned.
398 if (F->isDeclaration())
399 return cloneFunctionDecl(*GVsM, *F);
401 // Definitions in the original module (which we have emitted stubs
402 // for at this point) get turned into a constant alias to the stub
404 const DataLayout &DL = GVsM->getDataLayout();
405 std::string FName = mangle(F->getName(), DL);
406 auto StubSym = LD.StubsMgr->findStub(FName, false);
407 unsigned PtrBitWidth = DL.getPointerTypeSizeInBits(F->getType());
408 ConstantInt *StubAddr =
409 ConstantInt::get(GVsM->getContext(),
410 APInt(PtrBitWidth, StubSym.getAddress()));
411 Constant *Init = ConstantExpr::getCast(Instruction::IntToPtr,
412 StubAddr, F->getType());
413 return GlobalAlias::create(F->getFunctionType(),
414 F->getType()->getAddressSpace(),
415 F->getLinkage(), F->getName(),
422 // Clone the global variable initializers.
423 for (auto &GV : SrcM.globals())
424 if (!GV.isDeclaration())
425 moveGlobalVariableInitializer(GV, VMap, &Materializer);
427 // Clone the global alias initializers.
428 for (auto &A : SrcM.aliases()) {
429 auto *NewA = cast<GlobalAlias>(VMap[&A]);
430 assert(NewA && "Alias not cloned?");
431 Value *Init = MapValue(A.getAliasee(), VMap, RF_None, nullptr,
433 NewA->setAliasee(cast<Constant>(Init));
436 // Build a resolver for the globals module and add it to the base layer.
437 auto GVsResolver = createLambdaResolver(
438 [this, &LD, LMId](const std::string &Name) {
439 if (auto Sym = LD.StubsMgr->findStub(Name, false))
441 if (auto Sym = LD.findSymbol(BaseLayer, Name, false))
443 return LD.ExternalSymbolResolver->findSymbolInLogicalDylib(Name);
445 [&LD](const std::string &Name) {
446 return LD.ExternalSymbolResolver->findSymbol(Name);
449 auto GVsH = LD.ModuleAdder(BaseLayer, std::move(GVsM),
450 std::move(GVsResolver));
451 LD.BaseLayerHandles.push_back(GVsH);
454 static std::string mangle(StringRef Name, const DataLayout &DL) {
455 std::string MangledName;
457 raw_string_ostream MangledNameStream(MangledName);
458 Mangler::getNameWithPrefix(MangledNameStream, Name, DL);
464 extractAndCompile(LogicalDylib &LD,
465 typename LogicalDylib::SourceModuleHandle LMId,
467 Module &SrcM = LD.getSourceModule(LMId);
469 // If F is a declaration we must already have compiled it.
470 if (F.isDeclaration())
473 // Grab the name of the function being called here.
474 std::string CalledFnName = mangle(F.getName(), SrcM.getDataLayout());
476 auto Part = Partition(F);
477 auto PartH = emitPartition(LD, LMId, Part);
479 JITTargetAddress CalledAddr = 0;
480 for (auto *SubF : Part) {
481 std::string FnName = mangle(SubF->getName(), SrcM.getDataLayout());
482 auto FnBodySym = BaseLayer.findSymbolIn(PartH, FnName, false);
483 assert(FnBodySym && "Couldn't find function body.");
485 JITTargetAddress FnBodyAddr = FnBodySym.getAddress();
487 // If this is the function we're calling record the address so we can
488 // return it from this function.
490 CalledAddr = FnBodyAddr;
492 // Update the function body pointer for the stub.
493 if (auto EC = LD.StubsMgr->updatePointer(FnName, FnBodyAddr))
500 template <typename PartitionT>
501 BaseLayerModuleSetHandleT
502 emitPartition(LogicalDylib &LD,
503 typename LogicalDylib::SourceModuleHandle LMId,
504 const PartitionT &Part) {
505 Module &SrcM = LD.getSourceModule(LMId);
507 // Create the module.
508 std::string NewName = SrcM.getName();
509 for (auto *F : Part) {
511 NewName += F->getName();
514 auto M = llvm::make_unique<Module>(NewName, SrcM.getContext());
515 M->setDataLayout(SrcM.getDataLayout());
516 ValueToValueMapTy VMap;
518 auto Materializer = createLambdaMaterializer([this, &LD, &LMId, &M,
519 &VMap](Value *V) -> Value * {
520 if (auto *GV = dyn_cast<GlobalVariable>(V))
521 return cloneGlobalVariableDecl(*M, *GV);
523 if (auto *F = dyn_cast<Function>(V)) {
524 // Check whether we want to clone an available_externally definition.
525 if (!LD.getStubsToClone(LMId).count(F))
526 return cloneFunctionDecl(*M, *F);
528 // Ok - we want an inlinable stub. For that to work we need a decl
529 // for the stub pointer.
530 auto *StubPtr = createImplPointer(*F->getType(), *M,
531 F->getName() + "$stub_ptr", nullptr);
532 auto *ClonedF = cloneFunctionDecl(*M, *F);
533 makeStub(*ClonedF, *StubPtr);
534 ClonedF->setLinkage(GlobalValue::AvailableExternallyLinkage);
535 ClonedF->addFnAttr(Attribute::AlwaysInline);
539 if (auto *A = dyn_cast<GlobalAlias>(V)) {
540 auto *Ty = A->getValueType();
541 if (Ty->isFunctionTy())
542 return Function::Create(cast<FunctionType>(Ty),
543 GlobalValue::ExternalLinkage, A->getName(),
546 return new GlobalVariable(*M, Ty, false, GlobalValue::ExternalLinkage,
547 nullptr, A->getName(), nullptr,
548 GlobalValue::NotThreadLocal,
549 A->getType()->getAddressSpace());
555 // Create decls in the new module.
557 cloneFunctionDecl(*M, *F, &VMap);
559 // Move the function bodies.
561 moveFunctionBody(*F, VMap, &Materializer);
563 // Create memory manager and symbol resolver.
564 auto Resolver = createLambdaResolver(
565 [this, &LD, LMId](const std::string &Name) {
566 if (auto Sym = LD.findSymbol(BaseLayer, Name, false))
568 return LD.ExternalSymbolResolver->findSymbolInLogicalDylib(Name);
570 [this, &LD](const std::string &Name) {
571 return LD.ExternalSymbolResolver->findSymbol(Name);
574 return LD.ModuleAdder(BaseLayer, std::move(M), std::move(Resolver));
577 BaseLayerT &BaseLayer;
578 PartitioningFtor Partition;
579 CompileCallbackMgrT &CompileCallbackMgr;
580 IndirectStubsManagerBuilderT CreateIndirectStubsManager;
582 LogicalDylibList LogicalDylibs;
583 bool CloneStubsIntoPartitions;
586 } // End namespace orc.
587 } // End namespace llvm.
589 #endif // LLVM_EXECUTIONENGINE_ORC_COMPILEONDEMANDLAYER_H