[FuncSpec] Global ranking of specialisations

The `FunctionSpecialization` pass chooses specializations among the
opportunities presented by a single function and its calls,
progressively penalizing subsequent specialization attempts by
artificially increasing the cost of a specialization, depending on how
many specialization were applied before. Thus the chosen
specializations are sensitive to the order the functions appear in the
module and may be worse than others, had those others been considered
earlier.

This patch makes the `FunctionSpecialization` pass rank the
specializations globally, i.e.  choose the "best" specializations
among the all possible specializations in the module, for all
functions.

Since this involved quite a bit of redesign of the pass data
structures, this patch also carries:

  * removal of duplicate specializations

  * optimization of call sites update, by collecting per
    specialization the list of call sites that can be directly
    rewritten, without prior expensive check if the call constants and
    their positions match those of the specialized function.

A bit of a write-up up about the FuncSpec data structures and
operation:

Each potential function specialisation is kept in a single vector
(`AllSpecs` in `FunctionSpecializer::run`).  This vector is populated
by `FunctionSpecializer::findSpecializations`.

The `findSpecializations` member function has a local `DenseMap` to
eliminate duplicates - with each call to the current function,
`findSpecializations` builds a specialisation signature (`SpecSig`)
and looks it in the duplicates map. If the signature is present, the
function records the call to rewrite into the existing specialisation
instance.  If the signature is absent, it means we have a new
specialisation instance - the function calculates the gain and creates
a new entry in `AllSpecs`. Negative gain specialisation are ignored at
this point, unless forced.

The potential specialisations for a function form a contiguous range
in the `AllSpecs` [1]. This range is recorded in `SpecMap SM`, so we
can quickly find all specialisations for a function.

Once we have all the potential specialisations with their gains we
need to choose the best ones, which fit in the module specialisation
budget. This is done by using a max-heap (`std::make_heap`,
`std::push_heap`, etc) to find the best `NSpec` specialisations with a
single traversal of the `AllSpecs` vector. The heap itself is
contained with a small vector (`BestSpecs`) of indices into
`AllSpecs`, since elements of `AllSpecs` are a bit too heavy to
shuffle around.

Next the chosen specialisation are performed, that is, functions
cloned, `SCCPSolver` primed, and known call sites updated.

Then we run the `SCCPSolver` to propagate constants in the cloned
functions, after which we walk the calls of the original functions to
update them to call the specialised functions.

---

[1] This range may contain specialisation that were discarded and is
not ordered in any way. One alternative design is to keep a vector
indices of all specialisations for this function (which would
initially be, `i`, `i+1`, `i+2`, etc) and later sort them by gain,
pushing non-applied ones to the back. This has the potential to speed
`updateCallSites` up.

Reviewed By: ChuanqiXu, labrinea

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

Change-Id: I708851eb38f07c42066637085b833ca91b195998

Author: momchil-velikov

llvm/include/llvm/Transforms/IPO/FunctionSpecialization.h

@@ -60,18 +60,57 @@
 using namespace llvm;
 
 namespace llvm {
-// Bookkeeping struct to pass data from the analysis and profitability phase
-// to the actual transform helper functions.
-struct SpecializationInfo {
-  SmallVector<ArgInfo, 8> Args; // Stores the {formal,actual} argument pairs.
-  InstructionCost Gain;         // Profitability: Gain = Bonus - Cost.
-  Function *Clone;              // The definition of the specialized function.
+// Specialization signature, used to uniquely designate a specialization within
+// a function.
+struct SpecSig {
+  // Hashing support, used to distinguish between ordinary, empty, or tombstone
+  // keys.
+  unsigned Key = 0;
+  SmallVector<ArgInfo, 4> Args;
+
+  bool operator==(const SpecSig &Other) const {
+    if (Key != Other.Key || Args.size() != Other.Args.size())
+      return false;
+    for (size_t I = 0; I < Args.size(); ++I)
+      if (Args[I] != Other.Args[I])
+        return false;
+    return true;
+  }
+
+  friend hash_code hash_value(const SpecSig &S) {
+    return hash_combine(hash_value(S.Key),
+                        hash_combine_range(S.Args.begin(), S.Args.end()));
+  }
+};
+
+// Specialization instance.
+struct Spec {
+  // Original function.
+  Function *F;
+
+  // Cloned function, a specialized version of the original one.
+  Function *Clone = nullptr;
+
+  // Specialization signature.
+  SpecSig Sig;
+
+  // Profitability of the specialization.
+  InstructionCost Gain;
+
+  // List of call sites, matching this specialization.
+  SmallVector<CallBase *> CallSites;
+
+  Spec(Function *F, const SpecSig &S, InstructionCost G)
+      : F(F), Sig(S), Gain(G) {}
+  Spec(Function *F, const SpecSig &&S, InstructionCost G)
+      : F(F), Sig(S), Gain(G) {}
 };
 
-using CallSpecBinding = std::pair<CallBase *, SpecializationInfo>;
-// We are using MapVector because it guarantees deterministic iteration
-// order across executions.
-using SpecializationMap = SmallMapVector<CallBase *, SpecializationInfo, 8>;
+// Map of potential specializations for each function. The FunctionSpecializer
+// keeps the discovered specialisation opportunities for the module in a single
+// vector, where the specialisations of each function form a contiguous range.
+// This map's value is the beginning and the end of that range.
+using SpecMap = DenseMap<Function *, std::pair<unsigned, unsigned>>;
 
 class FunctionSpecializer {
 
@@ -137,18 +176,23 @@ class FunctionSpecializer {
   // Compute the code metrics for function \p F.
   CodeMetrics &analyzeFunction(Function *F);
 
-  /// 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.
+  /// @brief  Find potential specialization opportunities.
+  /// @param F Function to specialize
+  /// @param Cost Cost of specializing a function. Final gain is this cost
+  /// minus benefit
+  /// @param AllSpecs A vector to add potential specializations to.
+  /// @param SM  A map for a function's specialisation range
+  /// @return True, if any potential specializations were found
   bool findSpecializations(Function *F, InstructionCost Cost,
-                           SmallVectorImpl<CallSpecBinding> &WorkList);
+                           SmallVectorImpl<Spec> &AllSpecs, SpecMap &SM);
 
   bool isCandidateFunction(Function *F);
 
-  Function *createSpecialization(Function *F, CallSpecBinding &Specialization);
+  /// @brief Create a specialization of \p F and prime the SCCPSolver
+  /// @param F Function to specialize
+  /// @param S Which specialization to create
+  /// @return The new, cloned function
+  Function *createSpecialization(Function *F, const SpecSig &S);
 
   /// Compute and return the cost of specializing function \p F.
   InstructionCost getSpecializationCost(Function *F);
@@ -165,9 +209,11 @@ class FunctionSpecializer {
   /// have a constant value. Return that constant.
   Constant *getCandidateConstant(Value *V);
 
-  /// Redirects callsites of function \p F to its specialized copies.
-  void updateCallSites(Function *F,
-                       SmallVectorImpl<CallSpecBinding> &Specializations);
+  /// @brief Find and update calls to \p F, which match a specialization
+  /// @param F Orginal function
+  /// @param Begin Start of a range of possibly matching specialisations
+  /// @param End End of a range (exclusive) of possibly matching specialisations
+  void updateCallSites(Function *F, const Spec *Begin, const Spec *End);
 };
 } // namespace llvm
 

llvm/include/llvm/Transforms/Utils/SCCPSolver.h

@@ -52,7 +52,17 @@ struct ArgInfo {
   Argument *Formal; // The Formal argument being analysed.
   Constant *Actual; // A corresponding actual constant argument.
 
-  ArgInfo(Argument *F, Constant *A) : Formal(F), Actual(A){};
+  ArgInfo(Argument *F, Constant *A) : Formal(F), Actual(A) {}
+
+  bool operator==(const ArgInfo &Other) const {
+    return Formal == Other.Formal && Actual == Other.Actual;
+  }
+
+  bool operator!=(const ArgInfo &Other) const { return !(*this == Other); }
+
+  friend hash_code hash_value(const ArgInfo &A) {
+    return hash_combine(hash_value(A.Formal), hash_value(A.Actual));
+  }
 };
 
 class SCCPInstVisitor;