[CSSPGO] Overwrite branch weight annotated in previous pass.

[hoy]

Summary

Sample profile loader can be run in both LTO prelink and postlink. Currently the counts annoation in postilnk doesn't fully overwrite what's done in prelink. I'm adding a switch (-overwrite-existing-weights=1) to enable a full overwrite, which includes:

1. Clear old metadata for calls when their parent block has a zero count. This could be caused by prelink code duplication.

2. Clear indirect call metadata if somehow all the rest targets have a sum of zero count.

3. Overwrite branch weight for basic blocks.

With a CS profile, I was seeing #1 and #2 help reduce code size by preventing post-sample ICP and CGSCC inliner working on obsolete metadata, which come from a partial global inlining in prelink. It's not expected to work well for non-CS case with a less-accurate post-inline count quality.

It's worth calling out that some prelink optimizations can damage counts quality in an irreversible way. One example is the loop rotate optimization. Due to lack of exact loop entry count (profiling can only give loop iteration count and loop exit count), moving one iteration out of the loop body leaves the rest iteration count unknown. We had to turn off prelink loop rotate to achieve a better postlink counts quality. A even better postlink counts quality can be archived by turning off prelink CGSCC inlining which is not context-sensitive.

[wenlei]

Thanks for bringing this up - it makes sure we give full context profile precedence even when we had to use context-less profile in prelink.

When counts are accurate, I was seeing #1 and #2 help reduce code size by preventing post-sample ICP and CGSCC inliner working on obselete metedata.

I assume the above is from CSSPGO, right?

It'd be good to call out in the change description as to why postlink overwrite could be beneficial.

If we target AutoFDO as well, would be good to mention the impact on AutoFDO too.

llvm/lib/Transforms/IPO/SampleProfile.cpp

226

nit: keep the blank line before namespace?

1458–1459

Btw, this comment is no longer accurate, right? Based on the discussion in D100993, the factor is strictly for handling duplication now, not for scaling counts after ICP.

1482–1486

So we have special case for Sum == 0 because updateIDTMetaData with Sum 0 is reserved for marking promoted targets, right?

We may also reach Sum == 0 inside updateIDTMetaData after promoted target is removed, however it looks like annotateValueSite would simply bailout without removing existing metadata when new target list is empty.. so we may miss some overwrite still? Though for such cases, it needs to be updated even when OverwriteExistingWeights is not used.

[hoy]

In D102537#2761666, @wenlei wrote:

Thanks for bringing this up - it makes sure we give full context profile precedence even when we had to use context-less profile in prelink.

When counts are accurate, I was seeing #1 and #2 help reduce code size by preventing post-sample ICP and CGSCC inliner working on obselete metedata.

I assume the above is from CSSPGO, right?

It'd be good to call out in the change description as to why postlink overwrite could be beneficial.

Yes, they are for CSSPGO.

If we target AutoFDO as well, would be good to mention the impact on AutoFDO too.

I haven't tested for non-CS case. In theory it should help there but with a less-than-ideal counts quality, the change may not help in practice.

llvm/lib/Transforms/IPO/SampleProfile.cpp

226

done.

1458–1459

Good point. The proration is done based on a pre-ICP distribution factor.

1482–1486

So we have special case for Sum == 0 because updateIDTMetaData with Sum 0 is reserved for marking promoted targets, right?

Right, Sum == 0 for updateIDTMetaData is used to mark a certain target to be promoted already.

We may also reach Sum == 0 inside updateIDTMetaData after promoted target is removed, however it looks like annotateValueSite would simply bailout without removing existing metadata when new target list is empty.. so we may miss some overwrite still? Though for such cases, it needs to be updated even when OverwriteExistingWeights is not used.

This is true. When Sum reaches zero when promoted target counts are withdrawn, instead of removing the existing metadata, the current implementation just assigns the call instruction a new zero metadata. That should give equivalent information to subsequent passes.

[hoy]

Addressing Wenlei's comment.

[wenlei]

In D102537#2762431, @hoy wrote:

In D102537#2761666, @wenlei wrote:

Thanks for bringing this up - it makes sure we give full context profile precedence even when we had to use context-less profile in prelink.

When counts are accurate, I was seeing #1 and #2 help reduce code size by preventing post-sample ICP and CGSCC inliner working on obselete metedata.

I assume the above is from CSSPGO, right?

It'd be good to call out in the change description as to why postlink overwrite could be beneficial.

Yes, they are for CSSPGO.

If we target AutoFDO as well, would be good to mention the impact on AutoFDO too.

I haven't tested for non-CS case. In theory it should help there but with a less-than-ideal counts quality, the change may not help in practice.

For future reference, can you make it explicit in the change description: 1) what made this beneficial for CSSPGO (give post-link full context profile precedence over pre-link partial context profile, i.e. be more specific about obsolete metadata); 2) what is the expected impact on AutoFDO.

I think there's another caveat worth calling out, by overwriting all pre-link weights, we would lose the prof metadata adjustment done by prelink opt passes (e.g. prelink loop passes).

llvm/lib/Transforms/IPO/SampleProfile.cpp

1482–1486

Ok took another look, i was looking at the wrong overload of annotateValueSite. There's one that would bail out without setting metadata if value data is empty, but that's not the overload we call from updateIDTMetaData. So we should be good.

[hoy]

In D102537#2763883, @wenlei wrote:

In D102537#2762431, @hoy wrote:

In D102537#2761666, @wenlei wrote:

Thanks for bringing this up - it makes sure we give full context profile precedence even when we had to use context-less profile in prelink.

When counts are accurate, I was seeing #1 and #2 help reduce code size by preventing post-sample ICP and CGSCC inliner working on obselete metedata.

I assume the above is from CSSPGO, right?

It'd be good to call out in the change description as to why postlink overwrite could be beneficial.

Yes, they are for CSSPGO.

If we target AutoFDO as well, would be good to mention the impact on AutoFDO too.

I haven't tested for non-CS case. In theory it should help there but with a less-than-ideal counts quality, the change may not help in practice.

For future reference, can you make it explicit in the change description: 1) what made this beneficial for CSSPGO (give post-link full context profile precedence over pre-link partial context profile, i.e. be more specific about obsolete metadata); 2) what is the expected impact on AutoFDO.

I think there's another caveat worth calling out, by overwriting all pre-link weights, we would lose the prof metadata adjustment done by prelink opt passes (e.g. prelink loop passes).

Good point, summary updated.

[wmi]

If we target AutoFDO as well, would be good to mention the impact on AutoFDO too.

I haven't tested for non-CS case. In theory it should help there but with a less-than-ideal counts quality, the change may not help in practice.

AutoFDO doesn't have the distribution factor which CSSPGO+pseudo hook have so OverwriteExistingWeights may not help for AutoFDO.

llvm/lib/Transforms/IPO/SampleProfile.cpp

1483–1486

The second updateIDTMetaData call should be removed?

1496–1498

Remove MD_prof means compiler may use static heuristic to infer the callsite hotness, which can still be hot. This is different from marking it as cold.

If the profile is accurate, why not set the count to 0?

1545–1547

Update the comment to mention OverwriteExistingWeights.

[hoy]

llvm/lib/Transforms/IPO/SampleProfile.cpp

1483–1486

Good catch, kind of rebasing issue.

1496–1498

I see. Thanks for pointing it out. We want it to be treated as cold here. Should be using zero count for direct callsites. For indirect callsites, I'm inclined to removing the metadata which shouldn't trigger PGO ICP to do anything. What do you think?

1545–1547

Done.

[hoy]

Addressing Wei's comment.

[wmi]

LGTM.

llvm/lib/Transforms/IPO/SampleProfile.cpp

1496–1498

Removing metadata for indirect callsites makes sense to me.

[wenlei]

lgtm, thanks.

llvm/lib/Transforms/IPO/SampleProfile.cpp

cl::desc(
         "call promotion in proirity-based sample profile loader inlining."));
 
 static cl::opt<bool> CallsitePrioritizedInline(
     "sample-profile-prioritized-inline", cl::Hidden, cl::ZeroOrMore,
     cl::init(false),
     cl::desc("Use call site prioritized inlining for sample profile loader."
              "Currently only CSSPGO is supported."));
 
-
 static cl::opt<std::string> ProfileInlineReplayFile(
     "sample-profile-inline-replay", cl::init(""), cl::value_desc("filename"),
     cl::desc(
         "Optimization remarks file containing inline remarks to be replayed "
         "by inlining from sample profile loader."),
     cl::Hidden);
 
 static cl::opt<unsigned>
     MaxNumPromotions("sample-profile-icp-max-prom", cl::init(3), cl::Hidden,
                      cl::ZeroOrMore,
                      cl::desc("Max number of promotions for a single indirect "
                               "call callsite in sample profile loader"));
 
+static cl::opt<bool> OverwriteExistingWeights(
+    "overwrite-existing-weights", cl::Hidden, cl::init(false),
+    cl::desc("Ignore existing branch weights on IR and always overwrite."));
+
 namespace {
 
 using BlockWeightMap = DenseMap<const BasicBlock *, uint64_t>;
 using EquivalenceClassMap = DenseMap<const BasicBlock *, const BasicBlock *>;
 using Edge = std::pair<const BasicBlock *, const BasicBlock *>;
 using EdgeWeightMap = DenseMap<Edge, uint64_t>;
 using BlockEdgeMap =
     DenseMap<const BasicBlock *, SmallVector<const BasicBlock *, 8>>;
if (BlockWeights[BB]) {
             continue;
           const DILocation *DIL = DLoc;
           const FunctionSamples *FS = findFunctionSamples(I);
           if (!FS)
             continue;
           auto CallSite = FunctionSamples::getCallSiteIdentifier(DIL);
           auto T = FS->findCallTargetMapAt(CallSite);
           if (!T || T.get().empty())
             continue;
-          // Prorate the callsite counts to reflect what is already done to the
-          // callsite, such as ICP or calliste cloning.
           if (FunctionSamples::ProfileIsProbeBased) {
+            // Prorate the callsite counts based on the pre-ICP distribution
+            // factor to reflect what is already done to the callsite before
+            // ICP, such as calliste cloning.
             if (Optional<PseudoProbe> Probe = extractProbe(I)) {
               if (Probe->Factor < 1)
                 T = SampleRecord::adjustCallTargets(T.get(), Probe->Factor);
             }
           }
           SmallVector<InstrProfValueData, 2> SortedCallTargets =
               GetSortedValueDataFromCallTargets(T.get());
           uint64_t Sum = 0;
           for (const auto &C : T.get())
             Sum += C.second;
           // With CSSPGO all indirect call targets are counted torwards the
           // original indirect call site in the profile, including both
           // inlined and non-inlined targets.
           if (!FunctionSamples::ProfileIsCS) {
             if (const FunctionSamplesMap *M =
                     FS->findFunctionSamplesMapAt(CallSite)) {
               for (const auto &NameFS : *M)
                 Sum += NameFS.second.getEntrySamples();
             }
           }
-          if (!Sum)
-            continue;
+          if (Sum)
             updateIDTMetaData(I, SortedCallTargets, Sum);
+          else if (OverwriteExistingWeights)
+            I.setMetadata(LLVMContext::MD_prof, nullptr);
         } else if (!isa<IntrinsicInst>(&I)) {
           I.setMetadata(LLVMContext::MD_prof,
                         MDB.createBranchWeights(
                             {static_cast<uint32_t>(BlockWeights[BB])}));
         }
       }
+    } else if (OverwriteExistingWeights) {
+      // Set profile metadata (possibly annotated by LTO prelink) to zero or
+      // clear it for cold code.
+      for (auto &I : BB->getInstList()) {
+        if (isa<CallInst>(I) || isa<InvokeInst>(I)) {
+          if (cast<CallBase>(I).isIndirectCall())
+            I.setMetadata(LLVMContext::MD_prof, nullptr);
+          else
+            I.setMetadata(LLVMContext::MD_prof, MDB.createBranchWeights(0));
+        }
+      }
     }
+
     Instruction *TI = BB->getTerminator();
     if (TI->getNumSuccessors() == 1)
       continue;
     if (!isa<BranchInst>(TI) && !isa<SwitchInst>(TI) &&
         !isa<IndirectBrInst>(TI))
       continue;
 
     DebugLoc BranchLoc = TI->getDebugLoc();
for (unsigned I = 0; I < TI->getNumSuccessors(); ++I) {
         if (Weight > MaxWeight) {
           MaxWeight = Weight;
           MaxDestInst = Succ->getFirstNonPHIOrDbgOrLifetime();
         }
       }
     }
 
     uint64_t TempWeight;
     // Only set weights if there is at least one non-zero weight.
     // In any other case, let the analyzer set weights.
-    // Do not set weights if the weights are present. In ThinLTO, the profile
-    // annotation is done twice. If the first annotation already set the
-    // weights, the second pass does not need to set it.
-    if (MaxWeight > 0 && !TI->extractProfTotalWeight(TempWeight)) {
+    // Do not set weights if the weights are present unless under
+    // OverwriteExistingWeights. In ThinLTO, the profile annotation is done
+    // twice. If the first annotation already set the weights, the second pass
+    // does not need to set it. With OverwriteExistingWeights, Blocks with zero
+    // weight should have their existing metadata (possibly annotated by LTO
+    // prelink) cleared.
+    if (MaxWeight > 0 &&
+        (!TI->extractProfTotalWeight(TempWeight) || OverwriteExistingWeights)) {
       LLVM_DEBUG(dbgs() << "SUCCESS. Found non-zero weights.\n");
-      TI->setMetadata(LLVMContext::MD_prof,
-                      MDB.createBranchWeights(Weights));
+      TI->setMetadata(LLVMContext::MD_prof, MDB.createBranchWeights(Weights));
       ORE->emit([&]() {
         return OptimizationRemark(DEBUG_TYPE, "PopularDest", MaxDestInst)
                << "most popular destination for conditional branches at "
                << ore::NV("CondBranchesLoc", BranchLoc);
       });
     } else {
+      if (OverwriteExistingWeights) {
+        TI->setMetadata(LLVMContext::MD_prof, nullptr);
+        LLVM_DEBUG(dbgs() << "CLEARED. All branch weights are zero.\n");
+      } else {
         LLVM_DEBUG(dbgs() << "SKIPPED. All branch weights are zero.\n");
       }
     }
   }
+}
 
 /// Once all the branch weights are computed, we emit the MD_prof
 /// metadata on BB using the computed values for each of its branches.
 ///
 /// \param F The function to query.
 ///
 /// \returns true if \p F was modified. Returns false, otherwise.
 bool SampleProfileLoader::emitAnnotations(Function &F) {

llvm/test/Transforms/SampleProfile/branch.ll

 ; RUN: opt < %s -sample-profile -sample-profile-file=%S/Inputs/branch.prof | opt -analyze -branch-prob -enable-new-pm=0 | FileCheck %s
 ; RUN: opt < %s -passes=sample-profile -sample-profile-file=%S/Inputs/branch.prof | opt -passes='print<branch-prob>' -disable-output 2>&1 | FileCheck %s
+; RUN: opt < %s -passes=sample-profile -sample-profile-file=%S/Inputs/branch.prof -overwrite-existing-weights=1 | opt -passes='print<branch-prob>' -disable-output 2>&1 | FileCheck %s  --check-prefix=OVW
 
 ; Original C++ code for this test case:
 ;
 ; #include <stdio.h>
 ; #include <stdlib.h>
 
 ; int main(int argc, char *argv[]) {
 ;   if (argc < 2)

 
 for.cond:                                         ; preds = %for.inc, %if.then.2
   %6 = load i32, i32* %u, align 4, !dbg !46
   %7 = load i32, i32* %limit, align 4, !dbg !48
   %cmp5 = icmp slt i32 %6, %7, !dbg !49
   br i1 %cmp5, label %for.body, label %for.end, !dbg !50, !prof !80
 ; CHECK: edge for.cond -> for.body probability is 0x73333333 / 0x80000000 = 90.00%
 ; CHECK: edge for.cond -> for.end probability is 0x0ccccccd / 0x80000000 = 10.00%
+; OVW: edge for.cond -> for.body probability is 0x76b3f3be / 0x80000000 = 92.74% 
+; OVW: edge for.cond -> for.end probability is 0x094c0c42 / 0x80000000 = 7.26% 
 
 for.body:                                         ; preds = %for.cond
   call void @llvm.dbg.declare(metadata double* %x, metadata !51, metadata !17), !dbg !53
   %8 = load double, double* %s, align 8, !dbg !54
   store double %8, double* %x, align 8, !dbg !53
   %9 = load double, double* %x, align 8, !dbg !55
   %add = fadd double %9, 3.049000e+00, !dbg !56
   %10 = load i32, i32* %u, align 4, !dbg !57

llvm/test/Transforms/SampleProfile/pseudo-probe-profile-metadata-2.ll

+; RUN: opt < %s -passes=sample-profile -sample-profile-file=%S/Inputs/pseudo-probe-profile.prof -pass-remarks=sample-profile -S | FileCheck %s
+; RUN: opt < %s -passes=sample-profile -sample-profile-file=%S/Inputs/pseudo-probe-profile.prof -pass-remarks=sample-profile -overwrite-existing-weights=1 -S | FileCheck %s -check-prefix=OVW
+
+define dso_local i32 @foo(i32 %x, void (i32)* %f) #0 !dbg !4 !prof !10 {
+entry:
+  %retval = alloca i32, align 4
+  %x.addr = alloca i32, align 4
+  store i32 %x, i32* %x.addr, align 4
+  %0 = load i32, i32* %x.addr, align 4
+  %cmp = icmp eq i32 %0, 0
+  call void @llvm.pseudoprobe(i64 6699318081062747564, i64 1, i32 0, i64 -1)
+  br i1 %cmp, label %if.then, label %if.else, !prof !11
+
+if.then:
+  call void @llvm.pseudoprobe(i64 6699318081062747564, i64 2, i32 0, i64 -1)
+  ; CHECK: call {{.*}}, !dbg ![[#]], !prof ![[#PROF:]]
+  ; OVW: call {{.*}}, !dbg ![[#]], !prof ![[#PROF:]]
+  call void %f(i32 1), !dbg !13, !prof !16
+  store i32 1, i32* %retval, align 4
+  br label %return
+
+if.else:
+  call void @llvm.pseudoprobe(i64 6699318081062747564, i64 2, i32 0, i64 0)
+  ; CHECK: call {{.*}}, !dbg ![[#]], !prof  ![[#PROF]]
+  ;; The block should have a 0 weight. Check the profile metadata is dropped.
+  ; OVW-NOT: call {{.*}}, !dbg ![[#]], !prof
+  call void %f(i32 2), !dbg !15, !prof !16
+  store i32 2, i32* %retval, align 4
+  br label %return
+
+return:
+  call void @llvm.pseudoprobe(i64 6699318081062747564, i64 4, i32 0, i64 -1)
+  %1 = load i32, i32* %retval, align 4
+  ret i32 %1
+}
+
+; CHECK: ![[#PROF]] = !{!"VP", i32 0, i64 7, i64 9191153033785521275, i64 5, i64 -1069303473483922844, i64 2}
+; OVW: ![[#PROF]] = !{!"VP", i32 0, i64 7, i64 9191153033785521275, i64 5, i64 -1069303473483922844, i64 2}
+
+declare void @llvm.pseudoprobe(i64, i64, i32, i64) #0
+
+attributes #0 = {"use-sample-profile"}
+
+!llvm.module.flags = !{!0, !1}
+!llvm.pseudo_probe_desc = !{!2}
+
+!0 = !{i32 7, !"Dwarf Version", i32 4}
+!1 = !{i32 2, !"Debug Info Version", i32 3}
+!2 = !{i64 6699318081062747564, i64 563022570642068, !"foo", 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 100, i32 0}
+;; A discriminator of 186646575 which is 0x6f80057 in hexdecimal, stands for an indirect call probe
+;; with an index of 5 and probe factor of 1.0.
+!12 = !DILexicalBlockFile(scope: !4, file: !5, discriminator: 186646575)
+!13 = distinct !DILocation(line: 10, column: 11, scope: !12)
+;; A discriminator of 134217775 which is 0x6f80057 in hexdecimal, stands for an indirect call probe
+;; with an index of 5 and probe factor of 0.
+!14 = !DILexicalBlockFile(scope: !4, file: !5, discriminator: 134217775)
+!15 = distinct !DILocation(line: 10, column: 11, scope: !14)
+!16 = !{!"VP", i32 0, i64 7, i64 9191153033785521275, i64 5, i64 -1069303473483922844, i64 2}
+