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
22namespace llvm {
23
24std::optional<PseudoProbe>
25extractProbeFromDiscriminator(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(Value: Discriminator);
32 Probe.Type =
33 PseudoProbeDwarfDiscriminator::extractProbeType(Value: Discriminator);
34 Probe.Attr =
35 PseudoProbeDwarfDiscriminator::extractProbeAttributes(Value: Discriminator);
36 Probe.Factor =
37 PseudoProbeDwarfDiscriminator::extractProbeFactor(Value: Discriminator) /
38 (float)PseudoProbeDwarfDiscriminator::FullDistributionFactor;
39 Probe.Discriminator = 0;
40 return Probe;
41 }
42 }
43 return std::nullopt;
44}
45
46std::optional<PseudoProbe>
47extractProbeFromDiscriminator(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(DIL: DLoc);
53 return std::nullopt;
54}
55
56std::optional<PseudoProbe> extractProbe(const Instruction &Inst) {
57 if (const auto *II = dyn_cast<PseudoProbeInst>(Val: &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>(Val: &Inst) && !isa<IntrinsicInst>(Val: &Inst))
71 return extractProbeFromDiscriminator(Inst);
72
73 return std::nullopt;
74}
75
76void 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>(Val: &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(From: II->getFactor(), To: Builder.getInt64(C: IntFactor));
87 } else if (isa<CallBase>(Val: &Inst) && !isa<IntrinsicInst>(Val: &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(Value: Discriminator);
94 auto Type =
95 PseudoProbeDwarfDiscriminator::extractProbeType(Value: Discriminator);
96 auto Attr = PseudoProbeDwarfDiscriminator::extractProbeAttributes(
97 Value: Discriminator);
98 auto DwarfBaseDiscriminator =
99 PseudoProbeDwarfDiscriminator::extractDwarfBaseDiscriminator(
100 Value: Discriminator);
101 // Round small factors to 0 to avoid over-counting.
102 uint32_t IntFactor =
103 PseudoProbeDwarfDiscriminator::FullDistributionFactor;
104 if (Factor < 1)
105 IntFactor *= Factor;
106 uint32_t V = PseudoProbeDwarfDiscriminator::packProbeData(
107 Index, Type, Flags: Attr, Factor: IntFactor, DwarfBaseDiscriminator);
108 DIL = DIL->cloneWithDiscriminator(Discriminator: V);
109 Inst.setDebugLoc(DIL);
110 }
111 }
112 }
113}
114
115} // namespace llvm
116