From f131edf6fbe8e2ab7306aba72698daa6153ec91e Mon Sep 17 00:00:00 2001
From: Avi Kivity <avi@scylladb.com>
Date: Mon, 26 Aug 2024 17:56:45 +0300
Subject: [PATCH] [Instrumentation] Fix EdgeCounts vector size in
 SetBranchWeights (#99064)

(cherry picked from commit 46a4132e167aa44d8ec7776262ce2a0e6d47de59)
---
 .../Instrumentation/PGOInstrumentation.cpp    | 14 +++++--
 .../Coroutines/coro-pgo-setbranchweights.ll   | 42 +++++++++++++++++++
 2 files changed, 52 insertions(+), 4 deletions(-)
 create mode 100644 llvm/test/Transforms/Coroutines/coro-pgo-setbranchweights.ll

diff --git a/llvm/lib/Transforms/Instrumentation/PGOInstrumentation.cpp b/llvm/lib/Transforms/Instrumentation/PGOInstrumentation.cpp
index 1ce8f58c1aa140..4924d5a317478f 100644
--- a/llvm/lib/Transforms/Instrumentation/PGOInstrumentation.cpp
+++ b/llvm/lib/Transforms/Instrumentation/PGOInstrumentation.cpp
@@ -1625,11 +1625,17 @@ void PGOUseFunc::setBranchWeights() {
       continue;
 
     // We have a non-zero Branch BB.
-    unsigned Size = BBCountInfo.OutEdges.size();
-    SmallVector<uint64_t, 2> EdgeCounts(Size, 0);
+
+    // SuccessorCount can be greater than OutEdgesCount, because
+    // removed edges don't appear in OutEdges.
+    unsigned OutEdgesCount = BBCountInfo.OutEdges.size();
+    unsigned SuccessorCount = BB.getTerminator()->getNumSuccessors();
+    assert(OutEdgesCount <= SuccessorCount);
+
+    SmallVector<uint64_t, 2> EdgeCounts(SuccessorCount, 0);
     uint64_t MaxCount = 0;
-    for (unsigned s = 0; s < Size; s++) {
-      const PGOUseEdge *E = BBCountInfo.OutEdges[s];
+    for (unsigned It = 0; It < OutEdgesCount; It++) {
+      const PGOUseEdge *E = BBCountInfo.OutEdges[It];
       const BasicBlock *SrcBB = E->SrcBB;
       const BasicBlock *DestBB = E->DestBB;
       if (DestBB == nullptr)
diff --git a/llvm/test/Transforms/Coroutines/coro-pgo-setbranchweights.ll b/llvm/test/Transforms/Coroutines/coro-pgo-setbranchweights.ll
new file mode 100644
index 00000000000000..4f5f936606ca3f
--- /dev/null
+++ b/llvm/test/Transforms/Coroutines/coro-pgo-setbranchweights.ll
@@ -0,0 +1,42 @@
+; RUN: rm -rf %t && split-file %s %t
+
+; RUN: llvm-profdata merge %t/a.proftext -o %t/a.profdata
+; RUN: opt < %t/a.ll --passes=pgo-instr-use -pgo-test-profile-file=%t/a.profdata
+
+;--- a.ll
+target datalayout = "e-m:e-p270:32:32-p271:32:32-p272:64:64-i64:64-i128:128-f80:128-n8:16:32:64-S128"
+target triple = "x86_64-redhat-linux-gnu"
+
+define void @_bar() presplitcoroutine personality ptr null {
+  %1 = call token @llvm.coro.save(ptr null)
+  %2 = call i8 @llvm.coro.suspend(token none, i1 false)
+  switch i8 %2, label %5 [
+    i8 0, label %3
+    i8 1, label %4
+  ]
+
+3:                                                ; preds = %0
+  ret void
+
+4:                                                ; preds = %0
+  ret void
+
+5:                                                ; preds = %0
+  ret void
+}
+
+declare token @llvm.coro.save(ptr)
+
+declare i8 @llvm.coro.suspend(token, i1)
+
+;--- a.proftext
+# IR level Instrumentation Flag
+:ir
+
+_bar
+# Func Hash:
+1063705160175073211
+# Num Counters:
+2
+1
+0