[FuncSpec] Fix specialisation based on literals

[chill]

Summary

The FunctionSpecialization pass has support for specialising
functions, which are called with literal arguments. This functionality
is disabled by default and is enabled with the option
-function-specialization-for-literal-constant . There are a few
issues with the implementation, though:

• even with the default, the pass will still specialise based on floating-point literals

• even when it's enabled, the pass will specialise only for the i1 type (or i2 if all of the possible 4 values occur, or i3 if all of the possible 8 values occur, etc)

The reason for this is incorrect check of the lattice value of the
function formal parameter. The lattice value is overdefined when the
constant range of the possible arguments is the full set, and this is
the reason for the specialisation to trigger. However, if the set of
the possible arguments is not the full set, that must not prevent the
specialisation.

This patch changes the pass to NOT consider a formal parameter when
specialising a function if the lattice value for that parameter is:

• unknown or undef
• a constant
• a constant range with a single element

on the basis that specialisation is pointless for those cases.

Is also changes the criteria for picking up an actual argument to
specialise if the argument is:

• a LLVM IR constant
• has constant lattice value has constantrange lattice value with a single element.

[ChuanqiXu]

Neat! I really enjoy reading your work. Thanks!

llvm/lib/Transforms/IPO/FunctionSpecialization.cpp

712–715

This change looks not related to this diff? Although it is bad that the tests still pass after deleting this. But I prefer to remain this if this is not intended.

762–775

nit:

[chill]

llvm/lib/Transforms/IPO/FunctionSpecialization.cpp

712–715

Indeed, this change is just a refactoring, I moved it towards the beginning of isArgumentInteresting since it's a more logical place for this test to be.

I don't mind not applying this change here in this patch, I have another one coming that's just refactoring/optimisaiton(?) and I can include it there.

762–775

Ack.

[ChuanqiXu]

llvm/lib/Transforms/IPO/FunctionSpecialization.cpp

712–715

Oh, my oversight. Then it should be fine.

[labrinea]

Another nit. Can we omit the "func-spec-" prefix from the test names (as well as in D135862 and D135867)? They are already under the relevant directory which makes it obvious to infer what the test is about.

[chill]

In D135893#3857985, @labrinea wrote:

Another nit. Can we omit the "func-spec-" prefix from the test names (as well as in D135862 and D135867)? They are already under the relevant directory which makes it obvious to infer what the test is about.

Agree!

[chill]

There is a build bot failure with the test literal-const.ll. AFAICT, the issue is with the test itself - it's not general enough (or, unlikely, there's a non-determinism in the FunctionSpecialization pass);

[chill]

For now, temporarily disable the test: https://reviews.llvm.org/D136755

[chill]

-DLLVM_REVERSE_ITERATION:BOOL=ON

Our cloned function names depend in the order we try to specialise the functions, which is the order as obtained by for (Function &F : M) { ...
and if this order is reversed tests that mention cloned function names break.

[chill]

Updated the test like in D136755

[labrinea]

llvm/test/Transforms/FunctionSpecialization/literal-const.ll

87–93

Maybe add "define" before the specialized function names to make sure we match the definitions?

[ChuanqiXu]

LGTM if @labrinea's comment get addressed.

llvm/test/Transforms/FunctionSpecialization/literal-const.ll

87–93

Maybe add "define" before the specialized function names to make sure we match the definitions?

+1

llvm/lib/Transforms/IPO/FunctionSpecialization.cpp

private:
   /// would result in any significant optimization opportunities. If
   /// optimization opportunities exist, the constant values of \p A on which to
   /// specialize the function are collected in \p Constants.
   ///
   /// \returns true if the function should be specialized on the given
   /// argument.
   bool isArgumentInteresting(Argument *A,
                              SmallVectorImpl<CallArgBinding> &Constants) {
+
+    // No point in specialization if the argument is unused.
+    if (A->user_empty())
+      return false;
+
     // For now, don't attempt to specialize functions based on the values of
     // composite types.
-    if (!A->getType()->isSingleValueType() || A->user_empty())
+    Type *ArgTy = A->getType()                              ;
+    if (!ArgTy->isSingleValueType())
+      return false;
+
+    // Specialization of integer and floating point types needs to be explicitly enabled.
+    if (!EnableSpecializationForLiteralConstant &&
+        (ArgTy->isIntegerTy() || ArgTy->isFloatingPointTy()))
       return false;
 
-    // If the argument isn't overdefined, there's nothing to do. It should
-    // already be constant.
-    if (!Solver.getLatticeValueFor(A).isOverdefined()) {
+    // SCCP solver does not record an argument that will be constructed on
+    // stack.
+    if (A->hasByValAttr() && !A->getParent()->onlyReadsMemory())
+      return false;
+
+    // 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.
+    const ValueLatticeElement &LV = Solver.getLatticeValueFor(A);
+    if (LV.isUnknownOrUndef() || LV.isConstant() ||
+        (LV.isConstantRange() && LV.getConstantRange().isSingleElement())) {
       LLVM_DEBUG(dbgs() << "FnSpecialization: Nothing to do, argument "
-                        << A->getNameOrAsOperand()
-                        << " is already constant?\n");
+                        << A->getNameOrAsOperand() << " is already constant\n");
       return false;
     }
 
     // Collect the constant values that the argument can take on. If the
     // argument can't take on any constant values, we aren't going to
     // specialize the function. While it's possible to specialize the function
     // based on non-constant arguments, there's likely not much benefit to
     // constant propagation in doing so.
private:
   }
 
   /// Collect in \p Constants all the constant values that argument \p A can
   /// take on.
   void getPossibleConstants(Argument *A,
                             SmallVectorImpl<CallArgBinding> &Constants) {
     Function *F = A->getParent();
 
-    // SCCP solver does not record an argument that will be constructed on
-    // stack.
-    if (A->hasByValAttr() && !F->onlyReadsMemory())
-      return;
-
     // Iterate over all the call sites of the argument's parent function.
     for (User *U : F->users()) {
       if (!isa<CallInst>(U) && !isa<InvokeInst>(U))
         continue;
       auto &CS = *cast<CallBase>(U);
       // If the call site has attribute minsize set, that callsite won't be
       // specialized.
       if (CS.hasFnAttr(Attribute::MinSize))
for (User *U : F->users()) {
         // global values.
         if (!GV->isConstant() && !SpecializeOnAddresses)
           continue;
 
         if (!GV->getValueType()->isSingleValueType())
           continue;
       }
 
-      if (isa<Constant>(V) && (Solver.getLatticeValueFor(V).isConstant() ||
-                               EnableSpecializationForLiteralConstant))
-        Constants.push_back({&CS, cast<Constant>(V)});
+      // Select for possible specialisation arguments which are constants or
+      // are deduced to be constants or constant ranges with a single element.
+      Constant *C = dyn_cast<Constant>(V);
+      if (!C) {
+        const ValueLatticeElement &LV = Solver.getLatticeValueFor(V);
+        if (LV.isConstant())
+          C = LV.getConstant();
+        else if (LV.isConstantRange() &&
+                 LV.getConstantRange().isSingleElement()) {
+          assert(V->getType()->isIntegerTy() && "Non-integral constant range");
+          C = Constant::getIntegerValue(
+              V->getType(), *LV.getConstantRange().getSingleElement());
+        } else
+          continue;
+      }
+
+      Constants.push_back({&CS, C});
     }
   }
 
   /// Rewrite calls to function \p F to call function \p Clone instead.
   ///
   /// This function modifies calls to function \p F as long as the actual
   /// arguments match those in \p Args. Note that for recursive calls we
   /// need to compare against the cloned formal arguments.

llvm/test/Transforms/FunctionSpecialization/literal-const.ll

+; RUN: opt -S --passes=function-specialization \
+; RUN:        -force-function-specialization < %s | FileCheck %s -check-prefix CHECK-NOLIT
+; RUN: opt -S --passes=function-specialization \
+; RUN:        -function-specialization-for-literal-constant \
+; RUN:        -force-function-specialization < %s | FileCheck %s -check-prefix CHECK-LIT
+
+define i32 @f0(i32 noundef %x) {
+entry:
+  %call = tail call i32 @neg(i32 noundef %x, i1 noundef zeroext false)
+  ret i32 %call
+}
+
+define i32 @f1(i32 noundef %x) {
+entry:
+  %call = tail call i32 @neg(i32 noundef %x, i1 noundef zeroext true)
+  ret i32 %call
+}
+
+define i32 @g0(i32 noundef %x) {
+entry:
+  %call = tail call i32 @add(i32 noundef %x, i32 noundef 1)
+  ret i32 %call
+}
+
+define i32 @g1(i32 noundef %x) {
+entry:
+  %call = tail call i32 @add(i32 noundef %x, i32 noundef 2)
+  ret i32 %call
+}
+
+define float @h0(float noundef %x) {
+entry:
+  %call = tail call float @addf(float noundef %x, float noundef 1.000000e+00)
+  ret float %call
+}
+
+define float @h1(float noundef %x) {
+entry:
+  %call = tail call float @addf(float noundef %x, float noundef 2.000000e+00)
+  ret float %call
+}
+
+define internal i32 @neg(i32 noundef %x, i1 noundef zeroext %b) {
+entry:
+  %sub = sub nsw i32 0, %x
+  %cond = select i1 %b, i32 %sub, i32 %x
+  ret i32 %cond
+}
+
+define internal i32 @add(i32 noundef %x, i32 noundef %y) {
+entry:
+  %add = add nsw i32 %y, %x
+  ret i32 %add
+}
+
+define internal float @addf(float noundef %x, float noundef %y) {
+entry:
+  %add = fadd float %x, %y
+  ret float %add
+}
+
+
+; Check no functions were specialised.
+; CHECK-NOLIT-NOT: @neg.
+; CHECK-NOLIT-NOT: @add.
+; CHECK-NOLIT-NOT: @addf.
+
+; CHECK-LIT-LABEL: define i32 @f0
+; CHECK-LIT: call i32 @neg.[[#A:]]
+
+; CHECK-LIT-LABEL: define i32 @f1
+; CHECK-LIT: call i32 @neg.[[#B:]]
+
+; CHECK-LIT-LABEL: define i32 @g0
+; CHECK-LIT: call i32 @add.[[#C:]]
+
+; CHECK-LIT-LABEL: define i32 @g1
+; CHECK-LIT: call i32 @add.[[#D:]]
+
+; CHECK-LIT-LABEL: define float @h0
+; CHECK-LIT: call float @addf.[[#E:]]
+
+; CHECK-LIT-LABEL: define float @h1
+; CHECK-LIT: call float @addf.[[#F:]]
+
+; Check all of `neg`, `add`, and `addf` were specialised.
+; CHECK-LIT-DAG: @neg.[[#A]]
+; CHECK-LIT-DAG: @neg.[[#B]]
+; CHECK-LIT-DAG: @add.[[#C]]
+; CHECK-LIT-DAG: @add.[[#D]]
+; CHECK-LIT-DAG: @addf.[[#E]]
+; CHECK-LIT-DAG: @addf.[[#F]]