[llvm] [FunctionSpecializer] Do not mark function dead if any unexecutable call site exists (PR #154668)

[via llvm-commits]

https://github.com/XChy updated https://github.com/llvm/llvm-project/pull/154668

>From bc9c37a53479327a1c0acab540f8a4d7257c9936 Mon Sep 17 00:00:00 2001
From: XChy <xxs_chy at outlook.com>
Date: Thu, 21 Aug 2025 12:22:24 +0800
Subject: [PATCH 1/2] [FunctionnSpecializer] Do not mark function dead if any
 unexecutable callsite exists

---
 .../Transforms/IPO/FunctionSpecialization.h   |  2 +
 .../Transforms/IPO/FunctionSpecialization.cpp |  9 ++++
 llvm/lib/Transforms/IPO/SCCP.cpp              |  4 +-
 .../reachable-after-specialization.ll         | 41 +++++++++++++++++++
 4 files changed, 55 insertions(+), 1 deletion(-)
 create mode 100644 llvm/test/Transforms/FunctionSpecialization/reachable-after-specialization.ll

diff --git a/llvm/include/llvm/Transforms/IPO/FunctionSpecialization.h b/llvm/include/llvm/Transforms/IPO/FunctionSpecialization.h
index 1575afa50e198..99640d0fde690 100644
--- a/llvm/include/llvm/Transforms/IPO/FunctionSpecialization.h
+++ b/llvm/include/llvm/Transforms/IPO/FunctionSpecialization.h
@@ -270,6 +270,8 @@ class FunctionSpecializer {
     return InstCostVisitor(GetBFI, F, M.getDataLayout(), TTI, Solver);
   }
 
+  bool isDeadFunction(Function *F) { return FullySpecialized.contains(F); }
+
 private:
   Constant *getPromotableAlloca(AllocaInst *Alloca, CallInst *Call);
 
diff --git a/llvm/lib/Transforms/IPO/FunctionSpecialization.cpp b/llvm/lib/Transforms/IPO/FunctionSpecialization.cpp
index c876a47ef2129..86805547ca67d 100644
--- a/llvm/lib/Transforms/IPO/FunctionSpecialization.cpp
+++ b/llvm/lib/Transforms/IPO/FunctionSpecialization.cpp
@@ -844,6 +844,15 @@ void FunctionSpecializer::removeDeadFunctions() {
                       << F->getName() << "\n");
     if (FAM)
       FAM->clear(*F, F->getName());
+
+    // Remove all the callsites that were proven unreachable once, and replace
+    // them with poison.
+    for (User *U : make_early_inc_range(F->users())) {
+      assert((isa<CallInst>(U) || isa<InvokeInst>(U)) && "User of dead function must be call or invoke");
+      Instruction *CS = cast<Instruction>(U);
+      CS->replaceAllUsesWith(PoisonValue::get(CS->getType()));
+      CS->eraseFromParent();
+    }
     F->eraseFromParent();
   }
   FullySpecialized.clear();
diff --git a/llvm/lib/Transforms/IPO/SCCP.cpp b/llvm/lib/Transforms/IPO/SCCP.cpp
index d50de34dfa482..604fc37d8ab34 100644
--- a/llvm/lib/Transforms/IPO/SCCP.cpp
+++ b/llvm/lib/Transforms/IPO/SCCP.cpp
@@ -167,7 +167,9 @@ static bool runIPSCCP(
   // constants if we have found them to be of constant values.
   bool MadeChanges = false;
   for (Function &F : M) {
-    if (F.isDeclaration())
+    // Skip the dead functions marked by FunctionSpecializer, avoiding removing
+    // blocks in dead functions.
+    if (F.isDeclaration() || Specializer.isDeadFunction(&F))
       continue;
 
     SmallVector<BasicBlock *, 512> BlocksToErase;
diff --git a/llvm/test/Transforms/FunctionSpecialization/reachable-after-specialization.ll b/llvm/test/Transforms/FunctionSpecialization/reachable-after-specialization.ll
new file mode 100644
index 0000000000000..53641e265a3fc
--- /dev/null
+++ b/llvm/test/Transforms/FunctionSpecialization/reachable-after-specialization.ll
@@ -0,0 +1,41 @@
+; NOTE: Assertions have been autogenerated by utils/update_test_checks.py UTC_ARGS: --version 5
+; RUN: opt -passes=ipsccp  --funcspec-min-function-size=1 -S < %s | FileCheck %s
+
+define i32 @caller() {
+; CHECK-LABEL: define i32 @caller() {
+; CHECK-NEXT:  [[ENTRY:.*:]]
+; CHECK-NEXT:    [[CALL1:%.*]] = call i32 @callee.specialized.1(i32 1)
+; CHECK-NEXT:    [[CALL2:%.*]] = call i32 @callee.specialized.2(i32 0)
+; CHECK-NEXT:    [[COND:%.*]] = icmp eq i32 undef, 0
+; CHECK-NEXT:    br i1 [[COND]], label %[[COMMON_RET:.*]], label %[[IF_THEN:.*]]
+; CHECK:       [[COMMON_RET]]:
+; CHECK-NEXT:    ret i32 0
+; CHECK:       [[IF_THEN]]:
+; CHECK-NEXT:    ret i32 0
+;
+entry:
+  %call1 = call i32 @callee(i32 1)
+  %call2 = call i32 @callee(i32 0)
+  %cond = icmp eq i32 %call2, 0
+  br i1 %cond, label %common.ret, label %if.then
+
+common.ret:                                       ; preds = %entry
+  ret i32 0
+
+if.then:                                         ; preds = %entry
+  %unreachable_call = call i32 @callee(i32 2)
+  ret i32 %unreachable_call
+}
+
+define internal i32 @callee(i32 %ac) {
+entry:
+  br label %ai
+
+ai:                                               ; preds = %ai, %entry
+  %add = or i32 0, 0
+  %cond = icmp eq i32 %ac, 1
+  br i1 %cond, label %aj, label %ai
+
+aj:                                               ; preds = %ai
+  ret i32 0
+}

>From 1067b5c640ab1b10bb2d4eae3bea319dcf7417dc Mon Sep 17 00:00:00 2001
From: XChy <xxs_chy at outlook.com>
Date: Sun, 24 Aug 2025 19:11:53 +0800
Subject: [PATCH 2/2] format

---
 llvm/lib/Transforms/IPO/FunctionSpecialization.cpp | 3 ++-
 1 file changed, 2 insertions(+), 1 deletion(-)

diff --git a/llvm/lib/Transforms/IPO/FunctionSpecialization.cpp b/llvm/lib/Transforms/IPO/FunctionSpecialization.cpp
index 86805547ca67d..745041a0d5709 100644
--- a/llvm/lib/Transforms/IPO/FunctionSpecialization.cpp
+++ b/llvm/lib/Transforms/IPO/FunctionSpecialization.cpp
@@ -848,7 +848,8 @@ void FunctionSpecializer::removeDeadFunctions() {
     // Remove all the callsites that were proven unreachable once, and replace
     // them with poison.
     for (User *U : make_early_inc_range(F->users())) {
-      assert((isa<CallInst>(U) || isa<InvokeInst>(U)) && "User of dead function must be call or invoke");
+      assert((isa<CallInst>(U) || isa<InvokeInst>(U)) &&
+             "User of dead function must be call or invoke");
       Instruction *CS = cast<Instruction>(U);
       CS->replaceAllUsesWith(PoisonValue::get(CS->getType()));
       CS->eraseFromParent();