[MachineOutliner] AArch64: Avoid saving + restoring LR if possible

[paquette]

Summary

This commit allows the outliner to avoid saving + restoring the link register on AArch64 when it is dead within an entire class of candidates. It also changes the interface for the outliner slightly to facilitate this.

By doing this, it's possible to find more outlining candidates in AArch64-- in the case that a candidate doesn't require a save/restore, that candidate can be far shorter. This allows us to find length 2 and 3 candidates and outline them.

This is still conservative. For example, if a length-3 candidate appears 20 times, and one of those times requires a save + restore, the outliner will assume all of them require a save + restore. Thus, none will be outlined. This will be fixed in later patches; the purpose of this is just to get the main idea going.

This, on average, improves the code size of the --benchmarking only SingleSource and MultiSource tests by about 2.94% in comparison to the current outliner. Several tests have no improvements; this is likely due to the conservative nature of the current technique. 36 of the tests improve by more than 10%. 4 tests regress, likely due to competing with linker deduping. This can likely be fixed by disallowing outlining from linkonceodr functions.

These were compiled with LNT using -Oz -mno-red-zone -mllvm -enable-machine-outliner and --benchmarking-only.

Full code size results: https://hastebin.com/raw/soloropiwo

[davide]

Some comments inline.

lib/CodeGen/MachineOutliner.cpp

902

I wonder why this (and other variables) are size_t and not just unsigned?

lib/Target/AArch64/AArch64InstrInfo.cpp

4561–4569

This one seems a decent candidate for any_of, but the use of stepBackward() could make it more difficult.

4572–4584

Is your cost model for the call overhead documented somewhere? If not, can you add a comment here?

4643–4648

Have you considered rewriting this code in terms of llvm::any_of (or all_of) for clarity?

4649–4656

Can you please add a comment here? It's not obvious when it's safe to outline without understanding the whole picture.

[MatzeB]

include/llvm/Target/TargetInstrInfo.h

1565–1567

Is the vector supposed to be changed by this function? I'd use an ArrayRef<...> instead of std::vector<...> if it isn't.

1567–1568

I wonder why this parameter is called CandidateClass, from reading the description I wonder why there is a Class suffix.

lib/Target/AArch64/AArch64InstrInfo.cpp

4543

Not sure if I miss something, but I would expect that it is possible to check LR liveness exactly at the position where the call will be inserted rather than requiring the whole block to have LR unused?

[MatzeB]

lib/CodeGen/MachineOutliner.cpp

904–913

Seeing that both getOutliningCallOverhead and getOutliningFrameOverhead are now called in sequence and both perform some computations in common, this smells like maybe they should become one function now? (It certainly feels odd that currently you test LR-liveness twice for each candidate because of this split).

[paquette]

lib/CodeGen/MachineOutliner.cpp

902

Mostly for consistency with the definition of length + size. I was thinking that because an "overhead" in this sense is the length of a sequence of instructions, it'd be more consistent to use size_t.

OTOH, having a std::pair<size_t, unsigned> is kind of weird, and "overhead" need not be considered a "length". I'll swap size_t out for unsigned.

lib/Target/AArch64/AArch64InstrInfo.cpp

4543

Oh, I think you're right. I thought that you'd have to ensure that at least the sequence is safe, but the outliner doesn't ever touch LR. That is much simpler. :)

4572–4584

I'll add a comment for sure.

There should also be more comments about how the targets interact with the outliner in the main pass, so I'll do that too.

4643–4648

... I don't know a ton of C++ magic, so you have just changed my life.

[paquette]

Updated the diff to address feedback. There are lots of changes, but here are the highlights:

• I added a struct for MachineOutlinerInfo which acts as the interface between the target and the pass.
• Lots of comments
• Updated the remarks test to reflect the cost model change
• Merged together the frame overhead/call overhead stuff into getOutliningCandidateInfo, which returns a MachineOutlinerInfo struct.

[MatzeB]

lib/Target/AArch64/AArch64InstrInfo.cpp

4603–4616

This still checks for LR being available for the whole range between the end of the basic block and CallInsertionPt. I would expect it to be enough to stepBackward() till CallInsertionPt and only then check if LR is unused.

4637–4644

llvms version of all_of even supports ranges I think, so you could do: all_of(RepeatedSequenceLocs, DoesntNeedLRSave) here.

And for the record: I personally prefer imperative for loops (at least in cases where you have to specifically create a lambda for the predicate); but that is subjective so please do not change it back or the next reviewer will just ask for all_of again and we circle around ;-)

test/CodeGen/AArch64/machine-outliner.mir

85–112

This test can be simplified, look at http://llvm.org/docs/MIRLangRef.html#simplifying-mir-files for some tactics.

[paquette]

lib/Target/AArch64/AArch64InstrInfo.cpp

4603–4616

I actually tried that, and the outliner ended up erroneously pulling out a bunch of things without saving LR.

I tried something like

LRU.stepBackward(*CallInsertionPt);
if (!LRU.available(AArch64::LR))
    return false;

Is this incorrect?

[paquette]

• Simplified the MIR test a bit
• Fixed the loop for the LiveRegisterUnits stuff so that we only check if LR is available at the prospective call-site

[MatzeB]

Thanks, LGTM

[paquette]

Committed:
r314341 - [MachineOutliner] AArch64: Avoid saving + restoring LR if possible

include/llvm/Target/TargetInstrInfo.h

virtual void reMaterialize(MachineBasicBlock &MBB,
                              const TargetRegisterInfo &TRI) const;
 
   /// \brief Clones instruction or the whole instruction bundle \p Orig and
   /// insert into \p MBB before \p InsertBefore. The target may update operands
   /// that are required to be unique.
   ///
   /// \p Orig must not return true for MachineInstr::isNotDuplicable().
   virtual MachineInstr &duplicate(MachineBasicBlock &MBB,
-      MachineBasicBlock::iterator InsertBefore, const MachineInstr &Orig) const;
+                                  MachineBasicBlock::iterator InsertBefore,
+                                  const MachineInstr &Orig) const;
 
   /// This method must be implemented by targets that
   /// set the M_CONVERTIBLE_TO_3_ADDR flag.  When this flag is set, the target
   /// may be able to convert a two-address instruction into one or more true
   /// three-address instructions on demand.  This allows the X86 target (for
   /// example) to convert ADD and SHL instructions into LEA instructions if they
   /// would require register copies due to two-addressness.
   ///
public:
 
   /// True if the instruction is bound to the top of its basic block and no
   /// other instructions shall be inserted before it. This can be implemented
   /// to prevent register allocator to insert spills before such instructions.
   virtual bool isBasicBlockPrologue(const MachineInstr &MI) const {
     return false;
   }
 
-  /// \brief Returns the number of instructions that will be taken to call a
-  /// function defined by the sequence on the closed interval [ \p StartIt, \p
-  /// EndIt].
-  ///
-  /// \returns The number of instructions for the call in the first member,
-  /// and a target-defined unsigned describing what type of call to emit in the
-  /// second member.
-  virtual std::pair<size_t, unsigned>
-  getOutliningCallOverhead(MachineBasicBlock::iterator &StartIt,
-                           MachineBasicBlock::iterator &EndIt) const {
-    llvm_unreachable(
-        "Target didn't implement TargetInstrInfo::getOutliningCallOverhead!");
-  }
+  /// \brief Describes the number of instructions that it will take to call and
+  /// construct a frame for a given outlining candidate.
+  struct MachineOutlinerInfo {
+    /// Number of instructions to call an outlined function for this candidate.
+    unsigned CallOverhead;
+
+    /// \brief Number of instructions to construct an outlined function frame
+    /// for this candidate.
+    unsigned FrameOverhead;
+
+    /// \brief Represents the specific instructions that must be emitted to
+    /// construct a call to this candidate.
+    unsigned CallConstructionID;
+
+    /// \brief Represents the specific instructions that must be emitted to
+    /// construct a frame for this candidate's outlined function.
+    unsigned FrameConstructionID;
+
+    MachineOutlinerInfo() {}
+    MachineOutlinerInfo(unsigned CallOverhead, unsigned FrameOverhead,
+                        unsigned CallConstructionID,
+                        unsigned FrameConstructionID)
+        : CallOverhead(CallOverhead), FrameOverhead(FrameOverhead),
+          CallConstructionID(CallConstructionID),
+          FrameConstructionID(FrameConstructionID) {}
+  };
 
-  /// \brief Returns the number of instructions that will be taken to construct
-  /// an outlined function frame for a function defined on the closed interval
-  /// [ \p StartIt, \p EndIt].
-  ///
-  /// \returns The number of instructions for the frame in the first member
-  /// and a target-defined unsigned describing what type of frame to construct
-  /// in the second member.
-  virtual std::pair<size_t, unsigned> getOutliningFrameOverhead(
+  /// \brief Returns a \p MachineOutlinerInfo struct containing target-specific
+  /// information for a set of outlining candidates.
+  virtual MachineOutlinerInfo getOutlininingCandidateInfo(
       std::vector<
           std::pair<MachineBasicBlock::iterator, MachineBasicBlock::iterator>>
-          &CandidateClass) const {
+          &RepeatedSequenceLocs) const {
     llvm_unreachable(
-        "Target didn't implement TargetInstrInfo::getOutliningFrameOverhead!");
+        "Target didn't implement TargetInstrInfo::getOutliningOverhead!");
   }
 
   /// Represents how an instruction should be mapped by the outliner.
   /// \p Legal instructions are those which are safe to outline.
   /// \p Illegal instructions are those which cannot be outlined.
   /// \p Invisible instructions are instructions which can be outlined, but
   /// shouldn't actually impact the outlining result.
   enum MachineOutlinerInstrType { Legal, Illegal, Invisible };
 
   /// Returns how or if \p MI should be outlined.
   virtual MachineOutlinerInstrType getOutliningType(MachineInstr &MI) const {
     llvm_unreachable(
         "Target didn't implement TargetInstrInfo::getOutliningType!");
   }
 
   /// Insert a custom epilogue for outlined functions.
   /// This may be empty, in which case no epilogue or return statement will be
   /// emitted.
   virtual void insertOutlinerEpilogue(MachineBasicBlock &MBB,
                                       MachineFunction &MF,
-                                      unsigned FrameClass) const {
+                                      const MachineOutlinerInfo &MInfo) const {
     llvm_unreachable(
         "Target didn't implement TargetInstrInfo::insertOutlinerEpilogue!");
   }
 
   /// Insert a call to an outlined function into the program.
   /// Returns an iterator to the spot where we inserted the call. This must be
   /// implemented by the target.
   virtual MachineBasicBlock::iterator
   insertOutlinedCall(Module &M, MachineBasicBlock &MBB,
                      MachineBasicBlock::iterator &It, MachineFunction &MF,
-                     unsigned CallClass) const {
+                     const MachineOutlinerInfo &MInfo) const {
     llvm_unreachable(
         "Target didn't implement TargetInstrInfo::insertOutlinedCall!");
   }
 
   /// Insert a custom prologue for outlined functions.
   /// This may be empty, in which case no prologue will be emitted.
   virtual void insertOutlinerPrologue(MachineBasicBlock &MBB,
                                       MachineFunction &MF,
-                                      unsigned FrameClass) const {
+                                      const MachineOutlinerInfo &MInfo) const {
     llvm_unreachable(
         "Target didn't implement TargetInstrInfo::insertOutlinerPrologue!");
   }
 
   /// Return true if the function can safely be outlined from.
-  /// By default, this means that the function has no red zone.
+  /// A function \p MF is considered safe for outlining if an outlined function
+  /// produced from instructions in F will produce a program which produces the
+  /// same output for any set of given inputs.
   virtual bool isFunctionSafeToOutlineFrom(MachineFunction &MF) const {
     llvm_unreachable("Target didn't implement "
                      "TargetInstrInfo::isFunctionSafeToOutlineFrom!");
   }
 
 private:
   unsigned CallFrameSetupOpcode, CallFrameDestroyOpcode;
   unsigned CatchRetOpcode;

lib/CodeGen/MachineOutliner.cpp

 /// \file
 /// Replaces repeated sequences of instructions with function calls.
 ///
 /// This works by placing every instruction from every basic block in a
 /// suffix tree, and repeatedly querying that tree for repeated sequences of
 /// instructions. If a sequence of instructions appears often, then it ought
 /// to be beneficial to pull out into a function.
 ///
+/// The MachineOutliner communicates with a given target using hooks defined in
+/// TargetInstrInfo.h. The target supplies the outliner with information on how
+/// a specific sequence of instructions should be outlined. This information
+/// is used to deduce the number of instructions necessary to
+///
+/// * Create an outlined function
+/// * Call that outlined function
+///
+/// Targets must implement
+///   * getOutliningCandidateInfo
+///   * insertOutlinerEpilogue
+///   * insertOutlinedCall
+///   * insertOutlinerPrologue
+///   * isFunctionSafeToOutlineFrom
+///
+/// in order to make use of the MachineOutliner.
+///
 /// This was originally presented at the 2016 LLVM Developers' Meeting in the
 /// talk "Reducing Code Size Using Outlining". For a high-level overview of
 /// how this pass works, the talk is available on YouTube at
 ///
 /// https://www.youtube.com/watch?v=yorld-WSOeU
 ///
 /// The slides for the talk are available at
 ///

 /// \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;
+  unsigned StartIdx;
 
   /// The number of instructions in this \p Candidate.
-  size_t Len;
+  unsigned Len;
 
   /// The index of this \p Candidate's \p OutlinedFunction in the list of
   /// \p OutlinedFunctions.
-  size_t FunctionIdx;
+  unsigned FunctionIdx;
 
-  /// Target-defined unsigned defining how to emit a call for this candidate.
-  unsigned CallClass = 0;
+  /// Contains all target-specific information for this \p Candidate.
+  TargetInstrInfo::MachineOutlinerInfo MInfo;
 
   /// \brief The number of instructions that would be saved by outlining every
   /// candidate of this type.
   ///
   /// This is a fixed value which is not updated during the candidate pruning
   /// process. It is only used for deciding which candidate to keep if two
   /// candidates overlap. The true benefit is stored in the OutlinedFunction
   /// for some given candidate.
   unsigned Benefit = 0;
 
-  Candidate(size_t StartIdx, size_t Len, size_t FunctionIdx, unsigned CallClass)
-      : StartIdx(StartIdx), Len(Len), FunctionIdx(FunctionIdx),
-        CallClass(CallClass) {}
+  Candidate(unsigned StartIdx, unsigned Len, unsigned 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 numbefr assigned to this function which appears at the end of its name.
-  size_t Name;
+  /// A number assigned to this function which appears at the end of its name.
+  unsigned Name;
 
   /// The number of candidates for this OutlinedFunction.
-  size_t OccurrenceCount = 0;
+  unsigned 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;
 
-  /// Target-defined unsigned defining how to emit the frame for this function.
-  unsigned FrameClass = 0;
+  /// Contains all target-specific information for this \p OutlinedFunction.
+  TargetInstrInfo::MachineOutlinerInfo MInfo;
 
-  OutlinedFunction(size_t Name, size_t OccurrenceCount,
+  OutlinedFunction(unsigned Name, unsigned OccurrenceCount,
                    const std::vector<unsigned> &Sequence, unsigned Benefit,
-                   unsigned FrameClass)
+                   TargetInstrInfo::MachineOutlinerInfo &MInfo)
       : Name(Name), OccurrenceCount(OccurrenceCount), Sequence(Sequence),
-        Benefit(Benefit), FrameClass(FrameClass) {}
+        Benefit(Benefit), MInfo(MInfo) {}
 };
 
 /// Represents an undefined index in the suffix tree.
-const size_t EmptyIdx = -1;
+const unsigned 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.
 ///
 /// Children are represented as a map between unsigned integers and nodes. If
 /// a node N has a child M on unsigned integer k, then the mapping represented
struct SuffixTreeNode {
   /// represented by the current node, there is a way to reach another
   /// mapping by tacking that character on the end of the current string.
   DenseMap<unsigned, SuffixTreeNode *> Children;
 
   /// A flag set to false if the node has been pruned from the tree.
   bool IsInTree = true;
 
   /// The start index of this node's substring in the main string.
-  size_t StartIdx = EmptyIdx;
+  unsigned StartIdx = EmptyIdx;
 
   /// The end index of this node's substring in the main string.
   ///
   /// Every leaf node must have its \p EndIdx incremented at the end of every
   /// step in the construction algorithm. To avoid having to update O(N)
   /// nodes individually at the end of every step, the end index is stored
   /// as a pointer.
-  size_t *EndIdx = nullptr;
+  unsigned *EndIdx = nullptr;
 
   /// For leaves, the start index of the suffix represented by this node.
   ///
   /// For all other nodes, this is ignored.
-  size_t SuffixIdx = EmptyIdx;
+  unsigned SuffixIdx = EmptyIdx;
 
   /// \brief For internal nodes, a pointer to the internal node representing
   /// the same sequence with the first character chopped off.
   ///
   /// This acts as a shortcut in Ukkonen's algorithm. One of the things that
   /// Ukkonen's algorithm does to achieve linear-time construction is
   /// keep track of which node the next insert should be at. This makes each
   /// insert O(1), and there are a total of O(N) inserts. The suffix link
struct SuffixTreeNode {
 
   /// The parent of this node. Every node except for the root has a parent.
   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;
+  unsigned OccurrenceCount = 0;
 
   /// The length of the string formed by concatenating the edge labels from the
   /// root to this node.
-  size_t ConcatLen = 0;
+  unsigned 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 {
 
     // Is it the root? If so, it's the empty string so return 0.
     if (isRoot())
       return 0;
 
     assert(*EndIdx != EmptyIdx && "EndIdx is undefined!");
 
     // Size = the number of elements in the string.
     // For example, [0 1 2 3] has length 4, not 3. 3-0 = 3, so we have 3-0+1.
     return *EndIdx - StartIdx + 1;
   }
 
-  SuffixTreeNode(size_t StartIdx, size_t *EndIdx, SuffixTreeNode *Link,
+  SuffixTreeNode(unsigned StartIdx, unsigned *EndIdx, SuffixTreeNode *Link,
                  SuffixTreeNode *Parent)
       : StartIdx(StartIdx), EndIdx(EndIdx), Link(Link), Parent(Parent) {}
 
   SuffixTreeNode() {}
 };
 
 /// A data structure for fast substring queries.
 ///
private:
   /// Each internal node is guaranteed to never have its end index change
   /// during the construction algorithm; however, leaves must be updated at
   /// every step. Therefore, we need to store leaf end indices by reference
   /// to avoid updating O(N) leaves at every step of construction. Thus,
   /// every internal node must be allocated its own end index.
   BumpPtrAllocator InternalEndIdxAllocator;
 
   /// The end index of each leaf in the tree.
-  size_t LeafEndIdx = -1;
+  unsigned LeafEndIdx = -1;
 
   /// \brief Helper struct which keeps track of the next insertion point in
   /// Ukkonen's algorithm.
   struct ActiveState {
     /// The next node to insert at.
     SuffixTreeNode *Node;
 
     /// The index of the first character in the substring currently being added.
-    size_t Idx = EmptyIdx;
+    unsigned Idx = EmptyIdx;
 
     /// The length of the substring we have to add at the current step.
-    size_t Len = 0;
+    unsigned Len = 0;
   };
 
   /// \brief The point the next insertion will take place at in the
   /// construction algorithm.
   ActiveState Active;
 
   /// Allocate a leaf node and add it to the tree.
   ///
   /// \param Parent The parent of this node.
   /// \param StartIdx The start index of this node's associated string.
   /// \param Edge The label on the edge leaving \p Parent to this node.
   ///
   /// \returns A pointer to the allocated leaf node.
-  SuffixTreeNode *insertLeaf(SuffixTreeNode &Parent, size_t StartIdx,
+  SuffixTreeNode *insertLeaf(SuffixTreeNode &Parent, unsigned StartIdx,
                              unsigned Edge) {
 
     assert(StartIdx <= LeafEndIdx && "String can't start after it ends!");
 
     SuffixTreeNode *N = new (NodeAllocator.Allocate())
         SuffixTreeNode(StartIdx, &LeafEndIdx, nullptr, &Parent);
     Parent.Children[Edge] = N;
 
     return N;
   }
 
   /// Allocate an internal node and add it to the tree.
   ///
   /// \param Parent The parent of this node. Only null when allocating the root.
   /// \param StartIdx The start index of this node's associated string.
   /// \param EndIdx The end index of this node's associated string.
   /// \param Edge The label on the edge leaving \p Parent to this node.
   ///
   /// \returns A pointer to the allocated internal node.
-  SuffixTreeNode *insertInternalNode(SuffixTreeNode *Parent, size_t StartIdx,
-                                     size_t EndIdx, unsigned Edge) {
+  SuffixTreeNode *insertInternalNode(SuffixTreeNode *Parent, unsigned StartIdx,
+                                     unsigned EndIdx, unsigned Edge) {
 
     assert(StartIdx <= EndIdx && "String can't start after it ends!");
     assert(!(!Parent && StartIdx != EmptyIdx) &&
            "Non-root internal nodes must have parents!");
 
-    size_t *E = new (InternalEndIdxAllocator) size_t(EndIdx);
+    unsigned *E = new (InternalEndIdxAllocator) unsigned(EndIdx);
     SuffixTreeNode *N = new (NodeAllocator.Allocate())
         SuffixTreeNode(StartIdx, E, Root, Parent);
     if (Parent)
       Parent->Children[Edge] = N;
 
     return N;
   }
 
   /// \brief Set the suffix indices of the leaves to the start indices of their
   /// respective suffixes. Also stores each leaf in \p LeafVector at its
   /// 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) {
+  void setSuffixIndices(SuffixTreeNode &CurrNode, unsigned 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();
private:
   /// suffix tree is a valid tree.
   ///
   /// \param EndIdx The end index of the current prefix in the main string.
   /// \param SuffixesToAdd The number of suffixes that must be added
   /// to complete the suffix tree at the current phase.
   ///
   /// \returns The number of suffixes that have not been added at the end of
   /// this step.
-  unsigned extend(size_t EndIdx, size_t SuffixesToAdd) {
+  unsigned extend(unsigned EndIdx, unsigned SuffixesToAdd) {
     SuffixTreeNode *NeedsLink = nullptr;
 
     while (SuffixesToAdd > 0) {
 
       // Are we waiting to add anything other than just the last character?
       if (Active.Len == 0) {
         // If not, then say the active index is the end index.
         Active.Idx = EndIdx;
while (SuffixesToAdd > 0) {
           NeedsLink->Link = Active.Node;
           NeedsLink = nullptr;
         }
       } else {
         // There's a match with FirstChar, so look for the point in the tree to
         // insert a new node.
         SuffixTreeNode *NextNode = Active.Node->Children[FirstChar];
 
-        size_t SubstringLen = NextNode->size();
+        unsigned SubstringLen = NextNode->size();
 
         // Is the current suffix we're trying to insert longer than the size of
         // the child we want to move to?
         if (Active.Len >= SubstringLen) {
           // If yes, then consume the characters we've seen and move to the next
           // node.
           Active.Idx += SubstringLen;
           Active.Len -= SubstringLen;
public:
   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());
 
     // Keep track of the number of suffixes we have to add of the current
     // prefix.
-    size_t SuffixesToAdd = 0;
+    unsigned SuffixesToAdd = 0;
     Active.Node = Root;
 
     // Construct the suffix tree iteratively on each prefix of the string.
     // PfxEndIdx is the end index of the current prefix.
     // End is one past the last element in the string.
-    for (size_t PfxEndIdx = 0, End = Str.size(); PfxEndIdx < End; PfxEndIdx++) {
+    for (unsigned PfxEndIdx = 0, End = Str.size(); PfxEndIdx < End;
+         PfxEndIdx++) {
       SuffixesToAdd++;
       LeafEndIdx = PfxEndIdx; // Extend each of the leaves.
       SuffixesToAdd = extend(PfxEndIdx, SuffixesToAdd);
     }
 
     // Set the suffix indices of each leaf.
     assert(Root && "Root node can't be nullptr!");
     setSuffixIndices(*Root, 0);
struct MachineOutliner : public ModulePass {
   /// \param TII TargetInstrInfo for the target.
   /// \param Mapper Contains outlining mapping information.
   /// \param[out] CandidateList Filled with candidates representing each
   /// beneficial substring.
   /// \param[out] FunctionList Filled with a list of \p OutlinedFunctions each
   /// type of candidate.
   ///
   /// \returns The length of the longest candidate found.
-  size_t findCandidates(SuffixTree &ST, const TargetInstrInfo &TII,
+  unsigned findCandidates(SuffixTree &ST, const TargetInstrInfo &TII,
                           InstructionMapper &Mapper,
                           std::vector<Candidate> &CandidateList,
                           std::vector<OutlinedFunction> &FunctionList);
 
   /// \brief Replace the sequences of instructions represented by the
   /// \p Candidates in \p CandidateList with calls to \p MachineFunctions
   /// described in \p FunctionList.
   ///
   /// \param M The module we are outlining from.
   /// \param CandidateList A list of candidates to be outlined.
   /// \param FunctionList A list of functions to be inserted into the module.

 
 namespace llvm {
 ModulePass *createMachineOutlinerPass() { return new MachineOutliner(); }
 } // namespace llvm
 
 INITIALIZE_PASS(MachineOutliner, DEBUG_TYPE, "Machine Function Outliner", false,
                 false)
 
-size_t
+unsigned
 MachineOutliner::findCandidates(SuffixTree &ST, const TargetInstrInfo &TII,
                                 InstructionMapper &Mapper,
                                 std::vector<Candidate> &CandidateList,
                                 std::vector<OutlinedFunction> &FunctionList) {
-
   CandidateList.clear();
   FunctionList.clear();
-  size_t FnIdx = 0;
-  size_t MaxLen = 0;
+  unsigned FnIdx = 0;
+  unsigned MaxLen = 0;
 
   // FIXME: Visit internal nodes instead of leaves.
   for (SuffixTreeNode *Leaf : ST.LeafVector) {
     assert(Leaf && "Leaves in LeafVector cannot be null!");
     if (!Leaf->IsInTree)
       continue;
 
     assert(Leaf->Parent && "All leaves must have parents!");
     SuffixTreeNode &Parent = *(Leaf->Parent);
 
     // If it doesn't appear enough, or we already outlined from it, skip it.
     if (Parent.OccurrenceCount < 2 || Parent.isRoot() || !Parent.IsInTree)
       continue;
 
     // Figure out if this candidate is beneficial.
-    size_t StringLen = Leaf->ConcatLen - Leaf->size();
+    unsigned StringLen = Leaf->ConcatLen - (unsigned)Leaf->size();
 
     // Too short to be beneficial; skip it.
     // FIXME: This isn't necessarily true for, say, X86. If we factor in
     // instruction lengths we need more information than this.
     if (StringLen < 2)
       continue;
 
-    size_t CallOverhead = 0;
-    size_t SequenceOverhead = StringLen;
-
     // If this is a beneficial class of candidate, then every one is stored in
     // this vector.
     std::vector<Candidate> CandidatesForRepeatedSeq;
 
-    // Used for getOutliningFrameOverhead.
+    // Describes the start and end point of each candidate. This allows the
+    // target to infer some information about each occurrence of each repeated
+    // sequence.
     // FIXME: CandidatesForRepeatedSeq and this should be combined.
     std::vector<
         std::pair<MachineBasicBlock::iterator, MachineBasicBlock::iterator>>
-        CandidateClass;
+        RepeatedSequenceLocs;
 
     // Figure out the call overhead for each instance of the sequence.
     for (auto &ChildPair : Parent.Children) {
       SuffixTreeNode *M = ChildPair.second;
 
       if (M && M->IsInTree && M->isLeaf()) {
         // Each sequence is over [StartIt, EndIt].
         MachineBasicBlock::iterator StartIt = Mapper.InstrList[M->SuffixIdx];
         MachineBasicBlock::iterator EndIt =
             Mapper.InstrList[M->SuffixIdx + StringLen - 1];
 
-        // Get the overhead for calling a function for this sequence and any
-        // target-specified data for how to construct the call.
-        std::pair<size_t, unsigned> CallOverheadPair =
-            TII.getOutliningCallOverhead(StartIt, EndIt);
-        CallOverhead += CallOverheadPair.first;
-        CandidatesForRepeatedSeq.emplace_back(M->SuffixIdx, StringLen, FnIdx,
-                                              CallOverheadPair.second);
-        CandidateClass.emplace_back(std::make_pair(StartIt, EndIt));
+        CandidatesForRepeatedSeq.emplace_back(M->SuffixIdx, StringLen, FnIdx);
+        RepeatedSequenceLocs.emplace_back(std::make_pair(StartIt, EndIt));
 
         // Never visit this leaf again.
         M->IsInTree = false;
       }
     }
 
-    std::pair<size_t, unsigned> FrameOverheadPair =
-        TII.getOutliningFrameOverhead(CandidateClass);
-    size_t FrameOverhead = FrameOverheadPair.first;
-
-    size_t OutliningCost = CallOverhead + FrameOverhead + SequenceOverhead;
-    size_t NotOutliningCost = SequenceOverhead * Parent.OccurrenceCount;
+    unsigned SequenceOverhead = StringLen;
+    TargetInstrInfo::MachineOutlinerInfo MInfo =
+        TII.getOutlininingCandidateInfo(RepeatedSequenceLocs);
+
+    unsigned OutliningCost =
+        (MInfo.CallOverhead * Parent.OccurrenceCount) + MInfo.FrameOverhead;
+    unsigned NotOutliningCost = SequenceOverhead * Parent.OccurrenceCount;
 
     // Is it better to outline this candidate than not?
     if (NotOutliningCost <= OutliningCost) {
       // Outlining this candidate would take more instructions than not
       // outlining.
       // Emit a remark explaining why we didn't outline this candidate.
       std::pair<MachineBasicBlock::iterator, MachineBasicBlock::iterator> C =
-          CandidateClass[0];
+          RepeatedSequenceLocs[0];
       MachineOptimizationRemarkEmitter MORE(
           *(C.first->getParent()->getParent()), nullptr);
       MachineOptimizationRemarkMissed R(DEBUG_TYPE, "NotOutliningCheaper",
                                         C.first->getDebugLoc(),
                                         C.first->getParent());
       R << "Did not outline " << NV("Length", StringLen) << " instructions"
-        << " from " << NV("NumOccurrences", CandidateClass.size())
+        << " from " << NV("NumOccurrences", RepeatedSequenceLocs.size())
         << " locations."
         << " Instructions from outlining all occurrences ("
         << NV("OutliningCost", OutliningCost) << ")"
         << " >= Unoutlined instruction count ("
         << NV("NotOutliningCost", NotOutliningCost) << ")"
         << " (Also found at: ";
 
       // Tell the user the other places the candidate was found.
-      for (size_t i = 1, e = CandidateClass.size(); i < e; i++) {
+      for (unsigned i = 1, e = RepeatedSequenceLocs.size(); i < e; i++) {
         R << NV((Twine("OtherStartLoc") + Twine(i)).str(),
-                CandidateClass[i].first->getDebugLoc());
+                RepeatedSequenceLocs[i].first->getDebugLoc());
         if (i != e - 1)
           R << ", ";
       }
 
       R << ")";
       MORE.emit(R);
 
       // Move to the next candidate.
       continue;
     }
 
-    size_t Benefit = NotOutliningCost - OutliningCost;
+    unsigned Benefit = NotOutliningCost - OutliningCost;
 
     if (StringLen > MaxLen)
       MaxLen = StringLen;
 
     // At this point, the candidate class is seen as beneficial. Set their
     // benefit values and save them in the candidate list.
     for (Candidate &C : CandidatesForRepeatedSeq) {
       C.Benefit = Benefit;
+      C.MInfo = MInfo;
       CandidateList.push_back(C);
     }
 
     // Save the function for the new candidate sequence.
     std::vector<unsigned> CandidateSequence;
     for (unsigned i = Leaf->SuffixIdx; i < Leaf->SuffixIdx + StringLen; i++)
       CandidateSequence.push_back(ST.Str[i]);
 
     FunctionList.emplace_back(FnIdx, CandidatesForRepeatedSeq.size(),
-                              CandidateSequence, Benefit,
-                              FrameOverheadPair.second);
+                              CandidateSequence, Benefit, MInfo);
 
     // Move to the next function.
     FnIdx++;
     Parent.IsInTree = false;
   }
 
   return MaxLen;
 }
for (auto Sit = It + 1; Sit != Et; Sit++) {
         // If not, remove this candidate and move to the next one.
         assert(F2.OccurrenceCount > 0 &&
                "Can't remove OutlinedFunction with no occurrences!");
         F2.OccurrenceCount--;
         C2.InCandidateList = false;
         continue;
       }
 
-      size_t C2End = C2.StartIdx + C2.Len - 1;
+      unsigned C2End = C2.StartIdx + C2.Len - 1;
 
       // Do C1 and C2 overlap?
       //
       // Not overlapping:
       // 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
for (auto Sit = It + 1; Sit != Et; Sit++) {
 
         // C1 is better, so remove C2 and update C2's OutlinedFunction to
         // reflect the removal.
         assert(F2.OccurrenceCount > 0 &&
                "Can't remove OutlinedFunction with no occurrences!");
         F2.OccurrenceCount--;
 
         // Remove the call overhead from the removed sequence.
-        MachineBasicBlock::iterator StartIt = Mapper.InstrList[C2.StartIdx];
-        MachineBasicBlock::iterator EndIt =
-            Mapper.InstrList[C2.StartIdx + C2.Len - 1];
+        F2.Benefit += C2.MInfo.CallOverhead;
 
-        F2.Benefit += TII.getOutliningCallOverhead(StartIt, EndIt).first;
         // Add back one instance of the sequence.
-
         if (F2.Sequence.size() > F2.Benefit)
           F2.Benefit = 0;
         else
           F2.Benefit -= F2.Sequence.size();
 
         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 {
         // C2 is better, so remove C1 and update C1's OutlinedFunction to
         // reflect the removal.
         assert(F1.OccurrenceCount > 0 &&
                "Can't remove OutlinedFunction with no occurrences!");
         F1.OccurrenceCount--;
 
         // Remove the call overhead from the removed sequence.
-        MachineBasicBlock::iterator StartIt = Mapper.InstrList[C1.StartIdx];
-        MachineBasicBlock::iterator EndIt =
-            Mapper.InstrList[C1.StartIdx + C1.Len - 1];
-
-        F1.Benefit += TII.getOutliningCallOverhead(StartIt, EndIt).first;
+        F1.Benefit += C1.MInfo.CallOverhead;
 
         // Add back one instance of the sequence.
         if (F1.Sequence.size() > F1.Benefit)
           F1.Benefit = 0;
         else
           F1.Benefit -= F1.Sequence.size();
 
         C1.InCandidateList = false;

 
 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.
-  size_t MaxCandidateLen = 0;              // Length of the longest candidate.
+  unsigned MaxCandidateLen = 0;            // Length of the longest candidate.
 
   MaxCandidateLen =
       findCandidates(ST, TII, Mapper, CandidateList, FunctionList);
 
   // 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());
MachineOutliner::createOutlinedFunction(Module &M, const OutlinedFunction &OF,
   MachineFunction &MF = MMI.getOrCreateMachineFunction(*F);
   MachineBasicBlock &MBB = *MF.CreateMachineBasicBlock();
   const TargetSubtargetInfo &STI = MF.getSubtarget();
   const TargetInstrInfo &TII = *STI.getInstrInfo();
 
   // Insert the new function into the module.
   MF.insert(MF.begin(), &MBB);
 
-  TII.insertOutlinerPrologue(MBB, MF, OF.FrameClass);
+  TII.insertOutlinerPrologue(MBB, MF, OF.MInfo);
 
   // Copy over the instructions for the function using the integer mappings in
   // its sequence.
   for (unsigned Str : OF.Sequence) {
     MachineInstr *NewMI =
         MF.CloneMachineInstr(Mapper.IntegerInstructionMap.find(Str)->second);
     NewMI->dropMemRefs();
 
     // Don't keep debug information for outlined instructions.
     // FIXME: This means outlined functions are currently undebuggable.
     NewMI->setDebugLoc(DebugLoc());
     MBB.insert(MBB.end(), NewMI);
   }
 
-  TII.insertOutlinerEpilogue(MBB, MF, OF.FrameClass);
+  TII.insertOutlinerEpilogue(MBB, MF, OF.MInfo);
 
   return &MF;
 }
 
 bool MachineOutliner::outline(Module &M,
                               const ArrayRef<Candidate> &CandidateList,
                               std::vector<OutlinedFunction> &FunctionList,
                               InstructionMapper &Mapper) {
 
   bool OutlinedSomething = false;
-
   // Replace the candidates with calls to their respective outlined functions.
   for (const Candidate &C : CandidateList) {
 
     // Was the candidate removed during pruneOverlaps?
     if (!C.InCandidateList)
       continue;
 
     // If not, then look at its OutlinedFunction.
if (!OF.MF) {
       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, C.CallClass);
+    TII.insertOutlinedCall(M, *MBB, StartIt, *MF, C.MInfo);
     StartIt = Mapper.InstrList[C.StartIdx];
     MBB->erase(StartIt, EndIt);
 
     OutlinedSomething = true;
 
     // Statistics.
     NumOutlined++;
   }

lib/Target/AArch64/AArch64InstrInfo.h

public:
   decomposeMachineOperandsTargetFlags(unsigned TF) const override;
   ArrayRef<std::pair<unsigned, const char *>>
   getSerializableDirectMachineOperandTargetFlags() const override;
   ArrayRef<std::pair<unsigned, const char *>>
   getSerializableBitmaskMachineOperandTargetFlags() const override;
   ArrayRef<std::pair<MachineMemOperand::Flags, const char *>>
   getSerializableMachineMemOperandTargetFlags() const override;
 
+  bool
+  canOutlineWithoutLRSave(MachineBasicBlock::iterator &CallInsertionPt) const;
   bool isFunctionSafeToOutlineFrom(MachineFunction &MF) const override;
-  std::pair<size_t, unsigned>
-  getOutliningCallOverhead(MachineBasicBlock::iterator &StartIt,
-                           MachineBasicBlock::iterator &EndIt) const override;
-  std::pair<size_t, unsigned> getOutliningFrameOverhead(
+  MachineOutlinerInfo getOutlininingCandidateInfo(
       std::vector<
           std::pair<MachineBasicBlock::iterator, MachineBasicBlock::iterator>>
-          &CandidateClass) const override;
+          &RepeatedSequenceLocs) const override;
   AArch64GenInstrInfo::MachineOutlinerInstrType
   getOutliningType(MachineInstr &MI) const override;
   void insertOutlinerEpilogue(MachineBasicBlock &MBB, MachineFunction &MF,
-                              unsigned FrameClass) const override;
+                              const MachineOutlinerInfo &MInfo) const override;
   void insertOutlinerPrologue(MachineBasicBlock &MBB, MachineFunction &MF,
-                              unsigned FrameClass) const override;
+                              const MachineOutlinerInfo &MInfo) const override;
   MachineBasicBlock::iterator
   insertOutlinedCall(Module &M, MachineBasicBlock &MBB,
                      MachineBasicBlock::iterator &It, MachineFunction &MF,
-                     unsigned CallClass) const override;
+                     const MachineOutlinerInfo &MInfo) const override;
   /// Returns true if the instruction has a shift left that can be executed
   /// more efficiently.
   bool isExynosShiftLeftFast(const MachineInstr &MI) const;
   /// Returns true if the instruction has a shift by immediate that can be
   /// executed in one cycle less.
   bool isFalkorShiftExtFast(const MachineInstr &MI) const;
 
 private:

lib/Target/AArch64/AArch64InstrInfo.cpp

 #include "AArch64InstrInfo.h"
 #include "AArch64MachineFunctionInfo.h"
 #include "AArch64Subtarget.h"
 #include "MCTargetDesc/AArch64AddressingModes.h"
 #include "Utils/AArch64BaseInfo.h"
 #include "llvm/ADT/ArrayRef.h"
 #include "llvm/ADT/STLExtras.h"
 #include "llvm/ADT/SmallVector.h"
+#include "llvm/CodeGen/LiveRegUnits.h"
 #include "llvm/CodeGen/MachineBasicBlock.h"
 #include "llvm/CodeGen/MachineFrameInfo.h"
 #include "llvm/CodeGen/MachineFunction.h"
 #include "llvm/CodeGen/MachineInstr.h"
 #include "llvm/CodeGen/MachineInstrBuilder.h"
 #include "llvm/CodeGen/MachineMemOperand.h"
 #include "llvm/CodeGen/MachineOperand.h"
 #include "llvm/CodeGen/MachineRegisterInfo.h"

 ArrayRef<std::pair<MachineMemOperand::Flags, const char *>>
 AArch64InstrInfo::getSerializableMachineMemOperandTargetFlags() const {
   static const std::pair<MachineMemOperand::Flags, const char *> TargetFlags[] =
       {{MOSuppressPair, "aarch64-suppress-pair"},
        {MOStridedAccess, "aarch64-strided-access"}};
   return makeArrayRef(TargetFlags);
 }
 
-// Constants defining how certain sequences should be outlined.
-const unsigned MachineOutlinerDefaultFn = 0;
-const unsigned MachineOutlinerTailCallFn = 1;
-
-std::pair<size_t, unsigned> AArch64InstrInfo::getOutliningCallOverhead(
-    MachineBasicBlock::iterator &StartIt,
-    MachineBasicBlock::iterator &EndIt) const {
-  // Is this a tail-call?
-  if (EndIt->isTerminator()) {
-    // Yes, so we only have to emit a call. Return a cost of 1 + signify that
-    // this candidate should be tail-called.
-    return std::make_pair(1, MachineOutlinerTailCallFn);
-  }
-
-  // No, so save + restore LR.
-  return std::make_pair(3, MachineOutlinerDefaultFn);
-}
-
-std::pair<size_t, unsigned> AArch64InstrInfo::getOutliningFrameOverhead(
-    std::vector<std::pair<MachineBasicBlock::iterator,
-                          MachineBasicBlock::iterator>> &CandidateClass) const {
-
-  // Is the last instruction in this class a terminator?
-  if (CandidateClass[0].second->isTerminator())
-    return std::make_pair(0, MachineOutlinerTailCallFn);
+/// Constants defining how certain sequences should be outlined.
+/// This encompasses how an outlined function should be called, and what kind of
+/// frame should be emitted for that outlined function.
+///
+/// \p MachineOutlinerDefault implies that the function should be called with
+/// a save and restore of LR to the stack.
+///
+/// That is,
+///
+/// I1     Save LR                    OUTLINED_FUNCTION:
+/// I2 --> BL OUTLINED_FUNCTION       I1
+/// I3     Restore LR                 I2
+///                                   I3
+///                                   RET
+///
+/// * Call construction overhead: 3 (save + BL + restore)
+/// * Frame construction overhead: 1 (ret)
+/// * Requires stack fixups? Yes
+///
+/// \p MachineOutlinerTailCall implies that the function is being created from
+/// a sequence of instructions ending in a return.
+///
+/// That is,
+///
+/// I1                             OUTLINED_FUNCTION:
+/// I2 --> B OUTLINED_FUNCTION     I1
+/// RET                            I2
+///                                RET
+///
+/// * Call construction overhead: 1 (B)
+/// * Frame construction overhead: 0 (Return included in sequence)
+/// * Requires stack fixups? No
+///
+/// \p MachineOutlinerNoLRSave implies that the function should be called using
+/// a BL instruction, but doesn't require LR to be saved and restored. This
+/// happens when LR is known to be dead.
+///
+/// That is,
+///
+/// I1                                OUTLINED_FUNCTION:
+/// I2 --> BL OUTLINED_FUNCTION       I1
+/// I3                                I2
+///                                   I3
+///                                   RET
+///
+/// * Call construction overhead: 1 (BL)
+/// * Frame construction overhead: 1 (RET)
+/// * Requires stack fixups? No
+///
+enum MachineOutlinerClass {
+  MachineOutlinerDefault,  /// Emit a save, restore, call, and return.
+  MachineOutlinerTailCall, /// Only emit a branch.
+  MachineOutlinerNoLRSave  /// Emit a call and return.
+};
+
+bool AArch64InstrInfo::canOutlineWithoutLRSave(
+    MachineBasicBlock::iterator &CallInsertionPt) const {
+  // Was LR saved in the function containing this basic block?
+  MachineBasicBlock &MBB = *(CallInsertionPt->getParent());
+  LiveRegUnits LRU(getRegisterInfo());
+  LRU.addLiveOuts(MBB);
+
+  // Get liveness information from the end of the block to the end of the
+  // prospective outlined region.
+  std::for_each(MBB.rbegin(),
+               (MachineBasicBlock::reverse_iterator)CallInsertionPt,
+               [&LRU](MachineInstr &MI) {LRU.stepBackward(MI);}
+               );
+
+  // If the link register is available at this point, then we can safely outline
+  // the region without saving/restoring LR. Otherwise, we must emit a save and
+  // restore.
+  return LRU.available(AArch64::LR);
+}
+
+AArch64GenInstrInfo::MachineOutlinerInfo
+AArch64InstrInfo::getOutlininingCandidateInfo(
+    std::vector<
+        std::pair<MachineBasicBlock::iterator, MachineBasicBlock::iterator>>
+        &RepeatedSequenceLocs) const {
+
+  unsigned CallID = MachineOutlinerDefault;
+  unsigned FrameID = MachineOutlinerDefault;
+  unsigned NumInstrsForCall = 3;
+  unsigned NumInstrsToCreateFrame = 1;
+
+  auto DoesntNeedLRSave =
+      [this](std::pair<MachineBasicBlock::iterator, MachineBasicBlock::iterator>
+                 &I) { return canOutlineWithoutLRSave(I.second); };
+
+  // If the last instruction in any candidate is a terminator, then we should
+  // tail call all of the candidates.
+  if (RepeatedSequenceLocs[0].second->isTerminator()) {
+    CallID = MachineOutlinerTailCall;
+    FrameID = MachineOutlinerTailCall;
+    NumInstrsForCall = 1;
+    NumInstrsToCreateFrame = 0;
+  }
+
+  else if (std::all_of(RepeatedSequenceLocs.begin(), RepeatedSequenceLocs.end(),
+                       DoesntNeedLRSave)) {
+    CallID = MachineOutlinerNoLRSave;
+    FrameID = MachineOutlinerNoLRSave;
+    NumInstrsForCall = 1;
+    NumInstrsToCreateFrame = 1;
+  }
 
-  // No, so we have to add a return to the end.
-  return std::make_pair(1, MachineOutlinerDefaultFn);
+  return MachineOutlinerInfo(NumInstrsForCall, NumInstrsToCreateFrame, CallID,
+                             FrameID);
 }
 
 bool AArch64InstrInfo::isFunctionSafeToOutlineFrom(MachineFunction &MF) const {
-  return MF.getFunction()->hasFnAttribute(Attribute::NoRedZone);
+  // If MF has a red zone, then we ought not to outline from it, since outlined
+  // functions can modify/read from the stack.
+  // If MF's address is taken, then we don't want to outline from it either
+  // since we don't really know what the user is doing with it.
+  return MF.getFunction()->hasFnAttribute(Attribute::NoRedZone) &&
+         !MF.getFunction()->hasAddressTaken();
 }
 
 AArch64GenInstrInfo::MachineOutlinerInstrType
 AArch64InstrInfo::getOutliningType(MachineInstr &MI) const {
 
   MachineFunction *MF = MI.getParent()->getParent();
   AArch64FunctionInfo *FuncInfo = MF->getInfo<AArch64FunctionInfo>();
 
   // Don't outline LOHs.
if (MI.readsRegister(AArch64::W30, &getRegisterInfo()) ||
       MI.modifiesRegister(AArch64::W30, &getRegisterInfo()))
     return MachineOutlinerInstrType::Illegal;
 
   // Make sure none of the operands are un-outlinable.
   for (const MachineOperand &MOP : MI.operands()) {
     if (MOP.isCPI() || MOP.isJTI() || MOP.isCFIIndex() || MOP.isFI() ||
         MOP.isTargetIndex())
       return MachineOutlinerInstrType::Illegal;
+
+    // Don't outline anything that uses the link register.
+    if (MOP.isReg() && getRegisterInfo().regsOverlap(MOP.getReg(), AArch64::LR))
+      return MachineOutlinerInstrType::Illegal;
   }
 
   // Does this use the stack?
   if (MI.modifiesRegister(AArch64::SP, &RI) ||
       MI.readsRegister(AArch64::SP, &RI)) {
 
     // Is it a memory operation?
     if (MI.mayLoadOrStore()) {
for (MachineInstr &MI : MBB) {
     // We've pushed the return address to the stack, so add 16 to the offset.
     // This is safe, since we already checked if it would overflow when we
     // checked if this instruction was legal to outline.
     int64_t NewImm = (Offset + 16) / Scale;
     StackOffsetOperand.setImm(NewImm);
   }
 }
 
-void AArch64InstrInfo::insertOutlinerEpilogue(MachineBasicBlock &MBB,
-                                              MachineFunction &MF,
-                                              unsigned FrameClass) const {
+void AArch64InstrInfo::insertOutlinerEpilogue(
+    MachineBasicBlock &MBB, MachineFunction &MF,
+    const MachineOutlinerInfo &MInfo) const {
 
   // If this is a tail call outlined function, then there's already a return.
-  if (FrameClass == MachineOutlinerTailCallFn)
+  if (MInfo.FrameConstructionID == MachineOutlinerTailCall)
     return;
 
   // It's not a tail call, so we have to insert the return ourselves.
   MachineInstr *ret = BuildMI(MF, DebugLoc(), get(AArch64::RET))
                           .addReg(AArch64::LR, RegState::Undef);
   MBB.insert(MBB.end(), ret);
 
+  // Did we have to modify the stack by saving the link register?
+  if (MInfo.FrameConstructionID == MachineOutlinerNoLRSave)
+    return;
+
+  // We modified the stack.
   // Walk over the basic block and fix up all the stack accesses.
   fixupPostOutline(MBB);
 }
 
-void AArch64InstrInfo::insertOutlinerPrologue(MachineBasicBlock &MBB,
-                                              MachineFunction &MF,
-                                              unsigned FrameClass) const {}
+void AArch64InstrInfo::insertOutlinerPrologue(
+    MachineBasicBlock &MBB, MachineFunction &MF,
+    const MachineOutlinerInfo &MInfo) const {}
 
 MachineBasicBlock::iterator AArch64InstrInfo::insertOutlinedCall(
     Module &M, MachineBasicBlock &MBB, MachineBasicBlock::iterator &It,
-    MachineFunction &MF, unsigned CallClass) const {
+    MachineFunction &MF, const MachineOutlinerInfo &MInfo) const {
 
   // Are we tail calling?
-  if (CallClass == MachineOutlinerTailCallFn) {
+  if (MInfo.CallConstructionID == MachineOutlinerTailCall) {
     // If yes, then we can just branch to the label.
     It = MBB.insert(It, BuildMI(MF, DebugLoc(), get(AArch64::B))
                             .addGlobalAddress(M.getNamedValue(MF.getName())));
     return It;
   }
 
-  // We're not tail calling, so we have to save LR before the call and restore
-  // it after.
+  // Are we saving the link register?
+  if (MInfo.CallConstructionID == MachineOutlinerNoLRSave) {
+    // No, so just insert the call.
+    It = MBB.insert(It, BuildMI(MF, DebugLoc(), get(AArch64::BL))
+                            .addGlobalAddress(M.getNamedValue(MF.getName())));
+    return It;
+  }
+
+  // We have a default call. Save the link register.
   MachineInstr *STRXpre = BuildMI(MF, DebugLoc(), get(AArch64::STRXpre))
                               .addReg(AArch64::SP, RegState::Define)
                               .addReg(AArch64::LR)
                               .addReg(AArch64::SP)
                               .addImm(-16);
   It = MBB.insert(It, STRXpre);
   It++;
 

lib/Target/X86/X86InstrInfo.h

MachineInstr *optimizeLoadInstr(MachineInstr &MI,
                                   MachineInstr *&DefMI) const override;
 
   std::pair<unsigned, unsigned>
   decomposeMachineOperandsTargetFlags(unsigned TF) const override;
 
   ArrayRef<std::pair<unsigned, const char *>>
   getSerializableDirectMachineOperandTargetFlags() const override;
 
-  std::pair<size_t, unsigned>
-  getOutliningCallOverhead(MachineBasicBlock::iterator &StartIt,
-                           MachineBasicBlock::iterator &EndIt) const override;
-
-  std::pair<size_t, unsigned> getOutliningFrameOverhead(
+  virtual MachineOutlinerInfo getOutlininingCandidateInfo(
       std::vector<
           std::pair<MachineBasicBlock::iterator, MachineBasicBlock::iterator>>
-          &CandidateClass) const override;
+          &RepeatedSequenceLocs) const override;
 
   bool isFunctionSafeToOutlineFrom(MachineFunction &MF) const override;
 
   llvm::X86GenInstrInfo::MachineOutlinerInstrType
   getOutliningType(MachineInstr &MI) const override;
 
   void insertOutlinerEpilogue(MachineBasicBlock &MBB, MachineFunction &MF,
-                              unsigned FrameClass) const override;
+                              const MachineOutlinerInfo &MInfo) const override;
 
   void insertOutlinerPrologue(MachineBasicBlock &MBB, MachineFunction &MF,
-                              unsigned FrameClass) const override;
+                              const MachineOutlinerInfo &MInfo) const override;
 
   MachineBasicBlock::iterator
   insertOutlinedCall(Module &M, MachineBasicBlock &MBB,
                      MachineBasicBlock::iterator &It, MachineFunction &MF,
-                     unsigned CallClass) const override;
+                     const MachineOutlinerInfo &MInfo) const override;
 
 protected:
   /// Commutes the operands in the given instruction by changing the operands
   /// order and/or changing the instruction's opcode and/or the immediate value
   /// operand.
   ///
   /// The arguments 'CommuteOpIdx1' and 'CommuteOpIdx2' specify the operands
   /// to be commuted.

lib/Target/X86/X86InstrInfo.cpp

struct LDTLSCleanup : public MachineFunctionPass {
     }
   };
 }
 
 char LDTLSCleanup::ID = 0;
 FunctionPass*
 llvm::createCleanupLocalDynamicTLSPass() { return new LDTLSCleanup(); }
 
-// Constants defining how certain sequences should be outlined.
-const unsigned MachineOutlinerDefaultFn = 0;
-const unsigned MachineOutlinerTailCallFn = 1;
-
-std::pair<size_t, unsigned> X86InstrInfo::getOutliningCallOverhead(
-    MachineBasicBlock::iterator &StartIt,
-    MachineBasicBlock::iterator &EndIt) const {
-
-  // Is this a tail call? If it is, make note of it.
-  if (EndIt->isTerminator())
-    return std::make_pair(1, MachineOutlinerTailCallFn);
-
-  return std::make_pair(1, MachineOutlinerDefaultFn);
-}
-
-std::pair<size_t, unsigned> X86InstrInfo::getOutliningFrameOverhead(
-  std::vector<std::pair<MachineBasicBlock::iterator, 
-                        MachineBasicBlock::iterator>> &CandidateClass) const {
-  // Is this a tail-call?
-  // Is the last instruction in this class a terminator?
-  if (CandidateClass[0].second->isTerminator())
-    return std::make_pair(0, MachineOutlinerTailCallFn);
+/// Constants defining how certain sequences should be outlined.
+///
+/// \p MachineOutlinerDefault implies that the function is called with a call
+/// instruction, and a return must be emitted for the outlined function frame.
+///
+/// That is,
+///
+/// I1                                 OUTLINED_FUNCTION:
+/// I2 --> call OUTLINED_FUNCTION       I1
+/// I3                                  I2
+///                                     I3
+///                                     ret
+///
+/// * Call construction overhead: 1 (call instruction)
+/// * Frame construction overhead: 1 (return instruction)
+/// 
+/// \p MachineOutlinerTailCall implies that the function is being tail called.
+/// A jump is emitted instead of a call, and the return is already present in
+/// the outlined sequence. That is,
+///
+/// I1                                 OUTLINED_FUNCTION:
+/// I2 --> jmp OUTLINED_FUNCTION       I1
+/// ret                                I2
+///                                    ret
+///
+/// * Call construction overhead: 1 (jump instruction)
+/// * Frame construction overhead: 0 (don't need to return)
+///
+enum MachineOutlinerClass {
+  MachineOutlinerDefault,
+  MachineOutlinerTailCall
+};
 
-  // No, so we have to add a return to the end.
-  return std::make_pair(1, MachineOutlinerDefaultFn);
+X86GenInstrInfo::MachineOutlinerInfo
+X86InstrInfo::getOutlininingCandidateInfo(
+  std::vector<
+      std::pair<MachineBasicBlock::iterator, MachineBasicBlock::iterator>>
+      &RepeatedSequenceLocs) const {
+
+  if (RepeatedSequenceLocs[0].second->isTerminator())
+    return MachineOutlinerInfo(1, // Number of instructions to emit call.
+                               0, // Number of instructions to emit frame.
+                               MachineOutlinerTailCall, // Type of call.
+                               MachineOutlinerTailCall // Type of frame.
+                              );
+  
+  return MachineOutlinerInfo(1, 1, MachineOutlinerDefault,
+                             MachineOutlinerDefault);
 }
 
 bool X86InstrInfo::isFunctionSafeToOutlineFrom(MachineFunction &MF) const {
   return MF.getFunction()->hasFnAttribute(Attribute::NoRedZone);
 }
 
 X86GenInstrInfo::MachineOutlinerInstrType
 X86InstrInfo::getOutliningType(MachineInstr &MI) const {
if (MOP.isCPI() || MOP.isJTI() || MOP.isCFIIndex() || MOP.isFI() ||
         MOP.isTargetIndex())
       return MachineOutlinerInstrType::Illegal;
 
   return MachineOutlinerInstrType::Legal;
 }
 
 void X86InstrInfo::insertOutlinerEpilogue(MachineBasicBlock &MBB,
                                           MachineFunction &MF,
-                                          unsigned FrameClass) const {
-
+                                          const MachineOutlinerInfo &MInfo)
+                                          const {
   // If we're a tail call, we already have a return, so don't do anything.
-  if (FrameClass == MachineOutlinerTailCallFn)
+  if (MInfo.FrameConstructionID == MachineOutlinerTailCall)
     return;
 
   // We're a normal call, so our sequence doesn't have a return instruction.
   // Add it in.
   MachineInstr *retq = BuildMI(MF, DebugLoc(), get(X86::RETQ));
   MBB.insert(MBB.end(), retq);
 }
 
 void X86InstrInfo::insertOutlinerPrologue(MachineBasicBlock &MBB,
                                           MachineFunction &MF,
-                                          unsigned FrameClass) const {}
+                                          const MachineOutlinerInfo &MInfo)
+                                          const {}
 
 MachineBasicBlock::iterator
 X86InstrInfo::insertOutlinedCall(Module &M, MachineBasicBlock &MBB,
                                  MachineBasicBlock::iterator &It,
                                  MachineFunction &MF,
-                                 unsigned CallClass) const {
+                                 const MachineOutlinerInfo &MInfo) const {
   // Is it a tail call?
-  if (CallClass == MachineOutlinerTailCallFn) {
+  if (MInfo.CallConstructionID == MachineOutlinerTailCall) {
     // Yes, just insert a JMP.
     It = MBB.insert(It,
                   BuildMI(MF, DebugLoc(), get(X86::JMP_1))
                       .addGlobalAddress(M.getNamedValue(MF.getName())));
   } else {
     // No, insert a call.
     It = MBB.insert(It,
                   BuildMI(MF, DebugLoc(), get(X86::CALL64pcrel32))
                       .addGlobalAddress(M.getNamedValue(MF.getName())));
   }
 
   return It;
 }

test/CodeGen/AArch64/machine-outliner-remarks.ll

 ; RUN: llc %s -enable-machine-outliner -mtriple=aarch64-unknown-unknown -pass-remarks-missed=machine-outliner -o /dev/null 2>&1 | FileCheck %s
 ; CHECK: machine-outliner-remarks.ll:5:9:
 ; CHECK-SAME: Did not outline 2 instructions from 2 locations.
-; CHECK-SAME: Instructions from outlining all occurrences (9) >=
+; CHECK-SAME: Instructions from outlining all occurrences (7) >=
 ; CHECK-SAME: Unoutlined instruction count (4)
 ; CHECK-SAME: (Also found at: machine-outliner-remarks.ll:13:9)
 ; RUN: llc %s -enable-machine-outliner -mtriple=aarch64-unknown-unknown -o /dev/null -pass-remarks-missed=machine-outliner -pass-remarks-output=%t.yaml
 ; RUN: cat %t.yaml | FileCheck %s -check-prefix=YAML
 ; YAML: --- !Missed
 ; YAML-NEXT: Pass:            machine-outliner
 ; YAML-NEXT: Name:            NotOutliningCheaper
 ; YAML-NEXT: DebugLoc:        { File: machine-outliner-remarks.ll, Line: 5, Column: 9 }
 ; YAML-NEXT: Function:        dog
 ; YAML-NEXT: Args:            
 ; YAML-NEXT:   - String:          'Did not outline '
 ; YAML-NEXT:   - Length:          '2'
 ; YAML-NEXT:   - String:          ' instructions'
 ; YAML-NEXT:   - String:          ' from '
 ; YAML-NEXT:   - NumOccurrences:  '2'
 ; YAML-NEXT:   - String:          ' locations.'
 ; YAML-NEXT:   - String:          ' Instructions from outlining all occurrences ('
-; YAML-NEXT:   - OutliningCost:   '9'
+; YAML-NEXT:   - OutliningCost:   '7'
 ; YAML-NEXT:   - String:          ')'
 ; YAML-NEXT:   - String:          ' >= Unoutlined instruction count ('
 ; YAML-NEXT:   - NotOutliningCost: '4'
 ; YAML-NEXT:   - String:          ')'
 ; YAML-NEXT:   - String:          ' (Also found at: '
 ; YAML-NEXT:   - OtherStartLoc1:  'machine-outliner-remarks.ll:13:9'
 ; YAML-NEXT:     DebugLoc:        { File: machine-outliner-remarks.ll, Line: 13, Column: 9 }
 ; YAML-NEXT:   - String:          ')'

test/CodeGen/AArch64/machine-outliner.mir

 # RUN: llc -mtriple=aarch64--- -run-pass=machine-outliner %s -o - | FileCheck %s
 --- |
-  target triple = "aarch64---"
-  
   define i32 @main() #0 {
-  entry:
     ret i32 0
   }
   
-  attributes #0 = { noinline noredzone nounwind optnone ssp uwtable }
-
-# CHECK-LABEL: @OUTLINED_FUNCTION_0
+  define void @bar(i32 %a) #0 {
+    ret void
+  }
   
+  attributes #0 = { noinline noredzone "no-frame-pointer-elim"="true" }
 ...
 ---
 # This test ensures that we
 # - Create outlined functions
 # - Don't outline anything to do with LR or W30
+# - Save LR when it's not available
 #
 # CHECK-LABEL: name: main
-# CHECK: BL @OUTLINED_FUNCTION_0
-# CHECK: STRHHroW %w16, %x9, %w30, 1, 1
-# CHECK: %lr = ORRXri %xzr, 1
-# CHECK: BL @OUTLINED_FUNCTION_0
-# CHECK: STRHHroW %w16, %x9, %w30, 1, 1
-# CHECK: %lr = ORRXri %xzr, 1
-# CHECK: BL @OUTLINED_FUNCTION_0
-# CHECK: STRHHroW %w16, %x9, %w30, 1, 1
-# CHECK: %lr = ORRXri %xzr, 1
+# CHECK: BL @OUTLINED_FUNCTION_[[F0:[0-9]+]]
+# CHECK-NEXT: early-clobber %sp, %lr = LDRXpost %sp, 16
+# CHECK-NEXT: STRHHroW %w16, %x9, %w30, 1, 1
+# CHECK-NEXT: %lr = ORRXri %xzr, 1
+# CHECK: BL @OUTLINED_FUNCTION_[[F0]]
+# CHECK-NEXT: early-clobber %sp, %lr = LDRXpost %sp, 16
+# CHECK-NEXT: STRHHroW %w16, %x9, %w30, 1, 1
+# CHECK-NEXT: %lr = ORRXri %xzr, 1
+# CHECK: BL @OUTLINED_FUNCTION_[[F0]]
+# CHECK-NEXT: early-clobber %sp, %lr = LDRXpost %sp, 16
+# CHECK-NEXT: STRHHroW %w16, %x9, %w30, 1, 1
+# CHECK-NEXT: %lr = ORRXri %xzr, 1
 name:            main
-alignment:       2
-tracksRegLiveness: true
-frameInfo:       
-  stackSize:       16
-  maxAlignment:    4
-  maxCallFrameSize: 0
-
 body:             |
-  bb.0.entry:
+  bb.0:
     %sp = frame-setup SUBXri %sp, 16, 0
     %x9 = ORRXri %xzr, 1
     %w16 = ORRWri %wzr, 1
     %w30 = ORRWri %wzr, 1
     %lr = ORRXri %xzr, 1
 
     %w16 = ORRWri %wzr, 1
     %w16 = ORRWri %wzr, 1
bb.0:
     %w16 = ORRWri %wzr, 1
     %w16 = ORRWri %wzr, 1
     %w16 = ORRWri %wzr, 1
     %w16 = ORRWri %wzr, 1
     %w16 = ORRWri %wzr, 1
     STRHHroW %w16, %x9, %w30, 1, 1
     %lr = ORRXri %xzr, 1
 
-    %w5 = ORRWri %wzr, 1995
-
     %sp = ADDXri %sp, 16, 0
     RET undef %lr
+
+...
+---
+# This test ensures that we can avoid saving LR when it's available.
+# CHECK-LABEL: bb.1:
+# CHECK: BL @OUTLINED_FUNCTION_[[F1:[0-9]+]], implicit-def %lr, implicit %sp
+# CHECK-NEXT: %w17 = ORRWri %wzr, 2
+# CHECK-NEXT: BL @OUTLINED_FUNCTION_[[F1]], implicit-def %lr, implicit %sp
+# CHECK-NEXT: %w8 = ORRWri %wzr, 0
+name:            bar
+tracksRegLiveness: true
+body:             |
+  bb.0:
+    liveins: %w0, %lr, %w8
+    %sp = frame-setup SUBXri %sp, 32, 0
+    %fp = frame-setup ADDXri %sp, 16, 0
+
+  bb.1:
+    %w17 = ORRWri %wzr, 1
+    %w17 = ORRWri %wzr, 1
+    %w17 = ORRWri %wzr, 1
+    %w17 = ORRWri %wzr, 1
+    %w17 = ORRWri %wzr, 2
+    %w17 = ORRWri %wzr, 1
+    %w17 = ORRWri %wzr, 1
+    %w17 = ORRWri %wzr, 1
+    %w17 = ORRWri %wzr, 1
+    %w8 = ORRWri %wzr, 0
+  
+  bb.2:
+    %fp, %lr = LDPXi %sp, 2
+    %sp = ADDXri %sp, 32, 0
+    RET undef %lr
+
+...
+
+# CHECK-LABEL: name:            OUTLINED_FUNCTION_{{[0-9]}}
+# CHECK=LABEL: name:            OUTLINED_FUNCTION_{{[1-9]}}