[PseudoProbe] Fix cleanup for pseudo probe after annotation

[HaohaiWen]

When using -sample-profile-remove-probe, pseudo probe desc should
also be removed and dwarf discriminator for call instruction should
be restored.

[llvmbot]

@llvm/pr-subscribers-pgo

@llvm/pr-subscribers-llvm-transforms

Author: Haohai Wen (HaohaiWen)

Changes

When using -sample-profile-remove-probe, pseudo probe desc should
also be removed and dwarf discriminator for call instruction should
be restored.

---

Full diff: https://github.com/llvm/llvm-project/pull/119660.diff

2 Files Affected:

• (modified) llvm/lib/Transforms/IPO/SampleProfile.cpp (+20-4)
• (modified) llvm/test/Transforms/SampleProfile/pseudo-probe-profile.ll (+2)

diff --git a/llvm/lib/Transforms/IPO/SampleProfile.cpp b/llvm/lib/Transforms/IPO/SampleProfile.cpp
index b2fa66f2a6d379..603beb3b883d7f 100644
--- a/llvm/lib/Transforms/IPO/SampleProfile.cpp
+++ b/llvm/lib/Transforms/IPO/SampleProfile.cpp
@@ -529,7 +529,7 @@ class SampleProfileLoader final : public SampleProfileLoaderBaseImpl<Function> {
   void generateMDProfMetadata(Function &F);
   bool rejectHighStalenessProfile(Module &M, ProfileSummaryInfo *PSI,
                                   const SampleProfileMap &Profiles);
-  void removePseudoProbeInsts(Module &M);
+  void removePseudoProbeInstsDiscriminator(Module &M);
 
   /// Map from function name to Function *. Used to find the function from
   /// the function name. If the function name contains suffix, additional
@@ -2138,13 +2138,25 @@ bool SampleProfileLoader::rejectHighStalenessProfile(
   return false;
 }
 
-void SampleProfileLoader::removePseudoProbeInsts(Module &M) {
+void SampleProfileLoader::removePseudoProbeInstsDiscriminator(Module &M) {
   for (auto &F : M) {
     std::vector<Instruction *> InstsToDel;
     for (auto &BB : F) {
       for (auto &I : BB) {
         if (isa<PseudoProbeInst>(&I))
           InstsToDel.push_back(&I);
+        else if (isa<CallBase>(&I))
+          if (const DILocation *DIL = I.getDebugLoc().get()) {
+            // Restore dwarf discriminator for call.
+            unsigned Discriminator = DIL->getDiscriminator();
+            if (DILocation::isPseudoProbeDiscriminator(Discriminator)) {
+              std::optional<uint32_t> DwarfDiscriminator =
+                  PseudoProbeDwarfDiscriminator::extractDwarfBaseDiscriminator(
+                      Discriminator);
+              I.setDebugLoc(DIL->cloneWithDiscriminator(
+                  DwarfDiscriminator ? *DwarfDiscriminator : 0));
+            }
+          }
       }
     }
     for (auto *I : InstsToDel)
@@ -2224,8 +2236,12 @@ bool SampleProfileLoader::runOnModule(Module &M, ModuleAnalysisManager *AM,
          notInlinedCallInfo)
       updateProfileCallee(pair.first, pair.second.entryCount);
 
-  if (RemoveProbeAfterProfileAnnotation && FunctionSamples::ProfileIsProbeBased)
-    removePseudoProbeInsts(M);
+  if (RemoveProbeAfterProfileAnnotation &&
+      FunctionSamples::ProfileIsProbeBased) {
+    removePseudoProbeInstsDiscriminator(M);
+    if (auto *FuncInfo = M.getNamedMetadata(PseudoProbeDescMetadataName))
+      M.eraseNamedMetadata(FuncInfo);
+  }
 
   return retval;
 }
diff --git a/llvm/test/Transforms/SampleProfile/pseudo-probe-profile.ll b/llvm/test/Transforms/SampleProfile/pseudo-probe-profile.ll
index b52f93763d4926..66dbc49a1d210a 100644
--- a/llvm/test/Transforms/SampleProfile/pseudo-probe-profile.ll
+++ b/llvm/test/Transforms/SampleProfile/pseudo-probe-profile.ll
@@ -4,6 +4,8 @@
 ; RUN: opt < %t -passes=sample-profile -sample-profile-file=%S/Inputs/pseudo-probe-profile.prof -sample-profile-remove-probe -S | FileCheck %s -check-prefix=REMOVE-PROBE
 
 ; REMOVE-PROBE-NOT: call void @llvm.pseudoprobe
+; REMOVE-PROBE-NOT: !llvm.pseudo_probe_desc
+; REMOVE-PROBE:     !DILexicalBlockFile({{.*}}, discriminator: 0)
 
 define dso_local i32 @foo(i32 %x, ptr %f) #0 !dbg !4 {
 entry:

[HaohaiWen]

Is zero discriminator same effect as no discriminator in DILexicalBlockFile?

[wlei-llvm]

I think so, looks the encoding https://github.com/llvm/llvm-project/blob/main/llvm/include/llvm/IR/DebugInfoMetadata.h#L2101 0 means all the components are empty.

[WenleiHe]

The flag was used to avoid probe intrinsic hindering optimizations, so discriminator value etc wasn't important for that purpose. Curious what are you using this flag for?

[HaohaiWen]

The flag was used to avoid probe intrinsic hindering optimizations, so discriminator value etc wasn't important for that purpose. Curious what are you using this flag for?

.llvm.pseudo_probe_desc and discriminator will create .pseudo_probe and .pseudo_probe_desc sections which would increase the binary size of second round sampling PGO build (about hundreds of MBs for large project). .pseudo_probe_desc would also expose function name in binary.
Those discriminators and .llvm.pseudo_probe_desc is useless after sample profile loader, so I think they can be removed safely.

[wlei-llvm]

LGTM, thanks!

[wlei-llvm]

A friendly caveat: this is still not the original DwarfBaseDiscriminator, it's only keep 3 bits(#94506), though it can cover most of the cases(95%+), just not the exact same.

[HaohaiWen]

Thanks. I noticed it.

[wlei-llvm]

I think so, looks the encoding https://github.com/llvm/llvm-project/blob/main/llvm/include/llvm/IR/DebugInfoMetadata.h#L2101 0 means all the components are empty.