[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 08984a2eec296f13cd1c3335ca323d5555e38391 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.

A unit test, provided by @ellishg, is included.

Fixes #97962

(regression from ffd337b995658266d83e15630d656e57882bcf0d)
---
 .../Instrumentation/PGOInstrumentation.cpp    | 10 +++--
 .../Coroutines/coro-pgo-setbranchweights.ll   | 42 +++++++++++++++++++
 2 files changed, 48 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 41618194d12ed7..0387727dfbc2aa 100644
--- a/llvm/lib/Transforms/Instrumentation/PGOInstrumentation.cpp
+++ b/llvm/lib/Transforms/Instrumentation/PGOInstrumentation.cpp
@@ -1628,11 +1628,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)
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