[FnSpecialization] Allow specialization of indirect function calls exposed by specialization

Author: bababuck

llvm/lib/Transforms/IPO/FunctionSpecialization.cpp

@@ -234,23 +234,33 @@ Cost InstCostVisitor::getCodeSizeSavingsForUser(Instruction *User, Value *Use,
 
   Cost CodeSize = 0;
   auto isChainableCall = [&](Instruction *I) -> bool {
+    if (!CallUsers || !UseEdge)
+      return false;
     if (CallInst *CI = dyn_cast<CallInst>(I);
         CI && CI->getIntrinsicID() == llvm::Intrinsic::not_intrinsic) {
       LLVM_DEBUG(
           dbgs() << "FnSpecialization:   Found constant forwarded via a call "
                  << *C << "\n");
       Function *F = CI->getCalledFunction();
-      if (F && CallUsers && UseEdge) { // Avoid function pointers
+      if (F) { // Avoid function pointers
         unsigned Idx = CI->getArgOperandNo(UseEdge);
         LLVM_DEBUG(dbgs() << "FnSpecialization:   Function called: "
                           << F->getName() << " argument number: " << Idx
                           << "\n");
         (*CallUsers)[CI].first.push_back({Idx, C});
         (*CallUsers)[CI].second = F;
         return true;
+      } else if (Use == CI->getCalledOperand()) {
+        LLVM_DEBUG(dbgs() << "FnSpecialization:   Found call to constant "
+                             "function pointer.\n");
+        Function *CF = dyn_cast<Function>(C);
+        assert(CF && "Indirect call to a non-Function type");
+        (*CallUsers)[CI].second = CF;
       } else {
         LLVM_DEBUG(
             dbgs() << "FnSpecialization:   Could not find call function.\n");
+        unsigned Idx = CI->getArgOperandNo(UseEdge);
+        (*CallUsers)[CI].first.push_back({Idx, C});
       }
     }
     return false;

llvm/test/Transforms/FunctionSpecialization/compiler-crash-60191.ll

@@ -60,7 +60,7 @@ define i32 @f2(i32 %offset) {
 }
 
 ; Tests that `func` has been specialized and it didn't cause compiler crash.
+; CHECK-DAG: func.specialized.4
+; CHECK-DAG: func.specialized.5
 ; CHECK-DAG: func.specialized.1
-; CHECK-DAG: func.specialized.2
-; CHECK-DAG: func.specialized.3
 

llvm/test/Transforms/FunctionSpecialization/track-ptr-return.ll

@@ -48,9 +48,8 @@ entry:
 ; CHECK-NEXT:    [[OP1:%.*]] = call ptr @select_op.specialized.1(ptr @global_true)
 ; CHECK-NEXT:    [[OP2:%.*]] = call ptr @select_op.specialized.2(ptr @global_false)
 ; CHECK-NEXT:    [[C1:%.*]] = call i64 @compute.specialized.3(ptr @plus)
-; CHECK-NEXT:    [[C2:%.*]] = call i64 @compute.specialized.4(ptr @minus)
-; CHECK-NEXT:    [[ADD:%.*]] = add i64 [[C1]], [[C2]]
-; CHECK-NEXT:    ret i64 [[ADD]]
+; CHECK-NEXT:    [[C2:%.*]] = call i64 @compute.specialized.5(ptr @minus)
+; CHECK-NEXT:    ret i64 2
 ;
 ;
 ; CHECK-LABEL: define ptr @select_op(
@@ -87,15 +86,27 @@ entry:
 ; CHECK-LABEL: define internal i64 @compute.specialized.3(
 ; CHECK-SAME: ptr [[OP:%.*]]) {
 ; CHECK-NEXT:  [[ENTRY:.*:]]
-; CHECK-NEXT:    [[RES:%.*]] = call i64 @plus(i64 1)
-; CHECK-NEXT:    ret i64 [[RES]]
+; CHECK-NEXT:    [[RES:%.*]] = call i64 @plus.specialized.4(i64 1)
+; CHECK-NEXT:    ret i64 poison
 ;
 ;
-; CHECK-LABEL: define internal i64 @compute.specialized.4(
+; CHECK-LABEL: define internal i64 @plus.specialized.4(
+; CHECK-SAME: i64 [[X:%.*]]) {
+; CHECK-NEXT:  [[ENTRY:.*:]]
+; CHECK-NEXT:    ret i64 poison
+;
+;
+; CHECK-LABEL: define internal i64 @compute.specialized.5(
 ; CHECK-SAME: ptr [[OP:%.*]]) {
 ; CHECK-NEXT:  [[ENTRY:.*:]]
-; CHECK-NEXT:    [[RES:%.*]] = call i64 @minus(i64 1)
-; CHECK-NEXT:    ret i64 [[RES]]
+; CHECK-NEXT:    [[RES:%.*]] = call i64 @minus.specialized.6(i64 1)
+; CHECK-NEXT:    ret i64 poison
+;
+;
+; CHECK-LABEL: define internal i64 @minus.specialized.6(
+; CHECK-SAME: i64 [[X:%.*]]) {
+; CHECK-NEXT:  [[ENTRY:.*:]]
+; CHECK-NEXT:    ret i64 poison
 ;
 ;
 ; NOLIT-LABEL: define i64 @main() {