[PseudoProbe] Add PseudoProbeDescUpdatePass

Add PseudoProbeDescUpdatePass to update pseudo_probe_desc metadata for
functions generated by optimization transformation. Those functions may
contain pseudo probe discriminator so that llvm-profgen will try to get
its function name based on GUID and then crash.

Author: HaohaiWen

llvm/include/llvm/Transforms/IPO/SampleProfileProbe.h

@@ -140,5 +140,13 @@ class PseudoProbeUpdatePass : public PassInfoMixin<PseudoProbeUpdatePass> {
   PreservedAnalyses run(Module &M, ModuleAnalysisManager &AM);
 };
 
+// Update pseudo_probe_desc metadata for functions generated by optimization
+// transformation.
+class PseudoProbeDescUpdatePass
+    : public PassInfoMixin<PseudoProbeDescUpdatePass> {
+public:
+  PreservedAnalyses run(Module &M, ModuleAnalysisManager &AM);
+};
+
 } // end namespace llvm
 #endif // LLVM_TRANSFORMS_IPO_SAMPLEPROFILEPROBE_H

llvm/lib/Passes/PassBuilderPipelines.cpp

@@ -1581,9 +1581,12 @@ PassBuilder::buildPerModuleDefaultPipeline(OptimizationLevel Level,
   // Now add the optimization pipeline.
   MPM.addPass(buildModuleOptimizationPipeline(Level, LTOPhase));
 
-  if (PGOOpt && PGOOpt->PseudoProbeForProfiling &&
-      PGOOpt->Action == PGOOptions::SampleUse)
-    MPM.addPass(PseudoProbeUpdatePass());
+  if (PGOOpt && PGOOpt->PseudoProbeForProfiling) {
+    if (PGOOpt->Action == PGOOptions::SampleUse)
+      MPM.addPass(PseudoProbeUpdatePass());
+    else
+      MPM.addPass(PseudoProbeDescUpdatePass());
+  }
 
   // Emit annotation remarks.
   addAnnotationRemarksPass(MPM);
@@ -1651,9 +1654,12 @@ PassBuilder::buildThinLTOPreLinkDefaultPipeline(OptimizationLevel Level) {
   if (RunPartialInlining)
     MPM.addPass(PartialInlinerPass());
 
-  if (PGOOpt && PGOOpt->PseudoProbeForProfiling &&
-      PGOOpt->Action == PGOOptions::SampleUse)
-    MPM.addPass(PseudoProbeUpdatePass());
+  if (PGOOpt && PGOOpt->PseudoProbeForProfiling) {
+    if (PGOOpt->Action == PGOOptions::SampleUse)
+      MPM.addPass(PseudoProbeUpdatePass());
+    else
+      MPM.addPass(PseudoProbeDescUpdatePass());
+  }
 
   // Handle Optimizer{Early,Last}EPCallbacks added by clang on PreLink. Actual
   // optimization is going to be done in PostLink stage, but clang can't add

llvm/lib/Passes/PassRegistry.def

@@ -119,6 +119,7 @@ MODULE_PASS("print<inline-advisor>", InlineAdvisorAnalysisPrinterPass(dbgs()))
 MODULE_PASS("print<module-debuginfo>", ModuleDebugInfoPrinterPass(dbgs()))
 MODULE_PASS("pseudo-probe", SampleProfileProbePass(TM))
 MODULE_PASS("pseudo-probe-update", PseudoProbeUpdatePass())
+MODULE_PASS("pseudo-probe-desc-update", PseudoProbeDescUpdatePass())
 MODULE_PASS("recompute-globalsaa", RecomputeGlobalsAAPass())
 MODULE_PASS("rel-lookup-table-converter", RelLookupTableConverterPass())
 MODULE_PASS("rewrite-statepoints-for-gc", RewriteStatepointsForGC())

llvm/lib/Transforms/IPO/SampleProfileProbe.cpp

@@ -341,9 +341,7 @@ uint32_t SampleProfileProber::getCallsiteId(const Instruction *Call) const {
   return Iter == CallProbeIds.end() ? 0 : Iter->second;
 }
 
-void SampleProfileProber::instrumentOneFunc(Function &F, TargetMachine *TM) {
-  Module *M = F.getParent();
-  MDBuilder MDB(F.getContext());
+static StringRef getFuncName(const Function &F) {
   // Since the GUID from probe desc and inline stack are computed separately, we
   // need to make sure their names are consistent, so here also use the name
   // from debug info.
@@ -353,6 +351,13 @@ void SampleProfileProber::instrumentOneFunc(Function &F, TargetMachine *TM) {
     if (FName.empty())
       FName = SP->getName();
   }
+  return FName;
+}
+
+void SampleProfileProber::instrumentOneFunc(Function &F, TargetMachine *TM) {
+  Module *M = F.getParent();
+  MDBuilder MDB(F.getContext());
+  StringRef FName = getFuncName(F);
   uint64_t Guid = Function::getGUID(FName);
 
   // Assign an artificial debug line to a probe that doesn't come with a real
@@ -513,3 +518,31 @@ PreservedAnalyses PseudoProbeUpdatePass::run(Module &M,
   }
   return PreservedAnalyses::none();
 }
+
+PreservedAnalyses PseudoProbeDescUpdatePass::run(Module &M,
+                                                 ModuleAnalysisManager &AM) {
+  if (NamedMDNode *FuncInfo = M.getNamedMetadata(PseudoProbeDescMetadataName)) {
+    DenseSet<uint64_t> GuidSet;
+    for (const auto *Operand : FuncInfo->operands()) {
+      const auto *MD = cast<MDNode>(Operand);
+      auto *GUID = mdconst::dyn_extract<ConstantInt>(MD->getOperand(0));
+      GuidSet.insert(GUID->getZExtValue());
+    }
+
+    MDBuilder MDB(M.getContext());
+    for (Function &F : M) {
+      if (F.isDeclaration())
+        continue;
+      StringRef FName = getFuncName(F);
+      uint64_t Guid = Function::getGUID(FName);
+      if (GuidSet.insert(Guid).second) {
+        // Set those function's hash to 0 since it may vary depending on IR
+        // input to optimization transformation.
+        auto *MD = MDB.createPseudoProbeDesc(Guid, 0, FName);
+        FuncInfo->addOperand(MD);
+        LLVM_DEBUG(dbgs() << "New fixup pseudo_probe_desc MD: " << *MD << "\n");
+      }
+    }
+  }
+  return PreservedAnalyses::all();
+}

llvm/test/Transforms/SampleProfile/pseudo-probe-desc-update.ll

@@ -0,0 +1,35 @@
+; RUN: opt < %s -passes='pseudo-probe-desc-update' -S | FileCheck %s
+
+; CHECK:      !llvm.pseudo_probe_desc = !{!0, !1, !2, !3, !4, !5}
+; CHECK:      !0 = !{i64 -2012135647395072713, i64 4294967295, !"bar"}
+; CHECK-NEXT: !1 = !{i64 9204417991963109735, i64 72617220756, !"work"}
+; CHECK-NEXT: !2 = !{i64 6699318081062747564, i64 844700110938769, !"foo"}
+; CHECK-NEXT: !3 = !{i64 -2624081020897602054, i64 281563657672557, !"main"}
+; CHECK-NEXT: !4 = !{i64 6028998432455395745, i64 0, !"extract1"}
+; CHECK-NEXT: !5 = !{i64 -8314581669044049530, i64 0, !"_Zextract2v"}
+
+target triple = "x86_64-unknown-linux-gnu"
+
+define void @extract1() !dbg !1 {
+entry:
+  ret void
+}
+
+define void @extract2() !dbg !2 {
+entry:
+  ret void
+}
+
+!llvm.pseudo_probe_desc = !{!4, !5, !6, !7}
+!llvm.module.flags = !{!8, !9}
+
+!0 = distinct !DICompileUnit(language: DW_LANG_C11, file: !3, emissionKind: FullDebug,  splitDebugInlining: false, debugInfoForProfiling: true)
+!1 = distinct !DISubprogram(name: "extract1", unit: !0)
+!2 = distinct !DISubprogram(name: "extract2", linkageName: "_Zextract2v", unit: !0)
+!3 = !DIFile(filename: "foo.c", directory: "/home/test")
+!4 = !{i64 -2012135647395072713, i64 4294967295, !"bar"}
+!5 = !{i64 9204417991963109735, i64 72617220756, !"work"}
+!6 = !{i64 6699318081062747564, i64 844700110938769, !"foo"}
+!7 = !{i64 -2624081020897602054, i64 281563657672557, !"main"}
+!8 = !{i32 7, !"Dwarf Version", i32 5}
+!9 = !{i32 2, !"Debug Info Version", i32 3}