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 "LogicalDylib.h"
21 #include "llvm/ADT/STLExtras.h"
22 #include "llvm/ExecutionEngine/SectionMemoryManager.h"
23 #include "llvm/Transforms/Utils/Cloning.h"
28 #include "llvm/Support/Debug.h"
33 /// @brief Compile-on-demand layer.
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 = JITCompileCallbackManagerBase,
42 typename IndirectStubsMgrT = IndirectStubsManagerBase>
43 class CompileOnDemandLayer {
46 template <typename MaterializerFtor>
47 class LambdaMaterializer final : public ValueMaterializer {
49 LambdaMaterializer(MaterializerFtor M) : M(std::move(M)) {}
50 Value* materializeValueFor(Value *V) final {
57 template <typename MaterializerFtor>
58 LambdaMaterializer<MaterializerFtor>
59 createLambdaMaterializer(MaterializerFtor M) {
60 return LambdaMaterializer<MaterializerFtor>(std::move(M));
63 typedef typename BaseLayerT::ModuleSetHandleT BaseLayerModuleSetHandleT;
65 struct LogicalModuleResources {
66 std::shared_ptr<Module> SourceModule;
67 std::set<const Function*> StubsToClone;
68 std::unique_ptr<IndirectStubsMgrT> StubsMgr;
70 LogicalModuleResources() = default;
72 // Explicit move constructor to make MSVC happy.
73 LogicalModuleResources(LogicalModuleResources &&Other)
74 : SourceModule(std::move(Other.SourceModule)),
75 StubsToClone(std::move(Other.StubsToClone)),
76 StubsMgr(std::move(Other.StubsMgr)) {}
78 // Explicit move assignment to make MSVC happy.
79 LogicalModuleResources& operator=(LogicalModuleResources &&Other) {
80 SourceModule = std::move(Other.SourceModule);
81 StubsToClone = std::move(Other.StubsToClone);
82 StubsMgr = std::move(Other.StubsMgr);
85 JITSymbol findSymbol(StringRef Name, bool ExportedSymbolsOnly) {
86 if (Name.endswith("$stub_ptr") && !ExportedSymbolsOnly) {
87 assert(!ExportedSymbolsOnly && "Stubs are never exported");
88 return StubsMgr->findPointer(Name.drop_back(9));
90 return StubsMgr->findStub(Name, ExportedSymbolsOnly);
95 struct LogicalDylibResources {
96 typedef std::function<RuntimeDyld::SymbolInfo(const std::string&)>
98 SymbolResolverFtor ExternalSymbolResolver;
101 typedef LogicalDylib<BaseLayerT, LogicalModuleResources,
102 LogicalDylibResources> CODLogicalDylib;
104 typedef typename CODLogicalDylib::LogicalModuleHandle LogicalModuleHandle;
105 typedef std::list<CODLogicalDylib> LogicalDylibList;
109 /// @brief Handle to a set of loaded modules.
110 typedef typename LogicalDylibList::iterator ModuleSetHandleT;
112 /// @brief Module partitioning functor.
113 typedef std::function<std::set<Function*>(Function&)> PartitioningFtor;
115 /// @brief Builder for IndirectStubsManagers.
116 typedef std::function<std::unique_ptr<IndirectStubsMgrT>()>
117 IndirectStubsManagerBuilderT;
119 /// @brief Construct a compile-on-demand layer instance.
120 CompileOnDemandLayer(BaseLayerT &BaseLayer, PartitioningFtor Partition,
121 CompileCallbackMgrT &CallbackMgr,
122 IndirectStubsManagerBuilderT CreateIndirectStubsManager,
123 bool CloneStubsIntoPartitions = true)
124 : BaseLayer(BaseLayer), Partition(Partition),
125 CompileCallbackMgr(CallbackMgr),
126 CreateIndirectStubsManager(std::move(CreateIndirectStubsManager)),
127 CloneStubsIntoPartitions(CloneStubsIntoPartitions) {}
129 /// @brief Add a module to the compile-on-demand layer.
130 template <typename ModuleSetT, typename MemoryManagerPtrT,
131 typename SymbolResolverPtrT>
132 ModuleSetHandleT addModuleSet(ModuleSetT Ms,
133 MemoryManagerPtrT MemMgr,
134 SymbolResolverPtrT Resolver) {
136 assert(MemMgr == nullptr &&
137 "User supplied memory managers not supported with COD yet.");
139 LogicalDylibs.push_back(CODLogicalDylib(BaseLayer));
140 auto &LDResources = LogicalDylibs.back().getDylibResources();
142 LDResources.ExternalSymbolResolver =
143 [Resolver](const std::string &Name) {
144 return Resolver->findSymbol(Name);
147 // Process each of the modules in this module set.
149 addLogicalModule(LogicalDylibs.back(),
150 std::shared_ptr<Module>(std::move(M)));
152 return std::prev(LogicalDylibs.end());
155 /// @brief Remove the module represented by the given handle.
157 /// This will remove all modules in the layers below that were derived from
158 /// the module represented by H.
159 void removeModuleSet(ModuleSetHandleT H) {
160 LogicalDylibs.erase(H);
163 /// @brief Search for the given named symbol.
164 /// @param Name The name of the symbol to search for.
165 /// @param ExportedSymbolsOnly If true, search only for exported symbols.
166 /// @return A handle for the given named symbol, if it exists.
167 JITSymbol findSymbol(StringRef Name, bool ExportedSymbolsOnly) {
168 for (auto LDI = LogicalDylibs.begin(), LDE = LogicalDylibs.end();
170 if (auto Symbol = findSymbolIn(LDI, Name, ExportedSymbolsOnly))
172 return BaseLayer.findSymbol(Name, ExportedSymbolsOnly);
175 /// @brief Get the address of a symbol provided by this layer, or some layer
177 JITSymbol findSymbolIn(ModuleSetHandleT H, const std::string &Name,
178 bool ExportedSymbolsOnly) {
179 return H->findSymbol(Name, ExportedSymbolsOnly);
184 void addLogicalModule(CODLogicalDylib &LD, std::shared_ptr<Module> SrcM) {
186 // Bump the linkage and rename any anonymous/privote members in SrcM to
187 // ensure that everything will resolve properly after we partition SrcM.
188 makeAllSymbolsExternallyAccessible(*SrcM);
190 // Create a logical module handle for SrcM within the logical dylib.
191 auto LMH = LD.createLogicalModule();
192 auto &LMResources = LD.getLogicalModuleResources(LMH);
194 LMResources.SourceModule = SrcM;
196 // Create the GlobalValues module.
197 const DataLayout &DL = SrcM->getDataLayout();
198 auto GVsM = llvm::make_unique<Module>((SrcM->getName() + ".globals").str(),
200 GVsM->setDataLayout(DL);
202 // Create function stubs.
203 ValueToValueMapTy VMap;
205 typename IndirectStubsMgrT::StubInitsMap StubInits;
206 for (auto &F : *SrcM) {
207 // Skip declarations.
208 if (F.isDeclaration())
211 // Record all functions defined by this module.
212 if (CloneStubsIntoPartitions)
213 LMResources.StubsToClone.insert(&F);
215 // Create a callback, associate it with the stub for the function,
216 // and set the compile action to compile the partition containing the
218 auto CCInfo = CompileCallbackMgr.getCompileCallback(SrcM->getContext());
219 StubInits[mangle(F.getName(), DL)] =
220 std::make_pair(CCInfo.getAddress(),
221 JITSymbolBase::flagsFromGlobalValue(F));
222 CCInfo.setCompileAction(
223 [this, &LD, LMH, &F]() {
224 return this->extractAndCompile(LD, LMH, F);
228 LMResources.StubsMgr = CreateIndirectStubsManager();
229 auto EC = LMResources.StubsMgr->init(StubInits);
231 // FIXME: This should be propagated back to the user. Stub creation may
232 // fail for remote JITs.
233 assert(!EC && "Error generating stubs");
236 // Clone global variable decls.
237 for (auto &GV : SrcM->globals())
238 if (!GV.isDeclaration() && !VMap.count(&GV))
239 cloneGlobalVariableDecl(*GVsM, GV, &VMap);
242 for (auto &A : SrcM->aliases())
244 cloneGlobalAliasDecl(*GVsM, A, VMap);
246 // Now we need to clone the GV and alias initializers.
248 // Initializers may refer to functions declared (but not defined) in this
249 // module. Build a materializer to clone decls on demand.
250 auto Materializer = createLambdaMaterializer(
251 [this, &GVsM, &LMResources](Value *V) -> Value* {
252 if (auto *F = dyn_cast<Function>(V)) {
253 // Decls in the original module just get cloned.
254 if (F->isDeclaration())
255 return cloneFunctionDecl(*GVsM, *F);
257 // Definitions in the original module (which we have emitted stubs
258 // for at this point) get turned into a constant alias to the stub
260 const DataLayout &DL = GVsM->getDataLayout();
261 std::string FName = mangle(F->getName(), DL);
262 auto StubSym = LMResources.StubsMgr->findStub(FName, false);
263 unsigned PtrBitWidth = DL.getPointerTypeSizeInBits(F->getType());
264 ConstantInt *StubAddr =
265 ConstantInt::get(GVsM->getContext(),
266 APInt(PtrBitWidth, StubSym.getAddress()));
267 Constant *Init = ConstantExpr::getCast(Instruction::IntToPtr,
268 StubAddr, F->getType());
269 return GlobalAlias::create(F->getFunctionType(),
270 F->getType()->getAddressSpace(),
271 F->getLinkage(), F->getName(),
278 // Clone the global variable initializers.
279 for (auto &GV : SrcM->globals())
280 if (!GV.isDeclaration())
281 moveGlobalVariableInitializer(GV, VMap, &Materializer);
283 // Clone the global alias initializers.
284 for (auto &A : SrcM->aliases()) {
285 auto *NewA = cast<GlobalAlias>(VMap[&A]);
286 assert(NewA && "Alias not cloned?");
287 Value *Init = MapValue(A.getAliasee(), VMap, RF_None, nullptr,
289 NewA->setAliasee(cast<Constant>(Init));
292 // Build a resolver for the globals module and add it to the base layer.
293 auto GVsResolver = createLambdaResolver(
294 [&LD, LMH](const std::string &Name) {
295 auto &LMResources = LD.getLogicalModuleResources(LMH);
296 if (auto Sym = LMResources.StubsMgr->findStub(Name, false))
297 return RuntimeDyld::SymbolInfo(Sym.getAddress(), Sym.getFlags());
298 return LD.getDylibResources().ExternalSymbolResolver(Name);
300 [](const std::string &Name) {
301 return RuntimeDyld::SymbolInfo(nullptr);
304 std::vector<std::unique_ptr<Module>> GVsMSet;
305 GVsMSet.push_back(std::move(GVsM));
307 BaseLayer.addModuleSet(std::move(GVsMSet),
308 llvm::make_unique<SectionMemoryManager>(),
309 std::move(GVsResolver));
310 LD.addToLogicalModule(LMH, GVsH);
313 static std::string mangle(StringRef Name, const DataLayout &DL) {
314 std::string MangledName;
316 raw_string_ostream MangledNameStream(MangledName);
317 Mangler::getNameWithPrefix(MangledNameStream, Name, DL);
322 TargetAddress extractAndCompile(CODLogicalDylib &LD,
323 LogicalModuleHandle LMH,
325 auto &LMResources = LD.getLogicalModuleResources(LMH);
326 Module &SrcM = *LMResources.SourceModule;
328 // If F is a declaration we must already have compiled it.
329 if (F.isDeclaration())
332 // Grab the name of the function being called here.
333 std::string CalledFnName = mangle(F.getName(), SrcM.getDataLayout());
335 auto Part = Partition(F);
336 auto PartH = emitPartition(LD, LMH, Part);
338 TargetAddress CalledAddr = 0;
339 for (auto *SubF : Part) {
340 std::string FnName = mangle(SubF->getName(), SrcM.getDataLayout());
341 auto FnBodySym = BaseLayer.findSymbolIn(PartH, FnName, false);
342 assert(FnBodySym && "Couldn't find function body.");
344 TargetAddress FnBodyAddr = FnBodySym.getAddress();
346 // If this is the function we're calling record the address so we can
347 // return it from this function.
349 CalledAddr = FnBodyAddr;
351 // Update the function body pointer for the stub.
352 if (auto EC = LMResources.StubsMgr->updatePointer(FnName, FnBodyAddr))
359 template <typename PartitionT>
360 BaseLayerModuleSetHandleT emitPartition(CODLogicalDylib &LD,
361 LogicalModuleHandle LMH,
362 const PartitionT &Part) {
363 auto &LMResources = LD.getLogicalModuleResources(LMH);
364 Module &SrcM = *LMResources.SourceModule;
366 // Create the module.
367 std::string NewName = SrcM.getName();
368 for (auto *F : Part) {
370 NewName += F->getName();
373 auto M = llvm::make_unique<Module>(NewName, SrcM.getContext());
374 M->setDataLayout(SrcM.getDataLayout());
375 ValueToValueMapTy VMap;
377 auto Materializer = createLambdaMaterializer(
378 [this, &LMResources, &M, &VMap](Value *V) -> Value* {
379 if (auto *GV = dyn_cast<GlobalVariable>(V)) {
380 return cloneGlobalVariableDecl(*M, *GV);
381 } else if (auto *F = dyn_cast<Function>(V)) {
382 // Check whether we want to clone an available_externally definition.
383 if (LMResources.StubsToClone.count(F)) {
384 // Ok - we want an inlinable stub. For that to work we need a decl
385 // for the stub pointer.
386 auto *StubPtr = createImplPointer(*F->getType(), *M,
387 F->getName() + "$stub_ptr",
389 auto *ClonedF = cloneFunctionDecl(*M, *F);
390 makeStub(*ClonedF, *StubPtr);
391 ClonedF->setLinkage(GlobalValue::AvailableExternallyLinkage);
392 ClonedF->addFnAttr(Attribute::AlwaysInline);
396 return cloneFunctionDecl(*M, *F);
397 } else if (auto *A = dyn_cast<GlobalAlias>(V)) {
398 auto *PTy = cast<PointerType>(A->getType());
399 if (PTy->getElementType()->isFunctionTy())
400 return Function::Create(cast<FunctionType>(PTy->getElementType()),
401 GlobalValue::ExternalLinkage,
402 A->getName(), M.get());
404 return new GlobalVariable(*M, PTy->getElementType(), false,
405 GlobalValue::ExternalLinkage,
406 nullptr, A->getName(), nullptr,
407 GlobalValue::NotThreadLocal,
408 PTy->getAddressSpace());
414 // Create decls in the new module.
416 cloneFunctionDecl(*M, *F, &VMap);
418 // Move the function bodies.
420 moveFunctionBody(*F, VMap, &Materializer);
422 // Create memory manager and symbol resolver.
423 auto MemMgr = llvm::make_unique<SectionMemoryManager>();
424 auto Resolver = createLambdaResolver(
425 [this, &LD, LMH](const std::string &Name) {
426 if (auto Symbol = LD.findSymbolInternally(LMH, Name))
427 return RuntimeDyld::SymbolInfo(Symbol.getAddress(),
429 return LD.getDylibResources().ExternalSymbolResolver(Name);
431 [this, &LD, LMH](const std::string &Name) {
432 if (auto Symbol = LD.findSymbolInternally(LMH, Name))
433 return RuntimeDyld::SymbolInfo(Symbol.getAddress(),
435 return RuntimeDyld::SymbolInfo(nullptr);
437 std::vector<std::unique_ptr<Module>> PartMSet;
438 PartMSet.push_back(std::move(M));
439 return BaseLayer.addModuleSet(std::move(PartMSet), std::move(MemMgr),
440 std::move(Resolver));
443 BaseLayerT &BaseLayer;
444 PartitioningFtor Partition;
445 CompileCallbackMgrT &CompileCallbackMgr;
446 IndirectStubsManagerBuilderT CreateIndirectStubsManager;
448 LogicalDylibList LogicalDylibs;
449 bool CloneStubsIntoPartitions;
452 } // End namespace orc.
453 } // End namespace llvm.
455 #endif // LLVM_EXECUTIONENGINE_ORC_COMPILEONDEMANDLAYER_H