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

[Avi Kivity via llvm-commits]

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

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.

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

Fixes #97962

>From eae8cec3053e5ffb95d57463f5c44904c11c2899 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.

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

Fixes #97962
---
 llvm/lib/Transforms/Instrumentation/PGOInstrumentation.cpp | 4 +++-
 1 file changed, 3 insertions(+), 1 deletion(-)

diff --git a/llvm/lib/Transforms/Instrumentation/PGOInstrumentation.cpp b/llvm/lib/Transforms/Instrumentation/PGOInstrumentation.cpp
index 35b1bbf21be97..d32d00700bcb3 100644
--- a/llvm/lib/Transforms/Instrumentation/PGOInstrumentation.cpp
+++ b/llvm/lib/Transforms/Instrumentation/PGOInstrumentation.cpp
@@ -1624,7 +1624,9 @@ void PGOUseFunc::setBranchWeights() {
 
     // We have a non-zero Branch BB.
     unsigned Size = BBCountInfo.OutEdges.size();
-    SmallVector<uint64_t, 2> EdgeCounts(Size, 0);
+    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];