[MachineOutliner] Don't outline sequences where x16/x17/nzcv are live across

[paquette]

Summary

This is for bug 37573. It isn't safe to outline sequences of instructions where x16/x17/nzcv live across the sequence.

This teaches the outliner to check whether or not a specific canidate has x16/x17/nzcv live across it and discard the candidate in the case that that is true.

https://bugs.llvm.org/show_bug.cgi?id=37573

[efriedma]

It was mentioned on https://reviews.llvm.org/D47553 that outlining also needs special handling for atomic sequences (on AArch64, LDXR/STXR etc. can fail if there's an indirect branch between the two). I guess we don't model the exclusive lock as a register at the moment?

Looks fine otherwise.

lib/Target/AArch64/AArch64InstrInfo.cpp

4949

Please put the title of the document in question, instead of a link; links tend to go stale.

[efriedma]

Oh, oops, there's still an issue: this checks whether x16 and x17 which are live after the outlined sequence, but not whether x16 and x17 are live into the outlined sequence.

[paquette]

Updated patch to handle the live-in case. Updated the test to be more explicit as well.

There are now two LRUs maintained rather than one.

• LRUIn: live ins to a candidate
• LRUOut: live outs from a candidate

LRUIn handles the beginning of the block up to the beginning of the candidate. LRUOut handles the beginning of the candidate up to the end of the block.

[efriedma]

We should probably use this functionality to allow outlining instructions which access/modify SP, in cases where we don't need to spill LR. But that's probably better to do as a followup; it could have some unexpected effects on the heuristics.

include/llvm/CodeGen/MachineOutliner.h

163

This could get expensive with a long basic block; I guess we can try to add some sort of cache later if it becomes an issue.

lib/Target/AArch64/AArch64InstrInfo.cpp

4977

We don't care if it's live out: we can ignore the ABI rules because we know the exact destination of the "bx lr".

We could save/restore x16/x17 in theory, instead of giving up, but probably not worth implementing.

5012

I guess there isn't any practical difference between LRUOut and LRUIn here because we don't outline instructions other than calls which read/write LR?

[paquette]

include/llvm/CodeGen/MachineOutliner.h

163

If it'd be possible to infer candidate liveness from the entire block's liveness, we could wrap the candidate's MBB in a struct like

struct CandidateMBB
{
    MachineBasicBlock *MBB;
    LiveRegUnits LRU;
    ...
};

Then we wouldn't risk recalculating liveness on one basic block over and over again in the case of, say, overlapping candidates.

lib/Target/AArch64/AArch64InstrInfo.cpp

4977

We don't care if it's live out: we can ignore the ABI rules because we know the exact destination of the "bx lr".

Oh, good, much simpler then. :)

We could save/restore x16/x17 in theory, instead of giving up, but probably not worth implementing.

Yeah, I think most outlined sequences are too short to make saving/restoring x16/x17 worth it. However, if it turned out that just *one* candidate required the save/restore, then we could possibly salvage the candidate and still outline. I haven't checked to see if that's a common case at all though.

5012

Yep. It's worth a comment explaining that though. LRUOut was effectively what was being used before for LR as well.

(It might be worth it to try allowing instructions that read/write LR in general though, and then discarding candidates that use them if they aren't part of a prospective tail call.)

[paquette]

• Added a comment for saving/restoring LR
• Removed LRUOut, since it's not necessary here, and incurs unnecessary calculation
• LRUIn -> LRU

[efriedma]

include/llvm/CodeGen/MachineOutliner.h

163

I don't think you can do anything useful at a per-block level, unless you want to try to compute some sort of approximation for long blocks. (I wouldn't suggest the approximation approach, though; probably just as easy to keep around one LiveRegUnits for every 100 instructions or so.)

lib/Target/AArch64/AArch64InstrInfo.cpp

4967

Outdated comment.

5011

ADRP doesn't use LR, as far as I know.

[paquette]

Fixed comments.

[Gerolf]

lib/Target/AArch64/AArch64InstrInfo.cpp

4957

Could this be abstracted behind a target-specific interface? Other platforms like x86 etc will have to take care of similar issues and an interface would make this explicit. This could be done in a follow-up commit though.

[paquette]

Ping! :)

lib/Target/AArch64/AArch64InstrInfo.cpp

4957

Yeah, I think it would be a good idea to have targets implement their own version of MachineOutliner.h in general. Then functions like this could be holed away in there. I'll do that in a follow-up.

[efriedma]

include/llvm/CodeGen/MachineOutliner.h

152

This isn't using stepBackward correctly... you have to start from the end of the block to get correct results. I'm surprised this somehow works with your testcases; you probably need a bit better test coverage.

[paquette]

Fixed my mistake using LiveRegUnits. This just walks back across the entire MBB now, so that we know if a register is live across the candidate.

Also extended the test to handle some more edge-cases and complex behaviour, along with some explanation of what's tested.

[efriedma]

include/llvm/CodeGen/MachineOutliner.h

148

This is still not right; it's computing the values which are live at the beginning of the block, not the beginning of the outlined sequence.

You need to start at MBB->rbegin() (the end of the block), and end at ++(MachineBasicBlock::reverse_iterator)front() (the beginning of the outlined sequence).

Alternatively, I guess you could conservatively use LiveRegUnits::accumulate.

[paquette]

Updated diff to actually fix the LRU stuff (oops).

Also updated the test because the change makes a couple (safe) cases that worked with the (wrong) fix no longer work. Added a comment explaining what's going on to the test to keep note of the fact that there are a few more cases we can probably catch.

[efriedma]

Your results weren't making any sense, so I started debugging a bit... it looks like somehow the mechanism for removing candidates is broken. I'm seeing candidates which should have been erased in AArch64InstrInfo::getOutlininingCandidateInfo, showing up later in MachineOutliner::outline. I haven't tracked down exactly where the bug is, though; any ideas?

include/llvm/CodeGen/MachineOutliner.h

151

The addLiveIns call shouldn't be necessary; it's causing the code to reject valid candidates.

[efriedma]

Nevermind, the issue is my fault; should be (MachineBasicBlock::reverse_iterator)front(), not ++(MachineBasicBlock::reverse_iterator)front(). The resulting issue is tricky to debug because MachineInstrs are a circular linked list.

[paquette]

In D47655#1144273, @efriedma wrote:

Nevermind, the issue is my fault; should be (MachineBasicBlock::reverse_iterator)front(), not ++(MachineBasicBlock::reverse_iterator)front(). The resulting issue is tricky to debug because MachineInstrs are a circular linked list.

Ah, okay, I'll update the patch with that in a second.

In D47655#1144262, @efriedma wrote:

Your results weren't making any sense, so I started debugging a bit... it looks like somehow the mechanism for removing candidates is broken.

Even though this wasn't the issue, I think that I ought to go back and refactor/rewrite the candidate removal mechanism anyway. It's pretty messy, and clearly a source of confusion. I'll start poking at that.

[paquette]

Updated diff.

• ++(MachineBasicBlock::reverse_iterator)front()->(MachineBasicBlock::reverse_iterator)front()
• Removed addLiveIns from initLRU in the Candidate class
• Removed FIXME case from the test, because after the fix, this now works correctly

[efriedma]

LGTM

[paquette]

Thanks!

Committed in r335758 (https://reviews.llvm.org/rL335758).

include/llvm/CodeGen/MachineOutliner.h

 /// Contains all data structures shared between the outliner implemented in
 /// MachineOutliner.cpp and target implementations of the outliner.
 ///
 //===----------------------------------------------------------------------===//
 
 #ifndef LLVM_MACHINEOUTLINER_H
 #define LLVM_MACHINEOUTLINER_H
 
+#include "llvm/CodeGen/LiveRegUnits.h"
 #include "llvm/CodeGen/MachineFunction.h"
+#include "llvm/CodeGen/TargetRegisterInfo.h"
 
 namespace outliner {
 
 using namespace llvm;
 
 /// Represents how an instruction should be mapped by the outliner.
 /// \p Legal instructions are those which are safe to outline.
 /// \p LegalTerminator instructions are safe to outline, but only as the
public:
   unsigned FunctionIdx;
 
   /// Set to false if the candidate overlapped with another candidate.
   bool InCandidateList = true;
 
   /// Contains all target-specific information for this \p Candidate.
   TargetCostInfo TCI;
 
+  /// Contains physical register liveness information entering this
+  /// \p Candidate.
+  ///
+  /// This is optionally used by the target to calculate more fine-grained
+  /// cost model information.
+  LiveRegUnits LRU;
+
   /// Return the number of instructions in this Candidate.
   unsigned getLength() const { return Len; }
 
   /// Return the start index of this candidate.
   unsigned getStartIdx() const { return StartIdx; }
 
   /// Return the end index of this candidate.
   unsigned getEndIdx() const { return StartIdx + Len - 1; }
public:
   unsigned Benefit = 0;
 
   Candidate(unsigned StartIdx, unsigned Len,
             MachineBasicBlock::iterator &FirstInst,
             MachineBasicBlock::iterator &LastInst, MachineBasicBlock *MBB,
             unsigned FunctionIdx)
       : StartIdx(StartIdx), Len(Len), FirstInst(FirstInst), LastInst(LastInst),
         MBB(MBB), FunctionIdx(FunctionIdx) {}
+  Candidate() {}
 
   /// 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 getStartIdx() > RHS.getStartIdx();
   }
+
+  /// Compute register liveness up to the end of the repeated sequence.
+  /// Used by targets that need this information for cost model calculation.
+  /// If a target does not need this information, then this should not be
+  /// called.
+  void initLRU(const TargetRegisterInfo &TRI) {
+    // Next, calculate register liveness on entry to this candidate. Walk from
+    // the beginning of the candidate back to the beginning of MBB. Whichever
+    // registers are live in LRU are the candidate's live ins.
+    LRU.init(TRI);
+    LRU.addLiveOuts(*MBB);
+    std::for_each((MachineBasicBlock::reverse_iterator)front(), MBB->rend(),
+                  [this](MachineInstr &MI) { LRU.stepBackward(MI); });
+  }
 };
 
 /// The information necessary to create an outlined function for some
 /// class of candidate.
 struct OutlinedFunction {
 
 private:
   /// The number of candidates for this \p OutlinedFunction.
   unsigned OccurrenceCount = 0;
 
 public:
   std::vector<std::shared_ptr<Candidate>> Candidates;
 
   /// The actual outlined function created.
   /// This is initialized after we go through and create the actual function.
   MachineFunction *MF = nullptr;
 
   /// A number assigned to this function which appears at the end of its name.

lib/CodeGen/MachineOutliner.cpp

 #include "llvm/ADT/Statistic.h"
 #include "llvm/ADT/Twine.h"
 #include "llvm/CodeGen/MachineFunction.h"
 #include "llvm/CodeGen/MachineModuleInfo.h"
 #include "llvm/CodeGen/MachineOptimizationRemarkEmitter.h"
 #include "llvm/CodeGen/MachineRegisterInfo.h"
 #include "llvm/CodeGen/Passes.h"
 #include "llvm/CodeGen/TargetInstrInfo.h"
-#include "llvm/CodeGen/TargetRegisterInfo.h"
 #include "llvm/CodeGen/TargetSubtargetInfo.h"
 #include "llvm/IR/DIBuilder.h"
 #include "llvm/IR/IRBuilder.h"
 #include "llvm/IR/Mangler.h"
 #include "llvm/Support/Allocator.h"
 #include "llvm/Support/CommandLine.h"
 #include "llvm/Support/Debug.h"
 #include "llvm/Support/raw_ostream.h"
for (auto &ChildPair : Parent.Children) {
         }
       }
     }
 
     // We've found something we might want to outline.
     // Create an OutlinedFunction to store it and check if it'd be beneficial
     // to outline.
     TargetCostInfo TCI =
                       TII.getOutlininingCandidateInfo(CandidatesForRepeatedSeq);
     std::vector<unsigned> Seq;
     for (unsigned i = Leaf->SuffixIdx; i < Leaf->SuffixIdx + StringLen; i++)
       Seq.push_back(ST.Str[i]);
     OutlinedFunction OF(FunctionList.size(), CandidatesForRepeatedSeq.size(),
                         Seq, TCI);
     unsigned Benefit = OF.getBenefit();
 
     // Is it better to outline this candidate than not?

lib/Target/AArch64/AArch64InstrInfo.h

public:
   getSerializableDirectMachineOperandTargetFlags() const override;
   ArrayRef<std::pair<unsigned, const char *>>
   getSerializableBitmaskMachineOperandTargetFlags() const override;
   ArrayRef<std::pair<MachineMemOperand::Flags, const char *>>
   getSerializableMachineMemOperandTargetFlags() const override;
 
   /// AArch64 supports the MachineOutliner.
   bool useMachineOutliner() const override { return true; }
-  
-  bool
-  canOutlineWithoutLRSave(MachineBasicBlock::iterator &CallInsertionPt) const;
   bool isFunctionSafeToOutlineFrom(MachineFunction &MF,
                                    bool OutlineFromLinkOnceODRs) const override;
   outliner::TargetCostInfo getOutlininingCandidateInfo(
       std::vector<outliner::Candidate> &RepeatedSequenceLocs) const override;
   outliner::InstrType getOutliningType(MachineBasicBlock::iterator &MIT,
                                        unsigned Flags) const override;
   unsigned getMachineOutlinerMBBFlags(MachineBasicBlock &MBB) const override;
   void

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"
enum MachineOutlinerClass {
   MachineOutlinerThunk,    /// Emit a call and tail-call.
 };
 
 enum MachineOutlinerMBBFlags {
   LRUnavailableSomewhere = 0x2,
   HasCalls = 0x4
 };
 
-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);
-}
 
 outliner::TargetCostInfo AArch64InstrInfo::getOutlininingCandidateInfo(
     std::vector<outliner::Candidate> &RepeatedSequenceLocs) const {
   unsigned SequenceSize =
        std::accumulate(RepeatedSequenceLocs[0].front(),
-                      std::next(RepeatedSequenceLocs[0].back()), 0,
-                      [this](unsigned Sum, const MachineInstr &MI) {
-                        return Sum + getInstSizeInBytes(MI);
-                      });
+
+  // Compute liveness information for each candidate.
+  const TargetRegisterInfo &TRI = getRegisterInfo();
+  std::for_each(RepeatedSequenceLocs.begin(),
+                RepeatedSequenceLocs.end(),
+                [&TRI](outliner::Candidate &C){
+                  C.initLRU(TRI);
+                });
+
+  // According to the AArch64 Procedure Call Standard, the following are
+  // undefined on entry/exit from a function call:
+  //
+  // * Registers x16, x17, (and thus w16, w17)
+  // * Condition codes (and thus the NZCV register)
+  //
+  // Therefore, if we have a situation like this:
+  //
+  // I1 (defines, say, x16)
+  // I2
+  // ---- Outlining candidate --------
+  // I3
+  // ...
+  // ---- End outlining candidate ----
+  // Ik (uses x16)
+  //
+  // We cannot outline that candidate. The value of x16 at Ik is not guaranteed
+  // to be the same as in I1 once we insert a call instruction. Similarly, if
+  // x16 was defined inside the outlining candidate region, we cannot say that
+  // the value of x16 at Ik will be the same there.
+  //
+  // Thus, we must erase every candidate that violates any of these conditions.
+  auto CantGuaranteeValueAcrossCall = [](outliner::Candidate &C) {
+    LiveRegUnits LRU = C.LRU;
+    return (!LRU.available(AArch64::W16) || !LRU.available(AArch64::W17) ||
+            !LRU.available(AArch64::NZCV));
+  };
+
+  // Erase every candidate that violates the restrictions above. (It could be
+  // true that we have viable candidates, so it's not worth bailing out in the
+  // case that, say, 1 out of 20 candidates violate the restructions.)
+  RepeatedSequenceLocs.erase(
+      std::remove_if(RepeatedSequenceLocs.begin(), RepeatedSequenceLocs.end(),
+                    CantGuaranteeValueAcrossCall),
+      RepeatedSequenceLocs.end());
+
+  // At this point, we have only "safe" candidates to outline. Figure out frame
+  // + call instruction information.
   unsigned CallID = MachineOutlinerDefault;
   unsigned FrameID = MachineOutlinerDefault;
   unsigned NumBytesForCall = 12;
   unsigned NumBytesToCreateFrame = 4;
-
-  auto DoesntNeedLRSave = [this](outliner::Candidate &I) {
-    return canOutlineWithoutLRSave(I.back());
-  };
-
   unsigned LastInstrOpcode = RepeatedSequenceLocs[0].back()->getOpcode();
 
   // If the last instruction in any candidate is a terminator, then we should
   // tail call all of the candidates.
   if (RepeatedSequenceLocs[0].back()->isTerminator()) {
     CallID = MachineOutlinerTailCall;
     FrameID = MachineOutlinerTailCall;
     NumBytesForCall = 4;
     NumBytesToCreateFrame = 0;
   }
 
   else if (LastInstrOpcode == AArch64::BL || LastInstrOpcode == AArch64::BLR) {
     // FIXME: Do we need to check if the code after this uses the value of LR?
     CallID = MachineOutlinerThunk;
     FrameID = MachineOutlinerThunk;
     NumBytesForCall = 4;
     NumBytesToCreateFrame = 0;
   }
 
-  else if (std::all_of(RepeatedSequenceLocs.begin(), RepeatedSequenceLocs.end(),
-                       DoesntNeedLRSave)) {
+  // Make sure that LR isn't live on entry to this candidate. The only
+  // instructions that use LR that could possibly appear in a repeated sequence
+  // are calls and ADRPs. Therefore, we only have to check and see if LR is
+  // dead on entry to (or exit from) some candidate.
+  else if (std::all_of(RepeatedSequenceLocs.begin(),
+                       RepeatedSequenceLocs.end(),
+                       [](outliner::Candidate &C) {
+                         return C.LRU.available(AArch64::LR);
+                         })) {
     CallID = MachineOutlinerNoLRSave;
     FrameID = MachineOutlinerNoLRSave;
     NumBytesForCall = 4;
     NumBytesToCreateFrame = 4;
   }
 
   // Check if the range contains a call. These require a save + restore of the
   // link register.

test/CodeGen/AArch64/machine-outliner-bad-register.mir

+# RUN: llc -mtriple=aarch64--- -run-pass=machine-outliner \
+# RUN: -verify-machineinstrs %s -o - | FileCheck %s
+
+# Ensure that we don't outline from regions where x16, x17, or nzcv are live
+# across the outlining candidate. These values are allowed to be clobbered by,
+# say, the linker, in the presence of function calls. Thus, we can't outline
+# these, since the insertion of the outlined call could change the values of
+# these registers.
+--- |
+  ; This function should be the only one with outlining candidates.
+  define void @should_outline(i32 %a) #0 {
+    ret void
+  }
+
+  define void @x16_live() #0 {
+    ret void
+  }
+
+  define void @x17_live() #0 {
+    ret void
+  }
+
+  define void @nzcv_live() #0 {
+    ret void
+  }
+  
+  attributes #0 = { noredzone }
+...
+---
+
+name:            should_outline
+tracksRegLiveness: true
+body:             |
+  bb.0:
+    ; CHECK-LABEL: bb.0:
+    ; CHECK: BL
+    liveins: $w8, $wzr
+    $w8 = ORRWri $wzr, 1
+    $w8 = ORRWri $wzr, 2
+    $w8 = ORRWri $wzr, 3
+    $w8 = ORRWri $wzr, 4
+  bb.1:
+    ; CHECK-LABEL: bb.1:
+    ; CHECK: BL
+    liveins: $w8, $wzr
+    $w8 = ORRWri $wzr, 1
+    $w8 = ORRWri $wzr, 2
+    $w8 = ORRWri $wzr, 3
+    $w8 = ORRWri $wzr, 4
+  bb.2:
+    RET undef $lr
+
+...
+---
+
+name:            x16_live
+tracksRegLiveness: true
+body:             |
+  bb.0:
+    ; CHECK-LABEL: bb.0:
+    ; CHECK-NOT: BL
+    liveins: $w8, $wzr
+    $w8 = ORRWri $wzr, 1
+    $w8 = ORRWri $wzr, 2
+    $w8 = ORRWri $wzr, 3
+    $w8 = ORRWri $wzr, 4
+  bb.1:
+    liveins: $x16
+    RET undef $lr
+
+...
+---
+
+name:            x17_live
+tracksRegLiveness: true
+body:             |
+  bb.0:
+    ; CHECK-LABEL: bb.0:
+    ; CHECK-NOT: BL
+    liveins: $w8, $wzr
+    $w8 = ORRWri $wzr, 1
+    $w8 = ORRWri $wzr, 2
+    $w8 = ORRWri $wzr, 3
+    $w8 = ORRWri $wzr, 4
+  bb.1:
+    liveins: $x17
+    RET undef $lr
+
+...
+---
+
+name:            nzcv_live
+tracksRegLiveness: true
+body:             |
+  bb.0:
+    liveins: $w8, $wzr
+    ; CHECK-LABEL: bb.0:
+    ; CHECK-NOT: BL
+    $w8 = ORRWri $wzr, 1
+    $w8 = ORRWri $wzr, 2
+    $w8 = ORRWri $wzr, 3
+    $w8 = ORRWri $wzr, 4
+  bb.1:
+    liveins: $nzcv
+    RET undef $lr