[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

883

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

lib/Target/AArch64/AArch64InstrInfo.cpp

4461–4469

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

4472–4484

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

4530–4535

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

4538–4541

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

1560–1561

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

1567–1569

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

lib/Target/AArch64/AArch64InstrInfo.cpp

4443

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

885–894

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

883

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

4443

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. :)

4472–4484

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.

4530–4535

... 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

4524–4534

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 ;-)

4578–4591

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.

test/CodeGen/AArch64/machine-outliner.mir

112–139

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

4578–4591

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

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].
+  /// function defined by a collection of closed intervals defining repeated
+  /// sequences.
+  ///
+  /// \param CandidateClass A collection of [Start,End] intervals definining the
+  /// start and end points of a repeated sequence.
   ///
   /// \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 {
+  virtual std::pair<size_t, unsigned> getOutliningCallOverhead(
+      std::vector<
+          std::pair<MachineBasicBlock::iterator, MachineBasicBlock::iterator>>
+          &CandidateClass) const {
+    llvm_unreachable(
+        "Target didn't implement TargetInstrInfo::getOutliningCallOverhead!");
+  }
+
+  /// \brief Returns the number of instructions taken to call a candidate with
+  /// a given \p CallClass.
+  virtual unsigned getOutliningCallOverhead(unsigned CallClass) const {
     llvm_unreachable(
         "Target didn't implement TargetInstrInfo::getOutliningCallOverhead!");
   }
 
   /// \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].
   ///

lib/CodeGen/MachineOutliner.cpp

for (SuffixTreeNode *Leaf : ST.LeafVector) {
     size_t StringLen = Leaf->ConcatLen - 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.
     // FIXME: CandidatesForRepeatedSeq and this should be combined.
     std::vector<
         std::pair<MachineBasicBlock::iterator, MachineBasicBlock::iterator>>
         CandidateClass;
 
     // 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);
+                                              0);
         CandidateClass.emplace_back(std::make_pair(StartIt, EndIt));
 
         // Never visit this leaf again.
         M->IsInTree = false;
       }
     }
 
+    size_t SequenceOverhead = StringLen;
+
+    // Get the overall cost of calling every candidate in this class, and the
+    // way that they should all be outlined.
+    std::pair<size_t, unsigned> CallOverheadPair =
+        TII.getOutliningCallOverhead(CandidateClass);
+    size_t CallOverhead = CallOverheadPair.first;
+
+    // Do the same thing for the frame for the candidates.
     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;
 
+    // Outlining would produce more instructions than not, so move on.
     if (NotOutliningCost <= OutliningCost)
       continue;
 
     size_t 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.CallClass = CallOverheadPair.second;
       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]);
 
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 += TII.getOutliningCallOverhead(StartIt, EndIt).first;
+        F2.Benefit += TII.getOutliningCallOverhead(C2.CallClass);
         // 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 += TII.getOutliningCallOverhead(C1.CallClass);
 
         // 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;

 }
 
 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.

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 &MBB) 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> getOutliningCallOverhead(
+      std::vector<
+          std::pair<MachineBasicBlock::iterator, MachineBasicBlock::iterator>>
+          &CandidateClass) const override;
+  unsigned getOutliningCallOverhead(unsigned CallClass) const override;
   std::pair<size_t, unsigned> getOutliningFrameOverhead(
       std::vector<
           std::pair<MachineBasicBlock::iterator, MachineBasicBlock::iterator>>
           &CandidateClass) const override;
   AArch64GenInstrInfo::MachineOutlinerInstrType
   getOutliningType(MachineInstr &MI) const override;
   void insertOutlinerEpilogue(MachineBasicBlock &MBB, MachineFunction &MF,
                               unsigned FrameClass) const override;

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"

 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;
+enum MachineOutlinerClass {
+  MachineOutlinerDefault,
+  MachineOutlinerTailCall,
+  MachineOutlinerNoLRSave
+};
+
+bool AArch64InstrInfo::canOutlineWithoutLRSave(MachineBasicBlock &MBB) const {
+  // Was LR saved in the function containing this basic block?
+  MachineFrameInfo &MFI = MBB.getParent()->getFrameInfo();
+  const std::vector<CalleeSavedInfo> &CSI = MFI.getCalleeSavedInfo();
+
+  bool CalleeSavedLR = false;
+  for (auto CalleeSave : CSI) {
+    if (CalleeSave.getReg() == AArch64::LR) {
+      CalleeSavedLR = true;
+      break;
+    }
+  }
+
+  if (!CalleeSavedLR)
+    return false;
+
+  LiveRegUnits LRU(getRegisterInfo());
+  LRU.addLiveOuts(MBB);
+  for (auto I = MBB.rbegin(); I != MBB.rend(); LRU.stepBackward(*I), ++I) {
+    // Is the link register dead at this point?
+    if (!LRU.available(AArch64::LR))
+      return false;
+  }
+
+  // The link register is dead throughout the entire block. We can skip saving
+  // and restoring it.
+  return true;
+}
+
+unsigned AArch64InstrInfo::getOutliningCallOverhead(unsigned CallClass) const {
+  unsigned Overhead = 1;
+
+  switch (CallClass) {
+  default:
+    llvm_unreachable("Unknown call class!");
+  case MachineOutlinerDefault:
+    Overhead = 3;
+    break;
+  case MachineOutlinerTailCall:
+  case MachineOutlinerNoLRSave:
+    break;
+  }
+
+  return Overhead;
+}
 
 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);
+    std::vector<std::pair<MachineBasicBlock::iterator,
+                          MachineBasicBlock::iterator>> &CandidateClass) const {
+
+  // If the last instruction in any candidate is a terminator, then we should
+  // tail call all of the candidates.
+  if (CandidateClass[0].second->isTerminator())
+    return std::make_pair(CandidateClass.size(), MachineOutlinerTailCall);
+
+  unsigned Type = MachineOutlinerNoLRSave;
+  unsigned NumInstrsToCallFn = 1;
+
+  // Check if it's possible to skip saving LR in every candidate.
+  // TODO: We should be able to get some more candidates by saving + restoring
+  // with a scratch register wherever possible.
+  // TODO: If two candidates can call the same outlined function (for example,
+  // neither uses the stack), we should allow them to belong to the same
+  // candidate class, but support different cost models.
+  for (const auto &Interval : CandidateClass) {
+    if (!canOutlineWithoutLRSave(*(Interval.first->getParent()))) {
+      Type = MachineOutlinerDefault;
+      NumInstrsToCallFn = 3;
+      break;
+    }
   }
 
-  // No, so save + restore LR.
-  return std::make_pair(3, MachineOutlinerDefaultFn);
+  return std::make_pair(CandidateClass.size() * NumInstrsToCallFn, Type);
 }
 
 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);
+    return std::make_pair(0, MachineOutlinerTailCall);
+
+  // No, so we have to add a return to the end. Check if this call is going to
+  // be NoLRSave.
+
+  // Are each of these NoLRSave calls?
+  for (const auto &Interval : CandidateClass) {
+    if (!canOutlineWithoutLRSave(*(Interval.first->getParent())))
+      return std::make_pair(1, MachineOutlinerDefault);
+  }
 
-  // No, so we have to add a return to the end.
-  return std::make_pair(1, MachineOutlinerDefaultFn);
+  return std::make_pair(1, MachineOutlinerNoLRSave);
 }
 
 bool AArch64InstrInfo::isFunctionSafeToOutlineFrom(MachineFunction &MF) const {
-  return MF.getFunction()->hasFnAttribute(Attribute::NoRedZone);
+  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.isPosition())
     return MachineOutlinerInstrType::Illegal;
 
   // Don't touch the link register or W30.
   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()) {
       unsigned Base;  // Filled with the base regiser of MI.
       int64_t Offset; // Filled with the offset of MI.
       unsigned DummyWidth;
 
       // Does it allow us to offset the base register and is the base SP?
       if (!getMemOpBaseRegImmOfsWidth(MI, Base, Offset, DummyWidth, &RI) ||
           Base != AArch64::SP)
         return MachineOutlinerInstrType::Illegal;
 
       // Find the minimum/maximum offset for this instruction and check if
void AArch64InstrInfo::fixupPostOutline(MachineBasicBlock &MBB) const {
   }
 }
 
 void AArch64InstrInfo::insertOutlinerEpilogue(MachineBasicBlock &MBB,
                                               MachineFunction &MF,
                                               unsigned FrameClass) const {
 
   // If this is a tail call outlined function, then there's already a return.
-  if (FrameClass == MachineOutlinerTailCallFn)
+  if (FrameClass == 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 (FrameClass == 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 {}
 
 MachineBasicBlock::iterator AArch64InstrInfo::insertOutlinedCall(
     Module &M, MachineBasicBlock &MBB, MachineBasicBlock::iterator &It,
     MachineFunction &MF, unsigned CallClass) const {
 
   // Are we tail calling?
-  if (CallClass == MachineOutlinerTailCallFn) {
+  if (CallClass == 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 (CallClass == 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> getOutliningCallOverhead(
+      std::vector<
+          std::pair<MachineBasicBlock::iterator, MachineBasicBlock::iterator>>
+          &CandidateClass) const override;
+
+  unsigned getOutliningCallOverhead(unsigned CallClass) const override;
 
   std::pair<size_t, unsigned> getOutliningFrameOverhead(
       std::vector<
           std::pair<MachineBasicBlock::iterator, MachineBasicBlock::iterator>>
           &CandidateClass) const override;
 
   bool isFunctionSafeToOutlineFrom(MachineFunction &MF) const override;
 

lib/Target/X86/X86InstrInfo.cpp

 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;
 
+unsigned X86InstrInfo::getOutliningCallOverhead(unsigned CallClass) const {
+  return 1;
+}
+
 std::pair<size_t, unsigned> X86InstrInfo::getOutliningCallOverhead(
-    MachineBasicBlock::iterator &StartIt,
-    MachineBasicBlock::iterator &EndIt) const {
+  std::vector<std::pair<MachineBasicBlock::iterator, MachineBasicBlock::iterator>>
+  &CandidateClass) const {
 
   // Is this a tail call? If it is, make note of it.
-  if (EndIt->isTerminator())
-    return std::make_pair(1, MachineOutlinerTailCallFn);
+ if (CandidateClass[0].second->isTerminator())
+    return std::make_pair(CandidateClass.size(), MachineOutlinerTailCallFn);
 
-  return std::make_pair(1, MachineOutlinerDefaultFn);
+  return std::make_pair(CandidateClass.size(), 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())

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 }
+  define void @bar(i32 %a) #0 {
+  entry:
+    %a.addr = alloca i32, align 4
+    %x = alloca i32, align 4
+    store i32 %a, i32* %a.addr, align 4
+    %0 = load i32, i32* %a.addr, align 4
+    %cmp = icmp eq i32 %0, 0
+    br i1 %cmp, label %if.then, label %if.end
+  
+  if.then:
+    store i32 1, i32* %x, align 4
+    br label %if.end
+  
+  if.end:
+    ret void
+  }
 
-# CHECK-LABEL: @OUTLINED_FUNCTION_0
+  declare void @llvm.stackprotector(i8*, i8**) #1
   
+  attributes #0 = { noinline noredzone nounwind optnone ssp uwtable "no-frame-pointer-elim"="true" "no-frame-pointer-elim-non-leaf" "stack-protector-buffer-size"="8" }
+  attributes #1 = { nounwind }
 ...
 ---
 # 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-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_0
-# CHECK: STRHHroW %w16, %x9, %w30, 1, 1
-# CHECK: %lr = ORRXri %xzr, 1
+# 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_0
-# CHECK: STRHHroW %w16, %x9, %w30, 1, 1
-# CHECK: %lr = ORRXri %xzr, 1
+# 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
 
bb.0.entry:
     %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.if.then:
+# CHECK: BL @OUTLINED_FUNCTION_1, implicit-def %lr, implicit %sp
+# CHECK-NEXT: %w17 = ORRWri %wzr, 2
+# CHECK-NEXT: BL @OUTLINED_FUNCTION_1, implicit-def %lr, implicit %sp
+# CHECK-NEXT: %w8 = ORRWri %wzr, 0
+name:            bar
+alignment:       2
+exposesReturnsTwice: false
+legalized:       false
+regBankSelected: false
+selected:        false
+tracksRegLiveness: true
+registers:       
+liveins:         
+  - { reg: '%w0', virtual-reg: '' }
+frameInfo:       
+  stackSize:       32
+  offsetAdjustment: 0
+  maxAlignment:    8
+  adjustsStack:    true
+  hasCalls:        true
+  stackProtector:  ''
+  maxCallFrameSize: 0
+fixedStack:      
+stack:           
+  - { id: 0, name: a.addr, type: default, offset: -20, size: 4, alignment: 4, 
+      stack-id: 0, callee-saved-register: '', local-offset: -4, di-variable: '', 
+      di-expression: '', di-location: '' }
+  - { id: 2, name: '', type: spill-slot, offset: -8, size: 8, alignment: 8, 
+      stack-id: 0, callee-saved-register: '%lr', di-variable: '', di-expression: '', 
+      di-location: '' }
+  - { id: 3, name: '', type: spill-slot, offset: -16, size: 8, alignment: 8, 
+      stack-id: 0, callee-saved-register: '%fp', di-variable: '', di-expression: '', 
+      di-location: '' }
+constants:       
+body:             |
+  bb.0.entry:
+    successors: %bb.2.if.end(0x40000000), %bb.1.if.then(0x40000000)
+    liveins: %w0, %lr
+  
+    %sp = frame-setup SUBXri %sp, 32, 0
+    frame-setup STPXi killed %fp, killed %lr, %sp, 2 :: (store 8 into %stack.3), (store 8 into %stack.2)
+    %fp = frame-setup ADDXri %sp, 16, 0
+    frame-setup CFI_INSTRUCTION def_cfa %w29, 16
+    frame-setup CFI_INSTRUCTION offset %w30, -8
+    frame-setup CFI_INSTRUCTION offset %w29, -16
+    STURWi killed %w0, %fp, -4 :: (store 4 into %stack.0.a.addr)
+    %w8 = LDURWi %fp, -4 :: (load 4 from %stack.0.a.addr)
+    CBNZW killed %w8, %bb.2.if.end
+  
+  bb.1.if.then:
+    successors: %bb.2.if.end(0x80000000)
+    %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.if.end:
+    %fp, %lr = LDPXi %sp, 2 :: (load 8 from %stack.3), (load 8 from %stack.2)
+    %sp = ADDXri %sp, 32, 0
+    RET undef %lr
+
+...
+
+# CHECK-LABEL: name:            OUTLINED_FUNCTION_0
+# CHECK=LABEL: name:            OUTLINED_FUNCTION_1
 No newline at end of file