[PseudoProbe] Make probe discriminator compatible with dwarf base discriminator

wlei-llvm · wlei-llvm · commit 6217a7bc073a · 2024-06-05T10:28:30.000-07:00
diff --git a/llvm/include/llvm/IR/DebugInfoMetadata.h b/llvm/include/llvm/IR/DebugInfoMetadata.h
@@ -2147,13 +2147,18 @@ 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);
+    // Extract the dwarf base discriminator if it's encoded in the pseudo probe
+    // discriminator. Otherwise, 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)) {
+      if (PseudoProbeDwarfDiscriminator::isDwarfBaseDiscriminatorEncoded(D))
+        return PseudoProbeDwarfDiscriminator::extractDwarfBaseDiscriminator(D);
+      else
+        return PseudoProbeDwarfDiscriminator::extractProbeIndex(D);
+    }
 
     if (IsFSDiscriminator)
       return getMaskedDiscriminator(D, getBaseDiscriminatorBits());
diff --git a/llvm/include/llvm/IR/PseudoProbe.h b/llvm/include/llvm/IR/PseudoProbe.h
@@ -44,26 +44,52 @@ struct PseudoProbeDwarfDiscriminator {
   // 32-bit integer which is organized as:
   //  [2:0] - 0x7, this is reserved for regular discriminator,
   //          see DWARF discriminator encoding rule
-  //  [18:3] - probe id
+  //  if the [28:28] bit is zero:
+  //    [18:3] for probe id.
+  //  else:
+  //    [15:3] for probe id, [18:16] for dwarf base discriminator.
   //  [25:19] - probe distribution factor
-  //  [28:26] - probe type, see PseudoProbeType
-  //  [31:29] - reserved for probe attributes
+  //  [27:26] - probe type, see PseudoProbeType
+  //  [28:28] - indicates whether dwarf base discriminator is encoded.
+  //  [30:29] - reserved for probe attributes
   static uint32_t packProbeData(uint32_t Index, uint32_t Type, uint32_t Flags,
-                                uint32_t Factor) {
+                                uint32_t Factor,
+                                uint32_t DwarfBaseDiscriminator) {
     assert(Index <= 0xFFFF && "Probe index too big to encode, exceeding 2^16");
-    assert(Type <= 0x7 && "Probe type too big to encode, exceeding 7");
+    assert(Type <= 0x3 && "Probe type too big to encode, exceeding 3");
     assert(Flags <= 0x7);
     assert(Factor <= 100 &&
            "Probe distribution factor too big to encode, exceeding 100");
-    return (Index << 3) | (Factor << 19) | (Type << 26) | 0x7;
+    uint32_t V = (Index << 3) | (Factor << 19) | (Type << 26) | 0x7;
+    // If both the probe id and dwarf base discriminator is small, the probe id
+    // space is shared with the dwarf base discriminator, this is to make the
+    // probe-based build to be compatible with the dwarf-based profile. Pack the
+    // dwarf base discriminator into [18:16] and set the [28:28] bit.
+    if (Index <= 0x1FFF && DwarfBaseDiscriminator <= 0x7)
+      V |= (1 << 28) | (DwarfBaseDiscriminator << 16);
+    return V;
   }
 
   static uint32_t extractProbeIndex(uint32_t Value) {
+    if (isDwarfBaseDiscriminatorEncoded(Value))
+      return (Value >> 3) & 0x1FFF;
     return (Value >> 3) & 0xFFFF;
   }
 
+  static uint32_t extractDwarfBaseDiscriminator(uint32_t Value) {
+    if (isDwarfBaseDiscriminatorEncoded(Value))
+      return (Value >> 16) & 0x7;
+    // Return an invalid value to indicate the dwarf base discriminator is not
+    // encoded.
+    return 0xFFFFFFFF;
+  }
+
+  static uint32_t isDwarfBaseDiscriminatorEncoded(uint32_t Value) {
+    return (Value >> 28) & 0x1;
+  }
+
   static uint32_t extractProbeType(uint32_t Value) {
-    return (Value >> 26) & 0x7;
+    return (Value >> 26) & 0x3;
   }
 
   static uint32_t extractProbeAttributes(uint32_t Value) {
diff --git a/llvm/lib/IR/PseudoProbe.cpp b/llvm/lib/IR/PseudoProbe.cpp
@@ -95,13 +95,21 @@ void setProbeDistributionFactor(Instruction &Inst, float Factor) {
             PseudoProbeDwarfDiscriminator::extractProbeType(Discriminator);
         auto Attr = PseudoProbeDwarfDiscriminator::extractProbeAttributes(
             Discriminator);
+        auto IsDwarfBaseDiscriminatorEncoded =
+            PseudoProbeDwarfDiscriminator::isDwarfBaseDiscriminatorEncoded(
+                Discriminator);
+        auto DwarfBaseDiscriminator =
+            PseudoProbeDwarfDiscriminator::extractDwarfBaseDiscriminator(
+                Discriminator);
         // Round small factors to 0 to avoid over-counting.
         uint32_t IntFactor =
             PseudoProbeDwarfDiscriminator::FullDistributionFactor;
         if (Factor < 1)
           IntFactor *= Factor;
         uint32_t V = PseudoProbeDwarfDiscriminator::packProbeData(
-            Index, Type, Attr, IntFactor);
+            Index, Type, Attr, IntFactor,
+            IsDwarfBaseDiscriminatorEncoded ? DwarfBaseDiscriminator
+                                            : 0xFFFFFFFF);
         DIL = DIL->cloneWithDiscriminator(V);
         Inst.setDebugLoc(DIL);
       }
diff --git a/llvm/lib/Transforms/IPO/SampleProfileProbe.cpp b/llvm/lib/Transforms/IPO/SampleProfileProbe.cpp
@@ -431,8 +431,8 @@ void SampleProfileProber::instrumentOneFunc(Function &F, TargetMachine *TM) {
       // and type of a callsite probe. This gets rid of the dependency on
       // plumbing a customized metadata through the codegen pipeline.
       uint32_t V = PseudoProbeDwarfDiscriminator::packProbeData(
-          Index, Type, 0,
-          PseudoProbeDwarfDiscriminator::FullDistributionFactor);
+          Index, Type, 0, PseudoProbeDwarfDiscriminator::FullDistributionFactor,
+          DIL->getBaseDiscriminator());
       DIL = DIL->cloneWithDiscriminator(V);
       Call->setDebugLoc(DIL);
     }
diff --git a/llvm/test/Transforms/SampleProfile/pseudo-probe-discriminator.ll b/llvm/test/Transforms/SampleProfile/pseudo-probe-discriminator.ll
@@ -64,8 +64,8 @@ attributes #1 = { "disable-tail-calls"="false" "less-precise-fpmad"="false" "fra
 
 
 ; PROBE: ![[CALL1]] = !DILocation(line: 4, column: 3, scope: ![[CALL1BLOCK:[0-9]+]])
-; PROBE: ![[CALL1BLOCK]] = !DILexicalBlockFile({{.*}} discriminator: 186646559)
+; PROBE: ![[CALL1BLOCK]] = !DILexicalBlockFile({{.*}} discriminator: 455147551)
 ; PROBE: ![[CALL2]] = !DILocation(line: 4, column: 9, scope: ![[CALL2BLOCK:[0-9]+]])
-; PROBE: ![[CALL2BLOCK]] = !DILexicalBlockFile({{.*}} discriminator: 186646567)
+; PROBE: ![[CALL2BLOCK]] = !DILexicalBlockFile({{.*}} discriminator: 455344167)
 ; PROBE: ![[INST]] = !DILocation(line: 4, column: 15, scope: ![[INSTBLOCK:[0-9]+]])
 ; PROBE: ![[INSTBLOCK]] = !DILexicalBlockFile({{.*}} discriminator: 4)
diff --git a/llvm/test/Transforms/SampleProfile/pseudo-probe-emit-inline.ll b/llvm/test/Transforms/SampleProfile/pseudo-probe-emit-inline.ll
@@ -42,13 +42,13 @@ define dso_local i32 @entry() !dbg !14 {
 ; CHECK-IL: ![[#SCOPE2:]] = distinct !DISubprogram(name: "foo"
 ; CHECK-IL: ![[#DL1]] = !DILocation(line: 3, column: 1,  scope: ![[#SCOPE1]], inlinedAt: ![[#INL1:]])
 ; CHECK-IL: ![[#INL1]] = distinct !DILocation(line: 7, column: 3, scope: ![[#BL1:]])
-;; A discriminator of 186646551 which is 0xb200017 in hexdecimal, stands for a direct call probe
+;; A discriminator of 455082007 which is 0x1b200017 in hexdecimal, stands for a direct call probe
 ;; with an index of 2 and a scale of 100%.
-; CHECK-IL: ![[#BL1]] = !DILexicalBlockFile(scope: ![[#SCOPE2]], file: !1, discriminator: 186646551)
+; CHECK-IL: ![[#BL1]] = !DILexicalBlockFile(scope: ![[#SCOPE2]], file: !1, discriminator: 455082007)
 ; CHECK-IL: ![[#SCOPE3:]] = distinct !DISubprogram(name: "entry"
 ; CHECK-IL: ![[#DL2]] = !DILocation(line: 7, column: 3,  scope: ![[#SCOPE2]], inlinedAt: ![[#INL2:]])
 ; CHECK-IL: ![[#INL2]] = distinct !DILocation(line: 11, column: 3, scope: ![[#BL2:]])
-; CHECK-IL: ![[#BL2]] = !DILexicalBlockFile(scope: ![[#SCOPE3]], file: !1, discriminator: 186646551)
+; CHECK-IL: ![[#BL2]] = !DILexicalBlockFile(scope: ![[#SCOPE3]], file: !1, discriminator: 455082007)
 ; CHECK-IL: ![[#DL3]] = !DILocation(line: 3, column: 1,  scope: ![[#SCOPE1]], inlinedAt: ![[#INL3:]])
 ; CHECK-IL: ![[#INL3]] = distinct !DILocation(line: 7, column: 3,  scope: ![[#BL1]], inlinedAt: ![[#INL2]])
 
diff --git a/llvm/test/Transforms/SampleProfile/pseudo-probe-emit.ll b/llvm/test/Transforms/SampleProfile/pseudo-probe-emit.ll
@@ -84,13 +84,13 @@ entry:
 ; CHECK-IL: ![[#REALLINE]] = !DILocation(line: 2, scope: ![[#DISC0:]])
 ; CHECK-IL: ![[#DISC0]] = !DILexicalBlockFile(scope: ![[#FOO]], file: ![[#]], discriminator: 0)
 ; CHECK-IL: ![[#PROBE0]] = !DILocation(line: 2, column: 20, scope: ![[#SCOPE0:]])
-;; A discriminator of 67108887 which is 0x7200017 in hexdecimal, stands for a direct call probe
+;; A discriminator of 387973143 which is 0x17200017 in hexdecimal, stands for a direct call probe
 ;; with an index of 2.
-; CHECK-IL: ![[#SCOPE0]] = !DILexicalBlockFile(scope: ![[#]], file: ![[#]], discriminator: 119537687)
+; CHECK-IL: ![[#SCOPE0]] = !DILexicalBlockFile(scope: ![[#]], file: ![[#]], discriminator: 387973143)
 ; CHECK-IL: ![[#PROBE1]] = !DILocation(line: 0, scope: ![[#SCOPE1:]])
-;; A discriminator of 186646559 which is 0xb20001f in hexdecimal, stands for a direct call probe
+;; A discriminator of 455082015 which is 0x1b20001f in hexdecimal, stands for a direct call probe
 ;; with an index of 3.
-; CHECK-IL: ![[#SCOPE1]] = !DILexicalBlockFile(scope: ![[#]], file: ![[#]], discriminator: 186646559)
+; CHECK-IL: ![[#SCOPE1]] = !DILexicalBlockFile(scope: ![[#]], file: ![[#]], discriminator: 455082015)
 
 ; Check the generation of .pseudo_probe_desc section
 ; CHECK-ASM: .section .pseudo_probe_desc,"G",@progbits,.pseudo_probe_desc_foo,comdat
diff --git a/llvm/test/Transforms/SampleProfile/pseudo-probe-inline.ll b/llvm/test/Transforms/SampleProfile/pseudo-probe-inline.ll
@@ -106,8 +106,6 @@ 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/Transforms/SampleProfile/pseudo-probe-profile.ll b/llvm/test/Transforms/SampleProfile/pseudo-probe-profile.ll
@@ -40,14 +40,14 @@ attributes #0 = {"use-sample-profile"}
 
 ; CHECK: ![[PD1]] = !{!"branch_weights", i32 8, i32 7}
 ; CHECK: ![[#PROBE1]] = !DILocation(line: 0, scope: ![[#SCOPE1:]])
-;; A discriminator of 119537695 which is 0x720001f in hexdecimal, stands for an indirect call probe
+;; A discriminator of 387973151 which is 0x1720001f in hexdecimal, stands for an indirect call probe
 ;; with an index of 3.
-; CHECK: ![[#SCOPE1]] = !DILexicalBlockFile(scope: ![[#]], file: ![[#]], discriminator: 119537695)
+; CHECK: ![[#SCOPE1]] = !DILexicalBlockFile(scope: ![[#]], file: ![[#]], discriminator: 387973151)
 ; CHECK: ![[PROF1]] = !{!"VP", i32 0, i64 7, i64 9191153033785521275, i64 5, i64 -1069303473483922844, i64 2}
-;; A discriminator of 119537711 which is 0x720002f in hexdecimal, stands for an indirect call probe
+;; A discriminator of 387973167 which is 0x1720002f in hexdecimal, stands for an indirect call probe
 ;; with an index of 5.
 ; CHECK: ![[#PROBE2]] = !DILocation(line: 0, scope: ![[#SCOPE2:]])
-; CHECK: ![[#SCOPE2]] = !DILexicalBlockFile(scope: ![[#]], file: ![[#]], discriminator: 119537711)
+; CHECK: ![[#SCOPE2]] = !DILexicalBlockFile(scope: ![[#]], file: ![[#]], discriminator: 387973167)
 ; CHECK: ![[PROF2]] = !{!"VP", i32 0, i64 6, i64 -1069303473483922844, i64 4, i64 9191153033785521275, i64 2}
 
 !llvm.module.flags = !{!9, !10}
diff --git a/llvm/test/Transforms/SampleProfile/pseudo-probe-verify.ll b/llvm/test/Transforms/SampleProfile/pseudo-probe-verify.ll
@@ -54,10 +54,10 @@ bb24:
   ret void
 }
 
-;; A discriminator of 186646583 which is 0xb200037 in hexdecimal, stands for a direct call probe
+;; A discriminator of 455082031 which is 0x1b200037 in hexdecimal, stands for a direct call probe
 ;; with an index of 6 and a scale of -1%.
 ; CHECK: ![[#PROBE6]] = !DILocation(line: 2, column: 20, scope: ![[#SCOPE:]])
-; CHECK: ![[#SCOPE]] = !DILexicalBlockFile(scope: ![[#]], file: ![[#]], discriminator: 186646575)
+; CHECK: ![[#SCOPE]] = !DILexicalBlockFile(scope: ![[#]], file: ![[#]], discriminator: 455082031)
 
 !llvm.dbg.cu = !{!0}
 !llvm.module.flags = !{!9, !10}
