[llvm] [PseudoProbe] Use probe id as the base dwarf discriminator for callsites (PR #65685)

[Hongtao Yu via llvm-commits]

https://github.com/htyu updated https://github.com/llvm/llvm-project/pull/65685:

>From 763ab7d3cc47934fb449528076f437008fb6e0f4 Mon Sep 17 00:00:00 2001
From: Hongtao Yu <hoy at fb.com>
Date: Wed, 6 Sep 2023 14:55:55 -0700
Subject: [PATCH] Return probe id for probe discriminator

---
 llvm/include/llvm/IR/DebugInfoMetadata.h                 | 9 +++++++++
 .../test/Transforms/SampleProfile/pseudo-probe-inline.ll | 2 ++
 2 files changed, 11 insertions(+)

diff --git a/llvm/include/llvm/IR/DebugInfoMetadata.h b/llvm/include/llvm/IR/DebugInfoMetadata.h
index b9f6d39a7491566..9beb514b87125af 100644
--- a/llvm/include/llvm/IR/DebugInfoMetadata.h
+++ b/llvm/include/llvm/IR/DebugInfoMetadata.h
@@ -22,6 +22,7 @@
 #include "llvm/ADT/iterator_range.h"
 #include "llvm/IR/Constants.h"
 #include "llvm/IR/Metadata.h"
+#include "llvm/IR/PseudoProbe.h"
 #include "llvm/Support/Casting.h"
 #include "llvm/Support/CommandLine.h"
 #include "llvm/Support/Discriminator.h"
@@ -2075,6 +2076,14 @@ class DILocation : public MDNode {
   static unsigned
   getBaseDiscriminatorFromDiscriminator(unsigned D,
                                         bool IsFSDiscriminator = false) {
+    // Return the probe id instead of zero for a pseudo probe discriminator.
+    // This should help differenciate callsites with same line numbers to
+    // achieve a decent AutoFDO profile under -fpseudo-probe-for-profiling,
+    // where the original callsite dwarf discriminator is overwritten by
+    // callsite probe information.
+    if (isPseudoProbeDiscriminator(D))
+      return PseudoProbeDwarfDiscriminator::extractProbeIndex(D);
+
     if (IsFSDiscriminator)
       return getMaskedDiscriminator(D, getBaseDiscriminatorBits());
     return getUnsignedFromPrefixEncoding(D);
diff --git a/llvm/test/Transforms/SampleProfile/pseudo-probe-inline.ll b/llvm/test/Transforms/SampleProfile/pseudo-probe-inline.ll
index 8157bebdc3777be..18cbd857d97bb2c 100644
--- a/llvm/test/Transforms/SampleProfile/pseudo-probe-inline.ll
+++ b/llvm/test/Transforms/SampleProfile/pseudo-probe-inline.ll
@@ -106,6 +106,8 @@ if.end:
 ;YAML-NEXT:    - Line:            '1'
 ;YAML-NEXT:    - String:          ':'
 ;YAML-NEXT:    - Column:          '11'
+;YAML-NEXT:    - String:          .
+;YAML-NEXT:    - Disc:            '2'
 ;YAML-NEXT:    - String:          ';'
 ;YAML-NEXT:  ...
 ;YAML:  --- !Analysis