[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 60726081abf5f0ff67b85c9d8740024b40ee3851 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
---
 .../Transforms/Instrumentation/PGOInstrumentation.cpp  | 10 ++++++----
 1 file changed, 6 insertions(+), 4 deletions(-)

diff --git a/llvm/lib/Transforms/Instrumentation/PGOInstrumentation.cpp b/llvm/lib/Transforms/Instrumentation/PGOInstrumentation.cpp
index 35b1bbf21be97..84b53d0e72fe4 100644
--- a/llvm/lib/Transforms/Instrumentation/PGOInstrumentation.cpp
+++ b/llvm/lib/Transforms/Instrumentation/PGOInstrumentation.cpp
@@ -1623,11 +1623,13 @@ 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++) {
-      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)