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