1 files changed, 16 insertions, 6 deletions
diff --git a/llvm/lib/IR/PseudoProbe.cpp b/llvm/lib/IR/PseudoProbe.cpp
index f3802af26a61..df5f78c51182 100644
--- a/llvm/lib/IR/PseudoProbe.cpp
+++ b/llvm/lib/IR/PseudoProbe.cpp
@@ -22,12 +22,8 @@ using namespace llvm;
 namespace llvm {
 
 std::optional<PseudoProbe>
-extractProbeFromDiscriminator(const Instruction &Inst) {
-  assert(isa<CallBase>(&Inst) && !isa<IntrinsicInst>(&Inst) &&
-         "Only call instructions should have pseudo probe encodes as their "
-         "Dwarf discriminators");
-  if (const DebugLoc &DLoc = Inst.getDebugLoc()) {
-    const DILocation *DIL = DLoc;
+extractProbeFromDiscriminator(const DILocation *DIL) {
+  if (DIL) {
     auto Discriminator = DIL->getDiscriminator();
     if (DILocation::isPseudoProbeDiscriminator(Discriminator)) {
       PseudoProbe Probe;
@@ -40,12 +36,23 @@ extractProbeFromDiscriminator(const Instruction &Inst) {
       Probe.Factor =
           PseudoProbeDwarfDiscriminator::extractProbeFactor(Discriminator) /
           (float)PseudoProbeDwarfDiscriminator::FullDistributionFactor;
+      Probe.Discriminator = 0;
       return Probe;
     }
   }
   return std::nullopt;
 }
 
+std::optional<PseudoProbe>
+extractProbeFromDiscriminator(const Instruction &Inst) {
+  assert(isa<CallBase>(&Inst) && !isa<IntrinsicInst>(&Inst) &&
+         "Only call instructions should have pseudo probe encodes as their "
+         "Dwarf discriminators");
+  if (const DebugLoc &DLoc = Inst.getDebugLoc())
+    return extractProbeFromDiscriminator(DLoc);
+  return std::nullopt;
+}
+
 std::optional<PseudoProbe> extractProbe(const Instruction &Inst) {
   if (const auto *II = dyn_cast<PseudoProbeInst>(&Inst)) {
     PseudoProbe Probe;
@@ -54,6 +61,9 @@ std::optional<PseudoProbe> extractProbe(const Instruction &Inst) {
     Probe.Attr = II->getAttributes()->getZExtValue();
     Probe.Factor = II->getFactor()->getZExtValue() /
                    (float)PseudoProbeFullDistributionFactor;
+    Probe.Discriminator = 0;
+    if (const DebugLoc &DLoc = Inst.getDebugLoc())
+      Probe.Discriminator = DLoc->getDiscriminator();
     return Probe;
   }