[llvm] cc7bb70 - [FuncSpec] Relax restrictions on candidates for specialisation

[Momchil Velikov via llvm-commits]

Author: Momchil Velikov
Date: 2023-04-20T12:06:22+01:00
New Revision: cc7bb7080fc8e6f4d217e7f9b971fbdbf091f9e7

URL: https://github.com/llvm/llvm-project/commit/cc7bb7080fc8e6f4d217e7f9b971fbdbf091f9e7
DIFF: https://github.com/llvm/llvm-project/commit/cc7bb7080fc8e6f4d217e7f9b971fbdbf091f9e7.diff

LOG: [FuncSpec] Relax restrictions on candidates for specialisation

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.

Reviewed By: labrinea, ChuanqiXu

Differential Revision: https://reviews.llvm.org/D148345

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

Modified: 
    llvm/lib/Transforms/IPO/FunctionSpecialization.cpp

Removed: 
    


################################################################################
diff  --git a/llvm/lib/Transforms/IPO/FunctionSpecialization.cpp b/llvm/lib/Transforms/IPO/FunctionSpecialization.cpp
index 4fc609e4405e4..207737f1185c3 100644
--- a/llvm/lib/Transforms/IPO/FunctionSpecialization.cpp
+++ b/llvm/lib/Transforms/IPO/FunctionSpecialization.cpp
@@ -498,15 +498,12 @@ bool FunctionSpecializer::findSpecializations(Function *F, InstructionCost Cost,
 }
 
 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;
@@ -533,6 +530,10 @@ bool FunctionSpecializer::isCandidateFunction(Function *F) {
 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.
@@ -676,6 +677,10 @@ bool FunctionSpecializer::isArgumentInteresting(Argument *A) {
   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.
@@ -761,7 +766,7 @@ 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) {
+  if (NCallsLeft == 0 && Solver.isArgumentTrackedFunction(F)) {
     Solver.markFunctionUnreachable(F);
     FullySpecialized.insert(F);
   }

diff  --git a/llvm/test/Transforms/FunctionSpecialization/non-argument-tracked.ll b/llvm/test/Transforms/FunctionSpecialization/non-argument-tracked.ll
new file mode 100644
index 0000000000000..2fb7205f7e33d
--- /dev/null
+++ b/llvm/test/Transforms/FunctionSpecialization/non-argument-tracked.ll
@@ -0,0 +1,64 @@
+; 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).
+
+ at 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]]() {