[llvm] [Instrumentation] Fix EdgeCounts vector size in SetBranchWeights (PR #99064)

[Avi Kivity via llvm-commits]

https://github.com/avikivity updated https://github.com/llvm/llvm-project/pull/99064

>From b7db364983a62943e2e9ef18479704e2f535056d Mon Sep 17 00:00:00 2001
From: Avi Kivity <avi at scylladb.com>
Date: Tue, 16 Jul 2024 19:38:53 +0300
Subject: [PATCH] [Instrumentation] Fix EdgeCounts vector size in
 SetBranchWeights

SetBranchWeights() calculates the size of the EdgeCounts vector
using OutEdges.Size(), but this is an under-estimate with coroutines.

Use the number of successors, as the vector will be indexed by
the result of the GetSuccessorNumber() function.

Rename the Size local, to make it clear what it refers to.

This crashes in 18.1, but somehow doesn't in main. Still, it looks
like the fix is applicable to both.

Fixes #97962
---
 .../lib/Transforms/Instrumentation/PGOInstrumentation.cpp | 8 +++++---
 1 file changed, 5 insertions(+), 3 deletions(-)

diff --git a/llvm/lib/Transforms/Instrumentation/PGOInstrumentation.cpp b/llvm/lib/Transforms/Instrumentation/PGOInstrumentation.cpp
index 35b1bbf21be97..5f9cd7b11d815 100644
--- a/llvm/lib/Transforms/Instrumentation/PGOInstrumentation.cpp
+++ b/llvm/lib/Transforms/Instrumentation/PGOInstrumentation.cpp
@@ -1623,10 +1623,12 @@ void PGOUseFunc::setBranchWeights() {
       continue;
 
     // We have a non-zero Branch BB.
-    unsigned Size = BBCountInfo.OutEdges.size();
-    SmallVector<uint64_t, 2> EdgeCounts(Size, 0);
+    unsigned OutEdgesCount = BBCountInfo.OutEdges.size();
+    unsigned SuccessorCount = BB.getTerminator()->getNumSuccessors();
+
+    SmallVector<uint64_t, 2> EdgeCounts(SuccessorCount, 0);
     uint64_t MaxCount = 0;
-    for (unsigned s = 0; s < Size; s++) {
+    for (unsigned s = 0; s < OutEdgesCount; s++) {
       const PGOUseEdge *E = BBCountInfo.OutEdges[s];
       const BasicBlock *SrcBB = E->SrcBB;
       const BasicBlock *DestBB = E->DestBB;