4 files changed, 114 insertions, 63 deletions
diff --git a/llvm/lib/CodeGen/MachineOutliner.cpp b/llvm/lib/CodeGen/MachineOutliner.cpp
index f742ac6b55c..d6ce08a0097 100644
--- a/llvm/lib/CodeGen/MachineOutliner.cpp
+++ b/llvm/lib/CodeGen/MachineOutliner.cpp
@@ -1306,12 +1306,6 @@ void MachineOutliner::populateMapper(InstructionMapper &Mapper, Module &M,
     if (F.empty())
       continue;
 
-    // Disable outlining from noreturn functions right now. Noreturn requires
-    // special handling for the case where what we are outlining could be a
-    // tail call.
-    if (F.hasFnAttribute(Attribute::NoReturn))
-      continue;
-
     // There's something in F. Check if it has a MachineFunction associated with
     // it.
     MachineFunction *MF = MMI.getMachineFunction(F);
diff --git a/llvm/lib/Target/AArch64/AArch64InstrInfo.cpp b/llvm/lib/Target/AArch64/AArch64InstrInfo.cpp
index abde3e50021..54f3f7c1013 100644
--- a/llvm/lib/Target/AArch64/AArch64InstrInfo.cpp
+++ b/llvm/lib/Target/AArch64/AArch64InstrInfo.cpp
@@ -5868,11 +5868,21 @@ outliner::OutlinedFunction AArch64InstrInfo::getOutliningCandidateInfo(
     unsigned NumBytesNoStackCalls = 0;
     std::vector<outliner::Candidate> CandidatesWithoutStackFixups;
 
+    // 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);
+
       // Is LR available? If so, we don't need a save.
-      if (C.LRU.available(AArch64::LR)) {
+      if (C.LRU.available(AArch64::LR) && !IsNoReturn) {
         NumBytesNoStackCalls += 4;
         C.setCallInfo(MachineOutlinerNoLRSave, 4);
         CandidatesWithoutStackFixups.push_back(C);
diff --git a/llvm/test/CodeGen/AArch64/machine-outliner-noreturn-save-lr.mir b/llvm/test/CodeGen/AArch64/machine-outliner-noreturn-save-lr.mir
new file mode 100644
index 00000000000..000752acb01
--- /dev/null
+++ b/llvm/test/CodeGen/AArch64/machine-outliner-noreturn-save-lr.mir
@@ -0,0 +1,103 @@
+# NOTE: Assertions have been autogenerated by utils/update_mir_test_checks.py
+# RUN: llc -mtriple=aarch64-unknown-unknown -run-pass=machine-outliner -verify-machineinstrs %s -o - | FileCheck %s
+
+# Verify that we save + restore LR around calls to outlined functions from
+# noreturn parents. LR isn't actually used, but debuggers need LR in this case
+# to get a backtrace.
+
+--- |
+  define void @save_lr_1() #0 { ret void }
+  define void @save_lr_2() #0 { ret void }
+  define void @save_lr_3() #0 { ret void }
+  define void @save_lr_4() #0 { ret void }
+  attributes #0 = { noredzone noreturn }
+...
+---
+
+name:            save_lr_1
+alignment:       4
+tracksRegLiveness: true
+frameInfo:
+  maxAlignment:    1
+  maxCallFrameSize: 0
+machineFunctionInfo: {}
+body:             |
+  bb.0:
+    liveins: $lr
+    ; CHECK-LABEL: name: save_lr_1
+    ; CHECK: liveins: $lr
+    ; CHECK: $x0 = ORRXrs $xzr, $lr, 0
+    ; CHECK: BL @OUTLINED_FUNCTION_0, implicit-def $lr, implicit $sp, implicit-def $lr, implicit-def $lr, implicit-def $w3, implicit-def $w4, implicit-def $w5, implicit-def $w6
+    ; CHECK: $lr = ORRXrs $xzr, $x0, 0
+    $w3 = ORRWri $wzr, 1
+    $w4 = ORRWri $wzr, 1
+    BRK 1
+    $w5 = ORRWri $wzr, 1
+    $w6 = ORRWri $wzr, 1
+...
+---
+name:            save_lr_2
+alignment:       4
+tracksRegLiveness: true
+frameInfo:
+  maxAlignment:    1
+  maxCallFrameSize: 0
+machineFunctionInfo: {}
+body:             |
+  bb.0:
+    liveins: $lr
+    ; CHECK-LABEL: name: save_lr_2
+    ; CHECK: liveins: $lr
+    ; CHECK: $x0 = ORRXrs $xzr, $lr, 0
+    ; CHECK: BL @OUTLINED_FUNCTION_0, implicit-def $lr, implicit $sp, implicit-def $lr, implicit-def $lr, implicit-def $w3, implicit-def $w4, implicit-def $w5, implicit-def $w6
+    ; CHECK: $lr = ORRXrs $xzr, $x0, 0
+    $w3 = ORRWri $wzr, 1
+    $w4 = ORRWri $wzr, 1
+    BRK 1
+    $w5 = ORRWri $wzr, 1
+    $w6 = ORRWri $wzr, 1
+...
+---
+name:            save_lr_3
+alignment:       4
+tracksRegLiveness: true
+frameInfo:
+  maxAlignment:    1
+  maxCallFrameSize: 0
+machineFunctionInfo: {}
+body:             |
+  bb.0:
+    liveins: $lr
+    ; CHECK-LABEL: name: save_lr_3
+    ; CHECK: liveins: $lr
+    ; CHECK: $x0 = ORRXrs $xzr, $lr, 0
+    ; CHECK: BL @OUTLINED_FUNCTION_0, implicit-def $lr, implicit $sp, implicit-def $lr, implicit-def $lr, implicit-def $w3, implicit-def $w4, implicit-def $w5, implicit-def $w6
+    ; CHECK: $lr = ORRXrs $xzr, $x0, 0
+    $w3 = ORRWri $wzr, 1
+    $w4 = ORRWri $wzr, 1
+    BRK 1
+    $w5 = ORRWri $wzr, 1
+    $w6 = ORRWri $wzr, 1
+...
+---
+name:            save_lr_4
+alignment:       4
+tracksRegLiveness: true
+frameInfo:
+  maxAlignment:    1
+  maxCallFrameSize: 0
+machineFunctionInfo: {}
+body:             |
+  bb.0:
+    liveins: $lr
+    ; CHECK-LABEL: name: save_lr_4
+    ; CHECK: liveins: $lr
+    ; CHECK: $x0 = ORRXrs $xzr, $lr, 0
+    ; CHECK: BL @OUTLINED_FUNCTION_0, implicit-def $lr, implicit $sp, implicit-def $lr, implicit-def $lr, implicit-def $w3, implicit-def $w4, implicit-def $w5, implicit-def $w6
+    ; CHECK: $lr = ORRXrs $xzr, $x0, 0
+    $w3 = ORRWri $wzr, 1
+    $w4 = ORRWri $wzr, 1
+    BRK 1
+    $w5 = ORRWri $wzr, 1
+    $w6 = ORRWri $wzr, 1
+...
diff --git a/llvm/test/CodeGen/AArch64/machine-outliner-noreturn.mir b/llvm/test/CodeGen/AArch64/machine-outliner-noreturn.mir
deleted file mode 100644
index 29166f9f12d..00000000000
--- a/llvm/test/CodeGen/AArch64/machine-outliner-noreturn.mir
+++ /dev/null
@@ -1,56 +0,0 @@
-# RUN: llc -mtriple=aarch64-unknown-unknown -run-pass=machine-outliner -verify-machineinstrs %s -o - | FileCheck %s
-
---- |
-  define void @foo() #0 { ret void }
-  define void @bar(i32 %a) #0 { ret void }
-  define void @baz(i32 %a) #0 { ret void }
-  attributes #0 = { noredzone noreturn }
-...
----
-
-# Temporarily disable outlining from noreturn functions. To do this, we need
-# to verify thst every function we want to outline from is noreturn.
-
-# CHECK-NOT: OUTLINED_FUNCTION
-
-name:            foo
-alignment:       4
-tracksRegLiveness: true
-frameInfo:
-  maxAlignment:    1
-  maxCallFrameSize: 0
-machineFunctionInfo: {}
-body:             |
-  bb.0:
-    $w3 = ORRWri $wzr, 1
-    $w4 = ORRWri $wzr, 1
-    BRK 1
-...
----
-name:            bar
-alignment:       4
-tracksRegLiveness: true
-frameInfo:
-  maxAlignment:    1
-  maxCallFrameSize: 0
-machineFunctionInfo: {}
-body:             |
-  bb.0:
-    $w3 = ORRWri $wzr, 1
-    $w4 = ORRWri $wzr, 1
-    BRK 1
-...
----
-name:            baz
-alignment:       4
-tracksRegLiveness: true
-frameInfo:
-  maxAlignment:    1
-  maxCallFrameSize: 0
-machineFunctionInfo: {}
-body:             |
-  bb.0:
-    $w3 = ORRWri $wzr, 1
-    $w4 = ORRWri $wzr, 1
-    BRK 1
-...