1#include "llvm/ExecutionEngine/Orc/ReOptimizeLayer.h"
2#include "llvm/ExecutionEngine/Orc/BindCallControllerHandlerSPS.h"
3#include "llvm/ExecutionEngine/Orc/LookupAndApply.h"
4#include "llvm/ExecutionEngine/Orc/Mangling.h"
5#include "llvm/ExecutionEngine/Orc/Shared/OrcRTBridge.h"
6
7using namespace llvm;
8using namespace orc;
9
10bool ReOptimizeLayer::ReOptMaterializationUnitState::tryStartReoptimize() {
11 std::unique_lock<std::mutex> Lock(Mutex);
12 if (Reoptimizing)
13 return false;
14
15 Reoptimizing = true;
16 return true;
17}
18
19void ReOptimizeLayer::ReOptMaterializationUnitState::reoptimizeSucceeded() {
20 std::unique_lock<std::mutex> Lock(Mutex);
21 assert(Reoptimizing && "Tried to mark unstarted reoptimization as done");
22 Reoptimizing = false;
23 CurVersion++;
24}
25
26void ReOptimizeLayer::ReOptMaterializationUnitState::reoptimizeFailed() {
27 std::unique_lock<std::mutex> Lock(Mutex);
28 assert(Reoptimizing && "Tried to mark unstarted reoptimization as done");
29 Reoptimizing = false;
30}
31
32static void orc_rt_lite_reoptimize_helper(
33 shared::CWrapperFunctionBuffer (*JITDispatch)(void *Ctx, void *Tag,
34 const char *Data,
35 size_t Size),
36 void *JITDispatchCtx, void *Tag, uint64_t MUID, uint32_t CurVersion) {
37 // Serialize the arguments into a WrapperFunctionBuffer and call dispatch.
38 using SPSArgs = shared::SPSArgList<uint64_t, uint32_t>;
39 auto ArgBytes =
40 shared::WrapperFunctionBuffer::allocate(Size: SPSArgs::size(Arg: MUID, Args: CurVersion));
41 shared::SPSOutputBuffer OB(ArgBytes.data(), ArgBytes.size());
42 if (!SPSArgs::serialize(OB, Arg: MUID, Args: CurVersion)) {
43 errs()
44 << "Reoptimization error: could not serialize reoptimization arguments";
45 abort();
46 }
47 shared::WrapperFunctionBuffer Buf{
48 JITDispatch(JITDispatchCtx, Tag, ArgBytes.data(), ArgBytes.size())};
49
50 if (const char *ErrMsg = Buf.getOutOfBandError()) {
51 errs() << "Reoptimization error: " << ErrMsg << "\naborting.\n";
52 abort();
53 }
54}
55
56Error ReOptimizeLayer::addOrcRTLiteSupport(JITDylib &PlatformJD,
57 const DataLayout &DL) {
58 auto Ctx = std::make_unique<LLVMContext>();
59 auto Mod = std::make_unique<Module>(args: "orc-rt-lite-reoptimize.ll", args&: *Ctx);
60 Mod->setDataLayout(DL);
61
62 IRBuilder<> Builder(*Mod);
63
64 // Create basic types portably
65 Type *VoidTy = Type::getVoidTy(C&: *Ctx);
66 Type *Int8Ty = Type::getInt8Ty(C&: *Ctx);
67 Type *Int32Ty = Type::getInt32Ty(C&: *Ctx);
68 Type *Int64Ty = Type::getInt64Ty(C&: *Ctx);
69 Type *VoidPtrTy = PointerType::getUnqual(C&: *Ctx);
70
71 // Helper function type: void (void*, void*, void*, uint64_t, uint32_t)
72 FunctionType *HelperFnTy = FunctionType::get(
73 Result: VoidTy, Params: {VoidPtrTy, VoidPtrTy, VoidPtrTy, Int64Ty, Int32Ty}, isVarArg: false);
74
75 // Define ReoptimizeTag with initializer = 0
76 GlobalVariable *ReoptimizeTag = new GlobalVariable(
77 *Mod, Int8Ty, false, GlobalValue::ExternalLinkage,
78 ConstantInt::get(Ty: Int8Ty, V: 0), "__orc_rt_reoptimize_tag");
79
80 // Define orc_rt_lite_reoptimize function: void (uint64_t, uint32_t)
81 FunctionType *ReOptimizeFnTy =
82 FunctionType::get(Result: VoidTy, Params: {Int64Ty, Int32Ty}, isVarArg: false);
83
84 Function *ReOptimizeFn =
85 Function::Create(Ty: ReOptimizeFnTy, Linkage: Function::ExternalLinkage,
86 N: "__orc_rt_reoptimize", M: Mod.get());
87
88 // Set parameter names
89 auto ArgIt = ReOptimizeFn->arg_begin();
90 Value *MUID = &*ArgIt++;
91 MUID->setName("MUID");
92 Value *CurVersion = &*ArgIt;
93 CurVersion->setName("CurVersion");
94
95 // Build function body
96 BasicBlock *Entry = BasicBlock::Create(Context&: *Ctx, Name: "entry", Parent: ReOptimizeFn);
97 Builder.SetInsertPoint(Entry);
98
99 ExecutorAddr JITDispatchSym, JITDispatchCtxSym;
100 if (auto Err =
101 lookupAndApply(JD&: ES.getBootstrapJITDylib(),
102 PrepareFns: {recordAddr(Name: rt::DispatchName, A: &JITDispatchSym),
103 recordAddr(Name: rt::DispatchCtxName, A: &JITDispatchCtxSym)}))
104 return Err;
105
106 Type *IntPtrTy = DL.getIntPtrType(C&: *Ctx);
107 Constant *JITDispatchPtr = ConstantExpr::getIntToPtr(
108 C: ConstantInt::get(Ty: IntPtrTy, V: JITDispatchSym.getValue()), Ty: VoidPtrTy);
109 Constant *JITDispatchCtxPtr = ConstantExpr::getIntToPtr(
110 C: ConstantInt::get(Ty: IntPtrTy, V: JITDispatchCtxSym.getValue()), Ty: VoidPtrTy);
111 Constant *HelperFnAddr = ConstantExpr::getIntToPtr(
112 C: ConstantInt::get(Ty: IntPtrTy, V: reinterpret_cast<uintptr_t>(
113 &orc_rt_lite_reoptimize_helper)),
114 Ty: PointerType::getUnqual(C&: *Ctx));
115
116 // Cast ReoptimizeTag to void*
117 Value *ReoptimizeTagPtr = Builder.CreatePointerCast(V: ReoptimizeTag, DestTy: VoidPtrTy);
118
119 // Call the helper function
120 Builder.CreateCall(
121 FTy: HelperFnTy, Callee: HelperFnAddr,
122 Args: {JITDispatchPtr, JITDispatchCtxPtr, ReoptimizeTagPtr, MUID, CurVersion});
123
124 // Return void
125 Builder.CreateRetVoid();
126
127 return BaseLayer.add(JD&: PlatformJD,
128 TSM: ThreadSafeModule(std::move(Mod), std::move(Ctx)));
129}
130
131Error ReOptimizeLayer::registerRuntimeFunctions(JITDylib &PlatformJD) {
132 using ReoptimizeSPSSig = shared::SPSError(uint64_t, uint32_t);
133 return ES.registerCallControllerHandlers(
134 JD&: PlatformJD, Hs: bindCallControllerHandlerSPS<ReoptimizeSPSSig>(
135 Name: SymbolNameSpec::c(Name: "__orc_rt_reoptimize_tag"), Instance: this,
136 Method: &ReOptimizeLayer::rt_reoptimize));
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