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