[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

6199

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

6127–6130

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.

6143

Which file?

6187–6190

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

6171–6172

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.

6178–6183

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

6212

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

6212

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

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

8

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

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

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

8

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

       FrameID = MachineOutlinerNoLRSave;
     } else {
       SetCandidateCallInfo(MachineOutlinerDefault, 12);
+
+      // Bugzilla ID: 46767
+      // TODO: Check if fixing up the stack more than once is safe so we can
+      // outline these.
+      //
+      // An outline resulting in a caller that requires stack fixups at the
+      // callsite to a callee that also requires stack fixups can happen when
+      // there are no available registers at the candidate callsite for a
+      // candidate that itself also has calls.
+      //
+      // In other words if function_containing_sequence in the following pseudo
+      // assembly requires that we save LR at the point of the call, but there
+      // are no available registers: in this case we save using SP and as a
+      // result the SP offsets requires stack fixups by multiples of 16.
+      //
+      // 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
+      //
+      // Because the code to handle more than one stack fixup does not
+      // currently have the proper checks for legality, these cases will assert
+      // in the AArch64 MachineOutliner. This is because the code to do this
+      // needs more hardening, testing, better checks that generated code is
+      // legal, etc and because it is only verified to handle a single pass of
+      // stack fixup.
+      //
+      // The assert happens in AArch64InstrInfo::buildOutlinedFrame to catch
+      // these cases until they are known to be handled. Bugzilla 46767 is
+      // referenced in comments at the assert site.
+      //
+      // To avoid asserting (or generating non-legal code on noassert builds)
+      // we remove all candidates which would need more than one stack fixup by
+      // pruning the cases where the candidate has calls while also having no
+      // available LR and having no available general purpise registers to copy
+      // LR to (ie one extra stack save/restore).
+      //
+      if (FlagsSetInAll & MachineOutlinerMBBFlags::HasCalls) {
+        erase_if(RepeatedSequenceLocs, [this](outliner::Candidate &C) {
+          return (std::any_of(
+                     C.front(), std::next(C.back()),
+                     [](const MachineInstr &MI) { return MI.isCall(); })) &&
+                 (!C.LRU.available(AArch64::LR) || !findRegisterToSaveLRTo(C));
+        });
+      }
     }
 
     // If we dropped all of the candidates, bail out here.

   if (std::any_of(MBB.instr_begin(), MBB.instr_end(), IsNonTailCall)) {
     // Fix up the instructions in the range, since we're going to modify the
     // stack.
+
+    // Bugzilla ID: 46767
+    // TODO: Check if fixing up twice is safe so we can outline these.
     assert(OF.FrameConstructionID != MachineOutlinerDefault &&
            "Can only fix up stack references once");
     fixupPostOutline(MBB);

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

+# RUN: llc -mtriple=aarch64--- -run-pass=machine-outliner \
+# RUN: -verify-machineinstrs %s -o - | FileCheck %s
+
+# CHECK-NOT: OUTLINED_FUNCTION
+
+--- |
+  @var_i32 = global i32 42
+  define void @f1() minsize noinline noredzone "frame-pointer"="all" "branch-target-enforcement" { ret void }
+  define void @f2() minsize noinline noredzone "frame-pointer"="all" "branch-target-enforcement" { ret void }
+  define void @f3() minsize noinline noredzone "frame-pointer"="all" "branch-target-enforcement" { ret void }
+  define void @f4() minsize noinline noredzone "frame-pointer"="all" "branch-target-enforcement" { ret void }
+...
+---
+name: f1
+tracksRegLiveness: true
+body: |
+  bb.0:
+  liveins: $lr, $x0, $x1, $x2, $x3, $x4, $x5, $x6, $x7, $x8, $x9, $x10, $x11, $x12, $x13, $x14, $x15, $x18, $x19, $x20, $x21, $x22, $x23, $x20, $x21, $x22, $x23, $x24, $x25, $x26, $x27, $x28, $fp
+    $x20, $x19 = LDPXi $sp, 11
+    $x20, $x19 = LDPXi $sp, 12
+    $x20, $x19 = LDPXi $sp, 13
+    $x20, $x19 = LDPXi $sp, 14
+    $x20, $x19 = LDPXi $sp, 18
+    $x20, $x19 = LDPXi $sp, 19
+    $x20, $x19 = LDPXi $sp, 20
+    $x20, $x19 = LDPXi $sp, 21
+    BLR $x20, implicit $sp
+  bb.2:
+  liveins: $lr, $x0, $x1, $x2, $x3, $x4, $x5, $x6, $x7, $x8, $x9, $x10, $x11, $x12, $x13, $x14, $x15, $x18, $x19, $x20, $x21, $x22, $x23, $x20, $x21, $x22, $x23, $x24, $x25, $x26, $x27, $x28, $fp
+    RET undef $lr
+...
+---
+name: f2
+tracksRegLiveness: true
+body: |
+  bb.0:
+  liveins: $lr, $x0, $x1, $x2, $x3, $x4, $x5, $x6, $x7, $x8, $x9, $x10, $x11, $x12, $x13, $x14, $x15, $x18, $x19, $x20, $x21, $x22, $x23, $x20, $x21, $x22, $x23, $x24, $x25, $x26, $x27, $x28, $fp
+    $x20, $x19 = LDPXi $sp, 11
+    $x20, $x19 = LDPXi $sp, 12
+    $x20, $x19 = LDPXi $sp, 13
+    $x20, $x19 = LDPXi $sp, 14
+    $x20, $x19 = LDPXi $sp, 18
+    $x20, $x19 = LDPXi $sp, 19
+    $x20, $x19 = LDPXi $sp, 20
+    $x20, $x19 = LDPXi $sp, 21
+    BLR $x20, implicit $sp
+  bb.2:
+  liveins: $lr, $x0, $x1, $x2, $x3, $x4, $x5, $x6, $x7, $x8, $x9, $x10, $x11, $x12, $x13, $x14, $x15, $x18, $x19, $x20, $x21, $x22, $x23, $x20, $x21, $x22, $x23, $x24, $x25, $x26, $x27, $x28, $fp
+    RET undef $lr
+...
+---
+name: f3
+tracksRegLiveness: true
+body: |
+  bb.0:
+  liveins: $lr, $x0, $x1, $x2, $x3, $x4, $x5, $x6, $x7, $x8, $x9, $x10, $x11, $x12, $x13, $x14, $x15, $x18, $x19, $x20, $x21, $x22, $x23, $x20, $x21, $x22, $x23, $x24, $x25, $x26, $x27, $x28, $fp
+    $x20, $x19 = LDPXi $sp, 11
+    $x20, $x19 = LDPXi $sp, 12
+    $x20, $x19 = LDPXi $sp, 13
+    $x20, $x19 = LDPXi $sp, 14
+    $x20, $x19 = LDPXi $sp, 18
+    $x20, $x19 = LDPXi $sp, 19
+    $x20, $x19 = LDPXi $sp, 20
+    $x20, $x19 = LDPXi $sp, 21
+    BLR $x20, implicit $sp
+  bb.2:
+  liveins: $lr, $x0, $x1, $x2, $x3, $x4, $x5, $x6, $x7, $x8, $x9, $x10, $x11, $x12, $x13, $x14, $x15, $x18, $x19, $x20, $x21, $x22, $x23, $x20, $x21, $x22, $x23, $x24, $x25, $x26, $x27, $x28, $fp
+    RET undef $lr
+...
+---
+name: f4
+tracksRegLiveness: true
+body: |
+  bb.0:
+    RET undef $lr

llvm/test/CodeGen/AArch64/machine-outliner-no-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: 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 }
+...
+---
+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

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