diff options
Diffstat (limited to 'lib/ExecutionEngine/Orc/CompileOnDemandLayer.cpp')
| -rw-r--r-- | lib/ExecutionEngine/Orc/CompileOnDemandLayer.cpp | 343 |
1 files changed, 343 insertions, 0 deletions
diff --git a/lib/ExecutionEngine/Orc/CompileOnDemandLayer.cpp b/lib/ExecutionEngine/Orc/CompileOnDemandLayer.cpp new file mode 100644 index 000000000000..d42e7b05ba67 --- /dev/null +++ b/lib/ExecutionEngine/Orc/CompileOnDemandLayer.cpp @@ -0,0 +1,343 @@ +//===----- CompileOnDemandLayer.cpp - Lazily emit IR on first call --------===// +// +// The LLVM Compiler Infrastructure +// +// This file is distributed under the University of Illinois Open Source +// License. See LICENSE.TXT for details. +// +//===----------------------------------------------------------------------===// + +#include "llvm/ExecutionEngine/Orc/CompileOnDemandLayer.h" +#include "llvm/Bitcode/BitcodeReader.h" +#include "llvm/Bitcode/BitcodeWriter.h" +#include "llvm/IR/Mangler.h" +#include "llvm/IR/Module.h" +#include "llvm/Support/raw_ostream.h" +#include "llvm/Transforms/Utils/Cloning.h" + +using namespace llvm; +using namespace llvm::orc; + +namespace { + +template <typename MaterializerFtor> +class LambdaValueMaterializer final : public ValueMaterializer { +public: + LambdaValueMaterializer(MaterializerFtor M) : M(std::move(M)) {} + + Value *materialize(Value *V) final { return M(V); } + +private: + MaterializerFtor M; +}; + +template <typename MaterializerFtor> +LambdaValueMaterializer<MaterializerFtor> +createLambdaValueMaterializer(MaterializerFtor M) { + return LambdaValueMaterializer<MaterializerFtor>(std::move(M)); +} +} // namespace + +static void extractAliases(MaterializationResponsibility &R, Module &M, + MangleAndInterner &Mangle) { + SymbolAliasMap Aliases; + + std::vector<GlobalAlias *> ModAliases; + for (auto &A : M.aliases()) + ModAliases.push_back(&A); + + for (auto *A : ModAliases) { + Constant *Aliasee = A->getAliasee(); + assert(A->hasName() && "Anonymous alias?"); + assert(Aliasee->hasName() && "Anonymous aliasee"); + std::string AliasName = A->getName(); + + Aliases[Mangle(AliasName)] = SymbolAliasMapEntry( + {Mangle(Aliasee->getName()), JITSymbolFlags::fromGlobalValue(*A)}); + + if (isa<Function>(Aliasee)) { + auto *F = cloneFunctionDecl(M, *cast<Function>(Aliasee)); + A->replaceAllUsesWith(F); + A->eraseFromParent(); + F->setName(AliasName); + } else if (isa<GlobalValue>(Aliasee)) { + auto *G = cloneGlobalVariableDecl(M, *cast<GlobalVariable>(Aliasee)); + A->replaceAllUsesWith(G); + A->eraseFromParent(); + G->setName(AliasName); + } + } + + R.replace(symbolAliases(std::move(Aliases))); +} + +static std::unique_ptr<Module> +extractAndClone(Module &M, LLVMContext &NewContext, StringRef Suffix, + function_ref<bool(const GlobalValue *)> ShouldCloneDefinition) { + SmallVector<char, 1> ClonedModuleBuffer; + + { + std::set<GlobalValue *> ClonedDefsInSrc; + ValueToValueMapTy VMap; + auto Tmp = CloneModule(M, VMap, [&](const GlobalValue *GV) { + if (ShouldCloneDefinition(GV)) { + ClonedDefsInSrc.insert(const_cast<GlobalValue *>(GV)); + return true; + } + return false; + }); + + for (auto *GV : ClonedDefsInSrc) { + // Delete the definition and bump the linkage in the source module. + if (isa<Function>(GV)) { + auto &F = *cast<Function>(GV); + F.deleteBody(); + F.setPersonalityFn(nullptr); + } else if (isa<GlobalVariable>(GV)) { + cast<GlobalVariable>(GV)->setInitializer(nullptr); + } else + llvm_unreachable("Unsupported global type"); + + GV->setLinkage(GlobalValue::ExternalLinkage); + } + + BitcodeWriter BCWriter(ClonedModuleBuffer); + + BCWriter.writeModule(*Tmp); + BCWriter.writeSymtab(); + BCWriter.writeStrtab(); + } + + MemoryBufferRef ClonedModuleBufferRef( + StringRef(ClonedModuleBuffer.data(), ClonedModuleBuffer.size()), + "cloned module buffer"); + + auto ClonedModule = + cantFail(parseBitcodeFile(ClonedModuleBufferRef, NewContext)); + ClonedModule->setModuleIdentifier((M.getName() + Suffix).str()); + return ClonedModule; +} + +static std::unique_ptr<Module> extractGlobals(Module &M, + LLVMContext &NewContext) { + return extractAndClone(M, NewContext, ".globals", [](const GlobalValue *GV) { + return isa<GlobalVariable>(GV); + }); +} + +namespace llvm { +namespace orc { + +class ExtractingIRMaterializationUnit : public IRMaterializationUnit { +public: + ExtractingIRMaterializationUnit(ExecutionSession &ES, + CompileOnDemandLayer2 &Parent, + std::unique_ptr<Module> M) + : IRMaterializationUnit(ES, std::move(M)), Parent(Parent) {} + + ExtractingIRMaterializationUnit(std::unique_ptr<Module> M, + SymbolFlagsMap SymbolFlags, + SymbolNameToDefinitionMap SymbolToDefinition, + CompileOnDemandLayer2 &Parent) + : IRMaterializationUnit(std::move(M), std::move(SymbolFlags), + std::move(SymbolToDefinition)), + Parent(Parent) {} + +private: + void materialize(MaterializationResponsibility R) override { + // FIXME: Need a 'notify lazy-extracting/emitting' callback to tie the + // extracted module key, extracted module, and source module key + // together. This could be used, for example, to provide a specific + // memory manager instance to the linking layer. + + auto RequestedSymbols = R.getRequestedSymbols(); + + // Extract the requested functions into a new module. + std::unique_ptr<Module> ExtractedFunctionsModule; + if (!RequestedSymbols.empty()) { + std::string Suffix; + std::set<const GlobalValue *> FunctionsToClone; + for (auto &Name : RequestedSymbols) { + auto I = SymbolToDefinition.find(Name); + assert(I != SymbolToDefinition.end() && I->second != nullptr && + "Should have a non-null definition"); + FunctionsToClone.insert(I->second); + Suffix += "."; + Suffix += *Name; + } + + std::lock_guard<std::mutex> Lock(SourceModuleMutex); + ExtractedFunctionsModule = + extractAndClone(*M, Parent.GetAvailableContext(), Suffix, + [&](const GlobalValue *GV) -> bool { + return FunctionsToClone.count(GV); + }); + } + + // Build a new ExtractingIRMaterializationUnit to delegate the unrequested + // symbols to. + SymbolFlagsMap DelegatedSymbolFlags; + IRMaterializationUnit::SymbolNameToDefinitionMap + DelegatedSymbolToDefinition; + for (auto &KV : SymbolToDefinition) { + if (RequestedSymbols.count(KV.first)) + continue; + DelegatedSymbolFlags[KV.first] = + JITSymbolFlags::fromGlobalValue(*KV.second); + DelegatedSymbolToDefinition[KV.first] = KV.second; + } + + if (!DelegatedSymbolFlags.empty()) { + assert(DelegatedSymbolFlags.size() == + DelegatedSymbolToDefinition.size() && + "SymbolFlags and SymbolToDefinition should have the same number " + "of entries"); + R.replace(llvm::make_unique<ExtractingIRMaterializationUnit>( + std::move(M), std::move(DelegatedSymbolFlags), + std::move(DelegatedSymbolToDefinition), Parent)); + } + + if (ExtractedFunctionsModule) + Parent.emitExtractedFunctionsModule(std::move(R), + std::move(ExtractedFunctionsModule)); + } + + void discard(const VSO &V, SymbolStringPtr Name) override { + // All original symbols were materialized by the CODLayer and should be + // final. The function bodies provided by M should never be overridden. + llvm_unreachable("Discard should never be called on an " + "ExtractingIRMaterializationUnit"); + } + + mutable std::mutex SourceModuleMutex; + CompileOnDemandLayer2 &Parent; +}; + +CompileOnDemandLayer2::CompileOnDemandLayer2( + ExecutionSession &ES, IRLayer &BaseLayer, JITCompileCallbackManager &CCMgr, + IndirectStubsManagerBuilder BuildIndirectStubsManager, + GetAvailableContextFunction GetAvailableContext) + : IRLayer(ES), BaseLayer(BaseLayer), CCMgr(CCMgr), + BuildIndirectStubsManager(std::move(BuildIndirectStubsManager)), + GetAvailableContext(std::move(GetAvailableContext)) {} + +Error CompileOnDemandLayer2::add(VSO &V, VModuleKey K, + std::unique_ptr<Module> M) { + return IRLayer::add(V, K, std::move(M)); +} + +void CompileOnDemandLayer2::emit(MaterializationResponsibility R, VModuleKey K, + std::unique_ptr<Module> M) { + auto &ES = getExecutionSession(); + assert(M && "M should not be null"); + + for (auto &GV : M->global_values()) + if (GV.hasWeakLinkage()) + GV.setLinkage(GlobalValue::ExternalLinkage); + + MangleAndInterner Mangle(ES, M->getDataLayout()); + + extractAliases(R, *M, Mangle); + + auto GlobalsModule = extractGlobals(*M, GetAvailableContext()); + + // Delete the bodies of any available externally functions, rename the + // rest, and build the compile callbacks. + std::map<SymbolStringPtr, std::pair<JITTargetAddress, JITSymbolFlags>> + StubCallbacksAndLinkages; + auto &TargetVSO = R.getTargetVSO(); + + for (auto &F : M->functions()) { + if (F.isDeclaration()) + continue; + + if (F.hasAvailableExternallyLinkage()) { + F.deleteBody(); + F.setPersonalityFn(nullptr); + continue; + } + + assert(F.hasName() && "Function should have a name"); + std::string StubUnmangledName = F.getName(); + F.setName(F.getName() + "$body"); + auto StubDecl = cloneFunctionDecl(*M, F); + StubDecl->setName(StubUnmangledName); + StubDecl->setPersonalityFn(nullptr); + StubDecl->setLinkage(GlobalValue::ExternalLinkage); + F.replaceAllUsesWith(StubDecl); + + auto StubName = Mangle(StubUnmangledName); + auto BodyName = Mangle(F.getName()); + if (auto CallbackAddr = CCMgr.getCompileCallback( + [BodyName, &TargetVSO, &ES]() -> JITTargetAddress { + if (auto Sym = lookup({&TargetVSO}, BodyName)) + return Sym->getAddress(); + else { + ES.reportError(Sym.takeError()); + return 0; + } + })) { + auto Flags = JITSymbolFlags::fromGlobalValue(F); + Flags &= ~JITSymbolFlags::Weak; + StubCallbacksAndLinkages[std::move(StubName)] = + std::make_pair(*CallbackAddr, Flags); + } else { + ES.reportError(CallbackAddr.takeError()); + R.failMaterialization(); + return; + } + } + + // Build the stub inits map. + IndirectStubsManager::StubInitsMap StubInits; + for (auto &KV : StubCallbacksAndLinkages) + StubInits[*KV.first] = KV.second; + + // Build the function-body-extracting materialization unit. + if (auto Err = R.getTargetVSO().define( + llvm::make_unique<ExtractingIRMaterializationUnit>(ES, *this, + std::move(M)))) { + ES.reportError(std::move(Err)); + R.failMaterialization(); + return; + } + + // Build the stubs. + // FIXME: Remove function bodies materialization unit if stub creation fails. + auto &StubsMgr = getStubsManager(TargetVSO); + if (auto Err = StubsMgr.createStubs(StubInits)) { + ES.reportError(std::move(Err)); + R.failMaterialization(); + return; + } + + // Resolve and finalize stubs. + SymbolMap ResolvedStubs; + for (auto &KV : StubCallbacksAndLinkages) { + if (auto Sym = StubsMgr.findStub(*KV.first, false)) + ResolvedStubs[KV.first] = Sym; + else + llvm_unreachable("Stub went missing"); + } + + R.resolve(ResolvedStubs); + + BaseLayer.emit(std::move(R), std::move(K), std::move(GlobalsModule)); +} + +IndirectStubsManager &CompileOnDemandLayer2::getStubsManager(const VSO &V) { + std::lock_guard<std::mutex> Lock(CODLayerMutex); + StubManagersMap::iterator I = StubsMgrs.find(&V); + if (I == StubsMgrs.end()) + I = StubsMgrs.insert(std::make_pair(&V, BuildIndirectStubsManager())).first; + return *I->second; +} + +void CompileOnDemandLayer2::emitExtractedFunctionsModule( + MaterializationResponsibility R, std::unique_ptr<Module> M) { + auto K = getExecutionSession().allocateVModule(); + BaseLayer.emit(std::move(R), std::move(K), std::move(M)); +} + +} // end namespace orc +} // end namespace llvm |