[FuncSpec] Relax restrictions on candidates for specialisation

[chill]

Summary

Allow a function to be specialised even if it has its address taken or
it's global. For such functions, consider all of the arguments as
overdefined. Don't delete the functions even if all the apparent calls
were redirected to specialised instances.

[labrinea]

I've checked the number of specializations created for CTMark with this patch. No changes there so I am not expecting a compile time hit. Approved.

[nikic]

In D148345#4272354, @labrinea wrote:

I've checked the number of specializations created for CTMark with this patch. No changes there so I am not expecting a compile time hit. Approved.

Looks like this assumption wasn't right. There is a compile-time regression (http://llvm-compile-time-tracker.com/compare.php?from=99d4c722e30816aca6698e20345bffab953d9d79&to=cc7bb7080fc8e6f4d217e7f9b971fbdbf091f9e7&stat=instructions:u) and code size changed on lencod (http://llvm-compile-time-tracker.com/compare.php?from=99d4c722e30816aca6698e20345bffab953d9d79&to=cc7bb7080fc8e6f4d217e7f9b971fbdbf091f9e7&stat=size-text).

[labrinea]

In D148345#4283318, @nikic wrote:

In D148345#4272354, @labrinea wrote:

I've checked the number of specializations created for CTMark with this patch. No changes there so I am not expecting a compile time hit. Approved.

Looks like this assumption wasn't right. There is a compile-time regression (http://llvm-compile-time-tracker.com/compare.php?from=99d4c722e30816aca6698e20345bffab953d9d79&to=cc7bb7080fc8e6f4d217e7f9b971fbdbf091f9e7&stat=instructions:u) and code size changed on lencod (http://llvm-compile-time-tracker.com/compare.php?from=99d4c722e30816aca6698e20345bffab953d9d79&to=cc7bb7080fc8e6f4d217e7f9b971fbdbf091f9e7&stat=size-text).

Right I see now. The compile time regression are without LTO, whereas I had checked only the LTO configuration. I can see three new specializations in lencod.

llvm/lib/Transforms/IPO/FunctionSpecialization.cpp

if (auto It = UM.find(S); It != UM.end()) {
       Found = true;
     }
   }
 
   return Found;
 }
 
 bool FunctionSpecializer::isCandidateFunction(Function *F) {
-  if (F->isDeclaration())
+  if (F->isDeclaration() || F->arg_empty())
     return false;
 
   if (F->hasFnAttribute(Attribute::NoDuplicate))
     return false;
 
-  if (!Solver.isArgumentTrackedFunction(F))
-    return false;
-
   // Do not specialize the cloned function again.
   if (Specializations.contains(F))
     return false;
 
   // If we're optimizing the function for size, we shouldn't specialize it.
   if (F->hasOptSize() ||
       shouldOptimizeForSize(F, nullptr, nullptr, PGSOQueryType::IRPass))
     return false;
bool FunctionSpecializer::isCandidateFunction(Function *F) {
   LLVM_DEBUG(dbgs() << "FnSpecialization: Try function: " << F->getName()
                     << "\n");
   return true;
 }
 
 Function *FunctionSpecializer::createSpecialization(Function *F, const SpecSig &S) {
   Function *Clone = cloneCandidateFunction(F);
 
+  // The original function does not neccessarily have internal linkage, but the
+  // clone must.
+  Clone->setLinkage(GlobalValue::InternalLinkage);
+
   // Initialize the lattice state of the arguments of the function clone,
   // marking the argument on which we specialized the function constant
   // with the given value.
   Solver.setLatticeValueForSpecializationArguments(Clone, S.Args);
 
   Solver.addArgumentTrackedFunction(Clone);
   Solver.markBlockExecutable(&Clone->front());
 
if (!Ty->isPointerTy() && (!SpecializeLiteralConstant ||
       (!Ty->isIntegerTy() && !Ty->isFloatingPointTy() && !Ty->isStructTy())))
     return false;
 
   // SCCP solver does not record an argument that will be constructed on
   // stack.
   if (A->hasByValAttr() && !A->getParent()->onlyReadsMemory())
     return false;
 
+  // For non-argument-tracked functions every argument is overdefined.
+  if (!Solver.isArgumentTrackedFunction(A->getParent()))
+    return true;
+
   // Check the lattice value and decide if we should attemt to specialize,
   // based on this argument. No point in specialization, if the lattice value
   // is already a constant.
   bool IsOverdefined = Ty->isStructTy()
     ? any_of(Solver.getStructLatticeValueFor(A), SCCPSolver::isOverdefined)
     : SCCPSolver::isOverdefined(Solver.getLatticeValueFor(A));
 
   LLVM_DEBUG(
for (CallBase *CS : ToUpdate) {
     }
 
     if (ShouldDecrementCount)
       --NCallsLeft;
   }
 
   // If the function has been completely specialized, the original function
   // is no longer needed. Mark it unreachable.
-  if (NCallsLeft == 0) {
+  if (NCallsLeft == 0 && Solver.isArgumentTrackedFunction(F)) {
     Solver.markFunctionUnreachable(F);
     FullySpecialized.insert(F);
   }
 }

llvm/test/Transforms/FunctionSpecialization/non-argument-tracked.ll

+; RUN: opt -S --passes=ipsccp,deadargelim -funcspec-for-literal-constant --force-specialization < %s | FileCheck %s
+
+; Test that all of `f0`, `f1`, and `f2` are specialised, even though `f0` has its address taken
+; and `f1` is with external linkage (`f2` was specialised anyway).
+
+@p = global ptr @f0
+
+; `f0` is kept even though all apparent calls are specialized
+; CHECK-LABEL: define internal i32 @f0(
+define internal i32 @f0(i32 %i) {
+  %v = add i32 %i, 1
+  ret i32 %v
+}
+
+; Likewise, `f1` is kept, because of the external linkage
+; CHECK-LABEL: define i32 @f1(
+define i32 @f1(i32 %i) {
+  %v = add i32 %i, 1
+  ret i32 %v
+}
+
+; `f2` is fully specialised.
+; CHECK-NOT: defined internal i32 @f2()
+define internal i32 @f2(i32 %i) {
+  %v = add i32 %i, 1
+  ret i32 %v
+}
+
+;; All calls are to specilisation instances.
+
+; CHECK-LABEL: define i32 @g0
+; CHECK:         [[U0:%.*]] = call i32 @f0.[[#A:]]()
+; CHECK-NEXT:    [[U1:%.*]] = call i32 @f1.[[#B:]]()
+; CHECK-NEXT:    [[U2:%.*]] = call i32 @f2.[[#C:]]()
+define i32 @g0(i32 %i) {
+  %u0 = call i32 @f0(i32 1)
+  %u1 = call i32 @f1(i32 2)
+  %u2 = call i32 @f2(i32 3)
+  %v0 = add i32 %u0, %u1
+  %v = add i32 %v0, %u2
+  ret i32 %v
+}
+
+; CHECK-LABEL: define i32 @g1
+; CHECK:         [[U0:%.*]] = call i32 @f0.[[#D:]]()
+; CHECK-NEXT:    [[U1:%.*]] = call i32 @f1.[[#E:]]()
+; CHECK-NEXT:    [[U2:%.*]] = call i32 @f2.[[#F:]]()
+define i32 @g1(i32 %i) {
+  %u0 = call i32 @f0(i32 2)
+  %u1 = call i32 @f1(i32 3)
+  %u2 = call i32 @f2(i32 4)
+  %v0 = add i32 %u0, %u1
+  %v = add i32 %v0, %u2
+  ret i32 %v
+}
+
+; All of the function are specialized and all clones are with internal linkage.
+
+; CHECK-DAG: define internal i32 @f0.[[#A]]() {
+; CHECK-DAG: define internal i32 @f1.[[#B]]() {
+; CHECK-DAG: define internal i32 @f2.[[#C]]() {
+; CHECK-DAG: define internal i32 @f0.[[#D]]() {
+; CHECK-DAG: define internal i32 @f1.[[#E]]() {
+; CHECK-DAG: define internal i32 @f2.[[#F]]() {