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