[llvm] 10f0c1a - [PGO] Ensure non-zero entry-count after `populateCounters` (#112029)

[via llvm-commits]

Author: Michael O'Farrell
Date: 2024-10-22T16:05:40-07:00
New Revision: 10f0c1aadd878be3bd7c586c4a2c7d7c76aee733

URL: https://github.com/llvm/llvm-project/commit/10f0c1aadd878be3bd7c586c4a2c7d7c76aee733
DIFF: https://github.com/llvm/llvm-project/commit/10f0c1aadd878be3bd7c586c4a2c7d7c76aee733.diff

LOG: [PGO] Ensure non-zero entry-count after `populateCounters` (#112029)

With sampled instrumentation (#69535), profile counts may appear corrupt
and `fixFuncEntryCount` may assert. In particular a function can have a
0 block count for its entry, while later blocks are non zero. This is
only likely to happen for colder functions, so it is reasonable to take
any action that does not crash. Here we simply bail from fixing the
entry count.

Added: 
    llvm/test/Transforms/PGOProfile/fix_entry_count_sampled.ll

Modified: 
    llvm/lib/Transforms/Instrumentation/PGOInstrumentation.cpp

Removed: 
    


################################################################################
diff  --git a/llvm/lib/Transforms/Instrumentation/PGOInstrumentation.cpp b/llvm/lib/Transforms/Instrumentation/PGOInstrumentation.cpp
index 919660e7a040f2..bceb6135cc1f92 100644
--- a/llvm/lib/Transforms/Instrumentation/PGOInstrumentation.cpp
+++ b/llvm/lib/Transforms/Instrumentation/PGOInstrumentation.cpp
@@ -1615,6 +1615,10 @@ void PGOUseFunc::populateCounters() {
     assert(BI->Count && "BB count is not valid");
   }
 #endif
+  // Now annotate select instructions.  This may fixup impossible block counts.
+  FuncInfo.SIVisitor.annotateSelects(this, &CountPosition);
+  assert(CountPosition == ProfileCountSize);
+
   uint64_t FuncEntryCount = *getBBInfo(&*F.begin()).Count;
   uint64_t FuncMaxCount = FuncEntryCount;
   for (auto &BB : F) {
@@ -1630,10 +1634,6 @@ void PGOUseFunc::populateCounters() {
   F.setEntryCount(ProfileCount(FuncEntryCount, Function::PCT_Real));
   markFunctionAttributes(FuncEntryCount, FuncMaxCount);
 
-  // Now annotate select instructions
-  FuncInfo.SIVisitor.annotateSelects(this, &CountPosition);
-  assert(CountPosition == ProfileCountSize);
-
   LLVM_DEBUG(FuncInfo.dumpInfo("after reading profile."));
 }
 
@@ -1742,8 +1742,13 @@ void SelectInstVisitor::annotateOneSelectInst(SelectInst &SI) {
   ++(*CurCtrIdx);
   uint64_t TotalCount = 0;
   auto BI = UseFunc->findBBInfo(SI.getParent());
-  if (BI != nullptr)
+  if (BI != nullptr) {
     TotalCount = *BI->Count;
+
+    // Fix the block count if it is impossible.
+    if (TotalCount < SCounts[0])
+      BI->Count = SCounts[0];
+  }
   // False Count
   SCounts[1] = (TotalCount > SCounts[0] ? TotalCount - SCounts[0] : 0);
   uint64_t MaxCount = std::max(SCounts[0], SCounts[1]);

diff  --git a/llvm/test/Transforms/PGOProfile/fix_entry_count_sampled.ll b/llvm/test/Transforms/PGOProfile/fix_entry_count_sampled.ll
new file mode 100644
index 00000000000000..91476d3a16b1cc
--- /dev/null
+++ b/llvm/test/Transforms/PGOProfile/fix_entry_count_sampled.ll
@@ -0,0 +1,42 @@
+; RUN: rm -rf %t && split-file %s %t
+
+; RUN: llvm-profdata merge %t/main.proftext -o %t/main.profdata
+; RUN: opt < %t/main.ll -passes=pgo-instr-use -pgo-test-profile-file=%t/main.profdata -S | FileCheck %s
+
+;--- main.ll
+
+; Instrumentation PGO sampling makes corrupt looking counters possible.  This
+; tests one extreme case:
+; Test loading zero profile counts for all instrumented blocks while the entry
+; block is not instrumented.  Additionally include a non-zero profile count for
+; a select instruction, which prevents short circuiting the PGO application.
+
+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_no_entry_block_counter(i32 %n) {
+; CHECK: define i32 @test_no_entry_block_counter(i32 %n)
+; CHECK-SAME: !prof ![[ENTRY_COUNT:[0-9]*]]
+entry:
+  %cmp = icmp slt i32 42, %n
+  br i1 %cmp, label %tail1, label %tail2
+tail1:
+  %ret = select i1 true, i32 %n, i32 42
+; CHECK:  %ret = select i1 true, i32 %n, i32 42
+; CHECK-SAME: !prof ![[BW_FOR_SELECT:[0-9]+]]
+  ret i32 %ret
+tail2:
+  ret i32 42
+}
+; CHECK: ![[ENTRY_COUNT]] = !{!"function_entry_count", i64 1}
+; CHECK: ![[BW_FOR_SELECT]] = !{!"branch_weights", i32 1, i32 0}
+
+;--- main.proftext
+:ir
+test_no_entry_block_counter
+431494656217155589
+3
+0
+0
+1
+