1#include "llvm/ExecutionEngine/Orc/ReOptimizeLayer.h"
2#include "llvm/ExecutionEngine/Orc/LookupAndApply.h"
3#include "llvm/ExecutionEngine/Orc/Mangling.h"
4#include "llvm/ExecutionEngine/Orc/Shared/OrcRTBridge.h"
5
6using namespace llvm;
7using namespace orc;
8
9bool ReOptimizeLayer::ReOptMaterializationUnitState::tryStartReoptimize() {
10 std::unique_lock<std::mutex> Lock(Mutex);
11 if (Reoptimizing)
12 return false;
13
14 Reoptimizing = true;
15 return true;
16}
17
18void ReOptimizeLayer::ReOptMaterializationUnitState::reoptimizeSucceeded() {
19 std::unique_lock<std::mutex> Lock(Mutex);
20 assert(Reoptimizing && "Tried to mark unstarted reoptimization as done");
21 Reoptimizing = false;
22 CurVersion++;
23}
24
25void ReOptimizeLayer::ReOptMaterializationUnitState::reoptimizeFailed() {
26 std::unique_lock<std::mutex> Lock(Mutex);
27 assert(Reoptimizing && "Tried to mark unstarted reoptimization as done");
28 Reoptimizing = false;
29}
30
31static void orc_rt_lite_reoptimize_helper(
32 shared::CWrapperFunctionBuffer (*JITDispatch)(void *Ctx, void *Tag,
33 const char *Data,
34 size_t Size),
35 void *JITDispatchCtx, void *Tag, uint64_t MUID, uint32_t CurVersion) {
36 // Serialize the arguments into a WrapperFunctionBuffer and call dispatch.
37 using SPSArgs = shared::SPSArgList<uint64_t, uint32_t>;
38 auto ArgBytes =
39 shared::WrapperFunctionBuffer::allocate(Size: SPSArgs::size(Arg: MUID, Args: CurVersion));
40 shared::SPSOutputBuffer OB(ArgBytes.data(), ArgBytes.size());
41 if (!SPSArgs::serialize(OB, Arg: MUID, Args: CurVersion)) {
42 errs()
43 << "Reoptimization error: could not serialize reoptimization arguments";
44 abort();
45 }
46 shared::WrapperFunctionBuffer Buf{
47 JITDispatch(JITDispatchCtx, Tag, ArgBytes.data(), ArgBytes.size())};
48
49 if (const char *ErrMsg = Buf.getOutOfBandError()) {
50 errs() << "Reoptimization error: " << ErrMsg << "\naborting.\n";
51 abort();
52 }
53}
54
55Error ReOptimizeLayer::addOrcRTLiteSupport(JITDylib &PlatformJD,
56 const DataLayout &DL) {
57 auto Ctx = std::make_unique<LLVMContext>();
58 auto Mod = std::make_unique<Module>(args: "orc-rt-lite-reoptimize.ll", args&: *Ctx);
59 Mod->setDataLayout(DL);
60
61 IRBuilder<> Builder(*Ctx);
62
63 // Create basic types portably
64 Type *VoidTy = Type::getVoidTy(C&: *Ctx);
65 Type *Int8Ty = Type::getInt8Ty(C&: *Ctx);
66 Type *Int32Ty = Type::getInt32Ty(C&: *Ctx);
67 Type *Int64Ty = Type::getInt64Ty(C&: *Ctx);
68 Type *VoidPtrTy = PointerType::getUnqual(C&: *Ctx);
69
70 // Helper function type: void (void*, void*, void*, uint64_t, uint32_t)
71 FunctionType *HelperFnTy = FunctionType::get(
72 Result: VoidTy, Params: {VoidPtrTy, VoidPtrTy, VoidPtrTy, Int64Ty, Int32Ty}, isVarArg: false);
73
74 // Define ReoptimizeTag with initializer = 0
75 GlobalVariable *ReoptimizeTag = new GlobalVariable(
76 *Mod, Int8Ty, false, GlobalValue::ExternalLinkage,
77 ConstantInt::get(Ty: Int8Ty, V: 0), "__orc_rt_reoptimize_tag");
78
79 // Define orc_rt_lite_reoptimize function: void (uint64_t, uint32_t)
80 FunctionType *ReOptimizeFnTy =
81 FunctionType::get(Result: VoidTy, Params: {Int64Ty, Int32Ty}, isVarArg: false);
82
83 Function *ReOptimizeFn =
84 Function::Create(Ty: ReOptimizeFnTy, Linkage: Function::ExternalLinkage,
85 N: "__orc_rt_reoptimize", M: Mod.get());
86
87 // Set parameter names
88 auto ArgIt = ReOptimizeFn->arg_begin();
89 Value *MUID = &*ArgIt++;
90 MUID->setName("MUID");
91 Value *CurVersion = &*ArgIt;
92 CurVersion->setName("CurVersion");
93
94 // Build function body
95 BasicBlock *Entry = BasicBlock::Create(Context&: *Ctx, Name: "entry", Parent: ReOptimizeFn);
96 Builder.SetInsertPoint(Entry);
97
98 ExecutorAddr JITDispatchSym, JITDispatchCtxSym;
99 if (auto Err =
100 lookupAndApply(JD&: ES.getBootstrapJITDylib(),
101 PrepareFns: {recordAddr(Name: rt::DispatchName, A: &JITDispatchSym),
102 recordAddr(Name: rt::DispatchCtxName, A: &JITDispatchCtxSym)}))
103 return Err;
104
105 Type *IntPtrTy = DL.getIntPtrType(C&: *Ctx);
106 Constant *JITDispatchPtr = ConstantExpr::getIntToPtr(
107 C: ConstantInt::get(Ty: IntPtrTy, V: JITDispatchSym.getValue()), Ty: VoidPtrTy);
108 Constant *JITDispatchCtxPtr = ConstantExpr::getIntToPtr(
109 C: ConstantInt::get(Ty: IntPtrTy, V: JITDispatchCtxSym.getValue()), Ty: VoidPtrTy);
110 Constant *HelperFnAddr = ConstantExpr::getIntToPtr(
111 C: ConstantInt::get(Ty: IntPtrTy, V: reinterpret_cast<uintptr_t>(
112 &orc_rt_lite_reoptimize_helper)),
113 Ty: PointerType::getUnqual(C&: *Ctx));
114
115 // Cast ReoptimizeTag to void*
116 Value *ReoptimizeTagPtr = Builder.CreatePointerCast(V: ReoptimizeTag, DestTy: VoidPtrTy);
117
118 // Call the helper function
119 Builder.CreateCall(
120 FTy: HelperFnTy, Callee: HelperFnAddr,
121 Args: {JITDispatchPtr, JITDispatchCtxPtr, ReoptimizeTagPtr, MUID, CurVersion});
122
123 // Return void
124 Builder.CreateRetVoid();
125
126 return BaseLayer.add(JD&: PlatformJD,
127 TSM: ThreadSafeModule(std::move(Mod), std::move(Ctx)));
128}
129
130Error ReOptimizeLayer::registerRuntimeFunctions(JITDylib &PlatformJD) {
131 ExecutionSession::JITDispatchHandlerAssociationMap WFs;
132 using ReoptimizeSPSSig = shared::SPSError(uint64_t, uint32_t);
133 WFs[Mangle("__orc_rt_reoptimize_tag")] =
134 ES.wrapAsyncWithSPS<ReoptimizeSPSSig>(Instance: this,
135 Method: &ReOptimizeLayer::rt_reoptimize);
136 return ES.registerJITDispatchHandlers(JD&: PlatformJD, WFs: std::move(WFs));
137}
138
139void ReOptimizeLayer::emit(std::unique_ptr<MaterializationResponsibility> R,
140 ThreadSafeModule TSM) {
141 auto &JD = R->getTargetJITDylib();
142
143 bool HasNonCallable = false;
144 for (auto &KV : R->getSymbols()) {
145 auto &Flags = KV.second;
146 if (!Flags.isCallable())
147 HasNonCallable = true;
148 }
149
150 if (HasNonCallable) {
151 BaseLayer.emit(R: std::move(R), TSM: std::move(TSM));
152 return;
153 }
154
155 auto &MUState = createMaterializationUnitState(TSM);
156
157 if (auto Err = R->withResourceKeyDo(F: [&](ResourceKey Key) {
158 registerMaterializationUnitResource(Key, State&: MUState);
159 })) {
160 ES.reportError(Err: std::move(Err));
161 R->failMaterialization();
162 return;
163 }
164
165 if (auto Err =
166 ProfilerFunc(*this, MUState.getID(), MUState.getCurVersion(), TSM)) {
167 ES.reportError(Err: std::move(Err));
168 R->failMaterialization();
169 return;
170 }
171
172 auto InitialDests =
173 emitMUImplSymbols(MUState, Version: MUState.getCurVersion(), JD, TSM: std::move(TSM));
174 if (!InitialDests) {
175 ES.reportError(Err: InitialDests.takeError());
176 R->failMaterialization();
177 return;
178 }
179
180 RSManager.emitRedirectableSymbols(MR: std::move(R), InitialDests: std::move(*InitialDests));
181}
182
183Error ReOptimizeLayer::reoptimizeIfCallFrequent(ReOptimizeLayer &Parent,
184 ReOptMaterializationUnitID MUID,
185 unsigned CurVersion,
186 ThreadSafeModule &TSM) {
187 return TSM.withModuleDo(F: [&](Module &M) -> Error {
188 Type *I64Ty = Type::getInt64Ty(C&: M.getContext());
189 GlobalVariable *Counter = new GlobalVariable(
190 M, I64Ty, false, GlobalValue::InternalLinkage,
191 Constant::getNullValue(Ty: I64Ty), "__orc_reopt_counter");
192 for (auto &F : M) {
193 if (F.isDeclaration())
194 continue;
195 auto &BB = F.getEntryBlock();
196 auto *IP = &*BB.getFirstInsertionPt();
197 IRBuilder<> IRB(IP);
198 Value *Threshold = ConstantInt::get(Ty: I64Ty, V: CallCountThreshold, IsSigned: true);
199 Value *Cnt = IRB.CreateLoad(Ty: I64Ty, Ptr: Counter);
200 // Use EQ to prevent further reoptimize calls.
201 Value *Cmp = IRB.CreateICmpEQ(LHS: Cnt, RHS: Threshold);
202 Value *Added = IRB.CreateAdd(LHS: Cnt, RHS: ConstantInt::get(Ty: I64Ty, V: 1));
203 (void)IRB.CreateStore(Val: Added, Ptr: Counter);
204 Instruction *SplitTerminator = SplitBlockAndInsertIfThen(Cond: Cmp, SplitBefore: IP, Unreachable: false);
205 createReoptimizeCall(M, IP&: *SplitTerminator, MUID, CurVersion);
206 }
207 return Error::success();
208 });
209}
210
211Expected<SymbolMap>
212ReOptimizeLayer::emitMUImplSymbols(ReOptMaterializationUnitState &MUState,
213 uint32_t Version, JITDylib &JD,
214 ThreadSafeModule TSM) {
215 DenseMap<SymbolStringPtr, SymbolStringPtr> RenamedMap;
216 cantFail(Err: TSM.withModuleDo(F: [&](Module &M) -> Error {
217 MangleAndInterner Mangle(ES, M.getDataLayout());
218 for (auto &F : M)
219 if (!F.isDeclaration()) {
220 std::string NewName =
221 (F.getName() + ".__def__." + Twine(Version)).str();
222 RenamedMap[Mangle(F.getName())] = Mangle(NewName);
223 F.setName(NewName);
224 }
225 return Error::success();
226 }));
227
228 auto RT = JD.createResourceTracker();
229 if (auto Err =
230 JD.define(MU: std::make_unique<BasicIRLayerMaterializationUnit>(
231 args&: BaseLayer, args: *getManglingOptions(), args: std::move(TSM)),
232 RT))
233 return Err;
234 MUState.setResourceTracker(RT);
235
236 SymbolLookupSet LookupSymbols;
237 for (auto [K, V] : RenamedMap)
238 LookupSymbols.add(Name: V);
239
240 auto ImplSymbols =
241 ES.lookup(SearchOrder: {{&JD, JITDylibLookupFlags::MatchAllSymbols}}, Symbols: LookupSymbols,
242 K: LookupKind::Static, RequiredState: SymbolState::Resolved);
243 if (auto Err = ImplSymbols.takeError())
244 return Err;
245
246 SymbolMap Result;
247 for (auto [K, V] : RenamedMap)
248 Result[K] = (*ImplSymbols)[V];
249
250 return Result;
251}
252
253void ReOptimizeLayer::rt_reoptimize(SendErrorFn SendResult,
254 ReOptMaterializationUnitID MUID,
255 uint32_t CurVersion) {
256 auto &MUState = getMaterializationUnitState(MUID);
257 if (CurVersion < MUState.getCurVersion() || !MUState.tryStartReoptimize()) {
258 SendResult(Error::success());
259 return;
260 }
261
262 ThreadSafeModule TSM = cloneToNewContext(TSMW: MUState.getThreadSafeModule());
263 auto OldRT = MUState.getResourceTracker();
264 auto &JD = OldRT->getJITDylib();
265
266 if (auto Err = ReOptFunc(*this, MUID, CurVersion + 1, OldRT, TSM)) {
267 ES.reportError(Err: std::move(Err));
268 MUState.reoptimizeFailed();
269 SendResult(Error::success());
270 return;
271 }
272
273 auto SymbolDests =
274 emitMUImplSymbols(MUState, Version: CurVersion + 1, JD, TSM: std::move(TSM));
275 if (!SymbolDests) {
276 ES.reportError(Err: SymbolDests.takeError());
277 MUState.reoptimizeFailed();
278 SendResult(Error::success());
279 return;
280 }
281
282 if (auto Err = RSManager.redirect(JD, NewDests: std::move(*SymbolDests))) {
283 ES.reportError(Err: std::move(Err));
284 MUState.reoptimizeFailed();
285 SendResult(Error::success());
286 return;
287 }
288
289 MUState.reoptimizeSucceeded();
290 SendResult(Error::success());
291}
292
293void ReOptimizeLayer::createReoptimizeCall(Module &M, Instruction &IP,
294 ReOptMaterializationUnitID MUID,
295 uint32_t CurVersion) {
296 Type *MUIDTy = IntegerType::get(C&: M.getContext(), NumBits: 64);
297 Type *VersionTy = IntegerType::get(C&: M.getContext(), NumBits: 32);
298 Function *ReoptimizeFunc = M.getFunction(Name: "__orc_rt_reoptimize");
299 if (!ReoptimizeFunc) {
300 std::vector<Type *> ArgTys = {MUIDTy, VersionTy};
301 FunctionType *FuncTy =
302 FunctionType::get(Result: Type::getVoidTy(C&: M.getContext()), Params: ArgTys, isVarArg: false);
303 ReoptimizeFunc = Function::Create(Ty: FuncTy, Linkage: GlobalValue::ExternalLinkage,
304 N: "__orc_rt_reoptimize", M: &M);
305 }
306 Constant *MUIDArg = ConstantInt::get(Ty: MUIDTy, V: MUID, IsSigned: false);
307 Constant *CurVersionArg = ConstantInt::get(Ty: VersionTy, V: CurVersion, IsSigned: false);
308 IRBuilder<> IRB(&IP);
309 (void)IRB.CreateCall(Callee: ReoptimizeFunc, Args: {MUIDArg, CurVersionArg});
310}
311
312ReOptimizeLayer::ReOptMaterializationUnitState &
313ReOptimizeLayer::createMaterializationUnitState(const ThreadSafeModule &TSM) {
314 std::unique_lock<std::mutex> Lock(Mutex);
315 ReOptMaterializationUnitID MUID = NextID;
316 MUStates.emplace(args&: MUID,
317 args: ReOptMaterializationUnitState(MUID, cloneToNewContext(TSMW: TSM)));
318 ++NextID;
319 return MUStates.at(k: MUID);
320}
321
322ReOptimizeLayer::ReOptMaterializationUnitState &
323ReOptimizeLayer::getMaterializationUnitState(ReOptMaterializationUnitID MUID) {
324 std::unique_lock<std::mutex> Lock(Mutex);
325 return MUStates.at(k: MUID);
326}
327
328void ReOptimizeLayer::registerMaterializationUnitResource(
329 ResourceKey Key, ReOptMaterializationUnitState &State) {
330 std::unique_lock<std::mutex> Lock(Mutex);
331 MUResources[Key].insert(V: State.getID());
332}
333
334Error ReOptimizeLayer::handleRemoveResources(JITDylib &JD, ResourceKey K) {
335 std::unique_lock<std::mutex> Lock(Mutex);
336 for (auto MUID : MUResources[K])
337 MUStates.erase(x: MUID);
338
339 MUResources.erase(Val: K);
340 return Error::success();
341}
342
343void ReOptimizeLayer::handleTransferResources(JITDylib &JD, ResourceKey DstK,
344 ResourceKey SrcK) {
345 std::unique_lock<std::mutex> Lock(Mutex);
346 MUResources[DstK].insert_range(R&: MUResources[SrcK]);
347 MUResources.erase(Val: SrcK);
348}
349