[CSSPGO] Do not import pseudo probe desc in thinLTO

[hoy]

Summary

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:

[wenlei]

Thanks for the change and cleanup. Do we see identical probe and probe_desc section for final binary for larger builds w/ and w/o this change?

[hoy]

In D105248#2851809, @wenlei wrote:

Thanks for the change and cleanup. Do we see identical probe and probe_desc section for final binary for larger builds w/ and w/o this change?

That's a good point. I have tried with a giant binary, and the generated probe sections are decodable. Haven't compared smaller workloads. Let me do that.

[hoy]

Updating D105248: [CSSPGO] Do not import pseudo probe desc in thinLTO

[hoy]

Verified that for SPEC, the generated profiles w/ and w/o this changes are very close. The runtime performance of pass2 has no difference two. Regarding the pseudo probe sections, the .pseudo_probe section is exactly the same for all benchmarks. The .pseudo_probe_desc sections are mostly the same, except for a couple benchmarks, for which the probe decoding gives exact result. There might be a glitch in setting the comdat group name for duplicated functions, like the ones with .llvm. suffix, or an issue of our internal thinlto dedup work.

[wmi]

-funique-internal-linkage-names is not enabled by default. Will -fpseudo-probe-for-profiling implicitly enable -funique-internal-linkage-names?

[wenlei]

In D105248#2855626, @wmi wrote:

-funique-internal-linkage-names is not enabled by default. Will -fpseudo-probe-for-profiling implicitly enable -funique-internal-linkage-names?

We still pass -funique-internal-linkage-names separately. I think the situation is similar to line-based AFDO.

[wenlei]

In D105248#2853977, @hoy wrote:

Verified that for SPEC, the generated profiles w/ and w/o this changes are very close. The runtime performance of pass2 has no difference two. Regarding the pseudo probe sections, the .pseudo_probe section is exactly the same for all benchmarks. The .pseudo_probe_desc sections are mostly the same, except for a couple benchmarks, for which the probe decoding gives exact result. There might be a glitch in setting the comdat group name for duplicated functions, like the ones with .llvm. suffix, or an issue of our internal thinlto dedup work.

Thanks for checking. For pseudo_probe_desc, do you know what is the actual difference, and why decoding leads to same output with different pseudo_probe_desc?

[hoy]

In D105248#2855763, @wenlei wrote:

In D105248#2853977, @hoy wrote:

Verified that for SPEC, the generated profiles w/ and w/o this changes are very close. The runtime performance of pass2 has no difference two. Regarding the pseudo probe sections, the .pseudo_probe section is exactly the same for all benchmarks. The .pseudo_probe_desc sections are mostly the same, except for a couple benchmarks, for which the probe decoding gives exact result. There might be a glitch in setting the comdat group name for duplicated functions, like the ones with .llvm. suffix, or an issue of our internal thinlto dedup work.

Thanks for checking. For pseudo_probe_desc, do you know what is the actual difference, and why decoding leads to same output with different pseudo_probe_desc?

It appeared due to a glitch of our internal thinlto comdat deduplication change. Only part of the duplicated comdat groups were removed. llvm-profgen can tolerate that by taking only the last duplicate.

[wenlei]

It appeared due to a glitch of our internal thinlto comdat deduplication change. Only part of the duplicated comdat groups were removed. llvm-profgen can tolerate that by taking only the last duplicate.

I think it'd be good to track it down why that happens - the expectation is to have identical probe_desc before/after this change. Though this change lgtm.

llvm/lib/CodeGen/AsmPrinter/AsmPrinter.cpp

if (!EmitCodeView || M.getDwarfVersion()) {
         Handlers.emplace_back(std::unique_ptr<DwarfDebug>(DD), DbgTimerName,
                               DbgTimerDescription, DWARFGroupName,
                               DWARFGroupDescription);
       }
     }
   }
 
   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);
   }
 
   switch (MAI->getExceptionHandlingType()) {
   case ExceptionHandling::None:
     // We may want to emit CFI for debug.
     LLVM_FALLTHROUGH;

llvm/lib/CodeGen/AsmPrinter/PseudoProbePrinter.h

 class AsmPrinter;
 class MCStreamer;
 class Module;
 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);
 
   // Unused.
   void setSymbolSize(const MCSymbol *Sym, uint64_t Size) override {}
   void endModule() override {}
   void beginFunction(const MachineFunction *MF) override {}
   void endFunction(const MachineFunction *MF) override {}
   void beginInstruction(const MachineInstr *MI) override {}
   void endInstruction() override {}
-  
-#ifndef NDEBUG
-  void dump() const;
-#endif
 };
 
 } // namespace llvm
 #endif // LLVM_LIB_CODEGEN_ASMPRINTER_PSEUDOPROBEPRINTER_H

llvm/lib/CodeGen/AsmPrinter/PseudoProbePrinter.cpp

 #include "llvm/IR/DebugInfoMetadata.h"
 #include "llvm/IR/Module.h"
 #include "llvm/IR/PseudoProbe.h"
 #include "llvm/MC/MCPseudoProbe.h"
 #include "llvm/MC/MCStreamer.h"
 
 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) {
   // Gather all the inlined-at nodes.
   // When it's done ReversedInlineStack looks like ([66, B], [88, A])
   // which means, Function A inlines function B at calliste with a probe id 88,
   // and B inlines C at probe 66 where C is represented by Guid.
   SmallVector<InlineSite, 8> ReversedInlineStack;
   auto *InlinedAt = DebugLoc ? DebugLoc->getInlinedAt() : nullptr;
   while (InlinedAt) {
     const DISubprogram *SP = InlinedAt->getScope()->getSubprogram();
     // Use linkage name for C++ if possible.
     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);
     InlinedAt = InlinedAt->getInlinedAt();
   }
 
   SmallVector<InlineSite, 8> InlineStack(ReversedInlineStack.rbegin(),
                                          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/Linker/IRMover.cpp

 #include "llvm/ADT/SetVector.h"
 #include "llvm/ADT/SmallString.h"
 #include "llvm/ADT/Triple.h"
 #include "llvm/IR/Constants.h"
 #include "llvm/IR/DebugInfo.h"
 #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"
 #include "llvm/Transforms/Utils/Cloning.h"
 #include <utility>
 using namespace llvm;
 
 //===----------------------------------------------------------------------===//

 
 /// Insert all of the named MDNodes in Src into the Dest module.
 void IRLinker::linkNamedMDNodes() {
   const NamedMDNode *SrcModFlags = SrcM->getModuleFlagsMetadata();
   for (const NamedMDNode &NMD : SrcM->named_metadata()) {
     // 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())
       DestNMD->addOperand(Mapper.mapMDNode(*Op));
   }
 }
 
 /// Merge the linker flags in Src into the Dest module.

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

+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}
 No newline at end of file

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

+; 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}