[MachineOutliner] Leaf Descendants

llvm/include/llvm/Support/SuffixTree.h

@@ -42,6 +42,9 @@ class SuffixTree {
   /// Each element is an integer representing an instruction in the module.
   ArrayRef<unsigned> Str;
 
+  /// Whether to consider leaf descendants or only leaf children.
+  bool OutlinerLeafDescendants;
+
   /// A repeated substring in the tree.
   struct RepeatedSubstring {
     /// The length of the string.
@@ -130,11 +133,27 @@ class SuffixTree {
   /// this step.
   unsigned extend(unsigned EndIdx, unsigned SuffixesToAdd);
 
+  /// This vector contains all leaf nodes of this suffix tree. These leaf nodes
+  /// are identified using post-order depth-first traversal, so that the order
+  /// of these leaf nodes in the vector matches the order of the leaves in the
+  /// tree from left to right if one were to draw the tree on paper.
+  std::vector<SuffixTreeLeafNode *> LeafNodes;
+
+  /// Perform a post-order depth-first traversal of the tree and perform two
+  /// tasks during the traversal. The first is to populate LeafNodes, adding
+  /// nodes in order of the traversal. The second is to keep track of the leaf
+  /// descendants of every internal node by assigning values to LeftLeafIndex
+  /// and RightLefIndex fields of SuffixTreeNode for all internal nodes.
+  void setLeafNodes();
+
 public:
   /// Construct a suffix tree from a sequence of unsigned integers.
   ///
   /// \param Str The string to construct the suffix tree for.
-  SuffixTree(const ArrayRef<unsigned> &Str);
+  /// \param OutlinerLeafDescendants Whether to consider leaf descendants or
+  /// only leaf children (used by Machine Outliner).
+  SuffixTree(const ArrayRef<unsigned> &Str,
+             bool OutlinerLeafDescendants = false);
 
   /// Iterator for finding all repeated substrings in the suffix tree.
   struct RepeatedSubstringIterator {
@@ -154,6 +173,12 @@ class SuffixTree {
     /// instruction lengths.
     const unsigned MinLength = 2;
 
+    /// Vector of leaf nodes of the suffix tree.
+    const std::vector<SuffixTreeLeafNode *> &LeafNodes;
+
+    /// Whether to consider leaf descendants or only leaf children.
+    bool OutlinerLeafDescendants = !LeafNodes.empty();
+
     /// Move the iterator to the next repeated substring.
     void advance();
 
@@ -179,7 +204,10 @@ class SuffixTree {
       return !(*this == Other);
     }
 
-    RepeatedSubstringIterator(SuffixTreeInternalNode *N) : N(N) {
+    RepeatedSubstringIterator(
+        SuffixTreeInternalNode *N,
+        const std::vector<SuffixTreeLeafNode *> &LeafNodes = {})
+        : N(N), LeafNodes(LeafNodes) {
       // Do we have a non-null node?
       if (!N)
         return;
@@ -191,7 +219,7 @@ class SuffixTree {
   };
 
   typedef RepeatedSubstringIterator iterator;
-  iterator begin() { return iterator(Root); }
+  iterator begin() { return iterator(Root, LeafNodes); }
   iterator end() { return iterator(nullptr); }
 };
 

llvm/include/llvm/Support/SuffixTreeNode.h

@@ -46,6 +46,17 @@ struct SuffixTreeNode {
   /// the root to this node.
   unsigned ConcatLen = 0;
 
+  /// These two indices give a range of indices for its leaf descendants.
+  /// Imagine drawing a tree on paper and assigning a unique index to each leaf
+  /// node in monotonically increasing order from left to right. This way of
+  /// numbering the leaf nodes allows us to associate a continuous range of
+  /// indices with each internal node. For example, if a node has leaf
+  /// descendants with indices i, i+1, ..., j, then its LeftLeafIdx is i and
+  /// its RightLeafIdx is j. These indices are for LeafNodes in the SuffixTree
+  /// class, which is constructed using post-order depth-first traversal.
+  unsigned LeftLeafIdx = EmptyIdx;
+  unsigned RightLeafIdx = EmptyIdx;
+
 public:
   // LLVM RTTI boilerplate.
   NodeKind getKind() const { return Kind; }
@@ -56,6 +67,18 @@ struct SuffixTreeNode {
   /// \returns the end index of this node.
   virtual unsigned getEndIdx() const = 0;
 
+  /// \return the index of this node's left most leaf node.
+  unsigned getLeftLeafIdx() const;
+
+  /// \return the index of this node's right most leaf node.
+  unsigned getRightLeafIdx() const;
+
+  /// Set the index of the left most leaf node of this node to \p Idx.
+  void setLeftLeafIdx(unsigned Idx);
+
+  /// Set the index of the right most leaf node of this node to \p Idx.
+  void setRightLeafIdx(unsigned Idx);
+
   /// Advance this node's StartIdx by \p Inc.
   void incrementStartIdx(unsigned Inc);
 
@@ -168,4 +191,4 @@ struct SuffixTreeLeafNode : SuffixTreeNode {
   virtual ~SuffixTreeLeafNode() = default;
 };
 } // namespace llvm
-#endif // LLVM_SUPPORT_SUFFIXTREE_NODE_H
+#endif // LLVM_SUPPORT_SUFFIXTREE_NODE_H

llvm/lib/CodeGen/MachineOutliner.cpp

@@ -121,6 +121,12 @@ static cl::opt<unsigned> OutlinerBenefitThreshold(
     cl::desc(
         "The minimum size in bytes before an outlining candidate is accepted"));
 
+static cl::opt<bool> OutlinerLeafDescendants(
+    "outliner-leaf-descendants", cl::init(true), cl::Hidden,
+    cl::desc("Consider all leaf descendants of internal nodes of the suffix "
+             "tree as candidates for outlining (if false, only leaf children "
+             "are considered)"));
+
 namespace {
 
 /// Maps \p MachineInstrs to unsigned integers and stores the mappings.
@@ -576,7 +582,7 @@ void MachineOutliner::emitOutlinedFunctionRemark(OutlinedFunction &OF) {
 void MachineOutliner::findCandidates(
     InstructionMapper &Mapper, std::vector<OutlinedFunction> &FunctionList) {
   FunctionList.clear();
-  SuffixTree ST(Mapper.UnsignedVec);
+  SuffixTree ST(Mapper.UnsignedVec, OutlinerLeafDescendants);
 
   // First, find all of the repeated substrings in the tree of minimum length
   // 2.

llvm/lib/Support/SuffixTree.cpp

@@ -11,9 +11,11 @@
 //===----------------------------------------------------------------------===//
 
 #include "llvm/Support/SuffixTree.h"
+#include "llvm/ADT/SmallPtrSet.h"
 #include "llvm/Support/Allocator.h"
 #include "llvm/Support/Casting.h"
 #include "llvm/Support/SuffixTreeNode.h"
+#include <stack>
 
 using namespace llvm;
 
@@ -26,7 +28,9 @@ static size_t numElementsInSubstring(const SuffixTreeNode *N) {
   return N->getEndIdx() - N->getStartIdx() + 1;
 }
 
-SuffixTree::SuffixTree(const ArrayRef<unsigned> &Str) : Str(Str) {
+SuffixTree::SuffixTree(const ArrayRef<unsigned> &Str,
+                       bool OutlinerLeafDescendants)
+    : Str(Str), OutlinerLeafDescendants(OutlinerLeafDescendants) {
   Root = insertRoot();
   Active.Node = Root;
 
@@ -46,6 +50,11 @@ SuffixTree::SuffixTree(const ArrayRef<unsigned> &Str) : Str(Str) {
   // Set the suffix indices of each leaf.
   assert(Root && "Root node can't be nullptr!");
   setSuffixIndices();
+
+  // Collect all leaf nodes of the suffix tree. And for each internal node,
+  // record the range of leaf nodes that are descendants of it.
+  if (OutlinerLeafDescendants)
+    setLeafNodes();
 }
 
 SuffixTreeNode *SuffixTree::insertLeaf(SuffixTreeInternalNode &Parent,
@@ -105,6 +114,71 @@ void SuffixTree::setSuffixIndices() {
   }
 }
 
+void SuffixTree::setLeafNodes() {
+  // A stack that keeps track of nodes to visit for post-order DFS traversal.
+  SmallVector<SuffixTreeNode *> ToVisit;
+  ToVisit.push_back(Root);
+
+  // This keeps track of the index of the next leaf node to be added to
+  // the LeafNodes vector of the suffix tree.
+  unsigned LeafCounter = 0;
+
+  // This keeps track of nodes whose children have been added to the stack.
+  // The value is a pair, representing a node's first and last children.
+  DenseMap<SuffixTreeInternalNode *,
+           std::pair<SuffixTreeNode *, SuffixTreeNode *>>
+      ChildrenMap;
+
+  // Traverse the tree in post-order.
+  while (!ToVisit.empty()) {
+    SuffixTreeNode *CurrNode = ToVisit.back();
+    ToVisit.pop_back();
+    if (auto *CurrInternalNode = dyn_cast<SuffixTreeInternalNode>(CurrNode)) {
+      // The current node is an internal node.
+      auto I = ChildrenMap.find(CurrInternalNode);
+      if (I == ChildrenMap.end()) {
+        // This is the first time we visit this node.
+        // Its children have not been added to the stack yet.
+        // We add current node back, and add its children to the stack.
+        // We keep track of the first and last children of the current node.
+        auto it = CurrInternalNode->Children.begin();
+        if (it != CurrInternalNode->Children.end()) {
+          ToVisit.push_back(CurrNode);
+          SuffixTreeNode *FirstChild = it->second;
+          SuffixTreeNode *LastChild = nullptr;
+          for (; it != CurrInternalNode->Children.end(); ++it) {
+            LastChild = it->second;
+            ToVisit.push_back(LastChild);
+          }
+          ChildrenMap[CurrInternalNode] = {FirstChild, LastChild};
+        }
+      } else {
+        // This is the second time we visit this node.
+        // All of its children have already been processed.
+        // Now, we can set its LeftLeafIdx and RightLeafIdx;
+        auto [FirstChild, LastChild] = I->second;
+        // Get the first child to use its RightLeafIdx.
+        // The first child is the first one added to the stack, so it is
+        // the last one to be processed. Hence, the leaf descendants
+        // of the first child are assigned the largest index numbers.
+        CurrNode->setRightLeafIdx(FirstChild->getRightLeafIdx());
+        // Get the last child to use its LeftLeafIdx.
+        CurrNode->setLeftLeafIdx(LastChild->getLeftLeafIdx());
+        assert(CurrNode->getLeftLeafIdx() <= CurrNode->getRightLeafIdx() &&
+               "LeftLeafIdx should not be larger than RightLeafIdx");
+      }
+    } else {
+      // The current node is a leaf node.
+      // We can simply set its LeftLeafIdx and RightLeafIdx.
+      CurrNode->setLeftLeafIdx(LeafCounter);
+      CurrNode->setRightLeafIdx(LeafCounter);
+      LeafCounter++;
+      auto *CurrLeafNode = cast<SuffixTreeLeafNode>(CurrNode);
+      LeafNodes.push_back(CurrLeafNode);
+    }
+  }
+}
+
 unsigned SuffixTree::extend(unsigned EndIdx, unsigned SuffixesToAdd) {
   SuffixTreeInternalNode *NeedsLink = nullptr;
 
@@ -241,30 +315,34 @@ void SuffixTree::RepeatedSubstringIterator::advance() {
     // it's too short, we'll quit.
     unsigned Length = Curr->getConcatLen();
 
-    // Iterate over each child, saving internal nodes for visiting, and
-    // leaf nodes' SuffixIdx in RepeatedSubstringStarts. Internal nodes
-    // represent individual strings, which may repeat.
-    for (auto &ChildPair : Curr->Children) {
+    // Iterate over each child, saving internal nodes for visiting.
+    // Internal nodes represent individual strings, which may repeat.
+    for (auto &ChildPair : Curr->Children)
       // Save all of this node's children for processing.
       if (auto *InternalChild =
-              dyn_cast<SuffixTreeInternalNode>(ChildPair.second)) {
+              dyn_cast<SuffixTreeInternalNode>(ChildPair.second))
         InternalNodesToVisit.push_back(InternalChild);
-        continue;
-      }
-
-      if (Length < MinLength)
-        continue;
 
-      // Have an occurrence of a potentially repeated string. Save it.
-      auto *Leaf = cast<SuffixTreeLeafNode>(ChildPair.second);
-      RepeatedSubstringStarts.push_back(Leaf->getSuffixIdx());
-    }
+    // If length of repeated substring is below threshold, then skip it.
+    if (Length < MinLength)
+      continue;
 
     // The root never represents a repeated substring. If we're looking at
     // that, then skip it.
     if (Curr->isRoot())
       continue;
 
+    // Collect leaf children or leaf descendants by OutlinerLeafDescendants.
+    if (!OutlinerLeafDescendants) {
+      for (auto &ChildPair : Curr->Children)
+        if (auto *Leaf = dyn_cast<SuffixTreeLeafNode>(ChildPair.second))
+          RepeatedSubstringStarts.push_back(Leaf->getSuffixIdx());
+    } else {
+      for (unsigned I = Curr->getLeftLeafIdx(); I <= Curr->getRightLeafIdx();
+           ++I)
+        RepeatedSubstringStarts.push_back(LeafNodes[I]->getSuffixIdx());
+    }
+
     // Do we have any repeated substrings?
     if (RepeatedSubstringStarts.size() < 2)
       continue;

llvm/lib/Support/SuffixTreeNode.cpp

@@ -38,3 +38,8 @@ unsigned SuffixTreeLeafNode::getEndIdx() const {
 
 unsigned SuffixTreeLeafNode::getSuffixIdx() const { return SuffixIdx; }
 void SuffixTreeLeafNode::setSuffixIdx(unsigned Idx) { SuffixIdx = Idx; }
+
+unsigned SuffixTreeNode::getLeftLeafIdx() const { return LeftLeafIdx; }
+unsigned SuffixTreeNode::getRightLeafIdx() const { return RightLeafIdx; }
+void SuffixTreeNode::setLeftLeafIdx(unsigned Idx) { LeftLeafIdx = Idx; }
+void SuffixTreeNode::setRightLeafIdx(unsigned Idx) { RightLeafIdx = Idx; }

llvm/test/CodeGen/AArch64/machine-outliner-cfi-tail-some.mir

@@ -1,5 +1,5 @@
 # NOTE: Assertions have been autogenerated by utils/update_mir_test_checks.py
-# RUN: llc -mtriple=aarch64-apple-unknown -run-pass=machine-outliner -verify-machineinstrs %s -o - | FileCheck %s
+# RUN: llc -mtriple=aarch64-apple-unknown -run-pass=machine-outliner -verify-machineinstrs -outliner-leaf-descendants=false %s -o - | FileCheck %s
 
 # Outlining CFI instructions is unsafe if we cannot outline all of the CFI
 # instructions from a function.  This shows that we choose not to outline the