[MachineOutliner][AArch64] Fix for noreturn functions

[plotfi]

Summary

Noreturn functions conservatively require save + restore LR.
When there is no available register, the stack is used to handle LR. If the block contains a call, the outlined function also needs to set a frame on the stack. This caused an assertion failure for OF.FrameConstructionID != MachineOutlinerDefault && "Can only fix up stack references once".
This checks whether a candidate is MachineOutlinerDefault while requiring the modification of the stack due to a call. If that is the case, outlining is bailed-out to avoid modifying the stack twice.

[plotfi]

@kyulee Did you by chance come across this on the MultiSource/Applications/Burg test on llvm-test-suite? I hit something like this with -flto=full and -Oz and have been working on debugging it.

[kyulee]

@plotfi I hit this in an internal workload (-Oz -flto=thin), but it won't surprise if the same issue is in llvm-test-suite.

[plotfi]

In D83464#2152352, @kyulee wrote:

@plotfi I hit this in an internal workload (-Oz -flto=thin), but it won't surprise if the same issue is in llvm-test-suite.

This is really nice. It does fix the issue I hit; I just checked.

[plotfi]

llvm/lib/Target/AArch64/AArch64InstrInfo.cpp

6162

From what I can tell, bailing here with a RepeatedSequenceLocs.clear(); return outliner::OutlinedFunction(); also works. It seems to be if the candidate does not have an available LR or is noreturn, doesn't have available registers, and isn't using using the SP is is probably highly likely to get set to MachineOutlinerDefault unless the candidate size doesn't exceed the 12 bytes of a standard LR/BL/restore LR.

[paquette]

llvm/lib/Target/AArch64/AArch64InstrInfo.cpp

6162

You don't always want to bail out here.

You only want to bail out here when the candidates contain a call, or IsNoReturn == true.

This just says that we'll save LR to the stack.

6198–6205

I think this should only happen if we have MachineOutlinerDefault *and* the candidate requires stack fixups.

In the code above:

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

I don't think you want to remove candidates like these.

Technically, the commit message is inaccurate here, because this is really checking for situations where we fix things up twice. noreturn functions are one such case.

[plotfi]

llvm/lib/Target/AArch64/AArch64InstrInfo.cpp

6162

Ah! Yeah, that makes sense. The case that I was hitting does have a call.

llvm/lib/Target/AArch64/AArch64InstrInfo.cpp

else {
     // 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 {
       SetCandidateCallInfo(MachineOutlinerDefault, 12);
     }
 
     // If we dropped all of the candidates, bail out here.
     if (RepeatedSequenceLocs.size() < 2) {
       RepeatedSequenceLocs.clear();
       return outliner::OutlinedFunction();
     }
   }
if (FlagsSetInAll & MachineOutlinerMBBFlags::HasCalls) {
 
     if (ModStackToSaveLR) {
       // We can't fix up the stack. Bail out.
       if (!AllStackInstrsSafe) {
         RepeatedSequenceLocs.clear();
         return outliner::OutlinedFunction();
       }
 
+      // We don't fix up the stack twice.
+      if (std::any_of(RepeatedSequenceLocs.begin(), RepeatedSequenceLocs.end(),
+                      [](const outliner::Candidate &C) {
+                        return C.CallConstructionID == MachineOutlinerDefault;
+                      })) {
+        RepeatedSequenceLocs.clear();
+        return outliner::OutlinedFunction();
+      }
+
       // Save + restore LR.
       NumBytesToCreateFrame += 8;
     }
   }
 
   // If we have CFI instructions, we can only outline if the outlined section
   // can be a tail call
   if (FrameID != MachineOutlinerTailCall && CFICount > 0)

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