[FunctionnSpecializer] Do not mark function dead if any unexecutable callsite exists

Author: XChy

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);
 

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();

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;

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
+}