[MachineOutliner][AArch64] WA for cases where AArch64 Outliner needs to do more than one stack fixup.

[plotfi]

Summary

My take on https://reviews.llvm.org/D83464

Basically it appears that in cases where we have noreturns or calls in an outliner candidate that has no LR, no available free regs, or any use of SP that we end up hitting stack fixup code with the MachineOutlinerDefault set as the FrameID. This triggers an assert of OF.FrameConstructionID != MachineOutlinerDefault && "Can only fix up stack references once" at llvm/lib/Target/AArch64/AArch64InstrInfo.cpp:6595 because a lot of the fixup code is not tested to handle fixing up more than once and needs some better checks for possible illegal code generated.

This diff has test cases to cover cases where you have noreturn functions as well as cases where you can hit this same behavior with a BLR terminator in your outliner candidate (but not have that counted as a thunk).

I have created a new diff instead of appending to D83464 because Phab wont let me.

[kyulee]

llvm/lib/Target/AArch64/AArch64InstrInfo.cpp

6126

Shouldn't FlagsSetInAll & MachineOutlinerMBBFlags::HasCalls cover?

[plotfi]

llvm/lib/Target/AArch64/AArch64InstrInfo.cpp

6126

Ah nice, I will use that. Thanks!

[plotfi]

Updated based on @kyulee's suggestion to use FlagsSetInAll & MachineOutlinerMBBFlags::HasCalls

[plotfi]

llvm/test/CodeGen/AArch64/machine-outliner-noreturn-no-stack.mir

2

will drop -run-pass=prologepilog for now, since D76570 has not been landed yet.

[plotfi]

@paquette gentle ping

[plotfi]

Filed a bugzilla task to track progress: https://bugs.llvm.org/show_bug.cgi?id=46767

[paquette]

llvm/lib/Target/AArch64/AArch64InstrInfo.cpp

6113–6115

It's actually that it hasn't been implemented, not that it's untested. If you fix up twice you can get wrong code, guaranteed, because you don't know if you'll overflow the instruction.

[plotfi]

llvm/lib/Target/AArch64/AArch64InstrInfo.cpp

6113–6115

Ahh! I will update the comment and the bugzilla task. Do you think the outliner should bail here if it is a Noreturn call or a call or just noreturn call?

[paquette]

llvm/lib/Target/AArch64/AArch64InstrInfo.cpp

6113–6115

I think what you want to do is bail as early as possible when you know

(1) The candidate will require stack fixups (it's noreturn, or MachineOutlinerDefault + SP is used in the sequence)

and

(2) The candidate contains calls

The (probably better) alternative is to remove those candidates from the candidate list rather than bail out.

At this point, you don't know if you're in the Default + SP unused case. You *do* know enough to handle the noreturn case at this point though.

[plotfi]

Solidified the cases where we want to bail out after speaking with @paquette some more

[plotfi]

llvm/lib/Target/AArch64/AArch64InstrInfo.cpp

6164–6169

I can put the above bail-out code here in case we want to only bail if the cases are true and we are sure there isn't enough space to save the LR. @paquette what do you think here?

[paquette]

llvm/lib/Target/AArch64/AArch64InstrInfo.cpp

6110

Why are we saving these instead of the candidates that we could potentially outline?

This part of the code is mostly intended to decide on whether or not we should discard candidates which require stack fixups. (It should probably be factored out into a function to make that clearer, but that's another patch for another day.)

I think that you probably want to drop these candidates from the main list specifically.

I think something like this should work (using the nice range-based erase_if and any_of from STLExtras.h)

if (FlagsSetInAll & MachineOutlinerMBBFlags::HasCalls) {
  // Remove all candidates which would require multiple stack fixups.
  erase_if(RepeatedSequenceLocs, [](Candidate &C) {
    if (C.LRU.available(AArch64::LR) && !IsNoReturn)
      return false;
    return any_of(C, [](const MachineInstr &MI) { return MI.isCall(); });
  });

  if (RepeatedSequenceLocs.size() < 2) {
    // We removed enough candidates that outlining would no longer be
    // beneficial. Bail out.
    RepeatedSequenceLocs.clear();
    return outliner::OutlinedFunction();
  }
}

6129–6132

Maybe an example would be useful here, rather than trying to explain everything in words:

// Fixing up the stack twice is not yet implemented.
// 
// E.g. cases like this are not yet supported:
//
// function_containing_sequence:
//   ...
//   save LR to SP <- Requires stack instr fixups in OUTLINED_FUNCTION_N
//   call OUTLINED_FUNCTION_N
//   restore LR from SP
//   ...
//
// OUTLINED_FUNCTION_N:
//   save LR to SP <- Requires stack instr fixups in OUTLINED_FUNCTION_N
//   ...
//   bl foo
//   restore LR from SP
//   ret
//
// This happens whenever function_containing_sequence
// - Is noreturn
// - Requires that we save LR at the point of the call, but there are no
//   available registers. In this case, we save using SP.
//
// TODO: Check if fixing up twice is safe so we can outline these.

6145

Which file?

6159–6162

I think what we really want to do here is remove the candidates which require multiple fixups from the list.

Let's say we have something like this:

some_noreturn_func:
  ...
  candidate_sequence_containing_call
  ...

dog:
  ...
  candidate_sequence_containing_call
  ...

cat:
  ...
  candidate_sequence_containing_call
  ...

fish:
  ...
  candidate_sequence_containing_call
  ...

The call imposes one round of stack fixups, since we save LR using SP inside the function frame.

Let's say that in dog, cat, and fish, we don't require fixups at the callsite.

For example, we could outline from dog like this:

dog:
  ...
  bl OUTLINED_FUNCTION_N
  ...

This is fine; we fix up the stack instructions once inside OUTLINED_FUNCTION_N.

This patch bails out when we have a situation like this. If you removed the candidate that appeared in some_noreturn_func from the list, however, you'd be able to outline from dog, cat, and fish safely.

So what I'd expect to end up with is, say, something like this:

some_noreturn_func:
  ...
  candidate_sequence_containing_call <-- Don't outline, noreturn means we need to fix up at the callsite.
  ...

dog:
  ...
  bl OUTLINED_FUNCTION_N
  ...

cat:
  ...
  save LR to a register <-- This is fine, it does not impose any stack fixups
  bl OUTLINED_FUNCTION_N
  restore LR from a register
  ...

fish:
  ...
  bl OUTLINED_FUNCTION_N
  ...

[plotfi]

Update based on @paquette's feedback

[paquette]

llvm/lib/Target/AArch64/AArch64InstrInfo.cpp

6179–6180

Now that I've thought about this a bit more, I don't think we should mention noreturn at all.

The only case we actually care about is the case in the example below. Whether or not, say, function_containing_sequence is noreturn is not relevant. The noreturn case only makes it so that we can never have:

function_containing_sequence:
  ...
  bl OUTLINED_FUNCTION_N
  ...

The actual problem we want to address here is MachineOutlinerDefault + requiring a save + restore to SP.

noreturn functions can force us into that case, but they don't disallow say,

function_containing_sequence:
  ...
  mov xn, lr
  bl OUTLINED_FUNCTION_N
  mov lr, xn
  ...

So there's only one case here: when at the call-site, and inside the function itself, we must modify the stack.

6186–6191

considering the comment I left above, you can probably just remove this

6220

maybe use llvm's any_of from STLExtras?

llvm/test/CodeGen/AArch64/machine-outliner-noreturn-no-stack.mir

2

Can you add a test for the non-noreturn case?

Also can you add a test which verifies that the stuff below works?

f1_doesnt_need_stack:
  ...
  seq containing call <- should be outlined
  ...
f2_doesnt_need_stack:
  ...
  seq containing call <- should be outlined
  ...
needs_stack_not_noreturn:
  ...
  seq containing call <- should not be outlined
  ...
needs_stack_noreturn:
  ...
  seq containing call <- should not be outlined
  ...

[plotfi]

llvm/lib/Target/AArch64/AArch64InstrInfo.cpp

6220

llvm::any_of requires a begin and end not a front and back :-/

[plotfi]

Adding BLR-based test case that explodes LR and all other registers to achieve the same assert without noreturn.

[paquette]

llvm/lib/Target/AArch64/AArch64InstrInfo.cpp

6235

typo: purpise

llvm/test/CodeGen/AArch64/machine-outliner-2fixup-blr-terminator.mir

7 ↗

(On Diff #284242)

I don't think you need this, since it isn't referenced at all in the MIR.

8 ↗

(On Diff #284242)

Very minor nit: it's probably easier to write

attributes #0 = {...}

This really doesn't matter so you don't have to change it, but in the future, it's probably a little cleaner for when you have multiple attributes.

Also do you need "frame-pointer"="all" and noinline here? At the very least, I'd expect that noinline shouldn't do anything this late in the pass pipeline.

llvm/test/CodeGen/AArch64/machine-outliner-no-noreturn-no-stack.mir

14–16 ↗

(On Diff #284242)

Nit: can this be consistent with the formatting for @stack_1 and @stack_2?

[plotfi]

llvm/lib/Target/AArch64/AArch64InstrInfo.cpp

6235

Oops, -> :set spell

llvm/test/CodeGen/AArch64/machine-outliner-2fixup-blr-terminator.mir

7 ↗

(On Diff #284242)

Ah yeah, you are right. I had a regular BL before that was using it, but not now.

8 ↗

(On Diff #284242)

Yeah, I was using clang -E to generate these earlier and the # was being problematic.

[paquette]

LGTM

[plotfi]

Closing, cause phab doesn't seem to be auto-closing.
Landed at https://github.com/llvm/llvm-project/commit/7bc03f55539f7f081daea5363f2e4845b2e75f57

llvm/lib/Target/AArch64/AArch64InstrInfo.cpp

outliner::OutlinedFunction AArch64InstrInfo::getOutliningCandidateInfo(
   }
 
   else {
     // We need to decide how to emit calls + frames. We can always emit the same
     // frame if we don't need to save to the stack. If we have to save to the
     // stack, then we need a different frame.
     unsigned NumBytesNoStackCalls = 0;
     std::vector<outliner::Candidate> CandidatesWithoutStackFixups;
+    bool DoCandidatesHaveCallsOrNoReturn =
+        FlagsSetInAll & MachineOutlinerMBBFlags::HasCalls;
 
     // Check if we have to save LR.
     for (outliner::Candidate &C : RepeatedSequenceLocs) {
       C.initLRU(TRI);
 
       // If we have a noreturn caller, then we're going to be conservative and
       // say that we have to save LR. If we don't have a ret at the end of the
       // block, then we can't reason about liveness accurately.
       //
       // FIXME: We can probably do better than always disabling this in
       // noreturn functions by fixing up the liveness info.
       bool IsNoReturn =
           C.getMF()->getFunction().hasFnAttribute(Attribute::NoReturn);
+      DoCandidatesHaveCallsOrNoReturn |= IsNoReturn;
 
       // Is LR available? If so, we don't need a save.
       if (C.LRU.available(AArch64::LR) && !IsNoReturn) {
         NumBytesNoStackCalls += 4;
         C.setCallInfo(MachineOutlinerNoLRSave, 4);
         CandidatesWithoutStackFixups.push_back(C);
       }
 
       // Is an unused register available? If so, we won't modify the stack, so
       // we can outline with the same frame type as those that don't save LR.
       else if (findRegisterToSaveLRTo(C)) {
         NumBytesNoStackCalls += 12;
         C.setCallInfo(MachineOutlinerRegSave, 12);
         CandidatesWithoutStackFixups.push_back(C);
       }
 
       // Is SP used in the sequence at all? If not, we don't have to modify
       // the stack, so we are guaranteed to get the same frame.
       else if (C.UsedInSequence.available(AArch64::SP)) {
         NumBytesNoStackCalls += 12;
         C.setCallInfo(MachineOutlinerDefault, 12);
         CandidatesWithoutStackFixups.push_back(C);
       }
 
       // If we outline this, we need to modify the stack. Pretend we don't
       // outline this by saving all of its bytes.
       else {
         NumBytesNoStackCalls += SequenceSize;
       }
     }
 
     // If there are no places where we have to save LR, then note that we
     // don't have to update the stack. Otherwise, give every candidate the
     // default call type, as long as it's safe to do so.
     if (!AllStackInstrsSafe ||
         NumBytesNoStackCalls <= RepeatedSequenceLocs.size() * 12) {
       RepeatedSequenceLocs = CandidatesWithoutStackFixups;
       FrameID = MachineOutlinerNoLRSave;
     } else {
+      if (DoCandidatesHaveCallsOrNoReturn) {
+        RepeatedSequenceLocs.clear();
+        return outliner::OutlinedFunction();
+      }
+
       SetCandidateCallInfo(MachineOutlinerDefault, 12);
     }
 
     // If we dropped all of the candidates, bail out here.
     if (RepeatedSequenceLocs.size() < 2) {
       RepeatedSequenceLocs.clear();
       return outliner::OutlinedFunction();
     }
   }
 
   // Does every candidate's MBB contain a call? If so, then we might have a call
   // in the range.
   if (FlagsSetInAll & MachineOutlinerMBBFlags::HasCalls) {
     // Check if the range contains a call. These require a save + restore of the
     // link register.
     bool ModStackToSaveLR = false;
     if (std::any_of(FirstCand.front(), FirstCand.back(),
                     [](const MachineInstr &MI) { return MI.isCall(); }))
       ModStackToSaveLR = true;
 
     // Handle the last instruction separately. If this is a tail call, then the
     // last instruction is a call. We don't want to save + restore in this case.
     // However, it could be possible that the last instruction is a call without
     // it being valid to tail call this sequence. We should consider this as
     // well.
     else if (FrameID != MachineOutlinerThunk &&
              FrameID != MachineOutlinerTailCall && FirstCand.back()->isCall())
       ModStackToSaveLR = true;
 
     if (ModStackToSaveLR) {
outliner::OutlinedFunction AArch64InstrInfo::getOutliningCandidateInfo(
   // can be a tail call
   if (FrameID != MachineOutlinerTailCall && CFICount > 0)
     return outliner::OutlinedFunction();
 
   return outliner::OutlinedFunction(RepeatedSequenceLocs, SequenceSize,
                                     NumBytesToCreateFrame, FrameID);
 }
 
 bool AArch64InstrInfo::isFunctionSafeToOutlineFrom(
     MachineFunction &MF, bool OutlineFromLinkOnceODRs) const {
   const Function &F = MF.getFunction();
 
   // Can F be deduplicated by the linker? If it can, don't outline from it.
   if (!OutlineFromLinkOnceODRs && F.hasLinkOnceODRLinkage())
     return false;
 
   // Don't outline from functions with section markings; the program could
   // expect that all the code is in the named section.
   // FIXME: Allow outlining from multiple functions with the same section
   // marking.
   if (F.hasSection())
     return false;
 
   // Outlining from functions with redzones is unsafe since the outliner may
   // modify the stack. Check if hasRedZone is true or unknown; if yes, don't
   // outline from it.
   AArch64FunctionInfo *AFI = MF.getInfo<AArch64FunctionInfo>();
   if (!AFI || AFI->hasRedZone().getValueOr(true))
     return false;
 
   // FIXME: Teach the outliner to generate/handle Windows unwind info.
   if (MF.getTarget().getMCAsmInfo()->usesWindowsCFI())

llvm/test/CodeGen/AArch64/machine-outliner-noreturn-no-stack.mir

+# RUN: llc -mtriple=arm64-apple-ios -run-pass=prologepilog -run-pass=machine-outliner -verify-machineinstrs  %s -o - | FileCheck %s
+
+# Noreturn functions conservatively need to save and restore lr.
+# When there is no available register, the stack is used at call site.
+# If the stack also needs to be set up for a call in the outlined function,
+# bail-out this case since we do not handle adjusting the stack twice.
+
+# CHECK-NOT: OUTLINED_FUNCTION
+
+--- |
+  @g = external global i32
+  define void @stack_1() #0 { ret void }
+  define void @stack_2() #0 { ret void }
+  define void @baz() {
+    ret void
+  }
+  attributes #0 = { noredzone noreturn }
+...
+---
+name:            stack_1
+tracksRegLiveness: true
+body:             |
+  bb.0:
+    liveins: $x4, $x0, $x1, $x2, $x3
+    $w8 = MOVZWi 259, 0
+    $x9 = ADRP target-flags(aarch64-page) @g
+    $x9 = ADDXri $x9, target-flags(aarch64-pageoff, aarch64-nc) @g, 0
+    STRXui  $x9, $x1, 0
+    STRHHui $w8, $x1, 8
+    $w8 = ORRWrs $wzr, $w4, 0, implicit-def $x8
+    STRXui $x8, $x3, 0
+    STPXi $x3, $xzr, $x2, 0
+    $w8 = MOVZWi 271, 0
+    STRHHui $w8, $x2, 8
+    $x8 = ORRXrs $xzr, $x0, 0
+    $x0 = ORRXrs $xzr, $x1, 0
+    $x1 = ORRXrs $xzr, $x2, 0
+    BL @baz, implicit-def dead $lr, implicit $sp, implicit $x8, implicit $x0, implicit $x1, implicit $x3, implicit $x4, implicit-def $sp, implicit-def $x5, implicit-def $x6, implicit-def $x7, implicit-def $x8,  implicit-def $x9, implicit-def $x10,  implicit-def $x11, implicit-def $x12, implicit-def $x13,  implicit-def $x14, implicit-def $x15, implicit-def $x18
+    BRK 1
+...
+---
+name:            stack_2
+tracksRegLiveness: true
+body:             |
+  bb.0:
+    liveins: $x4, $x0, $x1, $x2, $x3
+    $w8 = MOVZWi 259, 0
+    $x9 = ADRP target-flags(aarch64-page) @g
+    $x9 = ADDXri $x9, target-flags(aarch64-pageoff, aarch64-nc) @g, 0
+    STRXui $x9, $x1, 0
+    STRHHui $w8, $x1, 8
+    $w8 = ORRWrs $wzr, $w4, 0, implicit-def $x8
+    STRXui $x8, $x3, 0
+    STPXi $x3, $xzr, $x2, 0
+    $w8 = MOVZWi 271, 0
+    STRHHui $w8, $x2, 8
+    $x8 = ORRXrs $xzr, $x0, 0
+    $x0 = ORRXrs $xzr, $x1, 0
+    $x1 = ORRXrs $xzr, $x2, 0
+    BL @baz, implicit-def dead $lr, implicit $sp, implicit $x8, implicit $x0, implicit $x1, implicit $x3, implicit $x4, implicit-def $sp, implicit-def $x5, implicit-def $x6, implicit-def $x7, implicit-def $x8,  implicit-def $x9, implicit-def $x10,  implicit-def $x11, implicit-def $x12, implicit-def $x13,  implicit-def $x14, implicit-def $x15, implicit-def $x18
+    BRK 1
+...
+---
+name:            baz
+tracksRegLiveness: true
+body:             |
+  bb.0:
+    liveins: $w0, $lr, $w8
+    RET undef $lr