diff --git a/llvm/examples/Kaleidoscope/BuildingAJIT/Chapter3/KaleidoscopeJIT.h b/llvm/examples/Kaleidoscope/BuildingAJIT/Chapter3/KaleidoscopeJIT.h index fd0e081ff2b4bd..fee2d26e5d9254 100644 --- a/llvm/examples/Kaleidoscope/BuildingAJIT/Chapter3/KaleidoscopeJIT.h +++ b/llvm/examples/Kaleidoscope/BuildingAJIT/Chapter3/KaleidoscopeJIT.h @@ -21,6 +21,7 @@ #include "llvm/ExecutionEngine/Orc/ExecutionUtils.h" #include "llvm/ExecutionEngine/Orc/ExecutorProcessControl.h" #include "llvm/ExecutionEngine/Orc/IRCompileLayer.h" +#include "llvm/ExecutionEngine/Orc/IRPartitionLayer.h" #include "llvm/ExecutionEngine/Orc/IRTransformLayer.h" #include "llvm/ExecutionEngine/Orc/JITTargetMachineBuilder.h" #include "llvm/ExecutionEngine/Orc/RTDyldObjectLinkingLayer.h" @@ -48,6 +49,7 @@ class KaleidoscopeJIT { RTDyldObjectLinkingLayer ObjectLayer; IRCompileLayer CompileLayer; IRTransformLayer OptimizeLayer; + IRPartitionLayer IPLayer; CompileOnDemandLayer CODLayer; JITDylib &MainJD; @@ -68,8 +70,8 @@ class KaleidoscopeJIT { CompileLayer(*this->ES, ObjectLayer, std::make_unique(std::move(JTMB))), OptimizeLayer(*this->ES, CompileLayer, optimizeModule), - CODLayer(*this->ES, OptimizeLayer, - this->EPCIU->getLazyCallThroughManager(), + IPLayer(*this->ES, OptimizeLayer), + CODLayer(*this->ES, IPLayer, this->EPCIU->getLazyCallThroughManager(), [this] { return this->EPCIU->createIndirectStubsManager(); }), MainJD(this->ES->createBareJITDylib("
")) { MainJD.addGenerator( diff --git a/llvm/examples/SpeculativeJIT/SpeculativeJIT.cpp b/llvm/examples/SpeculativeJIT/SpeculativeJIT.cpp index 1659e5c5c8b40b..24ae953eca7a53 100644 --- a/llvm/examples/SpeculativeJIT/SpeculativeJIT.cpp +++ b/llvm/examples/SpeculativeJIT/SpeculativeJIT.cpp @@ -3,6 +3,7 @@ #include "llvm/ExecutionEngine/Orc/Core.h" #include "llvm/ExecutionEngine/Orc/ExecutionUtils.h" #include "llvm/ExecutionEngine/Orc/IRCompileLayer.h" +#include "llvm/ExecutionEngine/Orc/IRPartitionLayer.h" #include "llvm/ExecutionEngine/Orc/IndirectionUtils.h" #include "llvm/ExecutionEngine/Orc/JITTargetMachineBuilder.h" #include "llvm/ExecutionEngine/Orc/RTDyldObjectLinkingLayer.h" @@ -109,13 +110,14 @@ class SpeculativeJIT { IndirectStubsManagerBuilderFunction ISMBuilder, std::unique_ptr ProcessSymbolsGenerator) : ES(std::move(ES)), DL(std::move(DL)), - MainJD(this->ES->createBareJITDylib("
")), LCTMgr(std::move(LCTMgr)), + MainJD(this->ES->createBareJITDylib("
")), + LCTMgr(std::move(LCTMgr)), CompileLayer(*this->ES, ObjLayer, std::make_unique(std::move(JTMB))), S(Imps, *this->ES), SpeculateLayer(*this->ES, CompileLayer, S, Mangle, BlockFreqQuery()), - CODLayer(*this->ES, SpeculateLayer, *this->LCTMgr, - std::move(ISMBuilder)) { + IPLayer(*this->ES, SpeculateLayer), + CODLayer(*this->ES, IPLayer, *this->LCTMgr, std::move(ISMBuilder)) { MainJD.addGenerator(std::move(ProcessSymbolsGenerator)); this->CODLayer.setImplMap(&Imps); ExitOnErr(S.addSpeculationRuntime(MainJD, Mangle)); @@ -141,6 +143,7 @@ class SpeculativeJIT { Speculator S; RTDyldObjectLinkingLayer ObjLayer{*ES, createMemMgr}; IRSpeculationLayer SpeculateLayer; + IRPartitionLayer IPLayer; CompileOnDemandLayer CODLayer; }; diff --git a/llvm/include/llvm/ExecutionEngine/Orc/CompileOnDemandLayer.h b/llvm/include/llvm/ExecutionEngine/Orc/CompileOnDemandLayer.h index 41f2882c576e13..d492dfe29fba20 100644 --- a/llvm/include/llvm/ExecutionEngine/Orc/CompileOnDemandLayer.h +++ b/llvm/include/llvm/ExecutionEngine/Orc/CompileOnDemandLayer.h @@ -53,37 +53,15 @@ namespace llvm { namespace orc { class CompileOnDemandLayer : public IRLayer { - friend class PartitioningIRMaterializationUnit; - public: /// Builder for IndirectStubsManagers. using IndirectStubsManagerBuilder = std::function()>; - using GlobalValueSet = std::set; - - /// Partitioning function. - using PartitionFunction = - std::function(GlobalValueSet Requested)>; - - /// Off-the-shelf partitioning which compiles all requested symbols (usually - /// a single function at a time). - static std::optional - compileRequested(GlobalValueSet Requested); - - /// Off-the-shelf partitioning which compiles whole modules whenever any - /// symbol in them is requested. - static std::optional - compileWholeModule(GlobalValueSet Requested); - /// Construct a CompileOnDemandLayer. CompileOnDemandLayer(ExecutionSession &ES, IRLayer &BaseLayer, - LazyCallThroughManager &LCTMgr, - IndirectStubsManagerBuilder BuildIndirectStubsManager); - - /// Sets the partition function. - void setPartitionFunction(PartitionFunction Partition); - + LazyCallThroughManager &LCTMgr, + IndirectStubsManagerBuilder BuildIndirectStubsManager); /// Sets the ImplSymbolMap void setImplMap(ImplSymbolMap *Imp); @@ -110,22 +88,12 @@ class CompileOnDemandLayer : public IRLayer { PerDylibResources &getPerDylibResources(JITDylib &TargetD); - void cleanUpModule(Module &M); - - void expandPartition(GlobalValueSet &Partition); - - void emitPartition(std::unique_ptr R, - ThreadSafeModule TSM, - IRMaterializationUnit::SymbolNameToDefinitionMap Defs); - mutable std::mutex CODLayerMutex; IRLayer &BaseLayer; LazyCallThroughManager &LCTMgr; IndirectStubsManagerBuilder BuildIndirectStubsManager; PerDylibResourcesMap DylibResources; - PartitionFunction Partition = compileRequested; - SymbolLinkagePromoter PromoteSymbols; ImplSymbolMap *AliaseeImpls = nullptr; }; diff --git a/llvm/include/llvm/ExecutionEngine/Orc/IRPartitionLayer.h b/llvm/include/llvm/ExecutionEngine/Orc/IRPartitionLayer.h new file mode 100644 index 00000000000000..a2981fea61d3ea --- /dev/null +++ b/llvm/include/llvm/ExecutionEngine/Orc/IRPartitionLayer.h @@ -0,0 +1,85 @@ +//===- IRPartitionLayer.h - Partition IR module on lookup -------*- C++ -*-===// +// +// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions. +// See https://llvm.org/LICENSE.txt for license information. +// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception +// +//===----------------------------------------------------------------------===// +// +// JIT layer for breaking up modules into smaller submodules that only contains +// looked up symbols. +// +//===----------------------------------------------------------------------===// + +#ifndef LLVM_EXECUTIONENGINE_ORC_IRPARTITIONLAYER_H +#define LLVM_EXECUTIONENGINE_ORC_IRPARTITIONLAYER_H + +#include "llvm/ExecutionEngine/Orc/IndirectionUtils.h" +#include "llvm/ExecutionEngine/Orc/Layer.h" +#include "llvm/IR/Attributes.h" +#include "llvm/IR/Constant.h" +#include "llvm/IR/Constants.h" +#include "llvm/IR/DataLayout.h" +#include "llvm/IR/Function.h" +#include "llvm/IR/GlobalAlias.h" +#include "llvm/IR/GlobalValue.h" +#include "llvm/IR/GlobalVariable.h" +#include "llvm/IR/Instruction.h" +#include "llvm/IR/Mangler.h" +#include "llvm/IR/Module.h" +#include "llvm/IR/Type.h" + +namespace llvm { +namespace orc { + +/// A layer that breaks up IR modules into smaller submodules that only contains +/// looked up symbols. +class IRPartitionLayer : public IRLayer { + friend class PartitioningIRMaterializationUnit; + +public: + using GlobalValueSet = std::set; + + /// Partitioning function. + using PartitionFunction = + std::function(GlobalValueSet Requested)>; + + /// Construct a IRPartitionLayer. + IRPartitionLayer(ExecutionSession &ES, IRLayer &BaseLayer); + + /// Off-the-shelf partitioning which compiles all requested symbols (usually + /// a single function at a time). + static std::optional + compileRequested(GlobalValueSet Requested); + + /// Off-the-shelf partitioning which compiles whole modules whenever any + /// symbol in them is requested. + static std::optional + compileWholeModule(GlobalValueSet Requested); + + /// Sets the partition function. + void setPartitionFunction(PartitionFunction Partition); + + /// Emits the given module. This should not be called by clients: it will be + /// called by the JIT when a definition added via the add method is requested. + void emit(std::unique_ptr R, + ThreadSafeModule TSM) override; + +private: + void cleanUpModule(Module &M); + + void expandPartition(GlobalValueSet &Partition); + + void emitPartition(std::unique_ptr R, + ThreadSafeModule TSM, + IRMaterializationUnit::SymbolNameToDefinitionMap Defs); + + IRLayer &BaseLayer; + PartitionFunction Partition = compileRequested; + SymbolLinkagePromoter PromoteSymbols; +}; + +} // namespace orc +} // namespace llvm + +#endif diff --git a/llvm/include/llvm/ExecutionEngine/Orc/LLJIT.h b/llvm/include/llvm/ExecutionEngine/Orc/LLJIT.h index 2660b9f74f405a..a2364b4515f01b 100644 --- a/llvm/include/llvm/ExecutionEngine/Orc/LLJIT.h +++ b/llvm/include/llvm/ExecutionEngine/Orc/LLJIT.h @@ -18,6 +18,7 @@ #include "llvm/ExecutionEngine/Orc/CompileUtils.h" #include "llvm/ExecutionEngine/Orc/ExecutionUtils.h" #include "llvm/ExecutionEngine/Orc/IRCompileLayer.h" +#include "llvm/ExecutionEngine/Orc/IRPartitionLayer.h" #include "llvm/ExecutionEngine/Orc/IRTransformLayer.h" #include "llvm/ExecutionEngine/Orc/JITTargetMachineBuilder.h" #include "llvm/ExecutionEngine/Orc/ThreadSafeModule.h" @@ -271,9 +272,8 @@ class LLLazyJIT : public LLJIT { public: /// Sets the partition function. - void - setPartitionFunction(CompileOnDemandLayer::PartitionFunction Partition) { - CODLayer->setPartitionFunction(std::move(Partition)); + void setPartitionFunction(IRPartitionLayer::PartitionFunction Partition) { + IPLayer->setPartitionFunction(std::move(Partition)); } /// Returns a reference to the on-demand layer. @@ -293,6 +293,7 @@ class LLLazyJIT : public LLJIT { LLLazyJIT(LLLazyJITBuilderState &S, Error &Err); std::unique_ptr LCTMgr; + std::unique_ptr IPLayer; std::unique_ptr CODLayer; }; diff --git a/llvm/lib/ExecutionEngine/Orc/CMakeLists.txt b/llvm/lib/ExecutionEngine/Orc/CMakeLists.txt index 0ee056e9f63a19..c07e6293ad1464 100644 --- a/llvm/lib/ExecutionEngine/Orc/CMakeLists.txt +++ b/llvm/lib/ExecutionEngine/Orc/CMakeLists.txt @@ -26,6 +26,7 @@ add_llvm_component_library(LLVMOrcJIT IndirectionUtils.cpp IRCompileLayer.cpp IRTransformLayer.cpp + IRPartitionLayer.cpp JITTargetMachineBuilder.cpp LazyReexports.cpp Layer.cpp diff --git a/llvm/lib/ExecutionEngine/Orc/CompileOnDemandLayer.cpp b/llvm/lib/ExecutionEngine/Orc/CompileOnDemandLayer.cpp index 6448adaa0ceb36..9296bc2b389ab2 100644 --- a/llvm/lib/ExecutionEngine/Orc/CompileOnDemandLayer.cpp +++ b/llvm/lib/ExecutionEngine/Orc/CompileOnDemandLayer.cpp @@ -9,6 +9,7 @@ #include "llvm/ExecutionEngine/Orc/CompileOnDemandLayer.h" #include "llvm/ADT/Hashing.h" #include "llvm/ExecutionEngine/Orc/ExecutionUtils.h" +#include "llvm/ExecutionEngine/Orc/Layer.h" #include "llvm/IR/Mangler.h" #include "llvm/IR/Module.h" #include "llvm/Support/FormatVariadic.h" @@ -17,101 +18,6 @@ using namespace llvm; using namespace llvm::orc; -static ThreadSafeModule extractSubModule(ThreadSafeModule &TSM, - StringRef Suffix, - GVPredicate ShouldExtract) { - - auto DeleteExtractedDefs = [](GlobalValue &GV) { - // Bump the linkage: this global will be provided by the external module. - GV.setLinkage(GlobalValue::ExternalLinkage); - - // Delete the definition in the source module. - if (isa(GV)) { - auto &F = cast(GV); - F.deleteBody(); - F.setPersonalityFn(nullptr); - } else if (isa(GV)) { - cast(GV).setInitializer(nullptr); - } else if (isa(GV)) { - // We need to turn deleted aliases into function or variable decls based - // on the type of their aliasee. - auto &A = cast(GV); - Constant *Aliasee = A.getAliasee(); - assert(A.hasName() && "Anonymous alias?"); - assert(Aliasee->hasName() && "Anonymous aliasee"); - std::string AliasName = std::string(A.getName()); - - if (isa(Aliasee)) { - auto *F = cloneFunctionDecl(*A.getParent(), *cast(Aliasee)); - A.replaceAllUsesWith(F); - A.eraseFromParent(); - F->setName(AliasName); - } else if (isa(Aliasee)) { - auto *G = cloneGlobalVariableDecl(*A.getParent(), - *cast(Aliasee)); - A.replaceAllUsesWith(G); - A.eraseFromParent(); - G->setName(AliasName); - } else - llvm_unreachable("Alias to unsupported type"); - } else - llvm_unreachable("Unsupported global type"); - }; - - auto NewTSM = cloneToNewContext(TSM, ShouldExtract, DeleteExtractedDefs); - NewTSM.withModuleDo([&](Module &M) { - M.setModuleIdentifier((M.getModuleIdentifier() + Suffix).str()); - }); - - return NewTSM; -} - -namespace llvm { -namespace orc { - -class PartitioningIRMaterializationUnit : public IRMaterializationUnit { -public: - PartitioningIRMaterializationUnit(ExecutionSession &ES, - const IRSymbolMapper::ManglingOptions &MO, - ThreadSafeModule TSM, - CompileOnDemandLayer &Parent) - : IRMaterializationUnit(ES, MO, std::move(TSM)), Parent(Parent) {} - - PartitioningIRMaterializationUnit( - ThreadSafeModule TSM, Interface I, - SymbolNameToDefinitionMap SymbolToDefinition, - CompileOnDemandLayer &Parent) - : IRMaterializationUnit(std::move(TSM), std::move(I), - std::move(SymbolToDefinition)), - Parent(Parent) {} - -private: - void materialize(std::unique_ptr R) override { - Parent.emitPartition(std::move(R), std::move(TSM), - std::move(SymbolToDefinition)); - } - - void discard(const JITDylib &V, const 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; - CompileOnDemandLayer &Parent; -}; - -std::optional -CompileOnDemandLayer::compileRequested(GlobalValueSet Requested) { - return std::move(Requested); -} - -std::optional -CompileOnDemandLayer::compileWholeModule(GlobalValueSet Requested) { - return std::nullopt; -} - CompileOnDemandLayer::CompileOnDemandLayer( ExecutionSession &ES, IRLayer &BaseLayer, LazyCallThroughManager &LCTMgr, IndirectStubsManagerBuilder BuildIndirectStubsManager) @@ -119,13 +25,10 @@ CompileOnDemandLayer::CompileOnDemandLayer( LCTMgr(LCTMgr), BuildIndirectStubsManager(std::move(BuildIndirectStubsManager)) {} -void CompileOnDemandLayer::setPartitionFunction(PartitionFunction Partition) { - this->Partition = std::move(Partition); -} - void CompileOnDemandLayer::setImplMap(ImplSymbolMap *Imp) { this->AliaseeImpls = Imp; } + void CompileOnDemandLayer::emit( std::unique_ptr R, ThreadSafeModule TSM) { assert(TSM && "Null module"); @@ -138,10 +41,6 @@ void CompileOnDemandLayer::emit( SymbolAliasMap NonCallables; SymbolAliasMap Callables; - TSM.withModuleDo([&](Module &M) { - // First, do some cleanup on the module: - cleanUpModule(M); - }); for (auto &KV : R->getSymbols()) { auto &Name = KV.first; @@ -152,11 +51,10 @@ void CompileOnDemandLayer::emit( NonCallables[Name] = SymbolAliasMapEntry(Name, Flags); } - // Create a partitioning materialization unit and lodge it with the - // implementation dylib. + // Lodge symbols with the implementation dylib. if (auto Err = PDR.getImplDylib().define( - std::make_unique( - ES, *getManglingOptions(), std::move(TSM), *this))) { + std::make_unique( + BaseLayer, *getManglingOptions(), std::move(TSM)))) { ES.reportError(std::move(Err)); R->failMaterialization(); return; @@ -210,173 +108,3 @@ CompileOnDemandLayer::getPerDylibResources(JITDylib &TargetD) { return I->second; } - -void CompileOnDemandLayer::cleanUpModule(Module &M) { - for (auto &F : M.functions()) { - if (F.isDeclaration()) - continue; - - if (F.hasAvailableExternallyLinkage()) { - F.deleteBody(); - F.setPersonalityFn(nullptr); - continue; - } - } -} - -void CompileOnDemandLayer::expandPartition(GlobalValueSet &Partition) { - // Expands the partition to ensure the following rules hold: - // (1) If any alias is in the partition, its aliasee is also in the partition. - // (2) If any aliasee is in the partition, its aliases are also in the - // partiton. - // (3) If any global variable is in the partition then all global variables - // are in the partition. - assert(!Partition.empty() && "Unexpected empty partition"); - - const Module &M = *(*Partition.begin())->getParent(); - bool ContainsGlobalVariables = false; - std::vector GVsToAdd; - - for (const auto *GV : Partition) - if (isa(GV)) - GVsToAdd.push_back( - cast(cast(GV)->getAliasee())); - else if (isa(GV)) - ContainsGlobalVariables = true; - - for (auto &A : M.aliases()) - if (Partition.count(cast(A.getAliasee()))) - GVsToAdd.push_back(&A); - - if (ContainsGlobalVariables) - for (auto &G : M.globals()) - GVsToAdd.push_back(&G); - - for (const auto *GV : GVsToAdd) - Partition.insert(GV); -} - -void CompileOnDemandLayer::emitPartition( - std::unique_ptr R, ThreadSafeModule TSM, - IRMaterializationUnit::SymbolNameToDefinitionMap Defs) { - - // 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 &ES = getExecutionSession(); - GlobalValueSet RequestedGVs; - for (auto &Name : R->getRequestedSymbols()) { - if (Name == R->getInitializerSymbol()) - TSM.withModuleDo([&](Module &M) { - for (auto &GV : getStaticInitGVs(M)) - RequestedGVs.insert(&GV); - }); - else { - assert(Defs.count(Name) && "No definition for symbol"); - RequestedGVs.insert(Defs[Name]); - } - } - - /// Perform partitioning with the context lock held, since the partition - /// function is allowed to access the globals to compute the partition. - auto GVsToExtract = - TSM.withModuleDo([&](Module &M) { return Partition(RequestedGVs); }); - - // Take a 'None' partition to mean the whole module (as opposed to an empty - // partition, which means "materialize nothing"). Emit the whole module - // unmodified to the base layer. - if (GVsToExtract == std::nullopt) { - Defs.clear(); - BaseLayer.emit(std::move(R), std::move(TSM)); - return; - } - - // If the partition is empty, return the whole module to the symbol table. - if (GVsToExtract->empty()) { - if (auto Err = - R->replace(std::make_unique( - std::move(TSM), - MaterializationUnit::Interface(R->getSymbols(), - R->getInitializerSymbol()), - std::move(Defs), *this))) { - getExecutionSession().reportError(std::move(Err)); - R->failMaterialization(); - return; - } - return; - } - - // Ok -- we actually need to partition the symbols. Promote the symbol - // linkages/names, expand the partition to include any required symbols - // (i.e. symbols that can't be separated from our partition), and - // then extract the partition. - // - // FIXME: We apply this promotion once per partitioning. It's safe, but - // overkill. - auto ExtractedTSM = - TSM.withModuleDo([&](Module &M) -> Expected { - auto PromotedGlobals = PromoteSymbols(M); - if (!PromotedGlobals.empty()) { - - MangleAndInterner Mangle(ES, M.getDataLayout()); - SymbolFlagsMap SymbolFlags; - IRSymbolMapper::add(ES, *getManglingOptions(), - PromotedGlobals, SymbolFlags); - - if (auto Err = R->defineMaterializing(SymbolFlags)) - return std::move(Err); - } - - expandPartition(*GVsToExtract); - - // Submodule name is given by hashing the names of the globals. - std::string SubModuleName; - { - std::vector HashGVs; - HashGVs.reserve(GVsToExtract->size()); - for (const auto *GV : *GVsToExtract) - HashGVs.push_back(GV); - llvm::sort(HashGVs, [](const GlobalValue *LHS, const GlobalValue *RHS) { - return LHS->getName() < RHS->getName(); - }); - hash_code HC(0); - for (const auto *GV : HashGVs) { - assert(GV->hasName() && "All GVs to extract should be named by now"); - auto GVName = GV->getName(); - HC = hash_combine(HC, hash_combine_range(GVName.begin(), GVName.end())); - } - raw_string_ostream(SubModuleName) - << ".submodule." - << formatv(sizeof(size_t) == 8 ? "{0:x16}" : "{0:x8}", - static_cast(HC)) - << ".ll"; - } - - // Extract the requested partiton (plus any necessary aliases) and - // put the rest back into the impl dylib. - auto ShouldExtract = [&](const GlobalValue &GV) -> bool { - return GVsToExtract->count(&GV); - }; - - return extractSubModule(TSM, SubModuleName , ShouldExtract); - }); - - if (!ExtractedTSM) { - ES.reportError(ExtractedTSM.takeError()); - R->failMaterialization(); - return; - } - - if (auto Err = R->replace(std::make_unique( - ES, *getManglingOptions(), std::move(TSM), *this))) { - ES.reportError(std::move(Err)); - R->failMaterialization(); - return; - } - BaseLayer.emit(std::move(R), std::move(*ExtractedTSM)); -} - -} // end namespace orc -} // end namespace llvm diff --git a/llvm/lib/ExecutionEngine/Orc/IRPartitionLayer.cpp b/llvm/lib/ExecutionEngine/Orc/IRPartitionLayer.cpp new file mode 100644 index 00000000000000..9ad171beac7fe2 --- /dev/null +++ b/llvm/lib/ExecutionEngine/Orc/IRPartitionLayer.cpp @@ -0,0 +1,303 @@ +//===----- IRPartitionLayer.cpp - Partition IR module into submodules -----===// +// +// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions. +// See https://llvm.org/LICENSE.txt for license information. +// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception +// +//===----------------------------------------------------------------------===// + +#include "llvm/ExecutionEngine/Orc/IRPartitionLayer.h" +#include "llvm/ExecutionEngine/Orc/ExecutionUtils.h" +#include "llvm/ExecutionEngine/Orc/IndirectionUtils.h" + +using namespace llvm; +using namespace llvm::orc; + +static ThreadSafeModule extractSubModule(ThreadSafeModule &TSM, + StringRef Suffix, + GVPredicate ShouldExtract) { + + auto DeleteExtractedDefs = [](GlobalValue &GV) { + // Bump the linkage: this global will be provided by the external module. + GV.setLinkage(GlobalValue::ExternalLinkage); + + // Delete the definition in the source module. + if (isa(GV)) { + auto &F = cast(GV); + F.deleteBody(); + F.setPersonalityFn(nullptr); + } else if (isa(GV)) { + cast(GV).setInitializer(nullptr); + } else if (isa(GV)) { + // We need to turn deleted aliases into function or variable decls based + // on the type of their aliasee. + auto &A = cast(GV); + Constant *Aliasee = A.getAliasee(); + assert(A.hasName() && "Anonymous alias?"); + assert(Aliasee->hasName() && "Anonymous aliasee"); + std::string AliasName = std::string(A.getName()); + + if (isa(Aliasee)) { + auto *F = cloneFunctionDecl(*A.getParent(), *cast(Aliasee)); + A.replaceAllUsesWith(F); + A.eraseFromParent(); + F->setName(AliasName); + } else if (isa(Aliasee)) { + auto *G = cloneGlobalVariableDecl(*A.getParent(), + *cast(Aliasee)); + A.replaceAllUsesWith(G); + A.eraseFromParent(); + G->setName(AliasName); + } else + llvm_unreachable("Alias to unsupported type"); + } else + llvm_unreachable("Unsupported global type"); + }; + + auto NewTSM = cloneToNewContext(TSM, ShouldExtract, DeleteExtractedDefs); + NewTSM.withModuleDo([&](Module &M) { + M.setModuleIdentifier((M.getModuleIdentifier() + Suffix).str()); + }); + + return NewTSM; +} + +namespace llvm { +namespace orc { + +class PartitioningIRMaterializationUnit : public IRMaterializationUnit { +public: + PartitioningIRMaterializationUnit(ExecutionSession &ES, + const IRSymbolMapper::ManglingOptions &MO, + ThreadSafeModule TSM, + IRPartitionLayer &Parent) + : IRMaterializationUnit(ES, MO, std::move(TSM)), Parent(Parent) {} + + PartitioningIRMaterializationUnit( + ThreadSafeModule TSM, Interface I, + SymbolNameToDefinitionMap SymbolToDefinition, IRPartitionLayer &Parent) + : IRMaterializationUnit(std::move(TSM), std::move(I), + std::move(SymbolToDefinition)), + Parent(Parent) {} + +private: + void materialize(std::unique_ptr R) override { + Parent.emitPartition(std::move(R), std::move(TSM), + std::move(SymbolToDefinition)); + } + + void discard(const JITDylib &V, const 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"); + } + + IRPartitionLayer &Parent; +}; + +} // namespace orc +} // namespace llvm + +IRPartitionLayer::IRPartitionLayer(ExecutionSession &ES, IRLayer &BaseLayer) + : IRLayer(ES, BaseLayer.getManglingOptions()), BaseLayer(BaseLayer) {} + +void IRPartitionLayer::setPartitionFunction(PartitionFunction Partition) { + this->Partition = Partition; +} + +std::optional +IRPartitionLayer::compileRequested(GlobalValueSet Requested) { + return std::move(Requested); +} + +std::optional +IRPartitionLayer::compileWholeModule(GlobalValueSet Requested) { + return std::nullopt; +} + +void IRPartitionLayer::emit(std::unique_ptr R, + ThreadSafeModule TSM) { + assert(TSM && "Null module"); + + auto &ES = getExecutionSession(); + TSM.withModuleDo([&](Module &M) { + // First, do some cleanup on the module: + cleanUpModule(M); + }); + + // Create a partitioning materialization unit and pass the responsibility. + if (auto Err = R->replace(std::make_unique( + ES, *getManglingOptions(), std::move(TSM), *this))) { + ES.reportError(std::move(Err)); + R->failMaterialization(); + return; + } +} + +void IRPartitionLayer::cleanUpModule(Module &M) { + for (auto &F : M.functions()) { + if (F.isDeclaration()) + continue; + + if (F.hasAvailableExternallyLinkage()) { + F.deleteBody(); + F.setPersonalityFn(nullptr); + continue; + } + } +} + +void IRPartitionLayer::expandPartition(GlobalValueSet &Partition) { + // Expands the partition to ensure the following rules hold: + // (1) If any alias is in the partition, its aliasee is also in the partition. + // (2) If any aliasee is in the partition, its aliases are also in the + // partiton. + // (3) If any global variable is in the partition then all global variables + // are in the partition. + assert(!Partition.empty() && "Unexpected empty partition"); + + const Module &M = *(*Partition.begin())->getParent(); + bool ContainsGlobalVariables = false; + std::vector GVsToAdd; + + for (const auto *GV : Partition) + if (isa(GV)) + GVsToAdd.push_back( + cast(cast(GV)->getAliasee())); + else if (isa(GV)) + ContainsGlobalVariables = true; + + for (auto &A : M.aliases()) + if (Partition.count(cast(A.getAliasee()))) + GVsToAdd.push_back(&A); + + if (ContainsGlobalVariables) + for (auto &G : M.globals()) + GVsToAdd.push_back(&G); + + for (const auto *GV : GVsToAdd) + Partition.insert(GV); +} + +void IRPartitionLayer::emitPartition( + std::unique_ptr R, ThreadSafeModule TSM, + IRMaterializationUnit::SymbolNameToDefinitionMap Defs) { + + // 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 &ES = getExecutionSession(); + GlobalValueSet RequestedGVs; + for (auto &Name : R->getRequestedSymbols()) { + if (Name == R->getInitializerSymbol()) + TSM.withModuleDo([&](Module &M) { + for (auto &GV : getStaticInitGVs(M)) + RequestedGVs.insert(&GV); + }); + else { + assert(Defs.count(Name) && "No definition for symbol"); + RequestedGVs.insert(Defs[Name]); + } + } + + /// Perform partitioning with the context lock held, since the partition + /// function is allowed to access the globals to compute the partition. + auto GVsToExtract = + TSM.withModuleDo([&](Module &M) { return Partition(RequestedGVs); }); + + // Take a 'None' partition to mean the whole module (as opposed to an empty + // partition, which means "materialize nothing"). Emit the whole module + // unmodified to the base layer. + if (GVsToExtract == std::nullopt) { + Defs.clear(); + BaseLayer.emit(std::move(R), std::move(TSM)); + return; + } + + // If the partition is empty, return the whole module to the symbol table. + if (GVsToExtract->empty()) { + if (auto Err = + R->replace(std::make_unique( + std::move(TSM), + MaterializationUnit::Interface(R->getSymbols(), + R->getInitializerSymbol()), + std::move(Defs), *this))) { + getExecutionSession().reportError(std::move(Err)); + R->failMaterialization(); + return; + } + return; + } + + // Ok -- we actually need to partition the symbols. Promote the symbol + // linkages/names, expand the partition to include any required symbols + // (i.e. symbols that can't be separated from our partition), and + // then extract the partition. + // + // FIXME: We apply this promotion once per partitioning. It's safe, but + // overkill. + auto ExtractedTSM = TSM.withModuleDo([&](Module &M) + -> Expected { + auto PromotedGlobals = PromoteSymbols(M); + if (!PromotedGlobals.empty()) { + + MangleAndInterner Mangle(ES, M.getDataLayout()); + SymbolFlagsMap SymbolFlags; + IRSymbolMapper::add(ES, *getManglingOptions(), PromotedGlobals, + SymbolFlags); + + if (auto Err = R->defineMaterializing(SymbolFlags)) + return std::move(Err); + } + + expandPartition(*GVsToExtract); + + // Submodule name is given by hashing the names of the globals. + std::string SubModuleName; + { + std::vector HashGVs; + HashGVs.reserve(GVsToExtract->size()); + for (const auto *GV : *GVsToExtract) + HashGVs.push_back(GV); + llvm::sort(HashGVs, [](const GlobalValue *LHS, const GlobalValue *RHS) { + return LHS->getName() < RHS->getName(); + }); + hash_code HC(0); + for (const auto *GV : HashGVs) { + assert(GV->hasName() && "All GVs to extract should be named by now"); + auto GVName = GV->getName(); + HC = hash_combine(HC, hash_combine_range(GVName.begin(), GVName.end())); + } + raw_string_ostream(SubModuleName) + << ".submodule." + << formatv(sizeof(size_t) == 8 ? "{0:x16}" : "{0:x8}", + static_cast(HC)) + << ".ll"; + } + + // Extract the requested partiton (plus any necessary aliases) and + // put the rest back into the impl dylib. + auto ShouldExtract = [&](const GlobalValue &GV) -> bool { + return GVsToExtract->count(&GV); + }; + + return extractSubModule(TSM, SubModuleName, ShouldExtract); + }); + + if (!ExtractedTSM) { + ES.reportError(ExtractedTSM.takeError()); + R->failMaterialization(); + return; + } + + if (auto Err = R->replace(std::make_unique( + ES, *getManglingOptions(), std::move(TSM), *this))) { + ES.reportError(std::move(Err)); + R->failMaterialization(); + return; + } + BaseLayer.emit(std::move(R), std::move(*ExtractedTSM)); +} diff --git a/llvm/lib/ExecutionEngine/Orc/LLJIT.cpp b/llvm/lib/ExecutionEngine/Orc/LLJIT.cpp index d3dd3b6bedfb65..d3db89a2c3e9f1 100644 --- a/llvm/lib/ExecutionEngine/Orc/LLJIT.cpp +++ b/llvm/lib/ExecutionEngine/Orc/LLJIT.cpp @@ -1293,9 +1293,12 @@ LLLazyJIT::LLLazyJIT(LLLazyJITBuilderState &S, Error &Err) : LLJIT(S, Err) { return; } + // Create the IP Layer. + IPLayer = std::make_unique(*ES, *InitHelperTransformLayer); + // Create the COD layer. - CODLayer = std::make_unique( - *ES, *InitHelperTransformLayer, *LCTMgr, std::move(ISMBuilder)); + CODLayer = std::make_unique(*ES, *IPLayer, *LCTMgr, + std::move(ISMBuilder)); if (*S.SupportConcurrentCompilation) CODLayer->setCloneToNewContextOnEmit(true); diff --git a/llvm/tools/lli/lli.cpp b/llvm/tools/lli/lli.cpp index f21d9810cca6a9..4c023bcdd61658 100644 --- a/llvm/tools/lli/lli.cpp +++ b/llvm/tools/lli/lli.cpp @@ -30,6 +30,7 @@ #include "llvm/ExecutionEngine/Orc/EPCEHFrameRegistrar.h" #include "llvm/ExecutionEngine/Orc/EPCGenericRTDyldMemoryManager.h" #include "llvm/ExecutionEngine/Orc/ExecutionUtils.h" +#include "llvm/ExecutionEngine/Orc/IRPartitionLayer.h" #include "llvm/ExecutionEngine/Orc/JITTargetMachineBuilder.h" #include "llvm/ExecutionEngine/Orc/LLJIT.h" #include "llvm/ExecutionEngine/Orc/ObjectTransformLayer.h" @@ -1060,7 +1061,7 @@ int runOrcJIT(const char *ProgName) { } if (PerModuleLazy) - J->setPartitionFunction(orc::CompileOnDemandLayer::compileWholeModule); + J->setPartitionFunction(orc::IRPartitionLayer::compileWholeModule); auto IRDump = createIRDebugDumper(); J->getIRTransformLayer().setTransform(