[FuncSpec] Skip SCCP on blocks of dead functions and poison their callsites (#154668)

Fixes #153295.
For test case below:
```llvm
define i32 @caller() {
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
}
```
Before specialization, the SCCP solver determines that
`unreachable_call` is unexecutable, as the value of `callee` can only be
zero.
After specializing the call sites `call1` and `call2`, FnSpecializer
announces `callee` is a dead function since all executable call sites
are specialized. However, the unexecutable call sites can become
executable again after solving specialized calls.
In this testcase, `call2` is considered `Overdefined` after
specialization, making `cond` also `Overdefined`. Thus,
`unreachable_call` becomes executable.
This patch skips SCCP on the blocks in dead functions, and poisons the
call sites of dead functions.

1 files changed, 17 insertions, 4 deletions

diff --git a/llvm/lib/Transforms/IPO/FunctionSpecialization.cpp b/llvm/lib/Transforms/IPO/FunctionSpecialization.cpp
index 9196a01..a459a9e 100644
--- a/llvm/lib/Transforms/IPO/FunctionSpecialization.cpp
+++ b/llvm/lib/Transforms/IPO/FunctionSpecialization.cpp
@@ -838,14 +838,24 @@ bool FunctionSpecializer::run() {
 }
 
 void FunctionSpecializer::removeDeadFunctions() {
-  for (Function *F : FullySpecialized) {
+  for (Function *F : DeadFunctions) {
     LLVM_DEBUG(dbgs() << "FnSpecialization: Removing dead function "
                       << 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();
+  DeadFunctions.clear();
 }
 
 /// Clone the function \p F and remove the ssa_copy intrinsics added by
@@ -1206,8 +1216,11 @@ void FunctionSpecializer::updateCallSites(Function *F, const Spec *Begin,
 
   // If the function has been completely specialized, the original function
   // is no longer needed. Mark it unreachable.
-  if (NCallsLeft == 0 && Solver.isArgumentTrackedFunction(F)) {
+  // NOTE: If the address of a function is taken, we cannot treat it as dead
+  // function.
+  if (NCallsLeft == 0 && Solver.isArgumentTrackedFunction(F) &&
+      !F->hasAddressTaken()) {
     Solver.markFunctionUnreachable(F);
-    FullySpecialized.insert(F);
+    DeadFunctions.insert(F);
   }
 }