llvm/lib/IR/PseudoProbe.cpp

   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