1//===- PGOCtxProfLowering.cpp - Contextual PGO Instr. Lowering ------------===//
2//
3// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
4// See https://llvm.org/LICENSE.txt for license information.
5// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
6//
7//===----------------------------------------------------------------------===//
8//
9
10#include "llvm/Transforms/Instrumentation/PGOCtxProfLowering.h"
11#include "llvm/ADT/STLExtras.h"
12#include "llvm/Analysis/CFG.h"
13#include "llvm/Analysis/OptimizationRemarkEmitter.h"
14#include "llvm/IR/Analysis.h"
15#include "llvm/IR/Constants.h"
16#include "llvm/IR/DiagnosticInfo.h"
17#include "llvm/IR/GlobalValue.h"
18#include "llvm/IR/IRBuilder.h"
19#include "llvm/IR/InstrTypes.h"
20#include "llvm/IR/Instructions.h"
21#include "llvm/IR/IntrinsicInst.h"
22#include "llvm/IR/Module.h"
23#include "llvm/IR/PassManager.h"
24#include "llvm/ProfileData/CtxInstrContextNode.h"
25#include "llvm/ProfileData/InstrProf.h"
26#include "llvm/Support/CommandLine.h"
27#include <utility>
28
29using namespace llvm;
30
31#define DEBUG_TYPE "ctx-instr-lower"
32
33static cl::list<std::string> ContextRoots(
34 "profile-context-root", cl::Hidden,
35 cl::desc(
36 "A function name, assumed to be global, which will be treated as the "
37 "root of an interesting graph, which will be profiled independently "
38 "from other similar graphs."));
39
40bool PGOCtxProfLoweringPass::isCtxIRPGOInstrEnabled() {
41 return !ContextRoots.empty();
42}
43
44// the names of symbols we expect in compiler-rt. Using a namespace for
45// readability.
46namespace CompilerRtAPINames {
47static auto StartCtx = "__llvm_ctx_profile_start_context";
48static auto ReleaseCtx = "__llvm_ctx_profile_release_context";
49static auto GetCtx = "__llvm_ctx_profile_get_context";
50static auto ExpectedCalleeTLS = "__llvm_ctx_profile_expected_callee";
51static auto CallsiteTLS = "__llvm_ctx_profile_callsite";
52} // namespace CompilerRtAPINames
53
54namespace {
55// The lowering logic and state.
56class CtxInstrumentationLowerer final {
57 Module &M;
58 ModuleAnalysisManager &MAM;
59 Type *ContextNodeTy = nullptr;
60 StructType *FunctionDataTy = nullptr;
61
62 DenseSet<const Function *> ContextRootSet;
63 Function *StartCtx = nullptr;
64 Function *GetCtx = nullptr;
65 Function *ReleaseCtx = nullptr;
66 GlobalVariable *ExpectedCalleeTLS = nullptr;
67 GlobalVariable *CallsiteInfoTLS = nullptr;
68 Constant *CannotBeRootInitializer = nullptr;
69
70public:
71 CtxInstrumentationLowerer(Module &M, ModuleAnalysisManager &MAM);
72 // return true if lowering happened (i.e. a change was made)
73 bool lowerFunction(Function &F);
74};
75
76// llvm.instrprof.increment[.step] captures the total number of counters as one
77// of its parameters, and llvm.instrprof.callsite captures the total number of
78// callsites. Those values are the same for instances of those intrinsics in
79// this function. Find the first instance of each and return them.
80std::pair<uint32_t, uint32_t> getNumCountersAndCallsites(const Function &F) {
81 uint32_t NumCounters = 0;
82 uint32_t NumCallsites = 0;
83 for (const auto &BB : F) {
84 for (const auto &I : BB) {
85 if (const auto *Incr = dyn_cast<InstrProfIncrementInst>(Val: &I)) {
86 uint32_t V =
87 static_cast<uint32_t>(Incr->getNumCounters()->getZExtValue());
88 assert((!NumCounters || V == NumCounters) &&
89 "expected all llvm.instrprof.increment[.step] intrinsics to "
90 "have the same total nr of counters parameter");
91 NumCounters = V;
92 } else if (const auto *CSIntr = dyn_cast<InstrProfCallsite>(Val: &I)) {
93 uint32_t V =
94 static_cast<uint32_t>(CSIntr->getNumCounters()->getZExtValue());
95 assert((!NumCallsites || V == NumCallsites) &&
96 "expected all llvm.instrprof.callsite intrinsics to have the "
97 "same total nr of callsites parameter");
98 NumCallsites = V;
99 }
100#ifdef NDEBUG
101 if (NumCounters && NumCallsites)
102 return std::make_pair(x&: NumCounters, y&: NumCallsites);
103#endif
104 }
105 }
106 return {NumCounters, NumCallsites};
107}
108
109void emitUnsupportedRootError(const Function &F, StringRef Reason) {
110 F.getContext().emitError(ErrorStr: "[ctxprof] The function " + F.getName() +
111 " was indicated as context root but " + Reason +
112 ", which is not supported.");
113}
114} // namespace
115
116// set up tie-in with compiler-rt.
117// NOTE!!!
118// These have to match compiler-rt/lib/ctx_profile/CtxInstrProfiling.h
119CtxInstrumentationLowerer::CtxInstrumentationLowerer(Module &M,
120 ModuleAnalysisManager &MAM)
121 : M(M), MAM(MAM) {
122 auto *PointerTy = PointerType::get(C&: M.getContext(), AddressSpace: 0);
123 auto *SanitizerMutexType = Type::getInt8Ty(C&: M.getContext());
124 auto *I32Ty = Type::getInt32Ty(C&: M.getContext());
125 auto *I64Ty = Type::getInt64Ty(C&: M.getContext());
126
127#define _PTRDECL(_, __) PointerTy,
128#define _VOLATILE_PTRDECL(_, __) PointerTy,
129#define _CONTEXT_ROOT PointerTy,
130#define _MUTEXDECL(_) SanitizerMutexType,
131
132 FunctionDataTy = StructType::get(
133 Context&: M.getContext(), Elements: {CTXPROF_FUNCTION_DATA(_PTRDECL, _CONTEXT_ROOT,
134 _VOLATILE_PTRDECL, _MUTEXDECL)});
135#undef _PTRDECL
136#undef _CONTEXT_ROOT
137#undef _VOLATILE_PTRDECL
138#undef _MUTEXDECL
139
140#define _PTRDECL(_, __) Constant::getNullValue(PointerTy),
141#define _VOLATILE_PTRDECL(_, __) _PTRDECL(_, __)
142#define _MUTEXDECL(_) Constant::getNullValue(SanitizerMutexType),
143#define _CONTEXT_ROOT \
144 Constant::getIntegerValue( \
145 PointerTy, \
146 APInt(M.getDataLayout().getPointerTypeSizeInBits(PointerTy), 1U)),
147 CannotBeRootInitializer = ConstantStruct::get(
148 T: FunctionDataTy, V: {CTXPROF_FUNCTION_DATA(_PTRDECL, _CONTEXT_ROOT,
149 _VOLATILE_PTRDECL, _MUTEXDECL)});
150#undef _PTRDECL
151#undef _CONTEXT_ROOT
152#undef _VOLATILE_PTRDECL
153#undef _MUTEXDECL
154
155 // The Context header.
156 ContextNodeTy = StructType::get(Context&: M.getContext(), Elements: {
157 I64Ty, /*Guid*/
158 PointerTy, /*Next*/
159 I32Ty, /*NumCounters*/
160 I32Ty, /*NumCallsites*/
161 });
162
163 // Define a global for each entrypoint. We'll reuse the entrypoint's name
164 // as prefix. We assume the entrypoint names to be unique.
165 for (const auto &Fname : ContextRoots) {
166 if (const auto *F = M.getFunction(Name: Fname)) {
167 if (F->isDeclaration())
168 continue;
169 ContextRootSet.insert(V: F);
170 for (const auto &BB : *F)
171 for (const auto &I : BB)
172 if (const auto *CB = dyn_cast<CallBase>(Val: &I))
173 if (CB->isMustTailCall())
174 emitUnsupportedRootError(F: *F, Reason: "it features musttail calls");
175 }
176 }
177
178 // Declare the functions we will call.
179 StartCtx = cast<Function>(
180 Val: M.getOrInsertFunction(
181 Name: CompilerRtAPINames::StartCtx,
182 T: FunctionType::get(Result: PointerTy,
183 Params: {PointerTy, /*FunctionData*/
184 I64Ty, /*Guid*/ I32Ty,
185 /*NumCounters*/ I32Ty /*NumCallsites*/},
186 isVarArg: false))
187 .getCallee());
188 GetCtx = cast<Function>(
189 Val: M.getOrInsertFunction(Name: CompilerRtAPINames::GetCtx,
190 T: FunctionType::get(Result: PointerTy,
191 Params: {PointerTy, /*FunctionData*/
192 PointerTy, /*Callee*/
193 I64Ty, /*Guid*/
194 I32Ty, /*NumCounters*/
195 I32Ty}, /*NumCallsites*/
196 isVarArg: false))
197 .getCallee());
198 ReleaseCtx = cast<Function>(
199 Val: M.getOrInsertFunction(Name: CompilerRtAPINames::ReleaseCtx,
200 T: FunctionType::get(Result: Type::getVoidTy(C&: M.getContext()),
201 Params: {
202 PointerTy, /*FunctionData*/
203 },
204 isVarArg: false))
205 .getCallee());
206
207 // Declare the TLSes we will need to use.
208 CallsiteInfoTLS =
209 new GlobalVariable(M, PointerTy, false, GlobalValue::ExternalLinkage,
210 nullptr, CompilerRtAPINames::CallsiteTLS);
211 CallsiteInfoTLS->setThreadLocal(true);
212 CallsiteInfoTLS->setVisibility(llvm::GlobalValue::HiddenVisibility);
213 ExpectedCalleeTLS =
214 new GlobalVariable(M, PointerTy, false, GlobalValue::ExternalLinkage,
215 nullptr, CompilerRtAPINames::ExpectedCalleeTLS);
216 ExpectedCalleeTLS->setThreadLocal(true);
217 ExpectedCalleeTLS->setVisibility(llvm::GlobalValue::HiddenVisibility);
218}
219
220PreservedAnalyses PGOCtxProfLoweringPass::run(Module &M,
221 ModuleAnalysisManager &MAM) {
222 CtxInstrumentationLowerer Lowerer(M, MAM);
223 bool Changed = false;
224 for (auto &F : M)
225 Changed |= Lowerer.lowerFunction(F);
226 return Changed ? PreservedAnalyses::none() : PreservedAnalyses::all();
227}
228
229bool CtxInstrumentationLowerer::lowerFunction(Function &F) {
230 if (F.isDeclaration())
231 return false;
232
233 // Probably pointless to try to do anything here, unlikely to be
234 // performance-affecting.
235 if (!llvm::canReturn(F)) {
236 for (auto &BB : F)
237 for (auto &I : make_early_inc_range(Range&: BB))
238 if (isa<InstrProfCntrInstBase>(Val: &I))
239 I.eraseFromParent();
240 if (ContextRootSet.contains(V: &F))
241 emitUnsupportedRootError(F, Reason: "it does not return");
242 return true;
243 }
244
245 auto &FAM = MAM.getResult<FunctionAnalysisManagerModuleProxy>(IR&: M).getManager();
246 auto &ORE = FAM.getResult<OptimizationRemarkEmitterAnalysis>(IR&: F);
247
248 Value *Guid = nullptr;
249 auto [NumCounters, NumCallsites] = getNumCountersAndCallsites(F);
250
251 Value *Context = nullptr;
252 Value *RealContext = nullptr;
253
254 StructType *ThisContextType = nullptr;
255 Value *TheRootFunctionData = nullptr;
256 Value *ExpectedCalleeTLSAddr = nullptr;
257 Value *CallsiteInfoTLSAddr = nullptr;
258 const bool HasMusttail = [&F]() {
259 for (auto &BB : F)
260 for (auto &I : BB)
261 if (auto *CB = dyn_cast<CallBase>(Val: &I))
262 if (CB->isMustTailCall())
263 return true;
264 return false;
265 }();
266
267 if (HasMusttail && ContextRootSet.contains(V: &F)) {
268 F.getContext().emitError(
269 ErrorStr: "[ctx_prof] A function with musttail calls was explicitly requested as "
270 "root. That is not supported because we cannot instrument a return "
271 "instruction to release the context: " +
272 F.getName());
273 return false;
274 }
275 auto &Head = F.getEntryBlock();
276 for (auto &I : Head) {
277 // Find the increment intrinsic in the entry basic block.
278 if (auto *Mark = dyn_cast<InstrProfIncrementInst>(Val: &I)) {
279 assert(Mark->getIndex()->isZero());
280
281 IRBuilder<> Builder(Mark);
282 Guid = Builder.getInt64(C: cast<Function>(Val&: *Mark->getNameValue()).getGUID());
283 // The type of the context of this function is now knowable since we have
284 // NumCallsites and NumCounters. We declare it here because it's more
285 // convenient - we have the Builder.
286 ThisContextType = StructType::get(
287 Context&: F.getContext(),
288 Elements: {ContextNodeTy, ArrayType::get(ElementType: Builder.getInt64Ty(), NumElements: NumCounters),
289 ArrayType::get(ElementType: Builder.getPtrTy(), NumElements: NumCallsites)});
290 // Figure out which way we obtain the context object for this function -
291 // if it's an entrypoint, then we call StartCtx, otherwise GetCtx. In the
292 // former case, we also set TheRootFunctionData since we need to release
293 // it at the end (plus it can be used to know if we have an entrypoint or
294 // a regular function). Don't set a name, they end up taking a lot of
295 // space and we don't need them.
296
297 // Zero-initialize the FunctionData, except for functions that have
298 // musttail calls. There, we set the CtxRoot field to 1, which will be
299 // treated as a "can't be set as root".
300 TheRootFunctionData = new GlobalVariable(
301 M, FunctionDataTy, false, GlobalVariable::InternalLinkage,
302 HasMusttail ? CannotBeRootInitializer
303 : Constant::getNullValue(Ty: FunctionDataTy));
304
305 if (ContextRootSet.contains(V: &F)) {
306 Context = Builder.CreateCall(
307 Callee: StartCtx, Args: {TheRootFunctionData, Guid, Builder.getInt32(C: NumCounters),
308 Builder.getInt32(C: NumCallsites)});
309 ORE.emit(
310 RemarkBuilder: [&] { return OptimizationRemark(DEBUG_TYPE, "Entrypoint", &F); });
311 } else {
312 Context = Builder.CreateCall(Callee: GetCtx, Args: {TheRootFunctionData, &F, Guid,
313 Builder.getInt32(C: NumCounters),
314 Builder.getInt32(C: NumCallsites)});
315 ORE.emit(RemarkBuilder: [&] {
316 return OptimizationRemark(DEBUG_TYPE, "RegularFunction", &F);
317 });
318 }
319 // The context could be scratch.
320 auto *CtxAsInt = Builder.CreatePtrToInt(V: Context, DestTy: Builder.getInt64Ty());
321 if (NumCallsites > 0) {
322 // Figure out which index of the TLS 2-element buffers to use.
323 // Scratch context => we use index == 1. Real contexts => index == 0.
324 auto *Index = Builder.CreateAnd(LHS: CtxAsInt, RHS: Builder.getInt64(C: 1));
325 // The GEPs corresponding to that index, in the respective TLS.
326 ExpectedCalleeTLSAddr = Builder.CreateGEP(
327 Ty: PointerType::getUnqual(C&: F.getContext()),
328 Ptr: Builder.CreateThreadLocalAddress(Ptr: ExpectedCalleeTLS), IdxList: {Index});
329 CallsiteInfoTLSAddr = Builder.CreateGEP(
330 Ty: Builder.getInt32Ty(),
331 Ptr: Builder.CreateThreadLocalAddress(Ptr: CallsiteInfoTLS), IdxList: {Index});
332 }
333 // Because the context pointer may have LSB set (to indicate scratch),
334 // clear it for the value we use as base address for the counter vector.
335 // This way, if later we want to have "real" (not clobbered) buffers
336 // acting as scratch, the lowering (at least this part of it that deals
337 // with counters) stays the same.
338 RealContext = Builder.CreateIntToPtr(
339 V: Builder.CreateAnd(LHS: CtxAsInt, RHS: Builder.getInt64(C: -2)),
340 DestTy: PointerType::getUnqual(C&: F.getContext()));
341 I.eraseFromParent();
342 break;
343 }
344 }
345 if (!Context) {
346 ORE.emit(RemarkBuilder: [&] {
347 return OptimizationRemarkMissed(DEBUG_TYPE, "Skip", &F)
348 << "Function doesn't have instrumentation, skipping";
349 });
350 return false;
351 }
352
353 bool ContextWasReleased = false;
354 for (auto &BB : F) {
355 for (auto &I : llvm::make_early_inc_range(Range&: BB)) {
356 if (auto *Instr = dyn_cast<InstrProfCntrInstBase>(Val: &I)) {
357 IRBuilder<> Builder(Instr);
358 switch (Instr->getIntrinsicID()) {
359 case llvm::Intrinsic::instrprof_increment:
360 case llvm::Intrinsic::instrprof_increment_step: {
361 // Increments (or increment-steps) are just a typical load - increment
362 // - store in the RealContext.
363 auto *AsStep = cast<InstrProfIncrementInst>(Val: Instr);
364 auto *GEP = Builder.CreateGEP(
365 Ty: ThisContextType, Ptr: RealContext,
366 IdxList: {Builder.getInt32(C: 0), Builder.getInt32(C: 1), AsStep->getIndex()});
367 Builder.CreateStore(
368 Val: Builder.CreateAdd(LHS: Builder.CreateLoad(Ty: Builder.getInt64Ty(), Ptr: GEP),
369 RHS: AsStep->getStep()),
370 Ptr: GEP);
371 } break;
372 case llvm::Intrinsic::instrprof_callsite:
373 // callsite lowering: write the called value in the expected callee
374 // TLS we treat the TLS as volatile because of signal handlers and to
375 // avoid these being moved away from the callsite they decorate.
376 auto *CSIntrinsic = dyn_cast<InstrProfCallsite>(Val: Instr);
377 Builder.CreateStore(Val: CSIntrinsic->getCallee(), Ptr: ExpectedCalleeTLSAddr,
378 isVolatile: true);
379 // write the GEP of the slot in the sub-contexts portion of the
380 // context in TLS. Now, here, we use the actual Context value - as
381 // returned from compiler-rt - which may have the LSB set if the
382 // Context was scratch. Since the header of the context object and
383 // then the values are all 8-aligned (or, really, insofar as we care,
384 // they are even) - if the context is scratch (meaning, an odd value),
385 // so will the GEP. This is important because this is then visible to
386 // compiler-rt which will produce scratch contexts for callers that
387 // have a scratch context.
388 Builder.CreateStore(
389 Val: Builder.CreateGEP(Ty: ThisContextType, Ptr: Context,
390 IdxList: {Builder.getInt32(C: 0), Builder.getInt32(C: 2),
391 CSIntrinsic->getIndex()}),
392 Ptr: CallsiteInfoTLSAddr, isVolatile: true);
393 break;
394 }
395 I.eraseFromParent();
396 } else if (!HasMusttail && isa<ReturnInst>(Val: I)) {
397 // Remember to release the context if we are an entrypoint.
398 IRBuilder<> Builder(&I);
399 Builder.CreateCall(Callee: ReleaseCtx, Args: {TheRootFunctionData});
400 ContextWasReleased = true;
401 }
402 }
403 }
404 if (!HasMusttail && !ContextWasReleased)
405 F.getContext().emitError(
406 ErrorStr: "[ctx_prof] A function that doesn't have musttail calls was "
407 "instrumented but it has no `ret` "
408 "instructions above which to release the context: " +
409 F.getName());
410 return true;
411}
412
413PreservedAnalyses NoinlineNonPrevailing::run(Module &M,
414 ModuleAnalysisManager &MAM) {
415 bool Changed = false;
416 for (auto &F : M) {
417 if (F.isDeclaration())
418 continue;
419 if (F.hasFnAttribute(Kind: Attribute::NoInline))
420 continue;
421 if (!F.isWeakForLinker())
422 continue;
423
424 if (F.hasFnAttribute(Kind: Attribute::AlwaysInline))
425 F.removeFnAttr(Kind: Attribute::AlwaysInline);
426
427 F.addFnAttr(Kind: Attribute::NoInline);
428 Changed = true;
429 }
430 if (Changed)
431 return PreservedAnalyses::none();
432 return PreservedAnalyses::all();
433}
434