[CSSPGO] Do not import pseudo probe desc in thinLTO

Previously we reliedy on pseudo probe descriptors to look up precomputed GUID during probe emission for inlined probes. Since we are moving to always using unique linkage names, GUID for functions can be computed in place from dwarf names. This eliminates the need of importing pseudo probe descs in thinlto, since those descs should be emitted by the original modules.

This significantly reduces thinlto memory footprint in some extreme case where the number of imported modules for a single module is massive.

Test Plan:

Reviewed By: wenlei

Differential Revision: https://reviews.llvm.org/D105248

Author: htyu

llvm/lib/CodeGen/AsmPrinter/AsmPrinter.cpp

@@ -359,7 +359,7 @@ bool AsmPrinter::doInitialization(Module &M) {
   }
 
   if (M.getNamedMetadata(PseudoProbeDescMetadataName)) {
-    PP = new PseudoProbeHandler(this, &M);
+    PP = new PseudoProbeHandler(this);
     Handlers.emplace_back(std::unique_ptr<PseudoProbeHandler>(PP), PPTimerName,
                           PPTimerDescription, PPGroupName, PPGroupDescription);
   }

llvm/lib/CodeGen/AsmPrinter/PseudoProbePrinter.cpp

@@ -20,31 +20,6 @@
 
 using namespace llvm;
 
-#define DEBUG_TYPE "pseudoprobe"
-
-PseudoProbeHandler::~PseudoProbeHandler() = default;
-
-PseudoProbeHandler::PseudoProbeHandler(AsmPrinter *A, Module *M) : Asm(A) {
-  NamedMDNode *FuncInfo = M->getNamedMetadata(PseudoProbeDescMetadataName);
-  assert(FuncInfo && "Pseudo probe descriptors are missing");
-  for (const auto *Operand : FuncInfo->operands()) {
-    const auto *MD = cast<MDNode>(Operand);
-    auto GUID =
-        mdconst::dyn_extract<ConstantInt>(MD->getOperand(0))->getZExtValue();
-    auto Name = cast<MDString>(MD->getOperand(2))->getString();
-    // We may see pairs with same name but different GUIDs here in LTO mode, due
-    // to static same-named functions inlined from other modules into this
-    // module. Function profiles with the same name will be merged no matter
-    // whether they are collected on the same function. Therefore we just pick
-    // up the last <Name, GUID> pair here to represent the same-named function
-    // collection and all probes from the collection will be merged into a
-    // single profile eventually.
-    Names[Name] = GUID;
-  }
-
-  LLVM_DEBUG(dump());
-}
-
 void PseudoProbeHandler::emitPseudoProbe(uint64_t Guid, uint64_t Index,
                                          uint64_t Type, uint64_t Attr,
                                          const DILocation *DebugLoc) {
@@ -60,8 +35,7 @@ void PseudoProbeHandler::emitPseudoProbe(uint64_t Guid, uint64_t Index,
     auto Name = SP->getLinkageName();
     if (Name.empty())
       Name = SP->getName();
-    assert(Names.count(Name) && "Pseudo probe descriptor missing for function");
-    uint64_t CallerGuid = Names[Name];
+    uint64_t CallerGuid = Function::getGUID(Name);
     uint64_t CallerProbeId = PseudoProbeDwarfDiscriminator::extractProbeIndex(
         InlinedAt->getDiscriminator());
     ReversedInlineStack.emplace_back(CallerGuid, CallerProbeId);
@@ -72,13 +46,3 @@ void PseudoProbeHandler::emitPseudoProbe(uint64_t Guid, uint64_t Index,
                                          ReversedInlineStack.rend());
   Asm->OutStreamer->emitPseudoProbe(Guid, Index, Type, Attr, InlineStack);
 }
-
-#ifndef NDEBUG
-void PseudoProbeHandler::dump() const {
-  dbgs() << "\n=============================\n";
-  dbgs() << "\nFunction Name to GUID map:\n";
-  dbgs() << "\n=============================\n";
-  for (const auto &Item : Names)
-    dbgs() << "Func: " << Item.first << "   GUID: " << Item.second << "\n";
-}
-#endif

llvm/lib/CodeGen/AsmPrinter/PseudoProbePrinter.h

@@ -26,12 +26,9 @@ class DILocation;
 class PseudoProbeHandler : public AsmPrinterHandler {
   // Target of pseudo probe emission.
   AsmPrinter *Asm;
-  // Name to GUID map
-  DenseMap<StringRef, uint64_t> Names;
 
 public:
-  PseudoProbeHandler(AsmPrinter *A, Module *M);
-  ~PseudoProbeHandler() override;
+  PseudoProbeHandler(AsmPrinter *A) : Asm(A){};
 
   void emitPseudoProbe(uint64_t Guid, uint64_t Index, uint64_t Type,
                        uint64_t Attr, const DILocation *DebugLoc);
@@ -43,10 +40,6 @@ class PseudoProbeHandler : public AsmPrinterHandler {
   void endFunction(const MachineFunction *MF) override {}
   void beginInstruction(const MachineInstr *MI) override {}
   void endInstruction() override {}
-  
-#ifndef NDEBUG
-  void dump() const;
-#endif
 };
 
 } // namespace llvm

llvm/lib/Linker/IRMover.cpp

@@ -16,6 +16,7 @@
 #include "llvm/IR/DiagnosticPrinter.h"
 #include "llvm/IR/GVMaterializer.h"
 #include "llvm/IR/Intrinsics.h"
+#include "llvm/IR/PseudoProbe.h"
 #include "llvm/IR/TypeFinder.h"
 #include "llvm/Object/ModuleSymbolTable.h"
 #include "llvm/Support/Error.h"
@@ -1207,6 +1208,10 @@ void IRLinker::linkNamedMDNodes() {
     // Don't link module flags here. Do them separately.
     if (&NMD == SrcModFlags)
       continue;
+    // Don't import pseudo probe descriptors here for thinLTO. They will be
+    // emitted by the originating module.
+    if (IsPerformingImport && NMD.getName() == PseudoProbeDescMetadataName)
+      continue;
     NamedMDNode *DestNMD = DstM.getOrInsertNamedMetadata(NMD.getName());
     // Add Src elements into Dest node.
     for (const MDNode *Op : NMD.operands())

llvm/test/ThinLTO/X86/Inputs/pseudo-probe-desc-import.ll

@@ -0,0 +1,16 @@
+target datalayout = "e-m:e-p270:32:32-p271:32:32-p272:64:64-i64:64-f80:128-n8:16:32:64-S128"
+target triple = "x86_64-unknown-linux-gnu"
+
+define void @foo() {
+entry:
+  call void @llvm.pseudoprobe(i64 6699318081062747564, i64 1, i32 0, i64 -1)
+  ret void
+}
+
+declare void @llvm.pseudoprobe(i64, i64, i32, i64) #0
+
+attributes #0 = { inaccessiblememonly nounwind willreturn }
+
+!llvm.pseudo_probe_desc = !{!0}
+
+!0 = !{i64 6699318081062747564, i64 4294967295, !"foo", null}

llvm/test/ThinLTO/X86/pseudo-probe-desc-import.ll

@@ -0,0 +1,28 @@
+; RUN: opt -module-summary %s -o %t1.bc
+; RUN: opt -module-summary %p/Inputs/pseudo-probe-desc-import.ll -o %t2.bc
+; RUN: llvm-lto -thinlto-action=thinlink -o %t.index.bc %t1.bc %t2.bc
+
+; Don't import pseudo probe desc.
+; RUN: llvm-lto -thinlto-action=import %t1.bc -thinlto-index=%t.index.bc -o - | llvm-dis -o - | FileCheck %s
+; CHECK-NOT: {i64 6699318081062747564, i64 4294967295, !"foo"
+; CHECK: !{i64 -2624081020897602054, i64 281479271677951, !"main"
+
+target datalayout = "e-m:e-p270:32:32-p271:32:32-p272:64:64-i64:64-f80:128-n8:16:32:64-S128"
+target triple = "x86_64-unknown-linux-gnu"
+
+define i32 @main() {
+entry:
+  call void (...) @foo()
+  call void @llvm.pseudoprobe(i64 -2624081020897602054, i64 1, i32 0, i64 -1)
+  ret i32 0
+}
+
+declare void @foo(...)
+
+declare void @llvm.pseudoprobe(i64, i64, i32, i64) #0
+
+attributes #0 = { inaccessiblememonly nounwind willreturn }
+
+!llvm.pseudo_probe_desc = !{!0}
+
+!0 = !{i64 -2624081020897602054, i64 281479271677951, !"main", null}