181 changes: 181 additions & 0 deletions llvm/include/llvm/ExecutionEngine/Orc/ReOptimizeLayer.h
View file
Original file line number Diff line number Diff line change
@@ -0,0 +1,181 @@
//===- ReOptimizeLayer.h - Re-optimization layer interface ------*- 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
//
//===----------------------------------------------------------------------===//
//
// Re-optimization layer interface.
//
//===----------------------------------------------------------------------===//
#ifndef LLVM_EXECUTIONENGINE_ORC_REOPTIMIZELAYER_H
#define LLVM_EXECUTIONENGINE_ORC_REOPTIMIZELAYER_H

#include "llvm/ExecutionEngine/Orc/Core.h"
#include "llvm/ExecutionEngine/Orc/Layer.h"
#include "llvm/ExecutionEngine/Orc/Mangling.h"
#include "llvm/ExecutionEngine/Orc/RedirectionManager.h"
#include "llvm/ExecutionEngine/Orc/ThreadSafeModule.h"
#include "llvm/IR/IRBuilder.h"
#include "llvm/Transforms/Utils/BasicBlockUtils.h"
#include "llvm/Transforms/Utils/Cloning.h"

namespace llvm {
namespace orc {

class ReOptimizeLayer : public IRLayer, public ResourceManager {
public:
using ReOptMaterializationUnitID = uint64_t;

/// AddProfilerFunc will be called when ReOptimizeLayer emits the first
/// version of a materialization unit in order to inject profiling code and
/// reoptimization request code.
using AddProfilerFunc = unique_function<Error(
ReOptimizeLayer &Parent, ReOptMaterializationUnitID MUID,
unsigned CurVersion, ThreadSafeModule &TSM)>;

/// ReOptimizeFunc will be called when ReOptimizeLayer reoptimization of a
/// materialization unit was requested in order to reoptimize the IR module
/// based on profile data. OldRT is the ResourceTracker that tracks the old
/// function definitions. The OldRT must be kept alive until it can be
/// guaranteed that every invocation of the old function definitions has been
/// terminated.
using ReOptimizeFunc = unique_function<Error(
ReOptimizeLayer &Parent, ReOptMaterializationUnitID MUID,
unsigned CurVersion, ResourceTrackerSP OldRT, ThreadSafeModule &TSM)>;

ReOptimizeLayer(ExecutionSession &ES, DataLayout &DL, IRLayer &BaseLayer,
RedirectableSymbolManager &RM)
: IRLayer(ES, BaseLayer.getManglingOptions()), ES(ES), Mangle(ES, DL),
BaseLayer(BaseLayer), RSManager(RM), ReOptFunc(identity),
ProfilerFunc(reoptimizeIfCallFrequent) {}

void setReoptimizeFunc(ReOptimizeFunc ReOptFunc) {
this->ReOptFunc = std::move(ReOptFunc);
}

void setAddProfilerFunc(AddProfilerFunc ProfilerFunc) {
this->ProfilerFunc = std::move(ProfilerFunc);
}

/// Registers reoptimize runtime dispatch handlers to given PlatformJD. The
/// reoptimization request will not be handled if dispatch handler is not
/// registered by using this function.
Error reigsterRuntimeFunctions(JITDylib &PlatformJD);

/// 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;

static const uint64_t CallCountThreshold = 10;

/// Basic AddProfilerFunc that reoptimizes the function when the call count
/// exceeds CallCountThreshold.
static Error reoptimizeIfCallFrequent(ReOptimizeLayer &Parent,
ReOptMaterializationUnitID MUID,
unsigned CurVersion,
ThreadSafeModule &TSM);

static Error identity(ReOptimizeLayer &Parent,
ReOptMaterializationUnitID MUID, unsigned CurVersion,
ResourceTrackerSP OldRT, ThreadSafeModule &TSM) {
return Error::success();
}

// Create IR reoptimize request fucntion call.
static void createReoptimizeCall(Module &M, Instruction &IP,
GlobalVariable *ArgBuffer);

Error handleRemoveResources(JITDylib &JD, ResourceKey K) override;
void handleTransferResources(JITDylib &JD, ResourceKey DstK,
ResourceKey SrcK) override;

private:
class ReOptMaterializationUnitState {
public:
ReOptMaterializationUnitState() = default;
ReOptMaterializationUnitState(ReOptMaterializationUnitID ID,
ThreadSafeModule TSM)
: ID(ID), TSM(std::move(TSM)) {}
ReOptMaterializationUnitState(ReOptMaterializationUnitState &&Other)
: ID(Other.ID), TSM(std::move(Other.TSM)), RT(std::move(Other.RT)),
Reoptimizing(std::move(Other.Reoptimizing)),
CurVersion(Other.CurVersion) {}

ReOptMaterializationUnitID getID() { return ID; }

const ThreadSafeModule &getThreadSafeModule() { return TSM; }

ResourceTrackerSP getResourceTracker() {
std::unique_lock<std::mutex> Lock(Mutex);
return RT;
}

void setResourceTracker(ResourceTrackerSP RT) {
std::unique_lock<std::mutex> Lock(Mutex);
this->RT = RT;
}

uint32_t getCurVersion() {
std::unique_lock<std::mutex> Lock(Mutex);
return CurVersion;
}

bool tryStartReoptimize();
void reoptimizeSucceeded();
void reoptimizeFailed();

private:
std::mutex Mutex;
ReOptMaterializationUnitID ID;
ThreadSafeModule TSM;
ResourceTrackerSP RT;
bool Reoptimizing = false;
uint32_t CurVersion = 0;
};

using SPSReoptimizeArgList =
shared::SPSArgList<ReOptMaterializationUnitID, uint32_t>;
using SendErrorFn = unique_function<void(Error)>;

Expected<SymbolMap> emitMUImplSymbols(ReOptMaterializationUnitState &MUState,
uint32_t Version, JITDylib &JD,
ThreadSafeModule TSM);

void rt_reoptimize(SendErrorFn SendResult, ReOptMaterializationUnitID MUID,
uint32_t CurVersion);

static Expected<Constant *>
createReoptimizeArgBuffer(Module &M, ReOptMaterializationUnitID MUID,
uint32_t CurVersion);

ReOptMaterializationUnitState &
createMaterializationUnitState(const ThreadSafeModule &TSM);

void
registerMaterializationUnitResource(ResourceKey Key,
ReOptMaterializationUnitState &State);

ReOptMaterializationUnitState &
getMaterializationUnitState(ReOptMaterializationUnitID MUID);

ExecutionSession &ES;
MangleAndInterner Mangle;
IRLayer &BaseLayer;
RedirectableSymbolManager &RSManager;

ReOptimizeFunc ReOptFunc;
AddProfilerFunc ProfilerFunc;

std::mutex Mutex;
std::map<ReOptMaterializationUnitID, ReOptMaterializationUnitState> MUStates;
DenseMap<ResourceKey, DenseSet<ReOptMaterializationUnitID>> MUResources;
ReOptMaterializationUnitID NextID = 1;
};

} // namespace orc
} // namespace llvm

#endif
103 changes: 103 additions & 0 deletions llvm/include/llvm/ExecutionEngine/Orc/RedirectionManager.h
View file
Original file line number Diff line number Diff line change
@@ -0,0 +1,103 @@
//===- RedirectionManager.h - Redirection manager interface -----*- 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
//
//===----------------------------------------------------------------------===//
//
// Redirection manager interface that redirects a call to symbol to another.
//
//===----------------------------------------------------------------------===//

#ifndef LLVM_EXECUTIONENGINE_ORC_REDIRECTIONMANAGER_H
#define LLVM_EXECUTIONENGINE_ORC_REDIRECTIONMANAGER_H

#include "llvm/ExecutionEngine/Orc/Core.h"

namespace llvm {
namespace orc {

/// Base class for performing redirection of call to symbol to another symbol in
/// runtime.
class RedirectionManager {
public:
/// Symbol name to symbol definition map.
using SymbolAddrMap = DenseMap<SymbolStringPtr, ExecutorSymbolDef>;

virtual ~RedirectionManager() = default;
/// Change the redirection destination of given symbols to new destination
/// symbols.
virtual Error redirect(JITDylib &JD, const SymbolAddrMap &NewDests) = 0;

/// Change the redirection destination of given symbol to new destination
/// symbol.
virtual Error redirect(JITDylib &JD, SymbolStringPtr Symbol,
ExecutorSymbolDef NewDest) {
return redirect(JD, {{Symbol, NewDest}});
}

private:
virtual void anchor();
};

/// Base class for managing redirectable symbols in which a call
/// gets redirected to another symbol in runtime.
class RedirectableSymbolManager : public RedirectionManager {
public:
/// Create redirectable symbols with given symbol names and initial
/// desitnation symbol addresses.
Error createRedirectableSymbols(ResourceTrackerSP RT,
const SymbolMap &InitialDests);

/// Create a single redirectable symbol with given symbol name and initial
/// desitnation symbol address.
Error createRedirectableSymbol(ResourceTrackerSP RT, SymbolStringPtr Symbol,
ExecutorSymbolDef InitialDest) {
return createRedirectableSymbols(RT, {{Symbol, InitialDest}});
}

/// Emit redirectable symbol
virtual void
emitRedirectableSymbols(std::unique_ptr<MaterializationResponsibility> MR,
const SymbolMap &InitialDests) = 0;
};

/// RedirectableMaterializationUnit materializes redirectable symbol
/// by invoking RedirectableSymbolManager::emitRedirectableSymbols
class RedirectableMaterializationUnit : public MaterializationUnit {
public:
RedirectableMaterializationUnit(RedirectableSymbolManager &RM,
const SymbolMap &InitialDests)
: MaterializationUnit(convertToFlags(InitialDests)), RM(RM),
InitialDests(InitialDests) {}

StringRef getName() const override {
return "RedirectableSymbolMaterializationUnit";
}

void materialize(std::unique_ptr<MaterializationResponsibility> R) override {
RM.emitRedirectableSymbols(std::move(R), std::move(InitialDests));
}

void discard(const JITDylib &JD, const SymbolStringPtr &Name) override {
InitialDests.erase(Name);
}

private:
static MaterializationUnit::Interface
convertToFlags(const SymbolMap &InitialDests) {
SymbolFlagsMap Flags;
for (auto [K, V] : InitialDests)
Flags[K] = V.getFlags();
return MaterializationUnit::Interface(Flags, {});
}

RedirectableSymbolManager &RM;
SymbolMap InitialDests;
};

} // namespace orc
} // namespace llvm

#endif
4 changes: 4 additions & 0 deletions llvm/lib/ExecutionEngine/Orc/CMakeLists.txt
View file
Original file line number Diff line number Diff line change
Expand Up @@ -26,6 +26,7 @@ add_llvm_component_library(LLVMOrcJIT
IndirectionUtils.cpp
IRCompileLayer.cpp
IRTransformLayer.cpp
IRPartitionLayer.cpp
JITTargetMachineBuilder.cpp
LazyReexports.cpp
Layer.cpp
Expand All @@ -50,6 +51,9 @@ add_llvm_component_library(LLVMOrcJIT
ExecutorProcessControl.cpp
TaskDispatch.cpp
ThreadSafeModule.cpp
RedirectionManager.cpp
JITLinkRedirectableSymbolManager.cpp
ReOptimizeLayer.cpp
ADDITIONAL_HEADER_DIRS
${LLVM_MAIN_INCLUDE_DIR}/llvm/ExecutionEngine/Orc

Expand Down
282 changes: 5 additions & 277 deletions llvm/lib/ExecutionEngine/Orc/CompileOnDemandLayer.cpp
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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"
Expand All @@ -17,115 +18,17 @@
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)
: IRLayer(ES, BaseLayer.getManglingOptions()), BaseLayer(BaseLayer),
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");
Expand All @@ -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;
Expand All @@ -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;
Expand Down Expand Up @@ -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
303 changes: 303 additions & 0 deletions llvm/lib/ExecutionEngine/Orc/IRPartitionLayer.cpp
View file
Original file line number Diff line number Diff line change
@@ -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));
}
182 changes: 182 additions & 0 deletions llvm/lib/ExecutionEngine/Orc/JITLinkRedirectableSymbolManager.cpp
View file
Original file line number Diff line number Diff line change
@@ -0,0 +1,182 @@
//===-- JITLinkRedirectableSymbolManager.cpp - JITLink redirection in Orc -===//
//
// 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/JITLinkRedirectableSymbolManager.h"
#include "llvm/ExecutionEngine/Orc/Core.h"

#define DEBUG_TYPE "orc"

using namespace llvm;
using namespace llvm::orc;

void JITLinkRedirectableSymbolManager::emitRedirectableSymbols(
std::unique_ptr<MaterializationResponsibility> R,
const SymbolAddrMap &InitialDests) {
auto &ES = ObjLinkingLayer.getExecutionSession();
std::unique_lock<std::mutex> Lock(Mutex);
if (GetNumAvailableStubs() < InitialDests.size())
if (auto Err = grow(InitialDests.size() - GetNumAvailableStubs())) {
ES.reportError(std::move(Err));
R->failMaterialization();
return;
}

JITDylib &TargetJD = R->getTargetJITDylib();
SymbolMap NewSymbolDefs;
std::vector<SymbolStringPtr> Symbols;
for (auto &[K, V] : InitialDests) {
StubHandle StubID = AvailableStubs.back();
if (SymbolToStubs[&TargetJD].count(K)) {
ES.reportError(make_error<StringError>(
"Tried to create duplicate redirectable symbols",
inconvertibleErrorCode()));
R->failMaterialization();
return;
}
SymbolToStubs[&TargetJD][K] = StubID;
NewSymbolDefs[K] = JumpStubs[StubID];
NewSymbolDefs[K].setFlags(V.getFlags());
Symbols.push_back(K);
AvailableStubs.pop_back();
}

// FIXME: when this fails we can return stubs to the pool
if (auto Err = redirectInner(TargetJD, InitialDests)) {
ES.reportError(std::move(Err));
R->failMaterialization();
return;
}

if (auto Err = R->replace(absoluteSymbols(NewSymbolDefs))) {
ES.reportError(std::move(Err));
R->failMaterialization();
return;
}

auto Err = R->withResourceKeyDo([&](ResourceKey Key) {
TrackedResources[Key].insert(TrackedResources[Key].end(), Symbols.begin(),
Symbols.end());
});
if (Err) {
ES.reportError(std::move(Err));
R->failMaterialization();
return;
}
}

Error JITLinkRedirectableSymbolManager::redirect(
JITDylib &TargetJD, const SymbolAddrMap &NewDests) {
std::unique_lock<std::mutex> Lock(Mutex);
return redirectInner(TargetJD, NewDests);
}

Error JITLinkRedirectableSymbolManager::redirectInner(
JITDylib &TargetJD, const SymbolAddrMap &NewDests) {
std::vector<tpctypes::PointerWrite> PtrWrites;
for (auto &[K, V] : NewDests) {
if (!SymbolToStubs[&TargetJD].count(K))
return make_error<StringError>(
"Tried to redirect non-existent redirectalbe symbol",
inconvertibleErrorCode());
StubHandle StubID = SymbolToStubs[&TargetJD].at(K);
PtrWrites.push_back({StubPointers[StubID].getAddress(), V.getAddress()});
}
return ObjLinkingLayer.getExecutionSession()
.getExecutorProcessControl()
.getMemoryAccess()
.writePointers(PtrWrites);
}

Error JITLinkRedirectableSymbolManager::grow(unsigned Need) {
unsigned OldSize = JumpStubs.size();
unsigned NumNewStubs = alignTo(Need, StubBlockSize);
unsigned NewSize = OldSize + NumNewStubs;

JumpStubs.resize(NewSize);
StubPointers.resize(NewSize);
AvailableStubs.reserve(NewSize);

SymbolLookupSet LookupSymbols;
DenseMap<SymbolStringPtr, ExecutorSymbolDef *> NewDefsMap;

auto &ES = ObjLinkingLayer.getExecutionSession();
Triple TT = ES.getTargetTriple();
auto G = std::make_unique<jitlink::LinkGraph>(
"<INDIRECT STUBS>", TT, TT.isArch64Bit() ? 8 : 4,
TT.isLittleEndian() ? endianness::little : endianness::big,
jitlink::getGenericEdgeKindName);
auto &PointerSection =
G->createSection(StubPtrTableName, MemProt::Write | MemProt::Read);
auto &StubsSection =
G->createSection(JumpStubTableName, MemProt::Exec | MemProt::Read);

// FIXME: We can batch the stubs into one block and use address to access them
for (size_t I = OldSize; I < NewSize; I++) {
auto Pointer = AnonymousPtrCreator(*G, PointerSection, nullptr, 0);
if (auto Err = Pointer.takeError())
return Err;

StringRef PtrSymName = StubPtrSymbolName(I);
Pointer->setName(PtrSymName);
Pointer->setScope(jitlink::Scope::Default);
LookupSymbols.add(ES.intern(PtrSymName));
NewDefsMap[ES.intern(PtrSymName)] = &StubPointers[I];

auto Stub = PtrJumpStubCreator(*G, StubsSection, *Pointer);
if (auto Err = Stub.takeError())
return Err;

StringRef JumpStubSymName = JumpStubSymbolName(I);
Stub->setName(JumpStubSymName);
Stub->setScope(jitlink::Scope::Default);
LookupSymbols.add(ES.intern(JumpStubSymName));
NewDefsMap[ES.intern(JumpStubSymName)] = &JumpStubs[I];
}

if (auto Err = ObjLinkingLayer.add(JD, std::move(G)))
return Err;

auto LookupResult = ES.lookup(makeJITDylibSearchOrder(&JD), LookupSymbols);
if (auto Err = LookupResult.takeError())
return Err;

for (auto &[K, V] : *LookupResult)
*NewDefsMap.at(K) = V;

for (size_t I = OldSize; I < NewSize; I++)
AvailableStubs.push_back(I);

return Error::success();
}

Error JITLinkRedirectableSymbolManager::handleRemoveResources(
JITDylib &TargetJD, ResourceKey K) {
std::unique_lock<std::mutex> Lock(Mutex);
for (auto &Symbol : TrackedResources[K]) {
if (!SymbolToStubs[&TargetJD].count(Symbol))
return make_error<StringError>(
"Tried to remove non-existent redirectable symbol",
inconvertibleErrorCode());
AvailableStubs.push_back(SymbolToStubs[&TargetJD].at(Symbol));
SymbolToStubs[&TargetJD].erase(Symbol);
if (SymbolToStubs[&TargetJD].empty())
SymbolToStubs.erase(&TargetJD);
}
TrackedResources.erase(K);

return Error::success();
}

void JITLinkRedirectableSymbolManager::handleTransferResources(
JITDylib &TargetJD, ResourceKey DstK, ResourceKey SrcK) {
std::unique_lock<std::mutex> Lock(Mutex);
TrackedResources[DstK].insert(TrackedResources[DstK].end(),
TrackedResources[SrcK].begin(),
TrackedResources[SrcK].end());
TrackedResources.erase(SrcK);
}
7 changes: 5 additions & 2 deletions llvm/lib/ExecutionEngine/Orc/LLJIT.cpp
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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);
Expand Down
279 changes: 279 additions & 0 deletions llvm/lib/ExecutionEngine/Orc/ReOptimizeLayer.cpp
View file
Original file line number Diff line number Diff line change
@@ -0,0 +1,279 @@
#include "llvm/ExecutionEngine/Orc/ReOptimizeLayer.h"
#include "llvm/ExecutionEngine/Orc/Mangling.h"

using namespace llvm;
using namespace orc;

bool ReOptimizeLayer::ReOptMaterializationUnitState::tryStartReoptimize() {
std::unique_lock<std::mutex> Lock(Mutex);
if (Reoptimizing)
return false;

Reoptimizing = true;
return true;
}

void ReOptimizeLayer::ReOptMaterializationUnitState::reoptimizeSucceeded() {
std::unique_lock<std::mutex> Lock(Mutex);
assert(Reoptimizing && "Tried to mark unstarted reoptimization as done");
Reoptimizing = false;
CurVersion++;
}

void ReOptimizeLayer::ReOptMaterializationUnitState::reoptimizeFailed() {
std::unique_lock<std::mutex> Lock(Mutex);
assert(Reoptimizing && "Tried to mark unstarted reoptimization as done");
Reoptimizing = false;
}

Error ReOptimizeLayer::reigsterRuntimeFunctions(JITDylib &PlatformJD) {
ExecutionSession::JITDispatchHandlerAssociationMap WFs;
using ReoptimizeSPSSig = shared::SPSError(uint64_t, uint32_t);
WFs[Mangle("__orc_rt_reoptimize_tag")] =
ES.wrapAsyncWithSPS<ReoptimizeSPSSig>(this,
&ReOptimizeLayer::rt_reoptimize);
return ES.registerJITDispatchHandlers(PlatformJD, std::move(WFs));
}

void ReOptimizeLayer::emit(std::unique_ptr<MaterializationResponsibility> R,
ThreadSafeModule TSM) {
auto &JD = R->getTargetJITDylib();

bool HasNonCallable = false;
for (auto &KV : R->getSymbols()) {
auto &Flags = KV.second;
if (!Flags.isCallable())
HasNonCallable = true;
}

if (HasNonCallable) {
BaseLayer.emit(std::move(R), std::move(TSM));
return;
}

auto &MUState = createMaterializationUnitState(TSM);

if (auto Err = R->withResourceKeyDo([&](ResourceKey Key) {
registerMaterializationUnitResource(Key, MUState);
})) {
ES.reportError(std::move(Err));
R->failMaterialization();
return;
}

if (auto Err =
ProfilerFunc(*this, MUState.getID(), MUState.getCurVersion(), TSM)) {
ES.reportError(std::move(Err));
R->failMaterialization();
return;
}

auto InitialDests =
emitMUImplSymbols(MUState, MUState.getCurVersion(), JD, std::move(TSM));
if (!InitialDests) {
ES.reportError(InitialDests.takeError());
R->failMaterialization();
return;
}

RSManager.emitRedirectableSymbols(std::move(R), std::move(*InitialDests));
}

Error ReOptimizeLayer::reoptimizeIfCallFrequent(ReOptimizeLayer &Parent,
ReOptMaterializationUnitID MUID,
unsigned CurVersion,
ThreadSafeModule &TSM) {
return TSM.withModuleDo([&](Module &M) -> Error {
Type *I64Ty = Type::getInt64Ty(M.getContext());
GlobalVariable *Counter = new GlobalVariable(
M, I64Ty, false, GlobalValue::InternalLinkage,
Constant::getNullValue(I64Ty), "__orc_reopt_counter");
auto ArgBufferConst = createReoptimizeArgBuffer(M, MUID, CurVersion);
if (auto Err = ArgBufferConst.takeError())
return Err;
GlobalVariable *ArgBuffer =
new GlobalVariable(M, (*ArgBufferConst)->getType(), true,
GlobalValue::InternalLinkage, (*ArgBufferConst));
for (auto &F : M) {
if (F.isDeclaration())
continue;
auto &BB = F.getEntryBlock();
auto *IP = &*BB.getFirstInsertionPt();
IRBuilder<> IRB(IP);
Value *Threshold = ConstantInt::get(I64Ty, CallCountThreshold, true);
Value *Cnt = IRB.CreateLoad(I64Ty, Counter);
// Use EQ to prevent further reoptimize calls.
Value *Cmp = IRB.CreateICmpEQ(Cnt, Threshold);
Value *Added = IRB.CreateAdd(Cnt, ConstantInt::get(I64Ty, 1));
(void)IRB.CreateStore(Added, Counter);
Instruction *SplitTerminator = SplitBlockAndInsertIfThen(Cmp, IP, false);
createReoptimizeCall(M, *SplitTerminator, ArgBuffer);
}
return Error::success();
});
}

Expected<SymbolMap>
ReOptimizeLayer::emitMUImplSymbols(ReOptMaterializationUnitState &MUState,
uint32_t Version, JITDylib &JD,
ThreadSafeModule TSM) {
DenseMap<SymbolStringPtr, SymbolStringPtr> RenamedMap;
cantFail(TSM.withModuleDo([&](Module &M) -> Error {
MangleAndInterner Mangle(ES, M.getDataLayout());
for (auto &F : M)
if (!F.isDeclaration()) {
std::string NewName =
(F.getName() + ".__def__." + Twine(Version)).str();
RenamedMap[Mangle(F.getName())] = Mangle(NewName);
F.setName(NewName);
}
return Error::success();
}));

auto RT = JD.createResourceTracker();
if (auto Err =
JD.define(std::make_unique<BasicIRLayerMaterializationUnit>(
BaseLayer, *getManglingOptions(), std::move(TSM)),
RT))
return Err;
MUState.setResourceTracker(RT);

SymbolLookupSet LookupSymbols;
for (auto [K, V] : RenamedMap)
LookupSymbols.add(V);

auto ImplSymbols =
ES.lookup({{&JD, JITDylibLookupFlags::MatchAllSymbols}}, LookupSymbols,
LookupKind::Static, SymbolState::Resolved);
if (auto Err = ImplSymbols.takeError())
return Err;

SymbolMap Result;
for (auto [K, V] : RenamedMap)
Result[K] = (*ImplSymbols)[V];

return Result;
}

void ReOptimizeLayer::rt_reoptimize(SendErrorFn SendResult,
ReOptMaterializationUnitID MUID,
uint32_t CurVersion) {
auto &MUState = getMaterializationUnitState(MUID);
if (CurVersion < MUState.getCurVersion() || !MUState.tryStartReoptimize()) {
SendResult(Error::success());
return;
}

ThreadSafeModule TSM = cloneToNewContext(MUState.getThreadSafeModule());
auto OldRT = MUState.getResourceTracker();
auto &JD = OldRT->getJITDylib();

if (auto Err = ReOptFunc(*this, MUID, CurVersion + 1, OldRT, TSM)) {
ES.reportError(std::move(Err));
MUState.reoptimizeFailed();
SendResult(Error::success());
return;
}

auto SymbolDests =
emitMUImplSymbols(MUState, CurVersion + 1, JD, std::move(TSM));
if (!SymbolDests) {
ES.reportError(SymbolDests.takeError());
MUState.reoptimizeFailed();
SendResult(Error::success());
return;
}

if (auto Err = RSManager.redirect(JD, std::move(*SymbolDests))) {
ES.reportError(std::move(Err));
MUState.reoptimizeFailed();
SendResult(Error::success());
return;
}

MUState.reoptimizeSucceeded();
SendResult(Error::success());
}

Expected<Constant *> ReOptimizeLayer::createReoptimizeArgBuffer(
Module &M, ReOptMaterializationUnitID MUID, uint32_t CurVersion) {
size_t ArgBufferSize = SPSReoptimizeArgList::size(MUID, CurVersion);
std::vector<char> ArgBuffer(ArgBufferSize);
shared::SPSOutputBuffer OB(ArgBuffer.data(), ArgBuffer.size());
if (!SPSReoptimizeArgList::serialize(OB, MUID, CurVersion))
return make_error<StringError>("Could not serealize args list",
inconvertibleErrorCode());
return ConstantDataArray::get(M.getContext(), ArrayRef(ArgBuffer));
}

void ReOptimizeLayer::createReoptimizeCall(Module &M, Instruction &IP,
GlobalVariable *ArgBuffer) {
GlobalVariable *DispatchCtx =
M.getGlobalVariable("__orc_rt_jit_dispatch_ctx");
if (!DispatchCtx)
DispatchCtx = new GlobalVariable(M, PointerType::get(M.getContext(), 0),
false, GlobalValue::ExternalLinkage,
nullptr, "__orc_rt_jit_dispatch_ctx");
GlobalVariable *ReoptimizeTag =
M.getGlobalVariable("__orc_rt_reoptimize_tag");
if (!ReoptimizeTag)
ReoptimizeTag = new GlobalVariable(M, PointerType::get(M.getContext(), 0),
false, GlobalValue::ExternalLinkage,
nullptr, "__orc_rt_reoptimize_tag");
Function *DispatchFunc = M.getFunction("__orc_rt_jit_dispatch");
if (!DispatchFunc) {
std::vector<Type *> Args = {PointerType::get(M.getContext(), 0),
PointerType::get(M.getContext(), 0),
PointerType::get(M.getContext(), 0),
IntegerType::get(M.getContext(), 64)};
FunctionType *FuncTy =
FunctionType::get(Type::getVoidTy(M.getContext()), Args, false);
DispatchFunc = Function::Create(FuncTy, GlobalValue::ExternalLinkage,
"__orc_rt_jit_dispatch", &M);
}
size_t ArgBufferSizeConst =
SPSReoptimizeArgList::size(ReOptMaterializationUnitID{}, uint32_t{});
Constant *ArgBufferSize = ConstantInt::get(
IntegerType::get(M.getContext(), 64), ArgBufferSizeConst, false);
IRBuilder<> IRB(&IP);
(void)IRB.CreateCall(DispatchFunc,
{DispatchCtx, ReoptimizeTag, ArgBuffer, ArgBufferSize});
}

ReOptimizeLayer::ReOptMaterializationUnitState &
ReOptimizeLayer::createMaterializationUnitState(const ThreadSafeModule &TSM) {
std::unique_lock<std::mutex> Lock(Mutex);
ReOptMaterializationUnitID MUID = NextID;
MUStates.emplace(MUID,
ReOptMaterializationUnitState(MUID, cloneToNewContext(TSM)));
++NextID;
return MUStates.at(MUID);
}

ReOptimizeLayer::ReOptMaterializationUnitState &
ReOptimizeLayer::getMaterializationUnitState(ReOptMaterializationUnitID MUID) {
std::unique_lock<std::mutex> Lock(Mutex);
return MUStates.at(MUID);
}

void ReOptimizeLayer::registerMaterializationUnitResource(
ResourceKey Key, ReOptMaterializationUnitState &State) {
std::unique_lock<std::mutex> Lock(Mutex);
MUResources[Key].insert(State.getID());
}

Error ReOptimizeLayer::handleRemoveResources(JITDylib &JD, ResourceKey K) {
std::unique_lock<std::mutex> Lock(Mutex);
for (auto MUID : MUResources[K])
MUStates.erase(MUID);

MUResources.erase(K);
return Error::success();
}

void ReOptimizeLayer::handleTransferResources(JITDylib &JD, ResourceKey DstK,
ResourceKey SrcK) {
std::unique_lock<std::mutex> Lock(Mutex);
MUResources[DstK].insert(MUResources[SrcK].begin(), MUResources[SrcK].end());
MUResources.erase(SrcK);
}
24 changes: 24 additions & 0 deletions llvm/lib/ExecutionEngine/Orc/RedirectionManager.cpp
View file
Original file line number Diff line number Diff line change
@@ -0,0 +1,24 @@
//===---- RedirectionManager.cpp - Redirection manager interface in Orc ---===//
//
// 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/RedirectionManager.h"

#define DEBUG_TYPE "orc"

using namespace llvm;
using namespace llvm::orc;

void RedirectionManager::anchor() {}

Error RedirectableSymbolManager::createRedirectableSymbols(
ResourceTrackerSP RT, const SymbolMap &InitialDests) {
auto &JD = RT->getJITDylib();
return JD.define(
std::make_unique<RedirectableMaterializationUnit>(*this, InitialDests),
RT);
}
3 changes: 2 additions & 1 deletion llvm/tools/lli/lli.cpp
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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"
Expand Down Expand Up @@ -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(
Expand Down
2 changes: 2 additions & 0 deletions llvm/unittests/ExecutionEngine/Orc/CMakeLists.txt
View file
Original file line number Diff line number Diff line change
Expand Up @@ -43,6 +43,8 @@ add_llvm_unittest(OrcJITTests
TaskDispatchTest.cpp
ThreadSafeModuleTest.cpp
WrapperFunctionUtilsTest.cpp
JITLinkRedirectionManagerTest.cpp
ReOptimizeLayerTest.cpp

EXPORT_SYMBOLS
)
Expand Down
99 changes: 99 additions & 0 deletions llvm/unittests/ExecutionEngine/Orc/JITLinkRedirectionManagerTest.cpp
View file
Original file line number Diff line number Diff line change
@@ -0,0 +1,99 @@
#include "OrcTestCommon.h"
#include "llvm/ExecutionEngine/JITLink/JITLinkMemoryManager.h"
#include "llvm/ExecutionEngine/Orc/ExecutorProcessControl.h"
#include "llvm/ExecutionEngine/Orc/JITLinkRedirectableSymbolManager.h"
#include "llvm/ExecutionEngine/Orc/JITTargetMachineBuilder.h"
#include "llvm/ExecutionEngine/Orc/ObjectLinkingLayer.h"
#include "llvm/ExecutionEngine/Orc/Shared/ExecutorAddress.h"
#include "llvm/Testing/Support/Error.h"
#include "gtest/gtest.h"

using namespace llvm;
using namespace llvm::orc;
using namespace llvm::jitlink;

static int initialTarget() { return 42; }
static int middleTarget() { return 13; }
static int finalTarget() { return 53; }

class JITLinkRedirectionManagerTest : public testing::Test {
public:
~JITLinkRedirectionManagerTest() {
if (ES)
if (auto Err = ES->endSession())
ES->reportError(std::move(Err));
}

protected:
void SetUp() override {
auto JTMB = JITTargetMachineBuilder::detectHost();
// Bail out if we can not detect the host.
if (!JTMB) {
consumeError(JTMB.takeError());
GTEST_SKIP();
}

ES = std::make_unique<ExecutionSession>(
std::make_unique<UnsupportedExecutorProcessControl>(
nullptr, nullptr, JTMB->getTargetTriple().getTriple()));
JD = &ES->createBareJITDylib("main");
ObjLinkingLayer = std::make_unique<ObjectLinkingLayer>(
*ES, std::make_unique<InProcessMemoryManager>(16384));
DL = std::make_unique<DataLayout>(
cantFail(JTMB->getDefaultDataLayoutForTarget()));
}
JITDylib *JD{nullptr};
std::unique_ptr<ExecutionSession> ES;
std::unique_ptr<ObjectLinkingLayer> ObjLinkingLayer;
std::unique_ptr<DataLayout> DL;
};

TEST_F(JITLinkRedirectionManagerTest, BasicRedirectionOperation) {
auto RM = JITLinkRedirectableSymbolManager::Create(*ObjLinkingLayer, *JD);
// Bail out if we can not create
if (!RM) {
consumeError(RM.takeError());
GTEST_SKIP();
}

auto DefineTarget = [&](StringRef TargetName, ExecutorAddr Addr) {
SymbolStringPtr Target = ES->intern(TargetName);
cantFail(JD->define(std::make_unique<SimpleMaterializationUnit>(
SymbolFlagsMap({{Target, JITSymbolFlags::Exported}}),
[&](std::unique_ptr<MaterializationResponsibility> R) -> void {
// No dependencies registered, can't fail.
cantFail(
R->notifyResolved({{Target, {Addr, JITSymbolFlags::Exported}}}));
cantFail(R->notifyEmitted({}));
})));
return cantFail(ES->lookup({JD}, TargetName));
};

auto InitialTarget =
DefineTarget("InitialTarget", ExecutorAddr::fromPtr(&initialTarget));
auto MiddleTarget =
DefineTarget("MiddleTarget", ExecutorAddr::fromPtr(&middleTarget));
auto FinalTarget =
DefineTarget("FinalTarget", ExecutorAddr::fromPtr(&finalTarget));

auto RedirectableSymbol = ES->intern("RedirectableTarget");
EXPECT_THAT_ERROR(
(*RM)->createRedirectableSymbols(JD->getDefaultResourceTracker(),
{{RedirectableSymbol, InitialTarget}}),
Succeeded());
auto RTDef = cantFail(ES->lookup({JD}, RedirectableSymbol));

auto RTPtr = RTDef.getAddress().toPtr<int (*)()>();
auto Result = RTPtr();
EXPECT_EQ(Result, 42) << "Failed to call initial target";

EXPECT_THAT_ERROR((*RM)->redirect(*JD, {{RedirectableSymbol, MiddleTarget}}),
Succeeded());
Result = RTPtr();
EXPECT_EQ(Result, 13) << "Failed to call middle redirected target";

EXPECT_THAT_ERROR((*RM)->redirect(*JD, {{RedirectableSymbol, FinalTarget}}),
Succeeded());
Result = RTPtr();
EXPECT_EQ(Result, 53) << "Failed to call redirected target";
}
152 changes: 152 additions & 0 deletions llvm/unittests/ExecutionEngine/Orc/ReOptimizeLayerTest.cpp
View file
Original file line number Diff line number Diff line change
@@ -0,0 +1,152 @@
#include "llvm/ExecutionEngine/Orc/ReOptimizeLayer.h"
#include "OrcTestCommon.h"
#include "llvm/ExecutionEngine/JITLink/JITLinkMemoryManager.h"
#include "llvm/ExecutionEngine/Orc/CompileUtils.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/JITLinkRedirectableSymbolManager.h"
#include "llvm/ExecutionEngine/Orc/JITTargetMachineBuilder.h"
#include "llvm/ExecutionEngine/Orc/ObjectLinkingLayer.h"
#include "llvm/ExecutionEngine/Orc/ObjectTransformLayer.h"
#include "llvm/ExecutionEngine/Orc/Shared/ExecutorAddress.h"
#include "llvm/IR/IRBuilder.h"
#include "llvm/Support/CodeGen.h"
#include "llvm/TargetParser/Host.h"
#include "llvm/Testing/Support/Error.h"
#include "gtest/gtest.h"

using namespace llvm;
using namespace llvm::orc;
using namespace llvm::jitlink;

class ReOptimizeLayerTest : public testing::Test {
public:
~ReOptimizeLayerTest() {
if (ES)
if (auto Err = ES->endSession())
ES->reportError(std::move(Err));
}

protected:
void SetUp() override {
auto JTMB = JITTargetMachineBuilder::detectHost();
// Bail out if we can not detect the host.
if (!JTMB) {
consumeError(JTMB.takeError());
GTEST_SKIP();
}

auto EPC = SelfExecutorProcessControl::Create();
if (!EPC) {
consumeError(EPC.takeError());
GTEST_SKIP();
}
ES = std::make_unique<ExecutionSession>(std::move(*EPC));
JD = &ES->createBareJITDylib("main");
ObjLinkingLayer = std::make_unique<ObjectLinkingLayer>(
*ES, std::make_unique<InProcessMemoryManager>(16384));
DL = std::make_unique<DataLayout>(
cantFail(JTMB->getDefaultDataLayoutForTarget()));

auto TM = JTMB->createTargetMachine();
if (!TM) {
consumeError(TM.takeError());
GTEST_SKIP();
}
auto CompileFunction =
std::make_unique<TMOwningSimpleCompiler>(std::move(*TM));
CompileLayer = std::make_unique<IRCompileLayer>(*ES, *ObjLinkingLayer,
std::move(CompileFunction));
}

Error addIRModule(ResourceTrackerSP RT, ThreadSafeModule TSM) {
assert(TSM && "Can not add null module");

TSM.withModuleDo([&](Module &M) { M.setDataLayout(*DL); });

return ROLayer->add(std::move(RT), std::move(TSM));
}

JITDylib *JD{nullptr};
std::unique_ptr<ExecutionSession> ES;
std::unique_ptr<ObjectLinkingLayer> ObjLinkingLayer;
std::unique_ptr<IRCompileLayer> CompileLayer;
std::unique_ptr<ReOptimizeLayer> ROLayer;
std::unique_ptr<DataLayout> DL;
};

static Function *createRetFunction(Module *M, StringRef Name,
uint32_t ReturnCode) {
Function *Result = Function::Create(
FunctionType::get(Type::getInt32Ty(M->getContext()), {}, false),
GlobalValue::ExternalLinkage, Name, M);

BasicBlock *BB = BasicBlock::Create(M->getContext(), Name, Result);
IRBuilder<> Builder(M->getContext());
Builder.SetInsertPoint(BB);

Value *RetValue = ConstantInt::get(M->getContext(), APInt(32, ReturnCode));
Builder.CreateRet(RetValue);
return Result;
}

TEST_F(ReOptimizeLayerTest, BasicReOptimization) {
MangleAndInterner Mangle(*ES, *DL);

auto &EPC = ES->getExecutorProcessControl();
EXPECT_THAT_ERROR(JD->define(absoluteSymbols(
{{Mangle("__orc_rt_jit_dispatch"),
{EPC.getJITDispatchInfo().JITDispatchFunction,
JITSymbolFlags::Exported}},
{Mangle("__orc_rt_jit_dispatch_ctx"),
{EPC.getJITDispatchInfo().JITDispatchContext,
JITSymbolFlags::Exported}},
{Mangle("__orc_rt_reoptimize_tag"),
{ExecutorAddr(), JITSymbolFlags::Exported}}})),
Succeeded());

auto RM = JITLinkRedirectableSymbolManager::Create(*ObjLinkingLayer, *JD);
EXPECT_THAT_ERROR(RM.takeError(), Succeeded());

ROLayer = std::make_unique<ReOptimizeLayer>(*ES, *DL, *CompileLayer, **RM);
ROLayer->setReoptimizeFunc(
[&](ReOptimizeLayer &Parent,
ReOptimizeLayer::ReOptMaterializationUnitID MUID, unsigned CurVerison,
ResourceTrackerSP OldRT, ThreadSafeModule &TSM) {
TSM.withModuleDo([&](Module &M) {
for (auto &F : M) {
if (F.isDeclaration())
continue;
for (auto &B : F) {
for (auto &I : B) {
if (ReturnInst *Ret = dyn_cast<ReturnInst>(&I)) {
Value *RetValue =
ConstantInt::get(M.getContext(), APInt(32, 53));
Ret->setOperand(0, RetValue);
}
}
}
}
});
return Error::success();
});
EXPECT_THAT_ERROR(ROLayer->reigsterRuntimeFunctions(*JD), Succeeded());

ThreadSafeContext Ctx(std::make_unique<LLVMContext>());
auto M = std::make_unique<Module>("<main>", *Ctx.getContext());
M->setTargetTriple(sys::getProcessTriple());

(void)createRetFunction(M.get(), "main", 42);

EXPECT_THAT_ERROR(addIRModule(JD->getDefaultResourceTracker(),
ThreadSafeModule(std::move(M), std::move(Ctx))),
Succeeded());

auto Result = cantFail(ES->lookup({JD}, Mangle("main")));
auto FuncPtr = Result.getAddress().toPtr<int (*)()>();
for (size_t I = 0; I <= ReOptimizeLayer::CallCountThreshold; I++)
EXPECT_EQ(FuncPtr(), 42);
EXPECT_EQ(FuncPtr(), 53);
}