[FuncSpec][NFC] Refactor finding specialisation opportunities

[chill]

Summary

This patch reorders the traversal of function call sites and function
formal parameters to:

• do various argument feasibility checks (isArgumentInteresting ) only once per argument, i.e. doing N-args checks instead of N-calls x N-args checks.

• do hash table lookups only once per call site, i.e. N-calls lookups/inserts instead of N-call x N-args lookups/inserts.

[labrinea]

llvm/lib/Transforms/IPO/FunctionSpecialization.cpp

464–465

This confused me a bit and then I remembered it's a MapVector. LGTM!

[ChuanqiXu]

llvm/lib/Transforms/IPO/FunctionSpecialization.cpp

671–673

Why this check missed?

[chill]

llvm/lib/Transforms/IPO/FunctionSpecialization.cpp

671–673

Oops, that's accidental. I'll put it back.

[SjoerdMeijer]

Yeah, nice patch.

Just a few nits inlined when I was reading the patch.

According to my definition of NFC, this is not so NFC as things are done quite differently now. I appreciate the codegen should be the same, which is the definition of NFC that some use I believe.

Just out of curiousity, did you measure compile times improvements with this, which I guess is the reason for doing this?

llvm/lib/Transforms/IPO/FunctionSpecialization.cpp

429

Nit: you can also drop these brackets I think

434

Nit: not sure the coding standard says anything about comparing against zero, but personally it's a bit verbose for me and I prefer !Args.size().

671–673

Looks like we are missing a test case? :)

682

After reducing/changing the comments, the function name and comment are a bit out of sync now for me, because specialisation should not only "possible" (the new comment), but also profitable which is why "interesting" is in the function name. Perhaps restoring some words about the profitability explains the "interesting" part.

[chill]

Unfortunately, this patch exposes a latent issue in the pass - one call site matching two or more specialisations - as can be seen in the failure of llvm/test/Transforms/FunctionSpecialization/identical-specializations.ll.

llvm/lib/Transforms/IPO/FunctionSpecialization.cpp

671–673

Ack!

682

In fact, the old comment was out of sync - neither the old isArgumentInteresting nor the old getPossibleConstants did any cost/benefit analysis.

[chill]

llvm/lib/Transforms/IPO/FunctionSpecialization.cpp

671–673

Testcase added in D136184

[chill]

In D135968#3861974, @SjoerdMeijer wrote:

According to my definition of NFC, this is not so NFC as things are done quite differently now. I appreciate the codegen should be the same, which is the definition of NFC that some use I believe.

AFAIK, the working notion of NFC in LLVM is "does not change any testcase", which is a bit of a compromise, of course. 😁

I think it can try specialisations of the same cost in a different order, which has the potential of changing the codegen, however, it does not
violate any explicit guarantees we make in the pass.

Just out of curiousity, did you measure compile times improvements with this, which I guess is the reason for doing this?

No, not measured yet, but planning to.

[ChuanqiXu]

LGTM.

[chill]

In D135968#3865784, @chill wrote:

Just out of curiousity, did you measure compile times improvements with this, which I guess is the reason for doing this?

No, not measured yet, but planning to.

I did some measurements on sqlite3 (with the whole series ending in D136332), with -O3 and -O3 -mllvm -enable-function-specialization).
I got about 0.42% overall regression compared to not running the FunctionSpecialization pass and about 0.45% improvement when the pass is enabled in both compilers.
Comparing just the pass (via clang ... -ftime-report 2>&1 | grep FunctionSpecialization) shows about 2.2% improvement.

[labrinea]

This revision looks similar to 467796, which must be accidental. Can you update it as it was when it landed? (470754)

llvm/lib/Transforms/IPO/FunctionSpecialization.cpp

for (auto *F : Candidates) {
             dbgs() << "FnSpecialization: Invalid specialization cost.\n");
         continue;
       }
 
       LLVM_DEBUG(dbgs() << "FnSpecialization: Specialization cost for "
                         << F->getName() << " is " << Cost << "\n");
 
       SmallVector<CallSpecBinding, 8> Specializations;
-      if (!calculateGains(F, Cost, Specializations)) {
-        LLVM_DEBUG(dbgs() << "FnSpecialization: No possible constants found\n");
+      if (!findSpecializations(F, Cost, Specializations)) {
+        LLVM_DEBUG(
+            dbgs() << "FnSpecialization: No possible specializations found\n");
         continue;
       }
 
       Changed = true;
       for (auto &Entry : Specializations)
         specializeFunction(F, Entry.second, WorkList);
     }
 
private:
   }
 
   /// This function decides whether it's worthwhile to specialize function
   /// \p F based on the known constant values its arguments can take on. It
   /// only discovers potential specialization opportunities without actually
   /// applying them.
   ///
   /// \returns true if any specializations have been found.
-  bool calculateGains(Function *F, InstructionCost Cost,
+  bool findSpecializations(Function *F, InstructionCost Cost,
                            SmallVectorImpl<CallSpecBinding> &WorkList) {
+    // Get a list of interesting arguments.
+    SmallVector<Argument *, 4> Args;
+    for (Argument &Arg : F->args())
+      if (isArgumentInteresting(&Arg))
+        Args.push_back(&Arg);
+
+    if (!Args.size())
+      return false;
+
+    // Find all the call sites for the function.
     SpecializationMap Specializations;
-    // Determine if we should specialize the function based on the values the
-    // argument can take on. If specialization is not profitable, we continue
-    // on to the next argument.
-    for (Argument &FormalArg : F->args()) {
-      // Determine if this argument is interesting. If we know the argument can
-      // take on any constant values, they are collected in Constants.
-      SmallVector<CallArgBinding, 8> ActualArgs;
-      if (!isArgumentInteresting(&FormalArg, ActualArgs)) {
-        LLVM_DEBUG(dbgs() << "FnSpecialization: Argument "
-                          << FormalArg.getNameOrAsOperand()
-                          << " is not interesting\n");
+    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))
+        continue;
+
+      // If the parent of the call site will never be executed, we don't need
+      // to worry about the passed value.
+      if (!Solver.isBlockExecutable(CS.getParent()))
         continue;
+
+      // Examine arguments and create specialization candidates from call sites
+      // with constant arguments.
+      bool Added = false;
+      for (Argument *A : Args) {
+        Constant *C = getCandidateConstant(CS.getArgOperand(A->getArgNo()));
+        if (!C)
+          continue;
+
+        if (!Added) {
+          Specializations[&CS] = {{}, 0 - Cost};
+          Added = true;
         }
 
-      for (const auto &Entry : ActualArgs) {
-        CallBase *Call = Entry.first;
-        Constant *ActualArg = Entry.second;
-
-        auto I = Specializations.insert({Call, SpecializationInfo()});
-        SpecializationInfo &S = I.first->second;
-
-        if (I.second)
-          S.Gain = 0 - Cost;
-        S.Gain += getSpecializationBonus(&FormalArg, ActualArg);
-        S.Args.push_back({&FormalArg, ActualArg});
+        SpecializationInfo &S = Specializations.back().second;
+        S.Gain += getSpecializationBonus(A, C);
+        S.Args.push_back({A, C});
       }
+      Added = false;
     }
 
     // Remove unprofitable specializations.
     if (!ForceFunctionSpecialization)
       Specializations.remove_if(
           [](const auto &Entry) { return Entry.second.Gain <= 0; });
 
     // Clear the MapVector and return the underlying vector.
for (User *U : A->users()) {
 
       LLVM_DEBUG(dbgs() << "FnSpecialization:   Inlining bonus " << Bonus
                         << " for user " << *U << "\n");
     }
 
     return TotalCost + Bonus;
   }
 
-  /// Determine if we should specialize a function based on the incoming values
-  /// of the given argument.
-  ///
-  /// This function implements the goal-directed heuristic. It determines if
-  /// specializing the function based on the incoming values of argument \p A
-  /// 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) {
-
+  /// Determine if it is possible to specialise the function for constant values
+  /// of the formal parameter \p A.
+  bool isArgumentInteresting(Argument *A) {
     // 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.
-    Type *ArgTy = A->getType()                              ;
+    Type *ArgTy = A->getType();
     if (!ArgTy->isSingleValueType())
       return false;
 
-    // Specialization of integer and floating point types needs to be explicitly enabled.
+    // Specialization of integer and floating point types needs to be explicitly
+    // enabled.
     if (!EnableSpecializationForLiteralConstant &&
         (ArgTy->isIntegerTy() || ArgTy->isFloatingPointTy()))
       return false;
 
     // 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");
       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.
-    //
-    // TODO 1: currently it won't specialize if there are over the threshold of
-    // calls using the same argument, e.g foo(a) x 4 and foo(b) x 1, but it
-    // might be beneficial to take the occurrences into account in the cost
-    // model, so we would need to find the unique constants.
-    //
-    // TODO 2: this currently does not support constants, i.e. integer ranges.
-    //
-    getPossibleConstants(A, Constants);
-
-    if (Constants.empty())
-      return false;
-
-    LLVM_DEBUG(dbgs() << "FnSpecialization: Found interesting argument "
-                      << A->getNameOrAsOperand() << "\n");
     return true;
   }
 
-  /// 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();
-
-    // 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))
-        continue;
-
-      // If the parent of the call site will never be executed, we don't need
-      // to worry about the passed value.
-      if (!Solver.isBlockExecutable(CS.getParent()))
-        continue;
-
-      auto *V = CS.getArgOperand(A->getArgNo());
-      if (isa<PoisonValue>(V))
-        continue;
+  /// Check if the valuy \p V  (an actual argument) is a constant or can only
+  /// have a constant value. Return that constant.
+  Constant *getCandidateConstant(Value *V) {
+      return nullptr;
 
-      // TrackValueOfGlobalVariable only tracks scalar global variables.
-      if (auto *GV = dyn_cast<GlobalVariable>(V)) {
-        // Check if we want to specialize on the address of non-constant
-        // global values.
-        if (!GV->isConstant() && !SpecializeOnAddresses)
-          continue;
+        return nullptr;
 
-        if (!GV->getValueType()->isSingleValueType())
-          continue;
-      }
+        return nullptr;
 
-      // 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;
-      }
+    // Select for possible specialisation values that are constants or
+        return nullptr;
 
-      Constants.push_back({&CS, C});
-    }
+    LLVM_DEBUG(dbgs() << "FnSpecialization: Found interesting argument "
+                      << V->getNameOrAsOperand() << "\n");
+
+    return 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.
   ///