[llvm] [PGO] Gracefully handle zero entry-count in `fixFuncEntryCount` (PR #112029)

[Michael O'Farrell via llvm-commits]

https://github.com/mofarrell updated https://github.com/llvm/llvm-project/pull/112029

>From 4d1221eb68b11c5fbad13fdaf396585ebb324101 Mon Sep 17 00:00:00 2001
From: Michael O'Farrell <micpof at gmail.com>
Date: Tue, 8 Oct 2024 10:41:11 -0700
Subject: [PATCH 1/2] [PGO] Gracefully handle zero entry-count

With sampled instrumentation (#69535), profile counts can appear
corrupt.  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.
---
 llvm/lib/Transforms/Instrumentation/PGOInstrumentation.cpp | 6 ++----
 1 file changed, 2 insertions(+), 4 deletions(-)

diff --git a/llvm/lib/Transforms/Instrumentation/PGOInstrumentation.cpp b/llvm/lib/Transforms/Instrumentation/PGOInstrumentation.cpp
index dbe908bb5e72f3..8e72cc367574ab 100644
--- a/llvm/lib/Transforms/Instrumentation/PGOInstrumentation.cpp
+++ b/llvm/lib/Transforms/Instrumentation/PGOInstrumentation.cpp
@@ -1972,11 +1972,9 @@ static void fixFuncEntryCount(PGOUseFunc &Func, LoopInfo &LI,
                               BranchProbabilityInfo &NBPI) {
   Function &F = Func.getFunc();
   BlockFrequencyInfo NBFI(F, NBPI, LI);
-#ifndef NDEBUG
   auto BFIEntryCount = F.getEntryCount();
-  assert(BFIEntryCount && (BFIEntryCount->getCount() > 0) &&
-         "Invalid BFI Entrycount");
-#endif
+  if (!BFIEntryCount || BFIEntryCount->getCount() == 0)
+    return;
   auto SumCount = APFloat::getZero(APFloat::IEEEdouble());
   auto SumBFICount = APFloat::getZero(APFloat::IEEEdouble());
   for (auto &BBI : F) {

>From 4494a6d09b964f0036d794bbc0ef86d265625328 Mon Sep 17 00:00:00 2001
From: Michael O'Farrell <micpof at gmail.com>
Date: Wed, 16 Oct 2024 12:25:15 -0700
Subject: [PATCH 2/2] [PGO] Test graceful handling of zero profile counts

This crashed before making fixFuncEntryCount gracefully handle zero
counts.
---
 .../PGOProfile/fix_entry_count_sampled.ll     | 28 +++++++++++++++++++
 1 file changed, 28 insertions(+)
 create mode 100644 llvm/test/Transforms/PGOProfile/fix_entry_count_sampled.ll

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..99863e631287af
--- /dev/null
+++ b/llvm/test/Transforms/PGOProfile/fix_entry_count_sampled.ll
@@ -0,0 +1,28 @@
+; RUN: llvm-profdata merge %S/Inputs/fix_entry_count_sampled.proftext -o %t.profdata
+; RUN: opt < %s -passes=pgo-instr-use -pgo-test-profile-file=%t.profdata -S | FileCheck %s --check-prefix=USE
+
+; 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) {
+; USE: define i32 @test_no_entry_block_counter(i32 %n)
+; USE-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
+; USE:  %ret = select i1 true, i32 %n, i32 42
+; USE-SAME: !prof ![[BW_FOR_SELECT:[0-9]+]]
+  ret i32 %ret
+tail2:
+  ret i32 42
+}
+; USE: ![[ENTRY_COUNT]] = !{!"function_entry_count", i64 0}
+; USE: ![[BW_FOR_SELECT]] = !{!"branch_weights", i32 1, i32 0}