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