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