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 "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         auto DwarfBaseDiscriminator =
  99             PseudoProbeDwarfDiscriminator::extractDwarfBaseDiscriminator(
 100                 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, Attr, IntFactor, DwarfBaseDiscriminator);
 108         DIL = DIL->cloneWithDiscriminator(V);
 109         Inst.setDebugLoc(DIL);
 110       }
 111     }
 112   }
 113 }
 114
 115 } // namespace llvm