1//===- SampleProfileProbe.cpp - Pseudo probe Instrumentation -------------===//
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// This file implements the SampleProfileProber transformation.
10//
11//===----------------------------------------------------------------------===//
12
13#include "llvm/Transforms/IPO/SampleProfileProbe.h"
14#include "llvm/ADT/Statistic.h"
15#include "llvm/ADT/StringSet.h"
16#include "llvm/Analysis/BlockFrequencyInfo.h"
17#include "llvm/Analysis/EHUtils.h"
18#include "llvm/Analysis/LoopInfo.h"
19#include "llvm/IR/BasicBlock.h"
20#include "llvm/IR/DebugInfoMetadata.h"
21#include "llvm/IR/DiagnosticInfo.h"
22#include "llvm/IR/IRBuilder.h"
23#include "llvm/IR/Instruction.h"
24#include "llvm/IR/IntrinsicInst.h"
25#include "llvm/IR/MDBuilder.h"
26#include "llvm/IR/Module.h"
27#include "llvm/IR/PassInstrumentation.h"
28#include "llvm/IR/PseudoProbe.h"
29#include "llvm/ProfileData/SampleProf.h"
30#include "llvm/Support/CRC.h"
31#include "llvm/Support/CommandLine.h"
32#include "llvm/Target/TargetMachine.h"
33#include "llvm/Transforms/Utils/Instrumentation.h"
34#include "llvm/Transforms/Utils/ModuleUtils.h"
35#include <vector>
36
37using namespace llvm;
38#define DEBUG_TYPE "pseudo-probe"
39
40STATISTIC(ArtificialDbgLine,
41 "Number of probes that have an artificial debug line");
42
43static cl::opt<bool>
44 VerifyPseudoProbe("verify-pseudo-probe", cl::init(Val: false), cl::Hidden,
45 cl::desc("Do pseudo probe verification"));
46
47static cl::list<std::string> VerifyPseudoProbeFuncList(
48 "verify-pseudo-probe-funcs", cl::Hidden,
49 cl::desc("The option to specify the name of the functions to verify."));
50
51static cl::opt<bool>
52 UpdatePseudoProbe("update-pseudo-probe", cl::init(Val: true), cl::Hidden,
53 cl::desc("Update pseudo probe distribution factor"));
54
55static uint64_t getCallStackHash(const DILocation *DIL) {
56 uint64_t Hash = 0;
57 const DILocation *InlinedAt = DIL ? DIL->getInlinedAt() : nullptr;
58 while (InlinedAt) {
59 Hash ^= MD5Hash(Str: std::to_string(val: InlinedAt->getLine()));
60 Hash ^= MD5Hash(Str: std::to_string(val: InlinedAt->getColumn()));
61 auto Name = InlinedAt->getSubprogramLinkageName();
62 Hash ^= MD5Hash(Str: Name);
63 InlinedAt = InlinedAt->getInlinedAt();
64 }
65 return Hash;
66}
67
68static uint64_t computeCallStackHash(const Instruction &Inst) {
69 return getCallStackHash(DIL: Inst.getDebugLoc());
70}
71
72bool PseudoProbeVerifier::shouldVerifyFunction(const Function *F) {
73 // Skip function declaration.
74 if (F->isDeclaration())
75 return false;
76 // Skip function that will not be emitted into object file. The prevailing
77 // defintion will be verified instead.
78 if (F->hasAvailableExternallyLinkage())
79 return false;
80 // Do a name matching.
81 static const StringSet<> VerifyFuncNames(llvm::from_range,
82 VerifyPseudoProbeFuncList);
83 return VerifyFuncNames.empty() || VerifyFuncNames.contains(key: F->getName());
84}
85
86void PseudoProbeVerifier::registerCallbacks(PassInstrumentationCallbacks &PIC) {
87 if (VerifyPseudoProbe) {
88 PIC.registerAfterPassCallback(
89 C: [this](StringRef P, IRUnitRef IR, const PreservedAnalyses &) {
90 this->runAfterPass(PassID: P, IR);
91 });
92 }
93}
94
95// Callback to run after each transformation for the new pass manager.
96void PseudoProbeVerifier::runAfterPass(StringRef PassID, IRUnitRef IR) {
97 std::string Banner =
98 "\n*** Pseudo Probe Verification After " + PassID.str() + " ***\n";
99 dbgs() << Banner;
100 if (const auto *M = dyn_cast<Module>(Val&: IR))
101 runAfterPass(M);
102 else if (const auto *F = dyn_cast<Function>(Val&: IR))
103 runAfterPass(F);
104 else if (const auto *C = dyn_cast<LazyCallGraph::SCC>(Val&: IR))
105 runAfterPass(C);
106 else if (const auto *L = dyn_cast<Loop>(Val&: IR))
107 runAfterPass(L);
108 else
109 llvm_unreachable("Unknown IR unit");
110}
111
112void PseudoProbeVerifier::runAfterPass(const Module *M) {
113 for (const Function &F : *M)
114 runAfterPass(F: &F);
115}
116
117void PseudoProbeVerifier::runAfterPass(const LazyCallGraph::SCC *C) {
118 for (const LazyCallGraph::Node &N : *C)
119 runAfterPass(F: &N.getFunction());
120}
121
122void PseudoProbeVerifier::runAfterPass(const Function *F) {
123 if (!shouldVerifyFunction(F))
124 return;
125 ProbeFactorMap ProbeFactors;
126 for (const auto &BB : *F)
127 collectProbeFactors(BB: &BB, ProbeFactors);
128 verifyProbeFactors(F, ProbeFactors);
129}
130
131void PseudoProbeVerifier::runAfterPass(const Loop *L) {
132 const Function *F = L->getHeader()->getParent();
133 runAfterPass(F);
134}
135
136void PseudoProbeVerifier::collectProbeFactors(const BasicBlock *Block,
137 ProbeFactorMap &ProbeFactors) {
138 for (const auto &I : *Block) {
139 if (std::optional<PseudoProbe> Probe = extractProbe(Inst: I)) {
140 uint64_t Hash = computeCallStackHash(Inst: I);
141 ProbeFactors[{Probe->Id, Hash}] += Probe->Factor;
142 }
143 }
144}
145
146void PseudoProbeVerifier::verifyProbeFactors(
147 const Function *F, const ProbeFactorMap &ProbeFactors) {
148 bool BannerPrinted = false;
149 auto &PrevProbeFactors = FunctionProbeFactors[F->getName()];
150 for (const auto &I : ProbeFactors) {
151 float CurProbeFactor = I.second;
152 auto [It, Inserted] = PrevProbeFactors.try_emplace(Key: I.first);
153 if (!Inserted) {
154 float PrevProbeFactor = It->second;
155 if (std::abs(x: CurProbeFactor - PrevProbeFactor) >
156 DistributionFactorVariance) {
157 if (!BannerPrinted) {
158 dbgs() << "Function " << F->getName() << ":\n";
159 BannerPrinted = true;
160 }
161 dbgs() << "Probe " << I.first.first << "\tprevious factor "
162 << format(Fmt: "%0.2f", Vals: PrevProbeFactor) << "\tcurrent factor "
163 << format(Fmt: "%0.2f", Vals: CurProbeFactor) << "\n";
164 }
165 }
166
167 // Update
168 It->second = I.second;
169 }
170}
171
172SampleProfileProber::SampleProfileProber(Function &Func) : F(&Func) {
173 BlockProbeIds.clear();
174 CallProbeIds.clear();
175 LastProbeId = (uint32_t)PseudoProbeReservedId::Last;
176
177 DenseSet<BasicBlock *> BlocksToIgnore;
178 DenseSet<BasicBlock *> BlocksAndCallsToIgnore;
179 computeBlocksToIgnore(BlocksToIgnore, BlocksAndCallsToIgnore);
180
181 computeProbeId(BlocksToIgnore, BlocksAndCallsToIgnore);
182 computeCFGHash(BlocksToIgnore);
183}
184
185// Two purposes to compute the blocks to ignore:
186// 1. Reduce the IR size.
187// 2. Make the instrumentation(checksum) stable. e.g. the frondend may
188// generate unstable IR while optimizing nounwind attribute, some versions are
189// optimized with the call-to-invoke conversion, while other versions do not.
190// This discrepancy in probe ID could cause profile mismatching issues.
191// Note that those ignored blocks are either cold blocks or new split blocks
192// whose original blocks are instrumented, so it shouldn't degrade the profile
193// quality.
194void SampleProfileProber::computeBlocksToIgnore(
195 DenseSet<BasicBlock *> &BlocksToIgnore,
196 DenseSet<BasicBlock *> &BlocksAndCallsToIgnore) {
197 // Ignore the cold EH and unreachable blocks and calls.
198 computeEHOnlyBlocks(F&: *F, EHBlocks&: BlocksAndCallsToIgnore);
199 findUnreachableBlocks(BlocksToIgnore&: BlocksAndCallsToIgnore);
200
201 BlocksToIgnore.insert_range(R&: BlocksAndCallsToIgnore);
202
203 // Handle the call-to-invoke conversion case: make sure that the probe id and
204 // callsite id are consistent before and after the block split. For block
205 // probe, we only keep the head block probe id and ignore the block ids of the
206 // normal dests. For callsite probe, it's different to block probe, there is
207 // no additional callsite in the normal dests, so we don't ignore the
208 // callsites.
209 findInvokeNormalDests(InvokeNormalDests&: BlocksToIgnore);
210}
211
212// Unreachable blocks and calls are always cold, ignore them.
213void SampleProfileProber::findUnreachableBlocks(
214 DenseSet<BasicBlock *> &BlocksToIgnore) {
215 for (auto &BB : *F) {
216 if (&BB != &F->getEntryBlock() && pred_size(BB: &BB) == 0)
217 BlocksToIgnore.insert(V: &BB);
218 }
219}
220
221// In call-to-invoke conversion, basic block can be split into multiple blocks,
222// only instrument probe in the head block, ignore the normal dests.
223void SampleProfileProber::findInvokeNormalDests(
224 DenseSet<BasicBlock *> &InvokeNormalDests) {
225 for (auto &BB : *F) {
226 auto *TI = BB.getTerminator();
227 if (auto *II = dyn_cast<InvokeInst>(Val: TI)) {
228 auto *ND = II->getNormalDest();
229 InvokeNormalDests.insert(V: ND);
230
231 // The normal dest and the try/catch block are connected by an
232 // unconditional branch.
233 while (pred_size(BB: ND) == 1) {
234 auto *Pred = *pred_begin(BB: ND);
235 if (succ_size(BB: Pred) == 1) {
236 InvokeNormalDests.insert(V: Pred);
237 ND = Pred;
238 } else
239 break;
240 }
241 }
242 }
243}
244
245// The call-to-invoke conversion splits the original block into a list of block,
246// we need to compute the hash using the original block's successors to keep the
247// CFG Hash consistent. For a given head block, we keep searching the
248// succesor(normal dest or unconditional branch dest) to find the tail block,
249// the tail block's successors are the original block's successors.
250const Instruction *SampleProfileProber::getOriginalTerminator(
251 const BasicBlock *Head, const DenseSet<BasicBlock *> &BlocksToIgnore) {
252 auto *TI = Head->getTerminator();
253 if (auto *II = dyn_cast<InvokeInst>(Val: TI)) {
254 return getOriginalTerminator(Head: II->getNormalDest(), BlocksToIgnore);
255 } else if (succ_size(BB: Head) == 1 &&
256 BlocksToIgnore.contains(V: *succ_begin(BB: Head))) {
257 // Go to the unconditional branch dest.
258 return getOriginalTerminator(Head: *succ_begin(BB: Head), BlocksToIgnore);
259 }
260 return TI;
261}
262
263// Compute Hash value for the CFG: the lower 32 bits are CRC32 of the index
264// value of each BB in the CFG. The higher 32 bits record the number of edges
265// preceded by the number of indirect calls.
266// This is derived from FuncPGOInstrumentation<Edge, BBInfo>::computeCFGHash().
267void SampleProfileProber::computeCFGHash(
268 const DenseSet<BasicBlock *> &BlocksToIgnore) {
269 std::vector<uint8_t> Indexes;
270 JamCRC JC;
271 for (auto &BB : *F) {
272 if (BlocksToIgnore.contains(V: &BB))
273 continue;
274
275 auto *TI = getOriginalTerminator(Head: &BB, BlocksToIgnore);
276 for (unsigned I = 0, E = TI->getNumSuccessors(); I != E; ++I) {
277 auto *Succ = TI->getSuccessor(Idx: I);
278 auto Index = getBlockId(BB: Succ);
279 // Ingore ignored-block(zero ID) to avoid unstable checksum.
280 if (Index == 0)
281 continue;
282 for (int J = 0; J < 4; J++)
283 Indexes.push_back(x: (uint8_t)(Index >> (J * 8)));
284 }
285 }
286
287 JC.update(Data: Indexes);
288
289 FunctionHash = (uint64_t)CallProbeIds.size() << 48 |
290 (uint64_t)Indexes.size() << 32 | JC.getCRC();
291 // Reserve bit 60-63 for other information purpose.
292 FunctionHash &= 0x0FFFFFFFFFFFFFFF;
293 assert(FunctionHash && "Function checksum should not be zero");
294 LLVM_DEBUG(dbgs() << "\nFunction Hash Computation for " << F->getName()
295 << ":\n"
296 << " CRC = " << JC.getCRC() << ", Edges = "
297 << Indexes.size() << ", ICSites = " << CallProbeIds.size()
298 << ", Hash = " << FunctionHash << "\n");
299}
300
301void SampleProfileProber::computeProbeId(
302 const DenseSet<BasicBlock *> &BlocksToIgnore,
303 const DenseSet<BasicBlock *> &BlocksAndCallsToIgnore) {
304 LLVMContext &Ctx = F->getContext();
305 Module *M = F->getParent();
306
307 for (auto &BB : *F) {
308 if (!BlocksToIgnore.contains(V: &BB))
309 BlockProbeIds[&BB] = ++LastProbeId;
310
311 if (BlocksAndCallsToIgnore.contains(V: &BB))
312 continue;
313 for (auto &I : BB) {
314 if (!isa<CallBase>(Val: I) || isa<IntrinsicInst>(Val: &I))
315 continue;
316
317 // The current implementation uses the lower 16 bits of the discriminator
318 // so anything larger than 0xFFFF will be ignored.
319 if (LastProbeId >= 0xFFFF) {
320 std::string Msg = "Pseudo instrumentation incomplete for " +
321 std::string(F->getName()) + " because it's too large";
322 Ctx.diagnose(
323 DI: DiagnosticInfoSampleProfile(M->getName().data(), Msg, DS_Warning));
324 return;
325 }
326
327 CallProbeIds[&I] = ++LastProbeId;
328 }
329 }
330}
331
332uint32_t SampleProfileProber::getBlockId(const BasicBlock *BB) const {
333 auto I = BlockProbeIds.find(Val: const_cast<BasicBlock *>(BB));
334 return I == BlockProbeIds.end() ? 0 : I->second;
335}
336
337uint32_t SampleProfileProber::getCallsiteId(const Instruction *Call) const {
338 auto Iter = CallProbeIds.find(Val: const_cast<Instruction *>(Call));
339 return Iter == CallProbeIds.end() ? 0 : Iter->second;
340}
341
342void SampleProfileProber::instrumentOneFunc(Function &F, TargetMachine *TM) {
343 Module *M = F.getParent();
344 MDBuilder MDB(F.getContext());
345 // Since the GUID from probe desc and inline stack are computed separately, we
346 // need to make sure their names are consistent, so here also use the name
347 // from debug info.
348 StringRef FName = F.getName();
349 if (auto *SP = F.getSubprogram()) {
350 FName = SP->getLinkageName();
351 if (FName.empty())
352 FName = SP->getName();
353 }
354 uint64_t Guid = Function::getGUIDAssumingExternalLinkage(GlobalName: FName);
355
356 // Assign an artificial debug line to a probe that doesn't come with a real
357 // line. A probe not having a debug line will get an incomplete inline
358 // context. This will cause samples collected on the probe to be counted
359 // into the base profile instead of a context profile. The line number
360 // itself is not important though.
361 auto AssignDebugLoc = [&](Instruction *I) {
362 assert((isa<PseudoProbeInst>(I) || isa<CallBase>(I)) &&
363 "Expecting pseudo probe or call instructions");
364 if (!I->getDebugLoc()) {
365 if (auto *SP = F.getSubprogram()) {
366 auto DIL = DILocation::get(Context&: SP->getContext(), Line: 0, Column: 0, Scope: SP);
367 I->setDebugLoc(DIL);
368 ArtificialDbgLine++;
369 LLVM_DEBUG({
370 dbgs() << "\nIn Function " << F.getName()
371 << " Probe gets an artificial debug line\n";
372 I->dump();
373 });
374 }
375 }
376 };
377
378 // Probe basic blocks.
379 for (auto &I : BlockProbeIds) {
380 BasicBlock *BB = I.first;
381 uint32_t Index = I.second;
382 // Insert a probe before an instruction with a valid debug line number which
383 // will be assigned to the probe. The line number will be used later to
384 // model the inline context when the probe is inlined into other functions.
385 // Debug instructions, phi nodes and lifetime markers do not have an valid
386 // line number. Real instructions generated by optimizations may not come
387 // with a line number either.
388 auto HasValidDbgLine = [](Instruction *J) {
389 return !isa<PHINode>(Val: J) && !J->isLifetimeStartOrEnd() && J->getDebugLoc();
390 };
391
392 Instruction *J = &*BB->getFirstInsertionPt();
393 while (J != BB->getTerminator() && !HasValidDbgLine(J)) {
394 J = J->getNextNode();
395 }
396
397 // A pseudo probe must not be inserted between a `musttail` or
398 // `llvm.experimental.deoptimize` call and its following `ret`, as this
399 // produces invalid IR. Such a call is required to immediately precede the
400 // block's `ret`, so only that position needs to be checked. Insert the
401 // probe before the call instead.
402 if (auto *Ret = dyn_cast<ReturnInst>(Val: BB->getTerminator()))
403 if (auto *CI = dyn_cast_or_null<CallInst>(Val: Ret->getPrevNode()))
404 if ((CI->isMustTailCall() ||
405 CI->getIntrinsicID() == Intrinsic::experimental_deoptimize) &&
406 !J->comesBefore(Other: CI))
407 J = CI;
408
409 IRBuilder<> Builder(J);
410 assert(Builder.GetInsertPoint() != BB->end() &&
411 "Cannot get the probing point");
412 Function *ProbeFn =
413 llvm::Intrinsic::getOrInsertDeclaration(M, id: Intrinsic::pseudoprobe);
414 Value *Args[] = {Builder.getInt64(C: Guid), Builder.getInt64(C: Index),
415 Builder.getInt32(C: 0),
416 Builder.getInt64(C: PseudoProbeFullDistributionFactor)};
417 auto *Probe = Builder.CreateCall(Callee: ProbeFn, Args);
418 AssignDebugLoc(Probe);
419 // Reset the dwarf discriminator if the debug location comes with any. The
420 // discriminator field may be used by FS-AFDO later in the pipeline.
421 if (auto DIL = Probe->getDebugLoc()) {
422 if (DIL->getDiscriminator()) {
423 DIL = DIL->cloneWithDiscriminator(Discriminator: 0);
424 Probe->setDebugLoc(DIL);
425 }
426 }
427 }
428
429 // Probe both direct calls and indirect calls. Direct calls are probed so that
430 // their probe ID can be used as an call site identifier to represent a
431 // calling context.
432 for (auto &I : CallProbeIds) {
433 auto *Call = I.first;
434 uint32_t Index = I.second;
435 uint32_t Type = cast<CallBase>(Val: Call)->getCalledFunction()
436 ? (uint32_t)PseudoProbeType::DirectCall
437 : (uint32_t)PseudoProbeType::IndirectCall;
438 AssignDebugLoc(Call);
439 if (auto DIL = Call->getDebugLoc()) {
440 // Levarge the 32-bit discriminator field of debug data to store the ID
441 // and type of a callsite probe. This gets rid of the dependency on
442 // plumbing a customized metadata through the codegen pipeline.
443 uint32_t V = PseudoProbeDwarfDiscriminator::packProbeData(
444 Index, Type, Flags: 0, Factor: PseudoProbeDwarfDiscriminator::FullDistributionFactor,
445 DwarfBaseDiscriminator: DIL->getBaseDiscriminator());
446 DIL = DIL->cloneWithDiscriminator(Discriminator: V);
447 Call->setDebugLoc(DIL);
448 }
449 }
450
451 // Create module-level metadata that contains function info necessary to
452 // synthesize probe-based sample counts, which are
453 // - FunctionGUID
454 // - FunctionHash.
455 // - FunctionName
456 auto Hash = getFunctionHash();
457 auto *MD = MDB.createPseudoProbeDesc(GUID: Guid, Hash, FName);
458 auto *NMD = M->getNamedMetadata(Name: PseudoProbeDescMetadataName);
459 assert(NMD && "llvm.pseudo_probe_desc should be pre-created");
460 NMD->addOperand(M: MD);
461}
462
463PreservedAnalyses SampleProfileProbePass::run(Module &M,
464 ModuleAnalysisManager &AM) {
465 // Create the pseudo probe desc metadata beforehand.
466 // Note that modules with only data but no functions will require this to
467 // be set up so that they will be known as probed later.
468 M.getOrInsertNamedMetadata(Name: PseudoProbeDescMetadataName);
469
470 for (auto &F : M) {
471 if (F.isDeclaration())
472 continue;
473 SampleProfileProber ProbeManager(F);
474 ProbeManager.instrumentOneFunc(F, TM);
475 }
476
477 return PreservedAnalyses::none();
478}
479
480void PseudoProbeUpdatePass::runOnFunction(Function &F,
481 FunctionAnalysisManager &FAM) {
482 BlockFrequencyInfo &BFI = FAM.getResult<BlockFrequencyAnalysis>(IR&: F);
483 auto BBProfileCount = [&BFI](BasicBlock *BB) {
484 return BFI.getBlockProfileCount(BB).value_or(u: 0);
485 };
486
487 // Collect the sum of execution weight for each probe.
488 ProbeFactorMap ProbeFactors;
489 for (auto &Block : F) {
490 for (auto &I : Block) {
491 if (std::optional<PseudoProbe> Probe = extractProbe(Inst: I)) {
492 uint64_t Hash = computeCallStackHash(Inst: I);
493 ProbeFactors[{Probe->Id, Hash}] += BBProfileCount(&Block);
494 }
495 }
496 }
497
498 // Fix up over-counted probes.
499 for (auto &Block : F) {
500 for (auto &I : Block) {
501 if (std::optional<PseudoProbe> Probe = extractProbe(Inst: I)) {
502 uint64_t Hash = computeCallStackHash(Inst: I);
503 float Sum = ProbeFactors[{Probe->Id, Hash}];
504 if (Sum != 0)
505 setProbeDistributionFactor(Inst&: I, Factor: BBProfileCount(&Block) / Sum);
506 }
507 }
508 }
509}
510
511PreservedAnalyses PseudoProbeUpdatePass::run(Module &M,
512 ModuleAnalysisManager &AM) {
513 if (UpdatePseudoProbe) {
514 for (auto &F : M) {
515 if (F.isDeclaration())
516 continue;
517 FunctionAnalysisManager &FAM =
518 AM.getResult<FunctionAnalysisManagerModuleProxy>(IR&: M).getManager();
519 runOnFunction(F, FAM);
520 }
521 }
522 return PreservedAnalyses::none();
523}
524