[llvm] [PseudoProbe] Fix cleanup for pseudo probe after annotation (PR #119660)

[Lei Wang via llvm-commits]

================
@@ -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(
----------------
wlei-llvm wrote:

A friendly caveat: this is still not the original DwarfBaseDiscriminator, it's only keep 3 bits(https://github.com/llvm/llvm-project/pull/94506), though it can cover most of the cases(95%+), just not the exact same. 

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