5be1d4b0dc
fa7f7a4cab4 changed the jit-dispatch function signature used in the orc_rt_lite_reoptimize_helper function, but jit-dispatch still takes a raw data pointer and size argument. Should fix the bug in https://lab.llvm.org/buildbot/#/builders/169/builds/18319 and similar builds.
346 lines
12 KiB
C++
346 lines
12 KiB
C++
#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;
|
|
}
|
|
|
|
static void orc_rt_lite_reoptimize_helper(
|
|
shared::CWrapperFunctionBuffer (*JITDispatch)(void *Ctx, void *Tag,
|
|
const char *Data,
|
|
size_t Size),
|
|
void *JITDispatchCtx, void *Tag, uint64_t MUID, uint32_t CurVersion) {
|
|
// Serialize the arguments into a WrapperFunctionBuffer and call dispatch.
|
|
using SPSArgs = shared::SPSArgList<uint64_t, uint32_t>;
|
|
auto ArgBytes =
|
|
shared::WrapperFunctionBuffer::allocate(SPSArgs::size(MUID, CurVersion));
|
|
shared::SPSOutputBuffer OB(ArgBytes.data(), ArgBytes.size());
|
|
if (!SPSArgs::serialize(OB, MUID, CurVersion)) {
|
|
errs()
|
|
<< "Reoptimization error: could not serialize reoptimization arguments";
|
|
abort();
|
|
}
|
|
shared::WrapperFunctionBuffer Buf{
|
|
JITDispatch(JITDispatchCtx, Tag, ArgBytes.data(), ArgBytes.size())};
|
|
|
|
if (const char *ErrMsg = Buf.getOutOfBandError()) {
|
|
errs() << "Reoptimization error: " << ErrMsg << "\naborting.\n";
|
|
abort();
|
|
}
|
|
}
|
|
|
|
Error ReOptimizeLayer::addOrcRTLiteSupport(JITDylib &PlatformJD,
|
|
const DataLayout &DL) {
|
|
auto Ctx = std::make_unique<LLVMContext>();
|
|
auto Mod = std::make_unique<Module>("orc-rt-lite-reoptimize.ll", *Ctx);
|
|
Mod->setDataLayout(DL);
|
|
|
|
IRBuilder<> Builder(*Ctx);
|
|
|
|
// Create basic types portably
|
|
Type *VoidTy = Type::getVoidTy(*Ctx);
|
|
Type *Int8Ty = Type::getInt8Ty(*Ctx);
|
|
Type *Int32Ty = Type::getInt32Ty(*Ctx);
|
|
Type *Int64Ty = Type::getInt64Ty(*Ctx);
|
|
Type *VoidPtrTy = PointerType::getUnqual(*Ctx);
|
|
|
|
// Helper function type: void (void*, void*, void*, uint64_t, uint32_t)
|
|
FunctionType *HelperFnTy = FunctionType::get(
|
|
VoidTy, {VoidPtrTy, VoidPtrTy, VoidPtrTy, Int64Ty, Int32Ty}, false);
|
|
|
|
// Define ReoptimizeTag with initializer = 0
|
|
GlobalVariable *ReoptimizeTag = new GlobalVariable(
|
|
*Mod, Int8Ty, false, GlobalValue::ExternalLinkage,
|
|
ConstantInt::get(Int8Ty, 0), "__orc_rt_reoptimize_tag");
|
|
|
|
// Define orc_rt_lite_reoptimize function: void (uint64_t, uint32_t)
|
|
FunctionType *ReOptimizeFnTy =
|
|
FunctionType::get(VoidTy, {Int64Ty, Int32Ty}, false);
|
|
|
|
Function *ReOptimizeFn =
|
|
Function::Create(ReOptimizeFnTy, Function::ExternalLinkage,
|
|
"__orc_rt_reoptimize", Mod.get());
|
|
|
|
// Set parameter names
|
|
auto ArgIt = ReOptimizeFn->arg_begin();
|
|
Value *MUID = &*ArgIt++;
|
|
MUID->setName("MUID");
|
|
Value *CurVersion = &*ArgIt;
|
|
CurVersion->setName("CurVersion");
|
|
|
|
// Build function body
|
|
BasicBlock *Entry = BasicBlock::Create(*Ctx, "entry", ReOptimizeFn);
|
|
Builder.SetInsertPoint(Entry);
|
|
|
|
// Create absolute address constants
|
|
auto &JDI = PlatformJD.getExecutionSession()
|
|
.getExecutorProcessControl()
|
|
.getJITDispatchInfo();
|
|
|
|
Type *IntPtrTy = DL.getIntPtrType(*Ctx);
|
|
Constant *JITDispatchPtr = ConstantExpr::getIntToPtr(
|
|
ConstantInt::get(IntPtrTy, JDI.JITDispatchFunction.getValue()),
|
|
VoidPtrTy);
|
|
Constant *JITDispatchCtxPtr = ConstantExpr::getIntToPtr(
|
|
ConstantInt::get(IntPtrTy, JDI.JITDispatchContext.getValue()), VoidPtrTy);
|
|
Constant *HelperFnAddr = ConstantExpr::getIntToPtr(
|
|
ConstantInt::get(IntPtrTy, reinterpret_cast<uintptr_t>(
|
|
&orc_rt_lite_reoptimize_helper)),
|
|
PointerType::getUnqual(*Ctx));
|
|
|
|
// Cast ReoptimizeTag to void*
|
|
Value *ReoptimizeTagPtr = Builder.CreatePointerCast(ReoptimizeTag, VoidPtrTy);
|
|
|
|
// Call the helper function
|
|
Builder.CreateCall(
|
|
HelperFnTy, HelperFnAddr,
|
|
{JITDispatchPtr, JITDispatchCtxPtr, ReoptimizeTagPtr, MUID, CurVersion});
|
|
|
|
// Return void
|
|
Builder.CreateRetVoid();
|
|
|
|
return BaseLayer.add(PlatformJD,
|
|
ThreadSafeModule(std::move(Mod), std::move(Ctx)));
|
|
}
|
|
|
|
Error ReOptimizeLayer::registerRuntimeFunctions(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");
|
|
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, MUID, CurVersion);
|
|
}
|
|
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());
|
|
}
|
|
|
|
void ReOptimizeLayer::createReoptimizeCall(Module &M, Instruction &IP,
|
|
ReOptMaterializationUnitID MUID,
|
|
uint32_t CurVersion) {
|
|
Type *MUIDTy = IntegerType::get(M.getContext(), 64);
|
|
Type *VersionTy = IntegerType::get(M.getContext(), 32);
|
|
Function *ReoptimizeFunc = M.getFunction("__orc_rt_reoptimize");
|
|
if (!ReoptimizeFunc) {
|
|
std::vector<Type *> ArgTys = {MUIDTy, VersionTy};
|
|
FunctionType *FuncTy =
|
|
FunctionType::get(Type::getVoidTy(M.getContext()), ArgTys, false);
|
|
ReoptimizeFunc = Function::Create(FuncTy, GlobalValue::ExternalLinkage,
|
|
"__orc_rt_reoptimize", &M);
|
|
}
|
|
Constant *MUIDArg = ConstantInt::get(MUIDTy, MUID, false);
|
|
Constant *CurVersionArg = ConstantInt::get(VersionTy, CurVersion, false);
|
|
IRBuilder<> IRB(&IP);
|
|
(void)IRB.CreateCall(ReoptimizeFunc, {MUIDArg, CurVersionArg});
|
|
}
|
|
|
|
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_range(MUResources[SrcK]);
|
|
MUResources.erase(SrcK);
|
|
}
|