[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 83a55a2bca4aee1659b05f9ad224762e9306d607 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 counts for all its
blocks, while having some non-zero counters for select instrumentation.
This is only possible for colder functions, and a reasonable
modification to ensure the entry is non-zero (required by
`fixFuncEntryCounts`) is to set the counter to one.  This is only likely
to happen for colder functions, so it is reasonable to take any action
that does not crash.
---
 .../lib/Transforms/Instrumentation/PGOInstrumentation.cpp | 8 ++++++--
 1 file changed, 6 insertions(+), 2 deletions(-)

diff --git a/llvm/lib/Transforms/Instrumentation/PGOInstrumentation.cpp b/llvm/lib/Transforms/Instrumentation/PGOInstrumentation.cpp
index dbe908bb5e72f3..3fe1d236ae22b5 100644
--- a/llvm/lib/Transforms/Instrumentation/PGOInstrumentation.cpp
+++ b/llvm/lib/Transforms/Instrumentation/PGOInstrumentation.cpp
@@ -1624,8 +1624,12 @@ void PGOUseFunc::populateCounters() {
     FuncMaxCount = std::max(FuncMaxCount, *BI->Count);
   }
 
-  // Fix the obviously inconsistent entry count.
-  if (FuncMaxCount > 0 && FuncEntryCount == 0)
+  // Fix the obviously inconsistent entry count.  A function that has all zero
+  // counters will skip populating the counters so a minimum entry count of 1
+  // makes sense.  Note that the presence of a non-zero counter does not imply
+  // FuncMaxCount is greater than zero because the only non-zero counts could
+  // be for select instrumentation which is handled later.
+  if (FuncEntryCount == 0)
     FuncEntryCount = 1;
   F.setEntryCount(ProfileCount(FuncEntryCount, Function::PCT_Real));
   markFunctionAttributes(FuncEntryCount, FuncMaxCount);

>From abfdcb663983eaa8c7ce0d03a3c6440c5fc97e9a 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.
---
 .../Inputs/fix_entry_count_sampled.proftext   |  8 ++++++
 .../PGOProfile/fix_entry_count_sampled.ll     | 28 +++++++++++++++++++
 2 files changed, 36 insertions(+)
 create mode 100644 llvm/test/Transforms/PGOProfile/Inputs/fix_entry_count_sampled.proftext
 create mode 100644 llvm/test/Transforms/PGOProfile/fix_entry_count_sampled.ll

diff --git a/llvm/test/Transforms/PGOProfile/Inputs/fix_entry_count_sampled.proftext b/llvm/test/Transforms/PGOProfile/Inputs/fix_entry_count_sampled.proftext
new file mode 100644
index 00000000000000..8578f8be07fa8a
--- /dev/null
+++ b/llvm/test/Transforms/PGOProfile/Inputs/fix_entry_count_sampled.proftext
@@ -0,0 +1,8 @@
+:ir
+test_no_entry_block_counter
+431494656217155589
+3
+0
+0
+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..56d26178cd6265
--- /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 1}
+; USE: ![[BW_FOR_SELECT]] = !{!"branch_weights", i32 1, i32 0}