[FuncSpec] Decouple Cost and Benefit analysis, to sort candidates. NFC.

[SjoerdMeijer]

Summary

This looks like a big change/rewrite, but it is not as bad as it looks as it mostly is the same code that is refactored to decouple the cost and benefit analysis. The biggest change is top-level function specializeFunctions that now drives the transformation likes this:

specializeFunctions() {
  Cost = getSpecializationCost(F);
  calculateGains(F, Cost);
  specializeFunction(F);
}

while this is just a restructuring and decoupling of the cost and benefit analysis, this separation helps the actual functional change in calculateGains. We now sort the candidates based on the expected specialisation gain, which we didn't do before. For this, a book keeping struct ArgInfo was introduced. If we have a list of N candidates, but we only want specialise less than N as set by option -func-specialization-max-constants, we sort the list and discard the candidates that give the least benefit.

Given a formal argument, this change results in selecting the best actual argument(s). In a follow up, I want to go one step further and compare all functions and all arguments. But that will mostly build on top of this refactoring and change, and will be less change; this is enough change for now.

[snehasish]

llvm/lib/Transforms/IPO/FunctionSpecialization.cpp

98

typo "specialization"

396

nit: We could avoid calling getSpecializationBonus if ForceFunctionSpecialization is true.

438

Use a reference to avoid copies here?

467

nit: just break here instead of return?

llvm/test/Transforms/FunctionSpecialization/function-specialization4.ll

41 ↗

(On Diff #393230)

It wasn't clear to me what changed in the code that this function is now specialized. Is it possible to split this into a refactor without functionality change and a smaller patch with the behaviour change?

[ChuanqiXu]

Do we notice the score/code size change in SPEC after this patch?

llvm/lib/Transforms/IPO/FunctionSpecialization.cpp

97

For a class used only in a cpp file, it may be better to wrap it into anonymous namespace to make sure the linkage is internal.

438

We could use llvm::sort here to use range style sort.

443–450

A note: this change is not NFC (I'm fine with this change) @snehasish

452

How about Worklist.size() < ActualConstArg.size()?

465

It might be better to:

Clone->getArg(AI.Arg->getArgNo);

[SjoerdMeijer]

Thanks for the reviews!
I have replied to 2 comments inline while I work on addressing the other comments.

llvm/lib/Transforms/IPO/FunctionSpecialization.cpp

396

Good point, that will help compile-times. We still need to set Gain though, but what I will do is initialise Gain to the maximum value with ~0 if specialisation is forced.

llvm/test/Transforms/FunctionSpecialization/function-specialization4.ll

41 ↗

(On Diff #393230)

Okay, good point. I changed the behaviour how MaxConstantsThreshold is interpreted.

Before, when the number of candidates in the worklist exceeded MaxConstantsThreshold, no specialisation was done at all. I have changed that to specialise up to MaxConstantsThreshold candidates, which seems to me what you would expect from MaxConstantsThreshold. After sorting the candidates on the maximum gain, it was this code getting rid of the remaining candidates with the lower gains:

// Truncate the worklist to 'MaxConstantsThreshold' candidates if                               
// necessary.
if (Worklist.size() > MaxConstantsThreshold) {                                                  
  LLVM_DEBUG(dbgs() << "FnSpecialization: number of constants exceed "                          
                    << "the maximum number of constants threshold.\n"                           
                    << "Truncating worklist to " << MaxConstantsThreshold                       
                    << " candidates.\n");
  Worklist.erase(Worklist.begin() + MaxConstantsThreshold,
                 Worklist.end());
}

I included this change in behaviour here, because otherwise the sorting doesn't serve a purpose if we are going to keep all candidates anyway. But given that I am going to follow up on this, I will keep the sorting but move this change in behaviour in a follow up and dependent patch.

To finally answer your question about this test change: there are 2 candidates here, but -func-specialization-max-constants=1. With the old behaviour, this means no specialisation at all like I mentioned before. With the new interpretation, this specialises up to 1 candidate, which is what you see here.

[SjoerdMeijer]

This addresses the other comments. This should now be NFC, so have added this to the Title.
I will double check the SPEC score remains the same, and will follow up with the functional change that was present of the first revision of this patch which caused the change in the test.

[ChuanqiXu]

In fact, this might not be a NFC change since it would try to sort the candidates.I would love to see the score/code size change after D115509 applied.

[ChuanqiXu]

BTW, it might be better to add a test case to show the sort matters.

[SjoerdMeijer]

I can confirm that the SPEC score with this patch and also with D115509 remains the same, as expected and should be.

This is NFC as the resulting code should be the same, which is the definition of NFC I use. But yeah, I do agree of course that under the hood things work slightly different, although that will only have an affect in subsequent patches. Anyway, happy to remove NFC from the title.

BTW, it might be better to add a test case to show the sort matters.

So it shouldn't matter here, but it's a good point and I will add an test-case for this to D115509 as the affect will only be visible there.

[ChuanqiXu]

If the score remains the same, it looks good to me. BTW, I think the change of sort matters. Previously, we banned a lot of cases which could have been specialized due to the consideration of the cost for code size and compilation time. But if we could sort them and assume we could have a good enough cost model (the current one might be coarse), then we could relax the restrictions we made previously to make function specialization more useful.
Nit: format codes before committing

(I suggest to commit this with D115509 in one shot, but it doesn't really matter)

[snehasish]

lgtm

llvm/test/Transforms/FunctionSpecialization/function-specialization4.ll

41 ↗

(On Diff #393230)

Thanks for clarifying the change.

llvm/lib/Transforms/IPO/FunctionSpecialization.cpp

 
 // TODO: This needs checking to see the impact on compile-times, which is why
 // this is off by default for now.
 static cl::opt<bool> EnableSpecializationForLiteralConstant(
     "function-specialization-for-literal-constant", cl::init(false), cl::Hidden,
     cl::desc("Enable specialization of functions that take a literal constant "
              "as an argument."));
 
+namespace {
+// Bookkeeping struct to pass data from the analysis and profitability phase
+// to the actual transform helper functions.
+struct ArgInfo {
+  Function *Fn;         // The function to perform specialisation on.
+  Argument *Arg;        // The Formal argument being analysed.
+  Constant *Const;      // A corresponding actual constant argument.
+  InstructionCost Gain; // Profitability: Gain = Bonus - Cost.
+
+  // Flag if this will be a partial specialization, in which case we will need
+  // to keep the original function around in addition to the added
+  // specializations.
+  bool Partial = false;
+
+  ArgInfo(Function *F, Argument *A, Constant *C, InstructionCost G)
+      : Fn(F), Arg(A), Const(C), Gain(G){};
+};
+} // Anonymous namespace
+
 // Helper to check if \p LV is either a constant or a constant
 // range with a single element. This should cover exactly the same cases as the
 // old ValueLatticeElement::isConstant() and is intended to be used in the
 // transition to ValueLatticeElement.
 static bool isConstant(const ValueLatticeElement &LV) {
   return LV.isConstant() ||
          (LV.isConstantRange() && LV.getConstantRange().isSingleElement());
 }
public:
 
   /// Attempt to specialize functions in the module to enable constant
   /// propagation across function boundaries.
   ///
   /// \returns true if at least one function is specialized.
   bool
   specializeFunctions(SmallVectorImpl<Function *> &FuncDecls,
                       SmallVectorImpl<Function *> &CurrentSpecializations) {
-
-    // Attempt to specialize the argument-tracked functions.
     bool Changed = false;
     for (auto *F : FuncDecls) {
-      if (specializeFunction(F, CurrentSpecializations)) {
-        Changed = true;
-        LLVM_DEBUG(dbgs() << "FnSpecialization: Can specialize this func.\n");
-      } else {
+      if (!isCandidateFunction(F, CurrentSpecializations))
+        continue;
+
+      auto Cost = getSpecializationCost(F);
+      if (!Cost.isValid()) {
         LLVM_DEBUG(
-            dbgs() << "FnSpecialization: Cannot specialize this func.\n");
+            dbgs() << "FnSpecialization: Invalid specialisation cost.\n");
+        continue;
+      }
+
+      auto ConstArgs = calculateGains(F, Cost);
+      if (ConstArgs.empty()) {
+        LLVM_DEBUG(dbgs() << "FnSpecialization: no possible constants found\n");
+        continue;
+      }
+
+      for (auto &CA : ConstArgs) {
+        specializeFunction(CA, CurrentSpecializations);
+        Changed = true;
       }
     }
 
     for (auto *SpecializedFunc : CurrentSpecializations) {
       SpecializedFuncs.insert(SpecializedFunc);
 
       // Initialize the state of the newly created functions, marking them
       // argument-tracked and executable.
private:
   /// the SCCPSolver in the cloned version.
   Function *cloneCandidateFunction(Function *F) {
     ValueToValueMapTy EmptyMap;
     Function *Clone = CloneFunction(F, EmptyMap);
     removeSSACopy(*Clone);
     return Clone;
   }
 
-  /// This function decides whether to specialize function \p F based on the
-  /// known constant values its arguments can take on. Specialization is
-  /// performed on the first interesting argument. Specializations based on
-  /// additional arguments will be evaluated on following iterations of the
-  /// main IPSCCP solve loop. \returns true if the function is specialized and
-  /// false otherwise.
-  bool specializeFunction(Function *F,
-                          SmallVectorImpl<Function *> &Specializations) {
+  /// This function decides whether it's worthwhile to specialize function \p F
+  /// based on the known constant values its arguments can take on, i.e. it
+  /// calculates a gain and returns a list of actual arguments that are deemed
+  /// profitable to specialize. Specialization is performed on the first
+  /// interesting argument. Specializations based on additional arguments will
+  /// be evaluated on following iterations of the main IPSCCP solve loop.
+  SmallVector<ArgInfo> calculateGains(Function *F, InstructionCost Cost) {
+    SmallVector<ArgInfo> Worklist;
+    // 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()) {
+      LLVM_DEBUG(dbgs() << "FnSpecialization: Analysing arg: "
+                        << FormalArg.getName() << "\n");
+      // Determine if this argument is interesting. If we know the argument can
+      // take on any constant values, they are collected in Constants. If the
+      // argument can only ever equal a constant value in Constants, the
+      // function will be completely specialized, and the IsPartial flag will
+      // be set to false by isArgumentInteresting (that function only adds
+      // values to the Constants list that are deemed profitable).
+      bool IsPartial = true;
+      SmallVector<Constant *> ActualConstArg;
+      if (!isArgumentInteresting(&FormalArg, ActualConstArg, IsPartial)) {
+        LLVM_DEBUG(dbgs() << "FnSpecialization: Argument is not interesting\n");
+        continue;
+      }
+
+      for (auto *ActualArg : ActualConstArg) {
+        InstructionCost Gain =
+            ForceFunctionSpecialization
+                ? 1
+                : getSpecializationBonus(&FormalArg, ActualArg) - Cost;
+
+        if (Gain <= 0)
+          continue;
+        Worklist.push_back({F, &FormalArg, ActualArg, Gain});
+      }
+
+      if (Worklist.empty())
+        continue;
+
+      // Sort the candidates in descending order.
+      llvm::sort(Worklist,
+                 [](ArgInfo &L, ArgInfo &R) { return L.Gain > R.Gain; });
+
+      // TODO: truncate the worklist to 'MaxConstantsThreshold' candidates if
+      // necessary.
+      if (Worklist.size() > MaxConstantsThreshold) {
+        Worklist.clear();
+        continue;
+      }
+
+      if (IsPartial || Worklist.size() < ActualConstArg.size())
+        for (auto &ActualArg : Worklist)
+          ActualArg.Partial = true;
+
+      LLVM_DEBUG(dbgs() << "Sorted list of candidates by gain:\n";
+                 for (auto &C
+                      : Worklist) {
+                   dbgs() << "- Function = " << C.Fn->getName() << ", ";
+                   dbgs() << "FormalArg = " << C.Arg->getName() << ", ";
+                   dbgs() << "ActualArg = " << C.Const->getName() << ", ";
+                   dbgs() << "Gain = " << C.Gain << "\n";
+                 });
+
+      // FIXME: Only one argument per function.
+      break;
+    }
+    return Worklist;
+  }
 
+  bool isCandidateFunction(Function *F,
+                           SmallVectorImpl<Function *> &Specializations) {
     // Do not specialize the cloned function again.
     if (SpecializedFuncs.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;
 
     // Exit if the function is not executable. There's no point in specializing
     // a dead function.
     if (!Solver.isBlockExecutable(&F->getEntryBlock()))
       return false;
 
     // It wastes time to specialize a function which would get inlined finally.
     if (F->hasFnAttribute(Attribute::AlwaysInline))
       return false;
 
     LLVM_DEBUG(dbgs() << "FnSpecialization: Try function: " << F->getName()
                       << "\n");
-
-    // Determine if it would be profitable to create a specialization of the
-    // function where the argument takes on the given constant value. If so,
-    // add the constant to Constants.
-    auto FnSpecCost = getSpecializationCost(F);
-    if (!FnSpecCost.isValid()) {
-      LLVM_DEBUG(dbgs() << "FnSpecialization: Invalid specialisation cost.\n");
-      return false;
-    }
-
-    LLVM_DEBUG(dbgs() << "FnSpecialization: func specialisation cost: ";
-               FnSpecCost.print(dbgs()); dbgs() << "\n");
-
-    // 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 &A : F->args()) {
-      LLVM_DEBUG(dbgs() << "FnSpecialization: Analysing arg: " << A.getName()
-                        << "\n");
-      // True if this will be a partial specialization. We will need to keep
-      // the original function around in addition to the added specializations.
-      bool IsPartial = true;
-
-      // Determine if this argument is interesting. If we know the argument can
-      // take on any constant values, they are collected in Constants. If the
-      // argument can only ever equal a constant value in Constants, the
-      // function will be completely specialized, and the IsPartial flag will
-      // be set to false by isArgumentInteresting (that function only adds
-      // values to the Constants list that are deemed profitable).
-      SmallVector<Constant *, 4> Constants;
-      if (!isArgumentInteresting(&A, Constants, FnSpecCost, IsPartial)) {
-        LLVM_DEBUG(dbgs() << "FnSpecialization: Argument is not interesting\n");
-        continue;
-      }
+    return true;
 
-      assert(!Constants.empty() && "No constants on which to specialize");
-      LLVM_DEBUG(dbgs() << "FnSpecialization: Argument is interesting!\n"
-                        << "FnSpecialization: Specializing '" << F->getName()
-                        << "' on argument: " << A << "\n"
-                        << "FnSpecialization: Constants are:\n\n";
-                 for (unsigned I = 0; I < Constants.size(); ++I) dbgs()
-                 << *Constants[I] << "\n";
-                 dbgs() << "FnSpecialization: End of constants\n\n");
-
-      // Create a version of the function in which the argument is marked
-      // constant with the given value.
-      for (auto *C : Constants) {
-        // Clone the function. We leave the ValueToValueMap empty to allow
-        // IPSCCP to propagate the constant arguments.
-        Function *Clone = cloneCandidateFunction(F);
-        Argument *ClonedArg = Clone->arg_begin() + A.getArgNo();
+  void specializeFunction(ArgInfo &AI,
+                          SmallVectorImpl<Function *> &Specializations) {
+    Function *Clone = cloneCandidateFunction(AI.Fn);
+    Argument *ClonedArg = Clone->getArg(AI.Arg->getArgNo());
 
-        // Rewrite calls to the function so that they call the clone instead.
-        rewriteCallSites(F, Clone, *ClonedArg, C);
+    rewriteCallSites(AI.Fn, Clone, *ClonedArg, AI.Const);
 
-        // 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.markArgInFuncSpecialization(F, ClonedArg, C);
+    Solver.markArgInFuncSpecialization(AI.Fn, ClonedArg, AI.Const);
 
-        // Mark all the specialized functions
-        Specializations.push_back(Clone);
-        NbFunctionsSpecialized++;
-      }
 
-      // If the function has been completely specialized, the original function
-      // is no longer needed. Mark it unreachable.
-      if (!IsPartial)
-        Solver.markFunctionUnreachable(F);
-
-      // FIXME: Only one argument per function.
-      return true;
-    }
-
-    return false;
+    if (!AI.Partial)
+      Solver.markFunctionUnreachable(AI.Fn);
   }
 
-  /// Compute the cost of specializing function \p F.
+  /// Compute and return the cost of specializing function \p F.
   InstructionCost getSpecializationCost(Function *F) {
     // Compute the code metrics for the function.
     SmallPtrSet<const Value *, 32> EphValues;
     CodeMetrics::collectEphemeralValues(F, &(GetAC)(*F), EphValues);
     CodeMetrics Metrics;
     for (BasicBlock &BB : *F)
       Metrics.analyzeBasicBlock(&BB, (GetTTI)(*F), EphValues);
 
private:
   /// \p Constants represent the complete set of values that \p A can take on,
   /// the function will be completely specialized, and the \p IsPartial flag is
   /// set to false.
   ///
   /// \returns true if the function should be specialized on the given
   /// argument.
   bool isArgumentInteresting(Argument *A,
                              SmallVectorImpl<Constant *> &Constants,
-                             const InstructionCost &FnSpecCost,
                              bool &IsPartial) {
     // For now, don't attempt to specialize functions based on the values of
     // composite types.
     if (!A->getType()->isSingleValueType() || A->user_empty())
       return false;
 
     // If the argument isn't overdefined, there's nothing to do. It should
     // already be constant.
bool isArgumentInteresting(Argument *A,
     //
     // 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.
     //
-    SmallVector<Constant *, 4> PossibleConstants;
-    bool AllConstant = getPossibleConstants(A, PossibleConstants);
-    if (PossibleConstants.empty()) {
-      LLVM_DEBUG(dbgs() << "FnSpecialization: no possible constants found\n");
-      return false;
-    }
-    if (PossibleConstants.size() > MaxConstantsThreshold) {
-      LLVM_DEBUG(dbgs() << "FnSpecialization: number of constants found exceed "
-                        << "the maximum number of constants threshold.\n");
-      return false;
-    }
-
-    for (auto *C : PossibleConstants) {
-      LLVM_DEBUG(dbgs() << "FnSpecialization: Constant: " << *C << "\n");
-      if (ForceFunctionSpecialization) {
-        LLVM_DEBUG(dbgs() << "FnSpecialization: Forced!\n");
-        Constants.push_back(C);
-        continue;
-      }
-      if (getSpecializationBonus(A, C) > FnSpecCost) {
-        LLVM_DEBUG(dbgs() << "FnSpecialization: profitable!\n");
-        Constants.push_back(C);
-      } else {
-        LLVM_DEBUG(dbgs() << "FnSpecialization: not profitable\n");
-      }
-    }
-
-    // None of the constant values the argument can take on were deemed good
-    // candidates on which to specialize the function.
-    if (Constants.empty())
-      return false;
-
-    // This will be a partial specialization if some of the constants were
-    // rejected due to their profitability.
-    IsPartial = !AllConstant || PossibleConstants.size() != Constants.size();
-
+    IsPartial = !getPossibleConstants(A, Constants);
+    LLVM_DEBUG(dbgs() << "FnSpecialization: interesting arg: " << *A << "\n");
     return true;
   }
 
   /// Collect in \p Constants all the constant values that argument \p A can
   /// take on.
   ///
   /// \returns true if all of the values the argument can take on are constant
   /// (e.g., the argument's parent function cannot be called with an
for (User *U : F->users()) {
         continue;
 
       auto *V = CS.getArgOperand(A->getArgNo());
       if (isa<PoisonValue>(V))
         return false;
 
       // For now, constant expressions are fine but only if they are function
       // calls.
-      if (auto *CE =  dyn_cast<ConstantExpr>(V))
+      if (auto *CE = dyn_cast<ConstantExpr>(V))
         if (!isa<Function>(CE->getOperand(0)))
           return false;
 
       // 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())