[llvm] [CFGMST] Remove special case for entry block with no successors (PR #140804)

[Arthur Eubanks via llvm-commits]

https://github.com/aeubanks created https://github.com/llvm/llvm-project/pull/140804

This breaks in the case where there are unreachable blocks after an entry block with no successors, which don't get visited in CFGMST, causing crashes.

Fixes #135828.

>From cba58725bbdf18417e6a514bccce5887335c192b Mon Sep 17 00:00:00 2001
From: Arthur Eubanks <aeubanks at google.com>
Date: Tue, 20 May 2025 21:16:05 +0000
Subject: [PATCH] [CFGMST] Remove special case for entry block with no
 successors

This breaks in the case where there are unreachable blocks after an entry block with no successors, which don't get visited in CFGMST, causing crashes.

Fixes #135828.
---
 .../include/llvm/Transforms/Instrumentation/CFGMST.h |  6 ------
 .../PGOProfile/Inputs/unreachable-block.proftext     |  9 +++++++++
 llvm/test/Transforms/PGOProfile/unreachable-block.ll | 12 ++++++++++++
 3 files changed, 21 insertions(+), 6 deletions(-)
 create mode 100644 llvm/test/Transforms/PGOProfile/Inputs/unreachable-block.proftext
 create mode 100644 llvm/test/Transforms/PGOProfile/unreachable-block.ll

diff --git a/llvm/include/llvm/Transforms/Instrumentation/CFGMST.h b/llvm/include/llvm/Transforms/Instrumentation/CFGMST.h
index f6bf045f7de2c..6f0c63b888353 100644
--- a/llvm/include/llvm/Transforms/Instrumentation/CFGMST.h
+++ b/llvm/include/llvm/Transforms/Instrumentation/CFGMST.h
@@ -133,12 +133,6 @@ template <class Edge, class BBInfo> class CFGMST {
     LLVM_DEBUG(dbgs() << "  Edge: from fake node to " << Entry->getName()
                       << " w = " << EntryWeight << "\n");
 
-    // Special handling for single BB functions.
-    if (succ_empty(Entry)) {
-      addEdge(Entry, nullptr, EntryWeight);
-      return;
-    }
-
     static const uint32_t CriticalEdgeMultiplier = 1000;
 
     for (BasicBlock &BB : F) {
diff --git a/llvm/test/Transforms/PGOProfile/Inputs/unreachable-block.proftext b/llvm/test/Transforms/PGOProfile/Inputs/unreachable-block.proftext
new file mode 100644
index 0000000000000..9be174fb32bca
--- /dev/null
+++ b/llvm/test/Transforms/PGOProfile/Inputs/unreachable-block.proftext
@@ -0,0 +1,9 @@
+# IR level Instrumentation Flag
+:ir
+foo
+# Func Hash:
+742261418966908927
+# Num Counters:
+1
+# Counter Values:
+1
diff --git a/llvm/test/Transforms/PGOProfile/unreachable-block.ll b/llvm/test/Transforms/PGOProfile/unreachable-block.ll
new file mode 100644
index 0000000000000..bd55185649989
--- /dev/null
+++ b/llvm/test/Transforms/PGOProfile/unreachable-block.ll
@@ -0,0 +1,12 @@
+; RUN: llvm-profdata merge %S/Inputs/unreachable-block.proftext -o %t.profdata
+; RUN: opt < %s -passes=pgo-instr-use -pgo-test-profile-file=%t.profdata -S
+
+declare ptr @bar()
+
+define ptr @foo() {
+entry:
+  ret ptr null
+
+2:
+  ret ptr null
+}