From b43a1d1bd97816d32388e78fd5659f732edc82c9 Mon Sep 17 00:00:00 2001
From: Wolfgang Pieb <wolfgang_pieb@playstation.sony.com>
Date: Fri, 16 Sep 2022 09:40:14 -0700
Subject: [PATCH] [PGO] Do not create block count annotations when all weights
 are 0, avoiding an assertion.

A BB with a nonzero count, whose successor blocks all have 0 counts, could
cause an assertion. Don't create any branch weights in this case.

Reviewed By: xur

Differential Revision: https://reviews.llvm.org/D134203
---
 .../Instrumentation/PGOInstrumentation.cpp    | 16 ++++++-
 .../Instrumentation/PGOMemOPSizeOpt.cpp       |  3 +-
 .../PGOProfile/Inputs/maxcountzero.proftext   | 11 +++++
 .../Transforms/PGOProfile/maxcountzero.ll     | 47 +++++++++++++++++++
 4 files changed, 75 insertions(+), 2 deletions(-)
 create mode 100755 llvm/test/Transforms/PGOProfile/Inputs/maxcountzero.proftext
 create mode 100755 llvm/test/Transforms/PGOProfile/maxcountzero.ll

diff --git a/llvm/lib/Transforms/Instrumentation/PGOInstrumentation.cpp b/llvm/lib/Transforms/Instrumentation/PGOInstrumentation.cpp
index cd50c9756408..d627c3eb5d80 100644
--- a/llvm/lib/Transforms/Instrumentation/PGOInstrumentation.cpp
+++ b/llvm/lib/Transforms/Instrumentation/PGOInstrumentation.cpp
@@ -1689,7 +1689,21 @@ void PGOUseFunc::setBranchWeights() {
         MaxCount = EdgeCount;
       EdgeCounts[SuccNum] = EdgeCount;
     }
-    setProfMetadata(M, TI, EdgeCounts, MaxCount);
+
+    if (MaxCount)
+      setProfMetadata(M, TI, EdgeCounts, MaxCount);
+    else {
+      // A zero MaxCount can come about when we have a BB with a positive
+      // count, and whose successor blocks all have 0 count. This can happen
+      // when there is no exit block and the code exits via a noreturn function.
+      auto &Ctx = M->getContext();
+      Ctx.diagnose(DiagnosticInfoPGOProfile(
+        M->getName().data(),
+        Twine("Profile in ") + F.getName().str() +
+            Twine(" partially ignored") +
+            Twine(", possibly due to the lack of a return path."),
+        DS_Warning));
+    }
   }
 }
 
diff --git a/llvm/lib/Transforms/Instrumentation/PGOMemOPSizeOpt.cpp b/llvm/lib/Transforms/Instrumentation/PGOMemOPSizeOpt.cpp
index b11f16894669..07c03ee2049a 100644
--- a/llvm/lib/Transforms/Instrumentation/PGOMemOPSizeOpt.cpp
+++ b/llvm/lib/Transforms/Instrumentation/PGOMemOPSizeOpt.cpp
@@ -440,7 +440,8 @@ bool MemOPSizeOpt::perform(MemOp MO) {
   DTU.applyUpdates(Updates);
   Updates.clear();
 
-  setProfMetadata(Func.getParent(), SI, CaseCounts, MaxCount);
+  if (MaxCount)
+    setProfMetadata(Func.getParent(), SI, CaseCounts, MaxCount);
 
   LLVM_DEBUG(dbgs() << *BB << "\n");
   LLVM_DEBUG(dbgs() << *DefaultBB << "\n");
diff --git a/llvm/test/Transforms/PGOProfile/Inputs/maxcountzero.proftext b/llvm/test/Transforms/PGOProfile/Inputs/maxcountzero.proftext
new file mode 100755
index 000000000000..d7deb6a40d7b
--- /dev/null
+++ b/llvm/test/Transforms/PGOProfile/Inputs/maxcountzero.proftext
@@ -0,0 +1,11 @@
+# IR level Instrumentation Flag
+:ir
+foo
+# Func Hash:
+146835645050580305
+# Num Counters:
+3
+# Counter Values:
+0
+1
+0
diff --git a/llvm/test/Transforms/PGOProfile/maxcountzero.ll b/llvm/test/Transforms/PGOProfile/maxcountzero.ll
new file mode 100755
index 000000000000..1137b2abcef3
--- /dev/null
+++ b/llvm/test/Transforms/PGOProfile/maxcountzero.ll
@@ -0,0 +1,47 @@
+; Make sure a profile that is generated from a function without an exit node
+; does not cause an assertion. The profile consists of a non-zero count in a
+; basic block and 0 counts in all succcessor blocks. Expect a warning.
+
+; RUN: llvm-profdata merge %S/Inputs/maxcountzero.proftext -o %t.profdata
+; RUN: opt < %s -passes=pgo-instr-use -pgo-instrument-entry=false -pgo-test-profile-file=%t.profdata -S 2>&1 | FileCheck %s
+
+define void @bar(i32 noundef %s) {
+entry:
+  %cmp = icmp sgt i32 %s, 20
+  br i1 %cmp, label %if.then, label %if.end
+
+if.then:
+  call void @exit(i32 noundef 1)
+  unreachable
+
+if.end:
+  ret void
+}
+
+declare void @exit(i32 noundef)
+
+define void @foo(i32 noundef %n) {
+entry:
+  %sum = alloca i32, align 4
+  store volatile i32 %n, ptr %sum, align 4
+  %sum.0.sum.0. = load volatile i32, ptr %sum, align 4
+  call void @bar(i32 noundef %sum.0.sum.0.)
+  %cmp = icmp slt i32 %n, 10
+  br i1 %cmp, label %if.then, label %if.end
+
+if.then:
+  %sum.0.sum.0.1 = load volatile i32, ptr %sum, align 4
+  call void @bar(i32 noundef %sum.0.sum.0.1)
+  br label %if.end
+
+if.end:
+  br label %for.cond
+
+for.cond:
+  %sum.0.sum.0.2 = load volatile i32, ptr %sum, align 4
+  call void @bar(i32 noundef %sum.0.sum.0.2)
+  br label %for.cond
+}
+
+; CHECK: warning:{{.*}}Profile in foo partially ignored
+; CHECK: define