[PGO] Fix incorrect function entry count

[xur]

Summary

Function entry count might be zero after the profile counts
reset and before reentry to the function.

Zero profile entry count is very bad as the profile count from BFI
will be wrong.

A simple fix is to set the profile entry count to 1 if there are
non-zero profile counts in this function.

[davidxl]

llvm/lib/Transforms/Instrumentation/PGOInstrumentation.cpp

1144

add comment here about the root cause: counter cleared before dump.

If counter promotion is disabled, in what situation entry count can not be recovered with propagation?

[xur]

llvm/lib/Transforms/Instrumentation/PGOInstrumentation.cpp

1144

This is independent of count promotion. Usually the entry count is not promoted (unless it's inlined to a callsite inside a loop).

This code here is for the case when we instrument the entry BB. If the entry BB is instrumented, it cannot be computed from propagation. (MST does not have the redundancy to recover)

Typically, it happens this way: for a func with a long running loop, after we reset the counter, the entry counters of this func are also zeroed out. If we dump the counters in the middle of the loop, the entry count will still be 0.

If we promote the counters inside the loop, the counter value inside the loop will be 0 (because the updates are in the exit block -- never exercised). If we disable the promote, the counter value insides the loop are OK.

[davidxl]

llvm/lib/Transforms/Instrumentation/PGOInstrumentation.cpp

1144

ok. what if entry is not instrumented, but two post-dom successors? If they are zeroed out, is there a way to recover?

[davidxl]

llvm/lib/Transforms/Instrumentation/PGOInstrumentation.cpp

1144

Also look at the caller of this function -- even when AllZeros is true, this function is still called, so the entry count may be wrongly fixed up.

1345

Just this fix should be enough to fix the insanity, right?

[xur]

llvm/lib/Transforms/Instrumentation/PGOInstrumentation.cpp

1144

It will not fix here, but later. When we set the entry count.

1144

When AllZeros is set, the entry count we set here is NOOP.

The caller (annotateAllFunctions) has a shortcut when AllZero is 0. It sets call setEntryCount directly (to 0) and skip counter population.

Maybe I should return early here if AllZero is 0.

1345

This only fixes the entry for the func. It will not contribute to the count value population.

[davidxl]

llvm/lib/Transforms/Instrumentation/PGOInstrumentation.cpp

1144

Skipping allzero sounds good.

1345

ok

[xur]

llvm/lib/Transforms/Instrumentation/PGOInstrumentation.cpp

1144

I think the reason I did not return early here was to detect the error when there was inconsistent number of counters.

[davidxl]

lgtm

llvm/lib/Transforms/Instrumentation/PGOInstrumentation.cpp

bool PGOUseFunc::setInstrumentedCounts(
   FuncInfo.getInstrumentBBs(InstrumentBBs);
   unsigned NumCounters =
       InstrumentBBs.size() + FuncInfo.SIVisitor.getNumOfSelectInsts();
   // The number of counters here should match the number of counters
   // in profile. Return if they mismatch.
   if (NumCounters != CountFromProfile.size()) {
     return false;
   }
+  auto *FuncEntry = &*F.begin();
+
   // Set the profile count to the Instrumented BBs.
   uint32_t I = 0;
   for (BasicBlock *InstrBB : InstrumentBBs) {
     uint64_t CountValue = CountFromProfile[I++];
     UseBBInfo &Info = getBBInfo(InstrBB);
+    // If we reach here, we know that we have some nonzero count
+    // values in this function. The entry count should not be 0.
+    // Fix it if necessary.
+    if (InstrBB == FuncEntry && CountValue == 0)
+      CountValue = 1;
     Info.setBBInfoCount(CountValue);
   }
   ProfileCountSize = CountFromProfile.size();
   CountPosition = I;
 
   // Set the edge count and update the count of unknown edges for BBs.
   auto setEdgeCount = [this](PGOUseEdge *E, uint64_t Value) -> void {
     E->setEdgeCount(Value);
#ifndef NDEBUG
   for (auto &BB : F) {
     auto BI = findBBInfo(&BB);
     if (BI == nullptr)
       continue;
     assert(BI->CountValid && "BB count is not valid");
   }
 #endif
   uint64_t FuncEntryCount = getBBInfo(&*F.begin()).CountValue;
-  F.setEntryCount(ProfileCount(FuncEntryCount, Function::PCT_Real));
   uint64_t FuncMaxCount = FuncEntryCount;
   for (auto &BB : F) {
     auto BI = findBBInfo(&BB);
     if (BI == nullptr)
       continue;
     FuncMaxCount = std::max(FuncMaxCount, BI->CountValue);
   }
+
+  // Fix the obviously inconsistent entry count.
+  if (FuncMaxCount > 0 && FuncEntryCount == 0)
+    FuncEntryCount = 1;
+  F.setEntryCount(ProfileCount(FuncEntryCount, Function::PCT_Real));
   markFunctionAttributes(FuncEntryCount, FuncMaxCount);
 
   // Now annotate select instructions
   FuncInfo.SIVisitor.annotateSelects(F, this, &CountPosition);
   assert(CountPosition == ProfileCountSize);
 
   LLVM_DEBUG(FuncInfo.dumpInfo("after reading profile."));
 }

llvm/test/Transforms/PGOProfile/Inputs/fix_entry_count.proftext

+:ir
+:entry_first
+test_simple_for
[phone]
+2
+0
+96
+

llvm/test/Transforms/PGOProfile/fix_entry_count.ll

+; RUN: llvm-profdata merge %S/Inputs/fix_entry_count.proftext -o %t.profdata
+; RUN: opt < %s -pgo-instr-use -pgo-instrument-entry=true -pgo-test-profile-file=%t.profdata -S | FileCheck %s --check-prefix=USE
+; RUN: opt < %s -passes=pgo-instr-use -pgo-instrument-entry=true -pgo-test-profile-file=%t.profdata -S | FileCheck %s --check-prefix=USE
+
+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 @test_simple_for(i32 %n) {
+; USE: define i32 @test_simple_for(i32 %n)
+; USE-SAME: !prof ![[ENTRY_COUNT:[0-9]*]]
+entry:
+  br label %for.cond
+
+for.cond:
+  %i = phi i32 [ 0, %entry ], [ %inc1, %for.inc ]
+  %sum = phi i32 [ 1, %entry ], [ %inc, %for.inc ]
+  %cmp = icmp slt i32 %i, %n
+  br i1 %cmp, label %for.body, label %for.end
+; USE: br i1 %cmp, label %for.body, label %for.end
+; USE-SAME: !prof ![[BW_FOR_COND:[0-9]+]]
+
+for.body:
+  %inc = add nsw i32 %sum, 1
+  br label %for.inc
+
+for.inc:
+  %inc1 = add nsw i32 %i, 1
+  br label %for.cond
+
+for.end:
+  ret i32 %sum
+}
+; USE: ![[ENTRY_COUNT]] = !{!"function_entry_count", i64 1}
+; USE: ![[BW_FOR_COND]] = !{!"branch_weights", i32 96, i32 1}