[ORC] Introduce IRPartitionLayer for common partition functionality.

This commit is contained in:
Sunho Kim 2023-09-20 05:11:15 +09:00
parent 2c01b27858
commit 04af63b267
10 changed files with 417 additions and 322 deletions

View File

@ -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 @@ private:
RTDyldObjectLinkingLayer ObjectLayer;
IRCompileLayer CompileLayer;
IRTransformLayer OptimizeLayer;
IRPartitionLayer IPLayer;
CompileOnDemandLayer CODLayer;
JITDylib &MainJD;
@ -68,8 +70,8 @@ public:
CompileLayer(*this->ES, ObjectLayer,
std::make_unique<ConcurrentIRCompiler>(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("<main>")) {
MainJD.addGenerator(

View File

@ -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 @@ private:
IndirectStubsManagerBuilderFunction ISMBuilder,
std::unique_ptr<DynamicLibrarySearchGenerator> ProcessSymbolsGenerator)
: ES(std::move(ES)), DL(std::move(DL)),
MainJD(this->ES->createBareJITDylib("<main>")), LCTMgr(std::move(LCTMgr)),
MainJD(this->ES->createBareJITDylib("<main>")),
LCTMgr(std::move(LCTMgr)),
CompileLayer(*this->ES, ObjLayer,
std::make_unique<ConcurrentIRCompiler>(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 @@ private:
Speculator S;
RTDyldObjectLinkingLayer ObjLayer{*ES, createMemMgr};
IRSpeculationLayer SpeculateLayer;
IRPartitionLayer IPLayer;
CompileOnDemandLayer CODLayer;
};

View File

@ -53,37 +53,15 @@ namespace llvm {
namespace orc {
class CompileOnDemandLayer : public IRLayer {
friend class PartitioningIRMaterializationUnit;
public:
/// Builder for IndirectStubsManagers.
using IndirectStubsManagerBuilder =
std::function<std::unique_ptr<IndirectStubsManager>()>;
using GlobalValueSet = std::set<const GlobalValue *>;
/// Partitioning function.
using PartitionFunction =
std::function<std::optional<GlobalValueSet>(GlobalValueSet Requested)>;
/// Off-the-shelf partitioning which compiles all requested symbols (usually
/// a single function at a time).
static std::optional<GlobalValueSet>
compileRequested(GlobalValueSet Requested);
/// Off-the-shelf partitioning which compiles whole modules whenever any
/// symbol in them is requested.
static std::optional<GlobalValueSet>
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 @@ private:
PerDylibResources &getPerDylibResources(JITDylib &TargetD);
void cleanUpModule(Module &M);
void expandPartition(GlobalValueSet &Partition);
void emitPartition(std::unique_ptr<MaterializationResponsibility> 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;
};

View File

@ -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<const GlobalValue *>;
/// Partitioning function.
using PartitionFunction =
std::function<std::optional<GlobalValueSet>(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<GlobalValueSet>
compileRequested(GlobalValueSet Requested);
/// Off-the-shelf partitioning which compiles whole modules whenever any
/// symbol in them is requested.
static std::optional<GlobalValueSet>
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<MaterializationResponsibility> R,
ThreadSafeModule TSM) override;
private:
void cleanUpModule(Module &M);
void expandPartition(GlobalValueSet &Partition);
void emitPartition(std::unique_ptr<MaterializationResponsibility> R,
ThreadSafeModule TSM,
IRMaterializationUnit::SymbolNameToDefinitionMap Defs);
IRLayer &BaseLayer;
PartitionFunction Partition = compileRequested;
SymbolLinkagePromoter PromoteSymbols;
};
} // namespace orc
} // namespace llvm
#endif

View File

@ -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 @@ private:
LLLazyJIT(LLLazyJITBuilderState &S, Error &Err);
std::unique_ptr<LazyCallThroughManager> LCTMgr;
std::unique_ptr<IRPartitionLayer> IPLayer;
std::unique_ptr<CompileOnDemandLayer> CODLayer;
};

View File

@ -26,6 +26,7 @@ add_llvm_component_library(LLVMOrcJIT
IndirectionUtils.cpp
IRCompileLayer.cpp
IRTransformLayer.cpp
IRPartitionLayer.cpp
JITTargetMachineBuilder.cpp
LazyReexports.cpp
Layer.cpp

View File

@ -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<Function>(GV)) {
auto &F = cast<Function>(GV);
F.deleteBody();
F.setPersonalityFn(nullptr);
} else if (isa<GlobalVariable>(GV)) {
cast<GlobalVariable>(GV).setInitializer(nullptr);
} else if (isa<GlobalAlias>(GV)) {
// We need to turn deleted aliases into function or variable decls based
// on the type of their aliasee.
auto &A = cast<GlobalAlias>(GV);
Constant *Aliasee = A.getAliasee();
assert(A.hasName() && "Anonymous alias?");
assert(Aliasee->hasName() && "Anonymous aliasee");
std::string AliasName = std::string(A.getName());
if (isa<Function>(Aliasee)) {
auto *F = cloneFunctionDecl(*A.getParent(), *cast<Function>(Aliasee));
A.replaceAllUsesWith(F);
A.eraseFromParent();
F->setName(AliasName);
} else if (isa<GlobalVariable>(Aliasee)) {
auto *G = cloneGlobalVariableDecl(*A.getParent(),
*cast<GlobalVariable>(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<MaterializationResponsibility> 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::GlobalValueSet>
CompileOnDemandLayer::compileRequested(GlobalValueSet Requested) {
return std::move(Requested);
}
std::optional<CompileOnDemandLayer::GlobalValueSet>
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<MaterializationResponsibility> 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<PartitioningIRMaterializationUnit>(
ES, *getManglingOptions(), std::move(TSM), *this))) {
std::make_unique<BasicIRLayerMaterializationUnit>(
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<const GlobalValue *> GVsToAdd;
for (const auto *GV : Partition)
if (isa<GlobalAlias>(GV))
GVsToAdd.push_back(
cast<GlobalValue>(cast<GlobalAlias>(GV)->getAliasee()));
else if (isa<GlobalVariable>(GV))
ContainsGlobalVariables = true;
for (auto &A : M.aliases())
if (Partition.count(cast<GlobalValue>(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<MaterializationResponsibility> 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<PartitioningIRMaterializationUnit>(
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<ThreadSafeModule> {
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<const GlobalValue*> 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<size_t>(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<PartitioningIRMaterializationUnit>(
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

View File

@ -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<Function>(GV)) {
auto &F = cast<Function>(GV);
F.deleteBody();
F.setPersonalityFn(nullptr);
} else if (isa<GlobalVariable>(GV)) {
cast<GlobalVariable>(GV).setInitializer(nullptr);
} else if (isa<GlobalAlias>(GV)) {
// We need to turn deleted aliases into function or variable decls based
// on the type of their aliasee.
auto &A = cast<GlobalAlias>(GV);
Constant *Aliasee = A.getAliasee();
assert(A.hasName() && "Anonymous alias?");
assert(Aliasee->hasName() && "Anonymous aliasee");
std::string AliasName = std::string(A.getName());
if (isa<Function>(Aliasee)) {
auto *F = cloneFunctionDecl(*A.getParent(), *cast<Function>(Aliasee));
A.replaceAllUsesWith(F);
A.eraseFromParent();
F->setName(AliasName);
} else if (isa<GlobalVariable>(Aliasee)) {
auto *G = cloneGlobalVariableDecl(*A.getParent(),
*cast<GlobalVariable>(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<MaterializationResponsibility> 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::GlobalValueSet>
IRPartitionLayer::compileRequested(GlobalValueSet Requested) {
return std::move(Requested);
}
std::optional<IRPartitionLayer::GlobalValueSet>
IRPartitionLayer::compileWholeModule(GlobalValueSet Requested) {
return std::nullopt;
}
void IRPartitionLayer::emit(std::unique_ptr<MaterializationResponsibility> 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<PartitioningIRMaterializationUnit>(
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<const GlobalValue *> GVsToAdd;
for (const auto *GV : Partition)
if (isa<GlobalAlias>(GV))
GVsToAdd.push_back(
cast<GlobalValue>(cast<GlobalAlias>(GV)->getAliasee()));
else if (isa<GlobalVariable>(GV))
ContainsGlobalVariables = true;
for (auto &A : M.aliases())
if (Partition.count(cast<GlobalValue>(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<MaterializationResponsibility> 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<PartitioningIRMaterializationUnit>(
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<ThreadSafeModule> {
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<const GlobalValue *> 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<size_t>(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<PartitioningIRMaterializationUnit>(
ES, *getManglingOptions(), std::move(TSM), *this))) {
ES.reportError(std::move(Err));
R->failMaterialization();
return;
}
BaseLayer.emit(std::move(R), std::move(*ExtractedTSM));
}

View File

@ -1293,9 +1293,12 @@ LLLazyJIT::LLLazyJIT(LLLazyJITBuilderState &S, Error &Err) : LLJIT(S, Err) {
return;
}
// Create the IP Layer.
IPLayer = std::make_unique<IRPartitionLayer>(*ES, *InitHelperTransformLayer);
// Create the COD layer.
CODLayer = std::make_unique<CompileOnDemandLayer>(
*ES, *InitHelperTransformLayer, *LCTMgr, std::move(ISMBuilder));
CODLayer = std::make_unique<CompileOnDemandLayer>(*ES, *IPLayer, *LCTMgr,
std::move(ISMBuilder));
if (*S.SupportConcurrentCompilation)
CODLayer->setCloneToNewContextOnEmit(true);

View File

@ -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(