[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 dc98f92f15568c9809612643d0778a51fd4c98da 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 ++
 .../SampleProfile/pseudo-probe-inline.ll      |  2 +
 .../tools/llvm-profgen/inline-probe-afdo.test | 82 +++++++++++++++++++
 3 files changed, 93 insertions(+)
 create mode 100644 llvm/test/tools/llvm-profgen/inline-probe-afdo.test

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
diff --git a/llvm/test/tools/llvm-profgen/inline-probe-afdo.test b/llvm/test/tools/llvm-profgen/inline-probe-afdo.test
new file mode 100644
index 000000000000000..763bf3b0f8c394a
--- /dev/null
+++ b/llvm/test/tools/llvm-profgen/inline-probe-afdo.test
@@ -0,0 +1,82 @@
+; RUN: llvm-profgen --format=text --use-dwarf-correlation --ignore-stack-samples --perfscript=%S/Inputs/cs-preinline-cost.perfscript --binary=%S/Inputs/cs-preinline-cost.perfbin --output %t
+; RUN: FileCheck %s --input-file %t
+
+; CHECK:     main:947937:0
+; CHECK-NEXT:  2: 545
+; CHECK-NEXT:  3: 545
+; CHECK-NEXT:  5: 545
+; CHECK-NEXT:  7: 0
+; CHECK-NEXT:  65496: 545
+; CHECK-NEXT:  3.7: _Z3fooi:915794
+; CHECK-NEXT:   1: 545
+; CHECK-NEXT:   5: 545
+; CHECK-NEXT:   6: 272
+; CHECK-NEXT:   10: 273
+; CHECK-NEXT:   11: 180
+; CHECK-NEXT:   12: 6965
+; CHECK-NEXT:   13: 6965
+; CHECK-NEXT:   14: 6965
+; CHECK-NEXT:   15: 6965
+; CHECK-NEXT:   20: 182
+; CHECK-NEXT:   21: 6958
+; CHECK-NEXT:   22: 6958
+; CHECK-NEXT:   23: 6958
+; CHECK-NEXT:   24: 6958
+; CHECK-NEXT:   29: 272
+; CHECK-NEXT:   65529: 182
+; CHECK-NEXT:  4.8: _Z3fooi:16338
+; CHECK-NEXT:   1: 272
+; CHECK-NEXT:   6: 545
+
+
+
+
+; binary is built with the source below using the following command line:
+;   clang -O3 -g -fpseudo-probe-for-profiling test.cpp
+;
+;#include <stdio.h>
+;
+;volatile int state = 9000;
+;
+;int foo(int x) {
+;    if (x == 0) {
+;        return 7;
+;    }
+;
+;    if ((x & 1) == 0) {
+;        state--;
+;        return 9;
+;    }
+;
+;    if (state > 5000) {
+;        while (state > 5000) {
+;               for (int i = 50; i >= 0; i--) {
+;                state *= 6;
+;                state /= 7;
+;                state -= 1;
+;            }
+;        }
+;    }
+;    else {
+;        while (state < 5000) {
+;            for (int i = 50; i >= 0; i--) {
+;                state *= 6;
+;                state /= 5;
+;                state += 1;
+;            }
+;        }
+;    }
+;
+;    return state;
+;}
+;
+;volatile int cnt = 10000000;//10000000;
+;int main() {
+;    int r = 0;
+;    for (int i = 0; i < cnt; i++) {
+;      r += foo(i);
+;      r -= foo(i & (~1));
+;      r += foo(0);
+;    }
+;    return r;
+;}