1//===- PseudoProbe.cpp - Pseudo Probe Helpers -----------------------------===//
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 helpers to manipulate pseudo probe IR intrinsic
10// calls.
11//
12//===----------------------------------------------------------------------===//
13
14#include "llvm/IR/PseudoProbe.h"
15#include "llvm/IR/DebugInfoMetadata.h"
16#include "llvm/IR/IRBuilder.h"
17#include "llvm/IR/Instruction.h"
18#include "llvm/IR/IntrinsicInst.h"
19
20using namespace llvm;
21
22static std::optional<PseudoProbe>
23extractProbeFromDiscriminator(const DILocation *DIL) {
24 if (DIL) {
25 auto Discriminator = DIL->getDiscriminator();
26 if (DILocation::isPseudoProbeDiscriminator(Discriminator)) {
27 PseudoProbe Probe;
28 Probe.Id =
29 PseudoProbeDwarfDiscriminator::extractProbeIndex(Value: Discriminator);
30 Probe.Type =
31 PseudoProbeDwarfDiscriminator::extractProbeType(Value: Discriminator);
32 Probe.Attr =
33 PseudoProbeDwarfDiscriminator::extractProbeAttributes(Value: Discriminator);
34 Probe.Factor =
35 PseudoProbeDwarfDiscriminator::extractProbeFactor(Value: Discriminator) /
36 (float)PseudoProbeDwarfDiscriminator::FullDistributionFactor;
37 Probe.Discriminator = 0;
38 return Probe;
39 }
40 }
41 return std::nullopt;
42}
43
44static std::optional<PseudoProbe>
45extractProbeFromDiscriminator(const Instruction &Inst) {
46 assert(isa<CallBase>(&Inst) && !isa<IntrinsicInst>(&Inst) &&
47 "Only call instructions should have pseudo probe encodes as their "
48 "Dwarf discriminators");
49 if (const DebugLoc &DLoc = Inst.getDebugLoc())
50 return extractProbeFromDiscriminator(DIL: DLoc);
51 return std::nullopt;
52}
53
54std::optional<PseudoProbe> llvm::extractProbe(const Instruction &Inst) {
55 if (const auto *II = dyn_cast<PseudoProbeInst>(Val: &Inst)) {
56 PseudoProbe Probe;
57 Probe.Id = II->getIndex()->getZExtValue();
58 Probe.Type = (uint32_t)PseudoProbeType::Block;
59 Probe.Attr = II->getAttributes()->getZExtValue();
60 Probe.Factor = II->getFactor()->getZExtValue() /
61 (float)PseudoProbeFullDistributionFactor;
62 Probe.Discriminator = 0;
63 if (const DebugLoc &DLoc = Inst.getDebugLoc())
64 Probe.Discriminator = DLoc->getDiscriminator();
65 return Probe;
66 }
67
68 if (isa<CallBase>(Val: &Inst) && !isa<IntrinsicInst>(Val: &Inst))
69 return extractProbeFromDiscriminator(Inst);
70
71 return std::nullopt;
72}
73
74void llvm::setProbeDistributionFactor(Instruction &Inst, float Factor) {
75 assert(Factor >= 0 && Factor <= 1 &&
76 "Distribution factor must be in [0, 1.0]");
77 if (auto *II = dyn_cast<PseudoProbeInst>(Val: &Inst)) {
78 IRBuilder<> Builder(&Inst);
79 uint64_t IntFactor = PseudoProbeFullDistributionFactor;
80 if (Factor < 1)
81 IntFactor *= Factor;
82 auto OrigFactor = II->getFactor()->getZExtValue();
83 if (IntFactor != OrigFactor)
84 II->replaceUsesOfWith(From: II->getFactor(), To: Builder.getInt64(C: IntFactor));
85 } else if (isa<CallBase>(Val: &Inst) && !isa<IntrinsicInst>(Val: &Inst)) {
86 if (const DebugLoc &DLoc = Inst.getDebugLoc()) {
87 const DILocation *DIL = DLoc;
88 auto Discriminator = DIL->getDiscriminator();
89 if (DILocation::isPseudoProbeDiscriminator(Discriminator)) {
90 auto Index =
91 PseudoProbeDwarfDiscriminator::extractProbeIndex(Value: Discriminator);
92 auto Type =
93 PseudoProbeDwarfDiscriminator::extractProbeType(Value: Discriminator);
94 auto Attr = PseudoProbeDwarfDiscriminator::extractProbeAttributes(
95 Value: Discriminator);
96 auto DwarfBaseDiscriminator =
97 PseudoProbeDwarfDiscriminator::extractDwarfBaseDiscriminator(
98 Value: Discriminator);
99 // Round small factors to 0 to avoid over-counting.
100 uint32_t IntFactor =
101 PseudoProbeDwarfDiscriminator::FullDistributionFactor;
102 if (Factor < 1)
103 IntFactor *= Factor;
104 uint32_t V = PseudoProbeDwarfDiscriminator::packProbeData(
105 Index, Type, Flags: Attr, Factor: IntFactor, DwarfBaseDiscriminator);
106 DIL = DIL->cloneWithDiscriminator(Discriminator: V);
107 Inst.setDebugLoc(DIL);
108 }
109 }
110 }
111}
112