[MachineOutliner] AArch64: Handle instrs that use SP and will never need fixups

This commit does two things. Firstly, it adds a collection of flags which can
be passed along to the target to encode information about the MBB that an
instruction lives in to the outliner.

Second, it adds some of those flags to the AArch64 outliner in order to add
more stack instructions to the list of legal instructions that are handled
by the outliner. The two flags added check if

- There are calls in the MachineBasicBlock containing the instruction
- The link register is available in the entire block

If the link register is available and there are no calls, then a stack
instruction can always be outlined without fixups, regardless of what it is,
since in this case, the outliner will never modify the stack to create a
call or outlined frame.

The motivation for doing this was checking which instructions are most often
missed by the outliner. Instructions like, say

%sp<def> = ADDXri %sp, 32, 0; flags: FrameDestroy

are very common, but cannot be outlined in the case that the outliner might
modify the stack. This commit allows us to outline instructions like this.
  

llvm-svn: 322048

include/llvm/CodeGen/TargetInstrInfo.h

@@ -1607,11 +1607,18 @@ public:
   enum MachineOutlinerInstrType { Legal, Illegal, Invisible };
 
   /// Returns how or if \p MI should be outlined.
-  virtual MachineOutlinerInstrType getOutliningType(MachineInstr &MI) const {
+  virtual MachineOutlinerInstrType
+  getOutliningType(MachineBasicBlock::iterator &MIT, unsigned Flags) const {
     llvm_unreachable(
         "Target didn't implement TargetInstrInfo::getOutliningType!");
   }
 
+  /// \brief Returns target-defined flags defining properties of the MBB for
+  /// the outliner.
+  virtual unsigned getMachineOutlinerMBBFlags(MachineBasicBlock &MBB) const {
+    return 0x0;
+  }
+
   /// Insert a custom epilogue for outlined functions.
   /// This may be empty, in which case no epilogue or return statement will be
   /// emitted.

lib/CodeGen/MachineOutliner.cpp

@@ -720,11 +720,13 @@ struct InstructionMapper {
   void convertToUnsignedVec(MachineBasicBlock &MBB,
                             const TargetRegisterInfo &TRI,
                             const TargetInstrInfo &TII) {
+    unsigned Flags = TII.getMachineOutlinerMBBFlags(MBB);
+
     for (MachineBasicBlock::iterator It = MBB.begin(), Et = MBB.end(); It != Et;
          It++) {
 
       // Keep track of where this instruction is in the module.
-      switch (TII.getOutliningType(*It)) {
+      switch (TII.getOutliningType(It, Flags)) {
       case TargetInstrInfo::MachineOutlinerInstrType::Illegal:
         mapToIllegalUnsigned(It);
         break;
@@ -1237,7 +1239,6 @@ MachineOutliner::createOutlinedFunction(Module &M, const OutlinedFunction &OF,
   }
 
   TII.insertOutlinerEpilogue(MBB, MF, OF.MInfo);
-
   return &MF;
 }
 

lib/Target/AArch64/AArch64InstrInfo.cpp

@@ -4696,6 +4696,11 @@ enum MachineOutlinerClass {
   MachineOutlinerNoLRSave  /// Emit a call and return.
 };
 
+enum MachineOutlinerMBBFlags {
+  LRUnavailableSomewhere = 0x2,
+  HasCalls = 0x4
+};
+
 bool AArch64InstrInfo::canOutlineWithoutLRSave(
     MachineBasicBlock::iterator &CallInsertionPt) const {
   // Was LR saved in the function containing this basic block?
@@ -4785,10 +4790,36 @@ bool AArch64InstrInfo::isFunctionSafeToOutlineFrom(
   return true;
 }
 
-AArch64GenInstrInfo::MachineOutlinerInstrType
-AArch64InstrInfo::getOutliningType(MachineInstr &MI) const {
+unsigned
+AArch64InstrInfo::getMachineOutlinerMBBFlags(MachineBasicBlock &MBB) const {
+  unsigned Flags = 0x0;
+  // Check if there's a call inside this MachineBasicBlock. If there is, then
+  // set a flag.
+  if (std::any_of(MBB.begin(), MBB.end(),
+                  [](MachineInstr &MI) { return MI.isCall(); }))
+    Flags |= MachineOutlinerMBBFlags::HasCalls;
 
-  MachineFunction *MF = MI.getParent()->getParent();
+  // Check if LR is available through all of the MBB. If it's not, then set
+  // a flag.
+  LiveRegUnits LRU(getRegisterInfo());
+  LRU.addLiveOuts(MBB);
+
+  std::for_each(MBB.rbegin(),
+                MBB.rend(),
+                [&LRU](MachineInstr &MI) { LRU.accumulate(MI); });
+
+  if (!LRU.available(AArch64::LR)) 
+      Flags |= MachineOutlinerMBBFlags::LRUnavailableSomewhere;
+
+  return Flags;
+}
+
+AArch64GenInstrInfo::MachineOutlinerInstrType
+AArch64InstrInfo::getOutliningType(MachineBasicBlock::iterator &MIT,
+                                   unsigned Flags) const {
+  MachineInstr &MI = *MIT;
+  MachineBasicBlock *MBB = MI.getParent();
+  MachineFunction *MF = MBB->getParent();
   AArch64FunctionInfo *FuncInfo = MF->getInfo<AArch64FunctionInfo>();
 
   // Don't outline LOHs.
@@ -4857,6 +4888,7 @@ AArch64InstrInfo::getOutliningType(MachineInstr &MI) const {
     if (Callee->doesNotAccessMemory()) 
       return MachineOutlinerInstrType::Legal;
     
+
     // It accesses memory. Get the machine function for the callee to see if
     // it's safe to outline.
     MachineFunction *CalleeMF = MF->getMMI().getMachineFunction(*Callee);
@@ -4896,7 +4928,52 @@ AArch64InstrInfo::getOutliningType(MachineInstr &MI) const {
   // Does this use the stack?
   if (MI.modifiesRegister(AArch64::SP, &RI) ||
       MI.readsRegister(AArch64::SP, &RI)) {
+    // True if there is no chance that any outlined candidate from this range
+    // could require stack fixups. That is, both
+    // * LR is available in the range (No save/restore around call)
+    // * The range doesn't include calls (No save/restore in outlined frame)
+    // are true.
+    bool MightNeedStackFixUp =
+        (Flags & (MachineOutlinerMBBFlags::LRUnavailableSomewhere |
+                  MachineOutlinerMBBFlags::HasCalls));
 
+    // If this instruction is in a range where it *never* needs to be fixed
+    // up, then we can *always* outline it. This is true even if it's not
+    // possible to fix that instruction up.
+    //
+    // Why? Consider two equivalent instructions I1, I2 where both I1 and I2
+    // use SP. Suppose that I1 sits within a range that definitely doesn't
+    // need stack fixups, while I2 sits in a range that does.
+    //
+    // First, I1 can be outlined as long as we *never* fix up the stack in
+    // any sequence containing it. I1 is already a safe instruction in the
+    // original program, so as long as we don't modify it we're good to go.
+    // So this leaves us with showing that outlining I2 won't break our
+    // program.
+    //
+    // Suppose I1 and I2 belong to equivalent candidate sequences. When we
+    // look at I2, we need to see if it can be fixed up. Suppose I2, (and
+    // thus I1) cannot be fixed up. Then I2 will be assigned an unique
+    // integer label; thus, I2 cannot belong to any candidate sequence (a
+    // contradiction). Suppose I2 can be fixed up. Then I1 can be fixed up
+    // as well, so we're good. Thus, I1 is always safe to outline.
+    //
+    // This gives us two things: first off, it buys us some more instructions
+    // for our search space by deeming stack instructions illegal only when
+    // they can't be fixed up AND we might have to fix them up. Second off,
+    // This allows us to catch tricky instructions like, say,
+    // %xi = ADDXri %sp, n, 0. We can't safely outline these since they might
+    // be paired with later SUBXris, which might *not* end up being outlined.
+    // If we mess with the stack to save something, then an ADDXri messes with
+    // it *after*, then we aren't going to restore the right something from
+    // the stack if we don't outline the corresponding SUBXri first. ADDXris and
+    // SUBXris are extremely common in prologue/epilogue code, so supporting
+    // them in the outliner can be a pretty big win!
+    if (!MightNeedStackFixUp)
+      return MachineOutlinerInstrType::Legal;
+
+    // At this point, we have a stack instruction that we might need to fix
+    // up. We'll handle it if it's a load or store.
     if (MI.mayLoadOrStore()) {
       unsigned Base;  // Filled with the base regiser of MI.
       int64_t Offset; // Filled with the offset of MI.

lib/Target/AArch64/AArch64InstrInfo.h

@@ -359,7 +359,8 @@ public:
           std::pair<MachineBasicBlock::iterator, MachineBasicBlock::iterator>>
           &RepeatedSequenceLocs) const override;
   AArch64GenInstrInfo::MachineOutlinerInstrType
-  getOutliningType(MachineInstr &MI) const override;
+  getOutliningType(MachineBasicBlock::iterator &MIT, unsigned Flags) const override;
+  unsigned getMachineOutlinerMBBFlags(MachineBasicBlock &MBB) const override;
   void insertOutlinerEpilogue(MachineBasicBlock &MBB, MachineFunction &MF,
                               const MachineOutlinerInfo &MInfo) const override;
   void insertOutlinerPrologue(MachineBasicBlock &MBB, MachineFunction &MF,

lib/Target/X86/X86InstrInfo.cpp

@@ -10937,8 +10937,8 @@ bool X86InstrInfo::isFunctionSafeToOutlineFrom(MachineFunction &MF,
 }
 
 X86GenInstrInfo::MachineOutlinerInstrType
-X86InstrInfo::getOutliningType(MachineInstr &MI) const {
-
+X86InstrInfo::getOutliningType(MachineBasicBlock::iterator &MIT,  unsigned Flags) const {
+  MachineInstr &MI = *MIT;
   // Don't allow debug values to impact outlining type.
   if (MI.isDebugValue() || MI.isIndirectDebugValue())
     return MachineOutlinerInstrType::Invisible;

lib/Target/X86/X86InstrInfo.h

@@ -568,7 +568,7 @@ public:
                                    bool OutlineFromLinkOnceODRs) const override;
 
   llvm::X86GenInstrInfo::MachineOutlinerInstrType
-  getOutliningType(MachineInstr &MI) const override;
+  getOutliningType(MachineBasicBlock::iterator &MIT, unsigned Flags) const override;
 
   void insertOutlinerEpilogue(MachineBasicBlock &MBB, MachineFunction &MF,
                               const MachineOutlinerInfo &MInfo) const override;

test/CodeGen/AArch64/machine-outliner.mir

@@ -20,21 +20,29 @@
 # - Create outlined functions
 # - Don't outline anything to do with LR or W30
 # - Save LR when it's not available
+# - Don't outline stack instructions when we might need to save + restore
 #
 # CHECK-LABEL: name: main
+
 # CHECK: BL @OUTLINED_FUNCTION_[[F0:[0-9]+]]
 # CHECK-NEXT: early-clobber %sp, %lr = LDRXpost %sp, 16
+# CHECK-NEXT: %x16 = ADDXri %sp, 48, 0
 # 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: %x16 = ADDXri %sp, 48, 0
 # 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: %x16 = ADDXri %sp, 48, 0
 # CHECK-NEXT: STRHHroW %w16, %x9, %w30, 1, 1
 # CHECK-NEXT: %lr = ORRXri %xzr, 1
 name:            main
+tracksRegLiveness: true
 body:             |
   bb.0:
     %sp = frame-setup SUBXri %sp, 16, 0
@@ -50,9 +58,9 @@ body:             |
     %w16 = ORRWri %wzr, 1
     %w16 = ORRWri %wzr, 1
     %w16 = ORRWri %wzr, 1
+    %x16 = ADDXri %sp, 48, 0;
     STRHHroW %w16, %x9, %w30, 1, 1
     %lr = ORRXri %xzr, 1
-
     %w3 = ORRWri %wzr, 1993
 
     %x20, %x19 = LDPXi %sp, 10
@@ -62,6 +70,7 @@ body:             |
     %w16 = ORRWri %wzr, 1
     %w16 = ORRWri %wzr, 1
     %w16 = ORRWri %wzr, 1
+    %x16 = ADDXri %sp, 48, 0;
     STRHHroW %w16, %x9, %w30, 1, 1
     %lr = ORRXri %xzr, 1 
 
@@ -74,6 +83,7 @@ body:             |
     %w16 = ORRWri %wzr, 1
     %w16 = ORRWri %wzr, 1
     %w16 = ORRWri %wzr, 1
+    %x16 = ADDXri %sp, 48, 0;
     STRHHroW %w16, %x9, %w30, 1, 1
     %lr = ORRXri %xzr, 1
 
@@ -114,6 +124,20 @@ body:             |
     %w8 = ORRWri %wzr, 0
     
   bb.2:
+    %w15 = ORRWri %wzr, 1
+    %w15 = ORRWri %wzr, 1
+    %w15 = ORRWri %wzr, 1
+    %w15 = ORRWri %wzr, 1
+    %x15 = ADDXri %sp, 48, 0;
+    %w9 = ORRWri %wzr, 0
+    %w15 = ORRWri %wzr, 1
+    %w15 = ORRWri %wzr, 1
+    %w15 = ORRWri %wzr, 1
+    %w15 = ORRWri %wzr, 1
+    %x15 = ADDXri %sp, 48, 0;
+    %w8 = ORRWri %wzr, 0
+    
+  bb.3:
     %fp, %lr = LDPXi %sp, 2
     %sp = ADDXri %sp, 32, 0
     RET undef %lr