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