[CSSPGO] Update pseudo probe distribution factor based on inline context.

[hoy]

Summary

With prelink inlining, pseudo probes with same ID can come from different inline contexts. Such probes should not share samples and their factors should be fixed up separately.

I'm seeing 0.3% speedup for SPEC2017 overall. Benchmark 631.deepsjeng_s benefits the most, about 4%.

[wmi]

llvm/include/llvm/IR/PseudoProbe.h

85

InlineAt saves the pointer of DILocation which may become dangling easily, so the current intention is to use it in a single pass PseudoProbeUpdatePass. If it is only used in a single pass, can we use a separate map to save the information instead of adding a field in PseudoProbe, so the fields in struct PseudoProbe all contain persistent information.

[hoy]

llvm/include/llvm/IR/PseudoProbe.h

85

Good point! Yes, inside PseudoProbeUpdatePass we can use an auxiliary map to save InlineAt. However, there's the PseudoProbeVerifier pass that tracker probe duplication throughout the pipeline and we need a globally stable InlineAt there. Since that pass is for debugging only, maybe we can just use a computed a hash for inline stack string though it's a bit slow?

[wenlei]

llvm/lib/Transforms/IPO/SampleProfileProbe.cpp

121

is getValue a new member of PseudoProbe struct? didn't find the definition.

[hoy]

llvm/lib/Transforms/IPO/SampleProfileProbe.cpp

121

getValue is a method of Optional. It returns the underlying PseudoProbe object.

[wenlei]

llvm/lib/Transforms/IPO/SampleProfileProbe.cpp

121

ah, nvm. didn't look careful enough. :)

[wlei]

Try to understand this:

So we merge the samples from the same function but different context, does this cause the issue?

Such probes should not share samples and their factors should be fixed up separately.

So with that, when code duplication happens, we need to fix probe factor, so here it will do it separately by grouping the probe by different inline context. then each group work for one context inlining. is my understanding correct?

[hoy]

In D102429#2760347, @wlei wrote:

Try to understand this:

So we merge the samples from the same function but different context, does this cause the issue?

Exactly.

Such probes should not share samples and their factors should be fixed up separately.

So with that, when code duplication happens, we need to fix probe factor, so here it will do it separately by grouping the probe by different inline context. then each group work for one context inlining. is my understanding correct?

That's right. Probe in a function profile can have multiple copies due to code duplication. We just want to share the corresponding samples among these duplications, but not among all probes with same id but from different function profiles.

[wmi]

llvm/include/llvm/IR/PseudoProbe.h

85

Sounds good. Likely computing a hash for inline stack won't be very slow, then it can be used in both PseudoProbeUpdatePass and PseudoProbeVerifier.

[wlei]

In D102429#2760386, @hoy wrote:

In D102429#2760347, @wlei wrote:

Try to understand this:

So we merge the samples from the same function but different context, does this cause the issue?

Exactly.

Such probes should not share samples and their factors should be fixed up separately.

So with that, when code duplication happens, we need to fix probe factor, so here it will do it separately by grouping the probe by different inline context. then each group work for one context inlining. is my understanding correct?

That's right. Probe in a function profile can have multiple copies due to code duplication. We just want to share the corresponding samples among these duplications, but not among all probes with same id but from different function profiles.

I see, thanks for the helpful explanation.

[hoy]

Addressing Wei's comment.

[wmi]

LGTM.

[wenlei]

lgtm, with some nit on comments.

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

36–38

comment for the key, value details would be helpful

llvm/lib/Transforms/IPO/SampleProfileProbe.cpp

391

Not directly related to this change, but PseudoProbeUpdatePass can use some header comments describing the high level idea for distribution factor update.

The fields in IR/PseudoProbe.h/PseudoProbe can also use some comments.

[hoy]

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

36–38

Comment added.

llvm/lib/Transforms/IPO/SampleProfileProbe.cpp

391

Good point. Comments added in both places.

[hoy]

Addressing Wenle's comments.

llvm/include/llvm/IR/PseudoProbe.h

 //===- PseudoProbe.h - Pseudo Probe IR Helpers ------------------*- C++ -*-===//
 //
 // Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
 // See https://llvm.org/LICENSE.txt for license information.
 // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
 //
 //===----------------------------------------------------------------------===//
 //
 // Pseudo probe IR intrinsic and dwarf discriminator manipulation routines.
 //
 //===----------------------------------------------------------------------===//
 
 #ifndef LLVM_IR_PSEUDOPROBE_H
 #define LLVM_IR_PSEUDOPROBE_H
 
 #include "llvm/ADT/Optional.h"
+#include "llvm/IR/DebugInfoMetadata.h"
 #include <cassert>
 #include <cstdint>
 #include <limits>
 
 namespace llvm {
 
 class Instruction;
 class BasicBlock;
public:
   constexpr static uint8_t FullDistributionFactor = 100;
 };
 
 struct PseudoProbe {
   uint32_t Id;
   uint32_t Type;
   uint32_t Attr;
   float Factor;
+  DILocation *InlineAt;
 
   bool isDangling() const {
     return Attr & (uint32_t)PseudoProbeAttributes::Dangling;
   }
+
+  struct Hash {
+    uint64_t operator()(const PseudoProbe &Probe) const {
+      return std::hash<uint64_t>()(Probe.Id) ^
+             std::hash<DILocation *>()(Probe.InlineAt);
+    }
+  };
+
+  struct Equal {
+    bool operator()(const PseudoProbe &L, const PseudoProbe &R) const {
+      return L.Id == R.Id && L.Type == R.Type && L.Attr == R.Attr &&
+             L.InlineAt == R.InlineAt;
+    }
+  };
 };
 
 Optional<PseudoProbe> extractProbe(const Instruction &Inst);
 
 void setProbeDistributionFactor(Instruction &Inst, float Factor);
 
 bool moveAndDanglePseudoProbes(BasicBlock *From, Instruction *To);
 
 bool removeRedundantPseudoProbes(BasicBlock *Block);
 } // end namespace llvm
 
 #endif // LLVM_IR_PSEUDOPROBE_H

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

 
 namespace llvm {
 
 class Module;
 
 using namespace sampleprof;
 using BlockIdMap = std::unordered_map<BasicBlock *, uint32_t>;
 using InstructionIdMap = std::unordered_map<Instruction *, uint32_t>;
-using ProbeFactorMap = std::unordered_map<uint64_t, float>;
+using ProbeFactorMap = std::unordered_map<PseudoProbe, float, PseudoProbe::Hash,
+                                          PseudoProbe::Equal>;
+
 using FuncProbeFactorMap = StringMap<ProbeFactorMap>;
 
 enum class PseudoProbeReservedId { Invalid = 0, Last = Invalid };
 
 class PseudoProbeDescriptor {
   uint64_t FunctionGUID;
   uint64_t FunctionHash;
 

llvm/lib/IR/PseudoProbe.cpp

if (DILocation::isPseudoProbeDiscriminator(Discriminator)) {
           PseudoProbeDwarfDiscriminator::extractProbeIndex(Discriminator);
       Probe.Type =
           PseudoProbeDwarfDiscriminator::extractProbeType(Discriminator);
       Probe.Attr =
           PseudoProbeDwarfDiscriminator::extractProbeAttributes(Discriminator);
       Probe.Factor =
           PseudoProbeDwarfDiscriminator::extractProbeFactor(Discriminator) /
           (float)PseudoProbeDwarfDiscriminator::FullDistributionFactor;
+      Probe.InlineAt = DIL->getInlinedAt();
       return Probe;
     }
   }
   return None;
 }
 
 Optional<PseudoProbe> extractProbe(const Instruction &Inst) {
   if (const auto *II = dyn_cast<PseudoProbeInst>(&Inst)) {
     PseudoProbe Probe;
     Probe.Id = II->getIndex()->getZExtValue();
     Probe.Type = (uint32_t)PseudoProbeType::Block;
     Probe.Attr = II->getAttributes()->getZExtValue();
     Probe.Factor = II->getFactor()->getZExtValue() /
                    (float)PseudoProbeFullDistributionFactor;
+    Probe.InlineAt =
+        II->getDebugLoc() ? II->getDebugLoc()->getInlinedAt() : nullptr;
     return Probe;
   }
 
   if (isa<CallBase>(&Inst) && !isa<IntrinsicInst>(&Inst))
     return extractProbeFromDiscriminator(Inst);
 
   return None;
 }

llvm/lib/Transforms/IPO/SampleProfileProbe.cpp

void PseudoProbeVerifier::runAfterPass(const Loop *L) {
   const Function *F = L->getHeader()->getParent();
   runAfterPass(F);
 }
 
 void PseudoProbeVerifier::collectProbeFactors(const BasicBlock *Block,
                                               ProbeFactorMap &ProbeFactors) {
   for (const auto &I : *Block) {
     if (Optional<PseudoProbe> Probe = extractProbe(I))
-      ProbeFactors[Probe->Id] += Probe->Factor;
+      ProbeFactors[Probe.getValue()] += Probe->Factor;
   }
 }
 
 void PseudoProbeVerifier::verifyProbeFactors(
     const Function *F, const ProbeFactorMap &ProbeFactors) {
   bool BannerPrinted = false;
   auto &PrevProbeFactors = FunctionProbeFactors[F->getName()];
   for (const auto &I : ProbeFactors) {
     float CurProbeFactor = I.second;
     if (PrevProbeFactors.count(I.first)) {
       float PrevProbeFactor = PrevProbeFactors[I.first];
       if (std::abs(CurProbeFactor - PrevProbeFactor) >
           DistributionFactorVariance) {
         if (!BannerPrinted) {
           dbgs() << "Function " << F->getName() << ":\n";
           BannerPrinted = true;
         }
-        dbgs() << "Probe " << I.first << "\tprevious factor "
+        dbgs() << "Probe " << I.first.Id << "\tprevious factor "
                << format("%0.2f", PrevProbeFactor) << "\tcurrent factor "
                << format("%0.2f", CurProbeFactor) << "\n";
       }
     }
 
     // Update
     PrevProbeFactors[I.first] = I.second;
   }
if (F.isDeclaration())
       continue;
     SampleProfileProber ProbeManager(F, ModuleId);
     ProbeManager.instrumentOneFunc(F, TM);
   }
 
   return PreservedAnalyses::none();
 }
 
 void PseudoProbeUpdatePass::runOnFunction(Function &F,
                                           FunctionAnalysisManager &FAM) {
   BlockFrequencyInfo &BFI = FAM.getResult<BlockFrequencyAnalysis>(F);
   auto BBProfileCount = [&BFI](BasicBlock *BB) {
     return BFI.getBlockProfileCount(BB)
                ? BFI.getBlockProfileCount(BB).getValue()
                : 0;
   };
 
   // Collect the sum of execution weight for each probe.
   ProbeFactorMap ProbeFactors;
   for (auto &Block : F) {
     for (auto &I : Block) {
       if (Optional<PseudoProbe> Probe = extractProbe(I)) {
         // Do not count dangling probes since they are logically deleted and the
         // current block that a dangling probe resides in doesn't reflect the
         // execution count of the probe. The original samples of the probe will
         // be distributed among the rest probes if there are any, this is
         // less-than-deal but at least we don't lose any samples.
         if (!Probe->isDangling())
-          ProbeFactors[Probe->Id] += BBProfileCount(&Block);
+          ProbeFactors[Probe.getValue()] += BBProfileCount(&Block);
       }
     }
   }
 
   // Fix up over-counted probes.
   for (auto &Block : F) {
     for (auto &I : Block) {
       if (Optional<PseudoProbe> Probe = extractProbe(I)) {
         // Ignore danling probes since they are logically deleted and should do
         // not consume any profile samples in the subsequent profile annotation.
         if (!Probe->isDangling()) {
-          float Sum = ProbeFactors[Probe->Id];
+          float Sum = ProbeFactors[Probe.getValue()];
           if (Sum != 0)
             setProbeDistributionFactor(I, BBProfileCount(&Block) / Sum);
         }
       }
     }
   }
 }
 

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

+; RUN: opt < %s -passes='pseudo-probe-update' -S  | FileCheck %s
+
+declare i32 @f1()
+
+declare i32 @f2()
+
+declare void @f3()
+
+define i32 @foo(i1 %cond, i1 %cond2) !dbg !4 !prof !10 {
+  call void @llvm.pseudoprobe(i64 6699318081062747564, i64 1, i32 0, i64 -1)
+  br i1 %cond, label %T1, label %Merge, !prof !11
+
+T1:                                               ; preds = %0
+  %v1 = call i32 @f1(), !prof !12
+  %cond3 = icmp eq i32 %v1, 412
+  call void @llvm.pseudoprobe(i64 6699318081062747564, i64 2, i32 0, i64 -1)
+;; The distribution factor -8513881372706734080 stands for 53.85%, whic is from 7/6+7.
+; CHECK: call void @llvm.pseudoprobe(i64 [[#GUID:]], i64 4, i32 0, i64 -8513881372706734080)
+  call void @llvm.pseudoprobe(i64 6699318081062747564, i64 4, i32 0, i64 -1), !dbg !13
+;; Probe 7 has two copies, since they don't share the same inline context, they are not
+;; considered sharing samples, thus their distribution factors are not fixed up.
+; CHECK: call void @llvm.pseudoprobe(i64 [[#GUID]], i64 7, i32 0, i64 -1)
+  call void @llvm.pseudoprobe(i64 6699318081062747564, i64 7, i32 0, i64 -1), !dbg !13
+;; Similar to Probe 7, one copy of Probe 8 doesn't have inline context.
+; CHECK: call void @llvm.pseudoprobe(i64 [[#GUID]], i64 8, i32 0, i64 -1)
+  call void @llvm.pseudoprobe(i64 6699318081062747564, i64 8, i32 0, i64 -1), !dbg !13
+  br i1 %cond3, label %T2, label %F2, !prof !11
+
+Merge:                                            ; preds = %0
+  %v2 = call i32 @f2(), !prof !12
+  call void @llvm.pseudoprobe(i64 6699318081062747564, i64 3, i32 0, i64 -1)
+;; The distribution factor 8513881922462547968 stands for 46.25%, which is from 6/6+7.
+; CHECK: call void @llvm.pseudoprobe(i64 [[#GUID]], i64 4, i32 0, i64 8513881922462547968)
+  call void @llvm.pseudoprobe(i64 6699318081062747564, i64 4, i32 0, i64 8513881922462547968), !dbg !13
+; CHECK: call void @llvm.pseudoprobe(i64 [[#GUID]], i64 7, i32 0, i64 -1)
+  call void @llvm.pseudoprobe(i64 6699318081062747564, i64 7, i32 0, i64 -1), !dbg !18
+; CHECK: call void @llvm.pseudoprobe(i64 [[#GUID]], i64 8, i32 0, i64 -1)
+  call void @llvm.pseudoprobe(i64 6699318081062747564, i64 8, i32 0, i64 -1), !dbg !19 
+  br i1 %cond2, label %T2, label %F2, !prof !11
+
+T2:                                               ; preds = %Merge, %T1
+  %B1 = phi i32 [ %v1, %T1 ], [ %v2, %Merge ]
+  call void @f3(), !prof !12
+  call void @llvm.pseudoprobe(i64 6699318081062747564, i64 5, i32 0, i64 -1)
+  ret i32 %B1
+
+F2:                                               ; preds = %Merge, %T1
+  %B2 = phi i32 [ %v1, %T1 ], [ %v2, %Merge ]
+  call void @llvm.pseudoprobe(i64 6699318081062747564, i64 6, i32 0, i64 -1)
+  ret i32 %B2
+}
+
+; Function Attrs: inaccessiblememonly nounwind willreturn
+declare void @llvm.pseudoprobe(i64, i64, i32, i64) #0
+
+attributes #0 = { inaccessiblememonly nounwind willreturn }
+
+!llvm.module.flags = !{!0, !1}
+!llvm.pseudo_probe_desc = !{!2, !3}
+
+!0 = !{i32 7, !"Dwarf Version", i32 4}
+!1 = !{i32 2, !"Debug Info Version", i32 3}
+!2 = !{i64 6699318081062747564, i64 281479271677951, !"foo", null}
+!3 = !{i64 6468398850841090686, i64 138828622701, !"zen", null}
+!4 = distinct !DISubprogram(name: "foo", scope: !5, file: !5, line: 9, type: !6, scopeLine: 9, spFlags: DISPFlagDefinition, unit: !9)
+!5 = !DIFile(filename: "test.cpp", directory: "test")
+!6 = !DISubroutineType(types: !7)
+!7 = !{!8, !8}
+!8 = !DIBasicType(name: "int", size: 32, encoding: DW_ATE_signed)
+!9 = distinct !DICompileUnit(language: DW_LANG_C_plus_plus_14, file: !5, isOptimized: false, runtimeVersion: 0, emissionKind: NoDebug)
+!10 = !{!"function_entry_count", i64 14}
+!11 = !{!"branch_weights", i32 8, i32 7}
+!12 = !{!"branch_weights", i32 7}
+!13 = !DILocation(line: 39, column: 9, scope: !14, inlinedAt: !16)
+!14 = distinct !DILexicalBlock(scope: !15, file: !5, line: 39, column: 7)
+!15 = distinct !DISubprogram(name: "zen", scope: !5, file: !5, line: 37, type: !6, scopeLine: 38, spFlags: DISPFlagDefinition, unit: !9)
+!16 = distinct !DILocation(line: 10, column: 11, scope: !17)
+!17 = !DILexicalBlockFile(scope: !4, file: !5, discriminator: 186646551)
+!18 = !DILocation(line: 53, column: 3, scope: !15, inlinedAt: !19)
+!19 = !DILocation(line: 12, column: 3, scope: !4)
 No newline at end of file