[Outliner] Fix compile-time overhead for candidate choice

[paquette]

Summary

The candidate collection method in the outliner can cause some dramatic compile-time regressions on large tests.

Currently, it works like this:

1. Build a suffix tree over the instructions in the program
2. Query the tree for good candidates
3. Prune tree
4. Go to step 2 until no candidates are found.

This hinges off the assumption that pruning the tree based off of overlapping candidates will reduce the search space enough that subsequent queries will take a negligible amount of time. Unfortunately, this assumption is shaky at best. Certain programs don't actually have many overlapping candidates. For very large tests, this is a big problem.

The new candidate collection method works like this:

1. Build a suffix tree over the instructions in the program
2. Iterate over each leaf in the tree.
3. Visit the parent of each leaf. If the parent represents a beneficial string, save a candidate for it.

There are O(I) leaves, where I is the number of instructions in the program. Since we only walk it once, this is now guaranteed to be at worst O(I). Overlapping candidates are now handled entirely by the pruneOverlaps method that handled anything the suffix tree pruning didn't handle. Since that method is on average O(n), this behaves far better.

Improvements

Compile-time was measured by compiling the LLVM test suite for AArch64 using LNT. Clang with Oz and mno-red-zone was tested against clang with Oz, mno-red-zone, and -mllvm -enable-machine-outliner.

Using the old candidate collection method, the worst compile-time regression was a spectacular 284.07 seconds on MultiSource/Benchmarks/7zip/7zip-benchmark.test. This was followed up by MultiSource/Benchmarks/Bullet/bullet.test with a 257.96 second regression. On average, there was a 9.69 second compile-time regression and a median compile-time regression of 0.85 seconds.

Using the new candidate collection method, the worst compile-time regression was 0.542 seconds for MultiSource/Benchmarks/7zip/7zip-benchmark.test, followed up by 0.496 seconds for MultiSource/Benchmarks/Bullet/bullet.test. On average, there was a 0.01 second compile-time regression, with a median compile-time regression of 0.002s.

Changes to code size

This change doesn't impact the code size results. In fact, by collecting all potential candidates, we can probably make the outliner make better decisions for what to outline in the future.

[MatzeB]

I am confused:

The candidate collection method in the outliner can cause some dramatic code size regressions on large tests.

... vs ...

This change doesn't impact the code size results. In fact, by collecting all potential candidates, we can probably make the outliner make better decisions for what to outline in the future.

[paquette]

In D31145#705561, @MatzeB wrote:

I am confused:

The candidate collection method in the outliner can cause some dramatic code size regressions on large tests.

... vs ...

This change doesn't impact the code size results. In fact, by collecting all potential candidates, we can probably make the outliner make better decisions for what to outline in the future.

***compile time!

Too used to typing code size. This is all about compile time being awful. :)

[silvas]

Overlapping candidates are now handled entirely by the pruneOverlaps method that handled anything the suffix tree pruning didn't handle. Since that method is on average O(n), this behaves far better.

Why is it O(n) on average? It is two nested for loops with some bailouts that only affect the asymptotic runtime if candidates overlap with little-omega(1) other candidates. And it needs to be big-omega(sqrt(n)) if the runtime is to be reduced to even O(n^1.5).

Also, as a side note, pruneOverlaps seems to only consider the first occurrences. E.g. decreasing occurrence counts by 1. Two candidates may overlap in more than one place. This probably explains something I was seeing: if you print out the total benefit of each candidate as it is outlined, the total benefit that the outliner thinks it is getting from the outlining operations it does is more can be more than the total string length, which reflects incorrect cost modeling.

Using the new candidate collection method, the worst compile-time regression was 0.542 seconds for MultiSource/Benchmarks/7zip/7zip-benchmark.test, followed up by 0.496 seconds for MultiSource/Benchmarks/Bullet/bullet.test. On average, there was a 0.01 second compile-time regression, with a median compile-time regression of 0.002s.

This is a big improvement over the previous numbers, but it would be good to phrase it as relative numbers. E.g. does the outliner after this patch ever take more than 10% of compile time? That would still be unfortunate.

Also, you should try this with FullLTO. That will smoke out compile time problems better. If you save off the combined bitcode files for each exe you care about, then you can just run llc on the bitcode directly which should let you iterate faster.

lib/CodeGen/MachineOutliner.cpp

558

Why not just phrase this routine as a single tree walk visiting the internal nodes instead of doing this thing where you walk the leaf vector looking at parents?

580

This FIXME doesn't make sense. Every internal node is the parent (not just ancestor) of a leaf. In fact you are relying on this so that looking at parents of the leaves considers all internal nodes.

620

Should be *Parent, right? Leaves will always have 1 occurrence.

[silvas]

lib/CodeGen/MachineOutliner.cpp

580

Or the code is failing to consider all internal nodes with this patch, which seems like a bummer. Do you have data on how many internal nodes aren't parents of leaves (and have found the missed opportunities to be negligible)?

[paquette]

In D31145#705994, @silvas wrote:

Why is it O(n) on average? It is two nested for loops with some bailouts that only affect the asymptotic runtime if candidates overlap with little-omega(1) other candidates. And it needs to be big-omega(sqrt(n)) if the runtime is to be reduced to even O(n^1.5).

Compiling for AArch64, for the SingleSource and MultiSource tests in the test suite the average length of a candidate is 10 MachineInstrs. The maximum length of a candidate is 108, which appears in a test with 11251 MachineInstrs. So far, it appears to be linear time on average, because for programs with a lot of candidates, the max length of a candidate is usually dominated by the number of candidates. I should really prove the bound though. It's a little tricky because it's O(m * n) where m is the maximum length of a candidate.

Hand waving/brain dump: The tricky case is where m = NumInstructions/k for k > 2 and n ~ O(NumInstructions). As k gets larger, m gets smaller. When m = 2, we only have to compare against one other candidate for each candidate, so we're good.

Previously, since we did some overlap pruning in the tree, we knew that as k gets smaller, the possible number of candidates in the program got smaller. Now that I think about it, that's not really true anymore. I guess the point here would be to show that we do the same number of overlap checks as in the tree here. But then we aren't really talking about pruneOverlaps, so ¯\_(ツ)_/¯

In the meantime, I can make the outliner bail out entirely in the event that things spin out of control. From there, I could start working on maybe

• Putting stuff in an interval tree or...
• Adding some stronger checks in the same vein as the max length check to reliably bound the search space now that we don't have the tree pruning

Also, as a side note, pruneOverlaps seems to only consider the first occurrences. E.g. decreasing occurrence counts by 1. Two candidates may overlap in more than one place. This probably explains something I was seeing: if you print out the total benefit of each candidate as it is outlined, the total benefit that the outliner thinks it is getting from the outlining operations it does is more can be more than the total string length, which reflects incorrect cost modeling.

I don't quite understand what you mean here. There are multiple occurrences of one candidate that can overlap with other candidates in multiple places, yes, but the decrement is done on their respective OutlinedFunctions. Also, the benefit used after this function is stored in the OutlinedFunction. Candidate benefit is only used for the greedy choice to prevent candidates from "cancelling each other out". Is that what you're getting at? I didn't really explain this well in the code on a second glance.

[paquette]

lib/CodeGen/MachineOutliner.cpp

580

Yeah I was wrong here. I got too excited about the idea of getting rid of a bunch of nodes.

620

BenefitFn looks at the parent of the leaf, so it's fine. It should really be refactored to take in the parent though, since that's what the traversal is concerned with.

[MatzeB]

• This looks good to me.
• +1 to the benefit function getting the parent passed in.
• I'd replace the FIXMEs with TODO, as they are really bugs but possible improvements (as far as I could understand)
• Superficially the suffixtree operations look okay to me, but I leave the judgement to the experts (=you).

lib/CodeGen/MachineOutliner.cpp

611

Maybe use a reference as Parent cannot be nullptr.

[silvas]

In D31145#707827, @paquette wrote:

In D31145#705994, @silvas wrote:

Why is it O(n) on average? It is two nested for loops with some bailouts that only affect the asymptotic runtime if candidates overlap with little-omega(1) other candidates. And it needs to be big-omega(sqrt(n)) if the runtime is to be reduced to even O(n^1.5).

Compiling for AArch64, for the SingleSource and MultiSource tests in the test suite the average length of a candidate is 10 MachineInstrs. The maximum length of a candidate is 108, which appears in a test with 11251 MachineInstrs. So far, it appears to be linear time on average, because for programs with a lot of candidates, the max length of a candidate is usually dominated by the number of candidates. I should really prove the bound though. It's a little tricky because it's O(m * n) where m is the maximum length of a candidate.

Hand waving/brain dump: The tricky case is where m = NumInstructions/k for k > 2 and n ~ O(NumInstructions). As k gets larger, m gets smaller. When m = 2, we only have to compare against one other candidate for each candidate, so we're good.

Previously, since we did some overlap pruning in the tree, we knew that as k gets smaller, the possible number of candidates in the program got smaller. Now that I think about it, that's not really true anymore. I guess the point here would be to show that we do the same number of overlap checks as in the tree here. But then we aren't really talking about pruneOverlaps, so ¯\_(ツ)_/¯

Oh, I think I was getting really confused about Candidate representing a single occurrence of a string or the equivalence class for that string itself. E.g. there is a comment "an occurrence of this candidate", but "Candidate" is the "occurrence", and "OutlinedFunction" is the "candidate". Maybe call them "Occurrence" and "CandidateClass" or something?

Also, I think I was thinking that this was trying to emulate the previous serialized greedy-choice, prune/recalculate benefit, greedy-choice, prune/recalculate benefit, ... that it used to do (or maybe I was misunderstanding that too... why would pruneOverlaps have been needed in that case...).

Anyway, this seems fine, sorry for the confusion. I see that on the version you committed, you added the comment that it is O(MaxCandidateLen * CandidateList.size()). It might be good to put some concrete numbers to give readers a feel. E.g. MaxCandidateLen in practice is basically a constant, and CandidateList.size() is typically some fraction of the number of instructions.

In the meantime, I can make the outliner bail out entirely in the event that things spin out of control. From there, I could start working on maybe

• Putting stuff in an interval tree or...
• Adding some stronger checks in the same vein as the max length check to reliably bound the search space now that we don't have the tree pruning

Now that I understand it better, I don't think this really needs any improvement. There seem to be only a couple ways of improving the code size benefit of the outliner:

1. Make more instructions viable for outlining (and you've been doing plenty of work on this, e.g. stack fixups and tail calls)
2. Make the outlining decision process get a more optimal result (this is a combinatorial optimization problem)
3. Heuristics for making outlining results across TU's converge more (I don't think there's much to do for this TBH)

I think the original hope for the suffix tree was to make 2 a lot better, but I think that on practical inputs even simple approaches like the finding parents of leaf children + pruneOverlaps already get pretty close to optimal (for example, it seems that just looking at parents of leaf children doesn't hurt the result too much).

Do you have any other plans here to exploit the suffix tree. I think that after this patch I think all we use the suffix tree for is to find repeated substrings. We might want to replace that with single sort + linear scan. I.e. initialize a vector of pointers into each instruction, then sort them by the suffix starting at that instruction.

You might want to do an experiment where you make "is instruction safe to outline" always return true and see how much more code size improvement is achieved. That will give a bound on how much you can expect to gain by improving 1., which seems useful to know.

Also, as a side note, pruneOverlaps seems to only consider the first occurrences. E.g. decreasing occurrence counts by 1. Two candidates may overlap in more than one place. This probably explains something I was seeing: if you print out the total benefit of each candidate as it is outlined, the total benefit that the outliner thinks it is getting from the outlining operations it does is more can be more than the total string length, which reflects incorrect cost modeling.

I don't quite understand what you mean here. There are multiple occurrences of one candidate that can overlap with other candidates in multiple places, yes, but the decrement is done on their respective OutlinedFunctions. Also, the benefit used after this function is stored in the OutlinedFunction. Candidate benefit is only used for the greedy choice to prevent candidates from "cancelling each other out". Is that what you're getting at? I didn't really explain this well in the code on a second glance.

I think I was really confused here. Sorry for the noise.

[MatzeB]

This seems to have landed a while ago in r298648

lib/CodeGen/MachineOutliner.cpp

 
 using namespace llvm;
 
 STATISTIC(NumOutlined, "Number of candidates outlined");
 STATISTIC(FunctionsCreated, "Number of functions created");
 
 namespace {
 
+/// \brief An individual sequence of instructions to be replaced with a call to
+/// an outlined function.
+struct Candidate {
+
+  /// Set to false if the candidate overlapped with another candidate.
+  bool InCandidateList = true;
+
+  /// The start index of this \p Candidate.
+  size_t StartIdx;
+
+  /// The number of instructions in this \p Candidate.
+  size_t Len;
+
+  /// The index of this \p Candidate's \p OutlinedFunction in the list of
+  /// \p OutlinedFunctions.
+  size_t FunctionIdx;
+
+  /// The number of instructions saved by outlining all candidates of this type.
+  unsigned Benefit = 0;
+
+  Candidate(size_t StartIdx, size_t Len, size_t FunctionIdx)
+      : StartIdx(StartIdx), Len(Len), FunctionIdx(FunctionIdx) {}
+
+  Candidate() {}
+
+  /// \brief Used to ensure that \p Candidates are outlined in an order that
+  /// preserves the start and end indices of other \p Candidates.
+  bool operator<(const Candidate &RHS) const { return StartIdx > RHS.StartIdx; }
+};
+
+/// \brief The information necessary to create an outlined function for some
+/// class of candidate.
+struct OutlinedFunction {
+
+  /// The actual outlined function created.
+  /// This is initialized after we go through and create the actual function.
+  MachineFunction *MF = nullptr;
+
+  /// A number assigned to this function which appears at the end of its name.
+  size_t Name;
+
+  /// The number of times that this function has appeared.
+  size_t OccurrenceCount = 0;
+
+  /// \brief The sequence of integers corresponding to the instructions in this
+  /// function.
+  std::vector<unsigned> Sequence;
+
+  /// The number of instructions this function would save.
+  unsigned Benefit = 0;
+
+  bool IsTailCall = false;
+
+  OutlinedFunction(size_t Name, size_t OccurrenceCount,
+                   const std::vector<unsigned> &Sequence,
+                   unsigned Benefit, bool IsTailCall)
+      : Name(Name), OccurrenceCount(OccurrenceCount), Sequence(Sequence),
+        Benefit(Benefit), IsTailCall(IsTailCall)
+        {}
+};
+
 /// Represents an undefined index in the suffix tree.
 const size_t EmptyIdx = -1;
 
 /// A node in a suffix tree which represents a substring or suffix.
 ///
 /// Each node has either no children or at least two children, with the root
 /// being a exception in the empty tree.
 ///
struct SuffixTreeNode {
   SuffixTreeNode *Parent = nullptr;
 
   /// The number of times this node's string appears in the tree.
   ///
   /// This is equal to the number of leaf children of the string. It represents
   /// the number of suffixes that the node's string is a prefix of.
   size_t OccurrenceCount = 0;
 
+  /// The length of the string formed by concatenating the edge labels from the
+  /// root to this node.
+  size_t ConcatLen = 0;
+
   /// Returns true if this node is a leaf.
   bool isLeaf() const { return SuffixIdx != EmptyIdx; }
 
   /// Returns true if this node is the root of its owning \p SuffixTree.
   bool isRoot() const { return StartIdx == EmptyIdx; }
 
   /// Return the number of elements in the substring associated with this node.
   size_t size() const {
private:
   /// respective suffix index.
   ///
   /// \param[in] CurrNode The node currently being visited.
   /// \param CurrIdx The current index of the string being visited.
   void setSuffixIndices(SuffixTreeNode &CurrNode, size_t CurrIdx) {
 
     bool IsLeaf = CurrNode.Children.size() == 0 && !CurrNode.isRoot();
 
+    // Store the length of the concatenation of all strings from the root to
+    // this node.
+    if (!CurrNode.isRoot()) {
+      if (CurrNode.ConcatLen == 0)
+        CurrNode.ConcatLen = CurrNode.size();
+
+      if (CurrNode.Parent)
+       CurrNode.ConcatLen += CurrNode.Parent->ConcatLen;
+    }
+
     // Traverse the tree depth-first.
     for (auto &ChildPair : CurrNode.Children) {
       assert(ChildPair.second && "Node had a null child!");
       setSuffixIndices(*ChildPair.second,
                        CurrIdx + ChildPair.second->size());
     }
 
     // Is this node a leaf?
while (SuffixesToAdd > 0) {
         // Start the next phase at the next smallest suffix.
         Active.Node = Active.Node->Link;
       }
     }
 
     return SuffixesToAdd;
   }
 
-  /// \brief Return the start index and length of a string which maximizes a
-  /// benefit function by traversing the tree depth-first.
-  ///
-  /// Helper function for \p bestRepeatedSubstring.
-  ///
-  /// \param CurrNode The node currently being visited.
-  /// \param CurrLen Length of the current string.
-  /// \param[out] BestLen Length of the most beneficial substring.
-  /// \param[out] MaxBenefit Benefit of the most beneficial substring.
-  /// \param[out] BestStartIdx Start index of the most beneficial substring.
-  /// \param BenefitFn The function the query should return a maximum string
-  /// for.
-  void findBest(SuffixTreeNode &CurrNode, size_t CurrLen, size_t &BestLen,
-                size_t &MaxBenefit, size_t &BestStartIdx,
-                const std::function<unsigned(SuffixTreeNode &, size_t CurrLen)>
-                &BenefitFn) {
-
-    if (!CurrNode.IsInTree)
-      return;
-
-    // Can we traverse further down the tree?
-    if (!CurrNode.isLeaf()) {
-      // If yes, continue the traversal.
-      for (auto &ChildPair : CurrNode.Children) {
-        if (ChildPair.second && ChildPair.second->IsInTree)
-          findBest(*ChildPair.second, CurrLen + ChildPair.second->size(),
-                   BestLen, MaxBenefit, BestStartIdx, BenefitFn);
-      }
-    } else {
-      // We hit a leaf.
-      size_t StringLen = CurrLen - CurrNode.size();
-      unsigned Benefit = BenefitFn(CurrNode, StringLen);
-
-      // Did we do better than in the last step?
-      if (Benefit <= MaxBenefit)
-        return;
-
-      // We did better, so update the best string.
-      MaxBenefit = Benefit;
-      BestStartIdx = CurrNode.SuffixIdx;
-      BestLen = StringLen;
-    }
-  }
-
 public:
 
   unsigned operator[](const size_t i) const {
     return Str[i];
   }
 
-  /// \brief Return a substring of the tree with maximum benefit if such a
-  /// substring exists.
+  /// Find all repeated substrings that satisfy \p BenefitFn.
   ///
-  /// Clears the input vector and fills it with a maximum substring or empty.
+  /// If a substring appears at least twice, then it must be represented by
+  /// an internal node which appears in at least two suffixes. Each suffix is
+  /// represented by a leaf node. Thus, by iterating over the leaves and looking
+  /// at their parents, we can find every outlining candidate in the module
+  /// in O(I) time, where I is the number of instructions in the program.
+  ///
+  /// Effectively, what we have is something that looks like this, going from
+  /// the leaves in LeafVector up:
+  /// 
+  ///              |  |   |    |     |
+  ///             s1  s2  s3  s4     sk
+  ///             /\  |  /|\  /\     |
+  /// LeafVector [0 1 2 3 4 5 6 7...I-1]
+  ///
+  /// s1, s3, and s4 might be considered for outlining based off the result
+  /// from BenefitFn. For, say, s3, we would visit leaf 3 first and then visit
+  /// s3. We would decide s3 is beneficial, look at leaves 3, 4, and 5 and then
+  /// mark them as out of the tree. When, in the next step, we visit leaf 4,
+  /// we'd skip it since we've already removed it.
+  ///
+  /// FIXME: Say we decide to outline s3. Then we know we'll never visit
+  /// leaves 4 or 5 after that step. We also know that s3 has 3 children. We
+  /// should skip over the #children-1 leaves that we'll never visit.
+  ///
+  /// FIXME: We only have to look at the leaves and their parents in the tree.
+  /// We should see if it's possible to save some space in the tree by taking
+  /// advantage of this fact.
+  /// 
+  /// \param[out] CandidateList Filled with candidates representing each
+  /// beneficial substring. Cleared initially.
+  /// \param[out] FunctionList Filled with a list of \p OutlinedFunctions each
+  /// type of candidate. Cleared initially.
+  /// \param BenefitFn Assigns positive integers to strings which would be
+  /// beneficial candidates, and 0 to unbeneficial candidates.
   ///
-  /// \param[in,out] Best The most beneficial substring in the tree. Empty
-  /// if it does not exist.
-  /// \param BenefitFn The function the query should return a maximum string
-  /// for.
-  void bestRepeatedSubstring(std::vector<unsigned> &Best,
-                 const std::function<unsigned(SuffixTreeNode &, size_t CurrLen)>
-                 &BenefitFn) {
-    Best.clear();
-    size_t Length = 0;   // Becomes the length of the best substring.
-    size_t Benefit = 0;  // Becomes the benefit of the best substring.
-    size_t StartIdx = 0; // Becomes the start index of the best substring.
-    findBest(*Root, 0, Length, Benefit, StartIdx, BenefitFn);
-
-    for (size_t Idx = 0; Idx < Length; Idx++)
-      Best.push_back(Str[Idx + StartIdx]);
-  }
-
-  /// Perform a depth-first search for \p QueryString on the suffix tree.
-  ///
-  /// \param QueryString The string to search for.
-  /// \param CurrIdx The current index in \p QueryString that is being matched
-  /// against.
-  /// \param CurrNode The suffix tree node being searched in.
-  ///
-  /// \returns A \p SuffixTreeNode that \p QueryString appears in if such a
-  /// node exists, and \p nullptr otherwise.
-  SuffixTreeNode *findString(const std::vector<unsigned> &QueryString,
-                             size_t &CurrIdx, SuffixTreeNode *CurrNode) {
-
-    // The search ended at a nonexistent or pruned node. Quit.
-    if (!CurrNode || !CurrNode->IsInTree)
-      return nullptr;
-
-    unsigned Edge = QueryString[CurrIdx]; // The edge we want to move on.
-    SuffixTreeNode *NextNode = CurrNode->Children[Edge]; // Next node in query.
-
-    if (CurrNode->isRoot()) {
-      // If we're at the root we have to check if there's a child, and move to
-      // that child. Don't consume the character since \p Root represents the
-      // empty string.
-      if (NextNode && NextNode->IsInTree)
-        return findString(QueryString, CurrIdx, NextNode);
-      return nullptr;
-    }
-
-    size_t StrIdx = CurrNode->StartIdx;
-    size_t MaxIdx = QueryString.size();
-    bool ContinueSearching = false;
-
-    // Match as far as possible into the string. If there's a mismatch, quit.
-    for (; CurrIdx < MaxIdx; CurrIdx++, StrIdx++) {
-      Edge = QueryString[CurrIdx];
-
-      // We matched perfectly, but still have a remainder to search.
-      if (StrIdx > *(CurrNode->EndIdx)) {
-        ContinueSearching = true;
-        break;
-      }
-
-      if (Edge != Str[StrIdx])
-        return nullptr;
-    }
-
-    NextNode = CurrNode->Children[Edge];
-
-    // Move to the node which matches what we're looking for and continue
-    // searching.
-    if (ContinueSearching)
-      return findString(QueryString, CurrIdx, NextNode);
-
-    // We matched perfectly so we're done.
-    return CurrNode;
-  }
-
-  /// \brief Remove a node from a tree and all nodes representing proper
-  /// suffixes of that node's string.
-  ///
-  /// This is used in the outlining algorithm to reduce the number of
-  /// overlapping candidates
-  ///
-  /// \param N The suffix tree node to start pruning from.
-  /// \param Len The length of the string to be pruned.
-  ///
-  /// \returns True if this candidate didn't overlap with a previously chosen
-  /// candidate.
-  bool prune(SuffixTreeNode *N, size_t Len) {
+  /// \returns The length of the longest candidate found.
+  size_t findCandidates(std::vector<Candidate> &CandidateList,
+  std::vector<OutlinedFunction> &FunctionList,
+  const std::function<unsigned(SuffixTreeNode &, size_t, bool&)> &BenefitFn) {
 
-    bool NoOverlap = true;
-    std::vector<unsigned> IndicesToPrune;
+    // Make sure there's nothing in the candidate or function lists.
+    CandidateList.clear();
+    FunctionList.clear();
 
-    // Look at each of N's children.
-    for (auto &ChildPair : N->Children) {
-      SuffixTreeNode *M = ChildPair.second;
+    size_t FnId = 0; // The ID assigned to the next OutlinedFunction produced.
+    size_t MaxLen = 0; // Length of the longest candidate found.
 
-      // Is this a leaf child?
-      if (M && M->IsInTree && M->isLeaf()) {
-        // Save each leaf child's suffix indices and remove them from the tree.
-        IndicesToPrune.push_back(M->SuffixIdx);
-        M->IsInTree = false;
-      }
-    }
+    // Find repeated substrings by iterating over the leaves in the tree.
+    for (SuffixTreeNode *Leaf : LeafVector) {
 
-    // Remove each suffix we have to prune from the tree. Each of these will be
-    // I + some offset for I in IndicesToPrune and some offset < Len.
-    unsigned Offset = 1;
-    for (unsigned CurrentSuffix = 1; CurrentSuffix < Len; CurrentSuffix++) {
-      for (unsigned I : IndicesToPrune) {
-
-        unsigned PruneIdx = I + Offset;
-
-        // Is this index actually in the string?
-        if (PruneIdx < LeafVector.size()) {
-          // If yes, we have to try and prune it.
-          // Was the current leaf already pruned by another candidate?
-          if (LeafVector[PruneIdx]->IsInTree) {
-            // If not, prune it.
-            LeafVector[PruneIdx]->IsInTree = false;
-          } else {
-            // If yes, signify that we've found an overlap, but keep pruning.
-            NoOverlap = false;
-          }
+      // Did we already find this leaf in a previous step?
+      if (!Leaf->IsInTree)
+        continue;
 
-          // Update the parent of the current leaf's occurrence count.
-          SuffixTreeNode *Parent = LeafVector[PruneIdx]->Parent;
+      // We didn't, so look at its parent.
+      SuffixTreeNode *Parent = Leaf->Parent;
 
-          // Is the parent still in the tree?
-          if (Parent->OccurrenceCount > 0) {
-            Parent->OccurrenceCount--;
-            Parent->IsInTree = (Parent->OccurrenceCount > 1);
-          }
-        }
-      }
+      // Did we already visit the parent?
+      if (!Parent->IsInTree)
+        continue;
 
-      // Move to the next character in the string.
-      Offset++;
+      // We found something new, so let's see if it's a good candidate.
+      size_t StringLen = Parent->ConcatLen;
+      bool IsTailCall;
+      unsigned Benefit = BenefitFn(*Leaf, StringLen, IsTailCall);
+
+      // Would it save us any instructions if we outlined it?
+      if (Benefit < 1) {
+        // It's unbeneficial, so remove the parent and the leaf and move on.
+        Parent->IsInTree = false;
+        Leaf->IsInTree = false;
+        continue;
       }
 
-    // We know we can never outline anything which starts one index back from
-    // the indices we want to outline. This is because our minimum outlining
-    // length is always 2.
-    for (unsigned I : IndicesToPrune) {
-      if (I > 0) {
+      // It could save us some instructions, so let's save it.
+      // First, check if the parent's string is longer than MaxLen and update
+      // accordingly.
+      if (StringLen > MaxLen)
+        MaxLen = StringLen;
 
-        unsigned PruneIdx = I-1;
-        SuffixTreeNode *Parent = LeafVector[PruneIdx]->Parent;
+      unsigned OccCount = 0; // Number of times a string appears.
+      for (auto &ChildPair : Parent->Children) {
+        SuffixTreeNode *M = ChildPair.second;
 
-        // Was the leaf one index back from I already pruned?
-        if (LeafVector[PruneIdx]->IsInTree) {
-          // If not, prune it.
-          LeafVector[PruneIdx]->IsInTree = false;
-        } else {
-          // If yes, signify that we've found an overlap, but keep pruning.
-          NoOverlap = false;
-        }
+        // Is it a leaf child?
+        if (M && M->IsInTree && M->isLeaf()) {
+          // It is, so we have an occurrence. Save it as a candidate.
+          OccCount++;
+          CandidateList.emplace_back(M->SuffixIdx, StringLen, FnId);
+          CandidateList.back().Benefit = Benefit;
 
-        // Update the parent of the current leaf's occurrence count.
-        if (Parent->OccurrenceCount > 0) {
-          Parent->OccurrenceCount--;
-          Parent->IsInTree = (Parent->OccurrenceCount > 1);
-        }
+          // Never save M as a candidate again.
+          M->IsInTree = false;
         }
       }
 
-    // Finally, remove N from the tree and set its occurrence count to 0.
-    N->IsInTree = false;
-    N->OccurrenceCount = 0;
+      // Save the string associated with the parent and create an
+      // OutlinedFunction for it.
+      std::vector<unsigned> CandidateSequence;
+      for (unsigned i = Leaf->SuffixIdx; i < Leaf->SuffixIdx + StringLen; i++)
+        CandidateSequence.push_back(Str[i]);
 
-    return NoOverlap;
-  }
+      FunctionList.emplace_back(FnId, OccCount, CandidateSequence, Benefit,
+                                IsTailCall);
 
-  /// \brief Find each occurrence of of a string in \p QueryString and prune
-  /// their nodes.
-  ///
-  /// \param QueryString The string to search for.
-  /// \param[out] Occurrences The start indices of each occurrence.
-  ///
-  /// \returns Whether or not the occurrence overlaps with a previous candidate.
-  bool findOccurrencesAndPrune(const std::vector<unsigned> &QueryString,
-                               std::vector<size_t> &Occurrences) {
-    size_t Dummy = 0;
-    SuffixTreeNode *N = findString(QueryString, Dummy, Root);
-
-    if (!N || !N->IsInTree)
-      return false;
+      // We've gotten every candidate for this OutlinedFunction.
+      // Move to the next available ID.
+      FnId++;
 
-    // If this is an internal node, occurrences are the number of leaf children
-    // of the node.
-    for (auto &ChildPair : N->Children) {
-      SuffixTreeNode *M = ChildPair.second;
-
-      // Is it a leaf? If so, we have an occurrence.
-      if (M && M->IsInTree && M->isLeaf())
-        Occurrences.push_back(M->SuffixIdx);
+      // The parent can't offer any new candidates, so take it out of the tree.
+      Parent->IsInTree = false;
     }
 
-    // If we're in a leaf, then this node is the only occurrence.
-    if (N->isLeaf())
-      Occurrences.push_back(N->SuffixIdx);
-
-    return prune(N, QueryString.size());
+    return MaxLen;
   }
  
   /// Construct a suffix tree from a sequence of unsigned integers.
   ///
   /// \param Str The string to construct the suffix tree for.
   SuffixTree(const std::vector<unsigned> &Str) : Str(Str) {
     Root = insertInternalNode(nullptr, EmptyIdx, EmptyIdx, 0);
     Root->IsInTree = true;
     Active.Node = Root;
     LeafVector = std::vector<SuffixTreeNode*>(Str.size());
SuffixTree(const std::vector<unsigned> &Str) : Str(Str) {
     }
 
     // Set the suffix indices of each leaf.
     assert(Root && "Root node can't be nullptr!");
     setSuffixIndices(*Root, 0);
   }
 };
 
-/// \brief An individual sequence of instructions to be replaced with a call to
-/// an outlined function.
-struct Candidate {
-
-  /// Set to false if the candidate overlapped with another candidate.
-  bool InCandidateList = true;
-
-  /// The start index of this \p Candidate.
-  size_t StartIdx;
-
-  /// The number of instructions in this \p Candidate.
-  size_t Len;
-
-  /// The index of this \p Candidate's \p OutlinedFunction in the list of
-  /// \p OutlinedFunctions.
-  size_t FunctionIdx;
-
-  Candidate(size_t StartIdx, size_t Len, size_t FunctionIdx)
-      : StartIdx(StartIdx), Len(Len), FunctionIdx(FunctionIdx) {}
-
-  Candidate() {}
-
-  /// \brief Used to ensure that \p Candidates are outlined in an order that
-  /// preserves the start and end indices of other \p Candidates.
-  bool operator<(const Candidate &RHS) const { return StartIdx > RHS.StartIdx; }
-};
-
-/// \brief The information necessary to create an outlined function for some
-/// class of candidate.
-struct OutlinedFunction {
-
-  /// The actual outlined function created.
-  /// This is initialized after we go through and create the actual function.
-  MachineFunction *MF = nullptr;
-
-  /// A number assigned to this function which appears at the end of its name.
-  size_t Name;
-
-  /// The number of times that this function has appeared.
-  size_t OccurrenceCount = 0;
-
-  /// \brief The sequence of integers corresponding to the instructions in this
-  /// function.
-  std::vector<unsigned> Sequence;
-
-  /// The number of instructions this function would save.
-  unsigned Benefit = 0;
-
-  bool IsTailCall = false;
-
-  OutlinedFunction(size_t Name, size_t OccurrenceCount,
-                   const std::vector<unsigned> &Sequence,
-                   unsigned Benefit, bool IsTailCall)
-      : Name(Name), OccurrenceCount(OccurrenceCount), Sequence(Sequence),
-        Benefit(Benefit), IsTailCall(IsTailCall)
-        {}
-};
-
 /// \brief Maps \p MachineInstrs to unsigned integers and stores the mappings.
 struct InstructionMapper {
 
   /// \brief The next available integer to assign to a \p MachineInstr that
   /// cannot be outlined.
   ///
   /// Set to -3 for compatability with \p DenseMapInfo<unsigned>.
   unsigned IllegalInstrNumber = -3;
struct MachineOutliner : public ModulePass {
   ///
   /// \returns The length of the longest candidate found. 0 if there are none.
   unsigned buildCandidateList(std::vector<Candidate> &CandidateList,
                               std::vector<OutlinedFunction> &FunctionList,
                               SuffixTree &ST,
                               InstructionMapper &Mapper,
                               const TargetInstrInfo &TII);
 
-  /// \brief Remove any overlapping candidates that weren't handled by the
-  /// suffix tree's pruning method.
+  /// \brief Remove overlapping candidates from \p CandidateList, greedily
+  /// picking more beneficial candidates in the case of overlaps.
   ///
-  /// Pruning from the suffix tree doesn't necessarily remove all overlaps.
-  /// If a short candidate is chosen for outlining, then a longer candidate
-  /// which has that short candidate as a suffix is chosen, the tree's pruning
-  /// method will not find it. Thus, we need to prune before outlining as well.
+  /// FIXME: This is currently on average (and at best) linear-time in the
+  /// number of candidates and worst-case quadratic.
+  /// It really ought to be O(log n).
   ///
   /// \param[in,out] CandidateList A list of outlining candidates.
   /// \param[in,out] FunctionList A list of functions to be outlined.
-  /// \param MaxCandidateLen The length of the longest candidate.
+  /// \param MaxCandidateLen The length of the longest candidate. Used to bound
+  /// the number of comparisons between candidates.
   /// \param TII TargetInstrInfo for the module.
   void pruneOverlaps(std::vector<Candidate> &CandidateList,
                      std::vector<OutlinedFunction> &FunctionList,
-                     unsigned MaxCandidateLen,
+                     const unsigned MaxCandidateLen,
                      const TargetInstrInfo &TII);
 
   /// Construct a suffix tree on the instructions in \p M and outline repeated
   /// strings from that tree.
   bool runOnModule(Module &M) override;
 };
 
 } // Anonymous namespace.
 
 char MachineOutliner::ID = 0;
 
 namespace llvm {
 ModulePass *createMachineOutlinerPass() { return new MachineOutliner(); }
 }
 
 INITIALIZE_PASS(MachineOutliner, "machine-outliner",
                 "Machine Function Outliner", false, false)
 
 void MachineOutliner::pruneOverlaps(std::vector<Candidate> &CandidateList,
                                     std::vector<OutlinedFunction> &FunctionList,
-                                    unsigned MaxCandidateLen,
+                                    const unsigned MaxCandidateLen,
                                     const TargetInstrInfo &TII) {
 
   // Check for overlaps in the range. This is O(n^2) worst case, but we can
   // alleviate that somewhat by bounding our search space using the start
   // index of our first candidate and the maximum distance an overlapping
   // candidate could have from the first candidate.
   for (auto It = CandidateList.begin(), Et = CandidateList.end(); It != Et;
        It++) {
for (auto Sit = It + 1; Sit != Et; Sit++) {
       // High indices... [C1End ... C1Start][C2End ... C2Start] ...Low indices
       //
       // We sorted our candidate list so C2Start <= C1Start. We know that
       // C2End > C2Start since each candidate has length >= 2. Therefore, all we
       // have to check is C2End < C2Start to see if we overlap.
       if (C2End < C1.StartIdx)
         continue;
 
-      // C2 overlaps with C1. Because we pruned the tree already, the only way
-      // this can happen is if C1 is a proper suffix of C2. Thus, we must have
-      // found C1 first during our query, so it must have benefit greater or
-      // equal to C2. Greedily pick C1 as the candidate to keep and toss out C2.
-      DEBUG (
-            size_t C1End = C1.StartIdx + C1.Len - 1;
-            dbgs() << "- Found an overlap to purge.\n";
-            dbgs() << "--- C1 :[" << C1.StartIdx << ", " << C1End << "]\n";
-            dbgs() << "--- C2 :[" << C2.StartIdx << ", " << C2End << "]\n";
-            );
-
-      // Update the function's occurrence count and benefit to reflec that C2
-      // is being removed.
+      // Choose the better of C1 and C2 to keep.
+      // NOTE: We use C1 and C2's benefits instead of their function's benefits.
+      // This is because we're making a greedy local choice.
+      if (C1.Benefit >= C2.Benefit) {
+        // Removed one candidate, so decrement occurrence count.
         F2.OccurrenceCount--;
+
+        // Recompute F2's benefit. This means that F2 could eventually become
+        // nonbeneficial and be skipped over.
         F2.Benefit = TII.getOutliningBenefit(F2.Sequence.size(),
                                              F2.OccurrenceCount,
                                              F2.IsTailCall
                                              );
 
         // Mark C2 as not in the list.
         C2.InCandidateList = false;
-
         DEBUG (
-            dbgs() << "- Removed C2. \n";
-            dbgs() << "--- Num fns left for C2: " << F2.OccurrenceCount << "\n";
-            dbgs() << "--- C2's benefit: " << F2.Benefit << "\n";
-            );
+      } else {
+        F1.OccurrenceCount--;
+        F1.Benefit = TII.getOutliningBenefit(F1.Sequence.size(),
+                                             F1.OccurrenceCount,
+                                             F1.IsTailCall
+                                             );
+        C1.InCandidateList = false;
+        DEBUG (
+          dbgs() << "- Removed C1. \n";
+          dbgs() << "--- Num fns left for C1: " << F1.OccurrenceCount << "\n";
+          dbgs() << "--- C1's benefit: " << F1.Benefit << "\n";
+        );
+      }
     }
   }
 }
 
+
 unsigned
 MachineOutliner::buildCandidateList(std::vector<Candidate> &CandidateList,
                                     std::vector<OutlinedFunction> &FunctionList,
                                     SuffixTree &ST,
                                     InstructionMapper &Mapper,
                                     const TargetInstrInfo &TII) {
 
   std::vector<unsigned> CandidateSequence; // Current outlining candidate.
-  unsigned MaxCandidateLen = 0; // Length of the longest candidate.
+  size_t MaxCandidateLen = 0; // Length of the longest candidate.
 
   // Function for maximizing query in the suffix tree.
   // This allows us to define more fine-grained types of things to outline in
   // the target without putting target-specific info in the suffix tree.
   auto BenefitFn = [&TII, &ST, &Mapper](const SuffixTreeNode &Curr,
-                                        size_t StringLen) {
-
+                                        size_t StringLen,
+                                        bool &IsTailCall) {
     // Any leaf whose parent is the root only has one occurrence.
     if (Curr.Parent->isRoot())
       return 0u;
 
-    // Anything with length < 2 will never be beneficial on any target.
+    // Is it long enough that outlining it would save space?
+    // FIXME: This is an approximation. On some targets, the call instruction
+    // might be smaller than the instruction we're outlining. Thus, outlining
+    // a single instruction might save space. Similarly, the call might be
+    // bigger than the two instruction sequence. We need instruction sizes to
+    // truly be accurate here.
     if (StringLen < 2)
       return 0u;
 
-    size_t Occurrences = Curr.Parent->OccurrenceCount;
+    size_t Occurrences = Curr.Parent->Children.size();
 
     // Anything with fewer than 2 occurrences will never be beneficial on any
     // target.
     if (Occurrences < 2)
       return 0u;
 
     // Check if the last instruction in the sequence is a return.
     MachineInstr *LastInstr =
     Mapper.IntegerInstructionMap[ST[Curr.SuffixIdx + StringLen - 1]];
     assert(LastInstr && "Last instruction in sequence was unmapped!");
 
     // The only way a terminator could be mapped as legal is if it was safe to
     // tail call.
-    bool IsTailCall = LastInstr->isTerminator();
-
+    IsTailCall = LastInstr->isTerminator();
     return TII.getOutliningBenefit(StringLen, Occurrences, IsTailCall);
   };
 
-  // Repeatedly query the suffix tree for the substring that maximizes
-  // BenefitFn. Find the occurrences of that string, prune the tree, and store
-  // each occurrence as a candidate.
-  for (ST.bestRepeatedSubstring(CandidateSequence, BenefitFn);
-       CandidateSequence.size() > 1;
-       ST.bestRepeatedSubstring(CandidateSequence, BenefitFn)) {
-
-    std::vector<size_t> Occurrences;
-
-    bool GotNonOverlappingCandidate =
-        ST.findOccurrencesAndPrune(CandidateSequence, Occurrences);
-
-    // Is the candidate we found known to overlap with something we already
-    // outlined?
-    if (!GotNonOverlappingCandidate)
-      continue;
-
-    // Is this candidate the longest so far?
-    if (CandidateSequence.size() > MaxCandidateLen)
-      MaxCandidateLen = CandidateSequence.size();
-
-    MachineInstr *LastInstr =
-    Mapper.IntegerInstructionMap[CandidateSequence.back()];
-    assert(LastInstr && "Last instruction in sequence was unmapped!");
-
-    // The only way a terminator could be mapped as legal is if it was safe to
-    // tail call.
-    bool IsTailCall = LastInstr->isTerminator();
-
-    // Keep track of the benefit of outlining this candidate in its
-    // OutlinedFunction.
-    unsigned FnBenefit = TII.getOutliningBenefit(CandidateSequence.size(),
-                                                 Occurrences.size(),
-                                                 IsTailCall
-                                                 );
-
-    assert(FnBenefit > 0 && "Function cannot be unbeneficial!");
-
-    // Save an OutlinedFunction for this candidate.
-    FunctionList.emplace_back(
-        FunctionList.size(), // Number of this function.
-        Occurrences.size(),  // Number of occurrences.
-        CandidateSequence,   // Sequence to outline.
-        FnBenefit,           // Instructions saved by outlining this function.
-        IsTailCall           // Flag set if this function is to be tail called.
-        );
-
-    // Save each of the occurrences of the candidate so we can outline them.
-    for (size_t &Occ : Occurrences)
-      CandidateList.emplace_back(
-          Occ,                      // Starting idx in that MBB.
-          CandidateSequence.size(), // Candidate length.
-          FunctionList.size() - 1   // Idx of the corresponding function.
-          );
-
-    FunctionsCreated++;
-  }
+  // Find all of the candidates in the tree.
+  MaxCandidateLen = ST.findCandidates(CandidateList, FunctionList, BenefitFn);
 
   // Sort the candidates in decending order. This will simplify the outlining
   // process when we have to remove the candidates from the mapping by
   // allowing us to cut them out without keeping track of an offset.
   std::stable_sort(CandidateList.begin(), CandidateList.end());
 
   return MaxCandidateLen;
 }
for (const Candidate &C : CandidateList) {
 
     assert(EndIdx < Mapper.InstrList.size() && "Candidate out of bounds!");
     MachineBasicBlock::iterator EndIt = Mapper.InstrList[EndIdx];
     assert(EndIt != MBB->end() && "EndIt out of bounds!");
 
     EndIt++; // Erase needs one past the end index.
 
     // Does this candidate have a function yet?
-    if (!OF.MF)
+    if (!OF.MF) {
       OF.MF = createOutlinedFunction(M, OF, Mapper);
+      FunctionsCreated++;
+    }
 
     MachineFunction *MF = OF.MF;
     const TargetSubtargetInfo &STI = MF->getSubtarget();
     const TargetInstrInfo &TII = *STI.getInstrInfo();
 
     // Insert a call to the new function and erase the old sequence.
     TII.insertOutlinedCall(M, *MBB, StartIt, *MF, OF.IsTailCall);
     StartIt = Mapper.InstrList[C.StartIdx];