From c877e033760ac2e55cb777b4c6550a08684a3e9e Mon Sep 17 00:00:00 2001
From: Jessica Paquette <jpaquette@apple.com>
Date: Sat, 1 Dec 2018 21:24:06 +0000
Subject: [PATCH] [MachineOutliner][AArch64] Improve checks for stack
 instructions

If we know that we'll definitely save LR to a register, there's no reason to
pre-check whether or not a stack instruction is unsafe to fix up.

This makes it so that we check for that condition before mapping instructions.

This allows us to outline more, since we don't pessimise as many instructions.

Also update some tests, since we outline more.

llvm-svn: 348081
---
 lib/Target/AArch64/AArch64InstrInfo.cpp       | 24 +++++++++++++++-
 .../AArch64/machine-outliner-remarks.ll       | 14 +++++-----
 test/CodeGen/AArch64/machine-outliner.ll      |  1 +
 test/CodeGen/AArch64/machine-outliner.mir     | 28 +++++++++----------
 4 files changed, 45 insertions(+), 22 deletions(-)

diff --git a/lib/Target/AArch64/AArch64InstrInfo.cpp b/lib/Target/AArch64/AArch64InstrInfo.cpp
index 784a4b02e6a..c8d9d2a6d02 100644
--- a/lib/Target/AArch64/AArch64InstrInfo.cpp
+++ b/lib/Target/AArch64/AArch64InstrInfo.cpp
@@ -5284,7 +5284,29 @@ bool AArch64InstrInfo::isMBBSafeToOutlineFrom(MachineBasicBlock &MBB,
   if (any_of(MBB, [](MachineInstr &MI) { return MI.isCall(); }))
     Flags |= MachineOutlinerMBBFlags::HasCalls;
 
-  if (!LRU.available(AArch64::LR))
+  MachineFunction *MF = MBB.getParent();
+
+  // In the event that we outline, we may have to save LR. If there is an
+  // available register in the MBB, then we'll always save LR there. Check if
+  // this is true.
+  bool CanSaveLR = false;
+  const AArch64RegisterInfo *ARI = static_cast<const AArch64RegisterInfo *>(
+      MF->getSubtarget().getRegisterInfo());
+
+  // Check if there is an available register across the sequence that we can
+  // use.
+  for (unsigned Reg : AArch64::GPR64RegClass) {
+    if (!ARI->isReservedReg(*MF, Reg) && Reg != AArch64::LR &&
+        Reg != AArch64::X16 && Reg != AArch64::X17 && LRU.available(Reg)) {
+      CanSaveLR = true;
+      break;
+    }
+  }
+
+  // Check if we have a register we can save LR to, and if LR was used
+  // somewhere. If both of those things are true, then we need to evaluate the
+  // safety of outlining stack instructions later.
+  if (!CanSaveLR && !LRU.available(AArch64::LR))
     Flags |= MachineOutlinerMBBFlags::LRUnavailableSomewhere;
 
   return true;
diff --git a/test/CodeGen/AArch64/machine-outliner-remarks.ll b/test/CodeGen/AArch64/machine-outliner-remarks.ll
index 29872d9518a..19351262b82 100644
--- a/test/CodeGen/AArch64/machine-outliner-remarks.ll
+++ b/test/CodeGen/AArch64/machine-outliner-remarks.ll
@@ -1,10 +1,10 @@
 ; RUN: llc %s -enable-machine-outliner -mtriple=aarch64-unknown-unknown -pass-remarks=machine-outliner -pass-remarks-missed=machine-outliner -o /dev/null 2>&1 | FileCheck %s
 ; CHECK: <unknown>:0:0:
 ; CHECK-SAME: Did not outline 2 instructions from 2 locations.
-; CHECK-SAME: Bytes from outlining all occurrences (36) >=
+; CHECK-SAME: Bytes from outlining all occurrences (16) >=
 ; CHECK-SAME: Unoutlined instruction bytes (16)
 ; CHECK-SAME: (Also found at: <UNKNOWN LOCATION>)
-; CHECK: remark: <unknown>:0:0: Saved 20 bytes by outlining 12 instructions
+; CHECK: remark: <unknown>:0:0: Saved 48 bytes by outlining 14 instructions
 ; CHECK-SAME: from 2 locations. (Found at: <UNKNOWN LOCATION>,
 ; CHECK-SAME: <UNKNOWN LOCATION>)
 ; RUN: llc %s -enable-machine-outliner -mtriple=aarch64-unknown-unknown -o /dev/null -pass-remarks-missed=machine-outliner -pass-remarks-output=%t.yaml
@@ -16,7 +16,7 @@
 ; YAML-NEXT: Pass:            machine-outliner
 ; YAML-NEXT: Name:            NotOutliningCheaper
 ; YAML-NEXT: Function:
-; YAML-NEXT: Args:            
+; YAML-NEXT: Args:
 ; YAML-NEXT:   - String:          'Did not outline '
 ; YAML-NEXT:   - Length:          '2'
 ; YAML-NEXT:   - String:          ' instructions'
@@ -24,7 +24,7 @@
 ; YAML-NEXT:   - NumOccurrences:  '2'
 ; YAML-NEXT:   - String:          ' locations.'
 ; YAML-NEXT:   - String:          ' Bytes from outlining all occurrences ('
-; YAML-NEXT:   - OutliningCost:   '36'
+; YAML-NEXT:   - OutliningCost:   '16'
 ; YAML-NEXT:   - String:          ')'
 ; YAML-NEXT:   - String:          ' >= Unoutlined instruction bytes ('
 ; YAML-NEXT:   - NotOutliningCost: '16'
@@ -36,12 +36,12 @@
 ; YAML-NEXT: Pass:            machine-outliner
 ; YAML-NEXT: Name:            OutlinedFunction
 ; YAML-NEXT: Function:        OUTLINED_FUNCTION_0
-; YAML-NEXT: Args:            
+; YAML-NEXT: Args:
 ; YAML-NEXT:   - String:          'Saved '
-; YAML-NEXT:   - OutliningBenefit: '20'
+; YAML-NEXT:   - OutliningBenefit: '48'
 ; YAML-NEXT:   - String:          ' bytes by '
 ; YAML-NEXT:   - String:          'outlining '
-; YAML-NEXT:   - Length:          '12'
+; YAML-NEXT:   - Length:          '14'
 ; YAML-NEXT:   - String:          ' instructions '
 ; YAML-NEXT:   - String:          'from '
 ; YAML-NEXT:   - NumOccurrences:  '2'
diff --git a/test/CodeGen/AArch64/machine-outliner.ll b/test/CodeGen/AArch64/machine-outliner.ll
index 19be14d8d39..42d0a09f028 100644
--- a/test/CodeGen/AArch64/machine-outliner.ll
+++ b/test/CodeGen/AArch64/machine-outliner.ll
@@ -103,6 +103,7 @@ define void @dog() #0 {
 ; CHECK-NEXT: str     w8, [sp, #12]
 ; CHECK-NEXT: orr     w8, wzr, #0x6
 ; CHECK-NEXT: str     w8, [sp, #8]
+; CHECK-NEXT: add     sp, sp, #32
 ; CHECK-NEXT: ret
 
 attributes #0 = { noredzone "target-cpu"="cyclone" "target-features"="+sse" }
diff --git a/test/CodeGen/AArch64/machine-outliner.mir b/test/CodeGen/AArch64/machine-outliner.mir
index 880c28e3de6..625d6a82684 100644
--- a/test/CodeGen/AArch64/machine-outliner.mir
+++ b/test/CodeGen/AArch64/machine-outliner.mir
@@ -22,39 +22,39 @@
 # - Create outlined functions
 # - Don't outline anything to do with LR or W30
 # - Save LR when it's not available
-# - Don't outline stack instructions when we might need to save + restore
 # - Functions whose addresses are taken can still be outlined
 #
-# CHECK-LABEL: name: main
-
-# CHECK: BL @OUTLINED_FUNCTION_[[F0:[0-9]+]]
+# CHECK-LABEL: main
+# CHECK-LABEL: bb.1:
+# CHECK-DAG: BL @OUTLINED_FUNCTION_[[F0:[0-9]+]]
 # CHECK-NEXT: $lr = ORRXrs $xzr, $x[[REG:[0-9]+]], 0
-# CHECK-NEXT: $x12 = ADDXri $sp, 48, 0
 # CHECK-NEXT: STRHHroW $w12, $x9, $w30, 1, 1
 # CHECK-NEXT: $lr = ORRXri $xzr, 1
 
+# CHECK-DAG: bb.2
 # CHECK: BL @OUTLINED_FUNCTION_[[F0]]
 # CHECK-NEXT: $lr = ORRXrs $xzr, $x[[REG]], 0
-# CHECK-NEXT: $x12 = ADDXri $sp, 48, 0
 # CHECK-NEXT: STRHHroW $w12, $x9, $w30, 1, 1
 # CHECK-NEXT: $lr = ORRXri $xzr, 1
 
+# CHECK-DAG: bb.3
 # CHECK: BL @OUTLINED_FUNCTION_[[F0]]
 # CHECK-NEXT: $lr = ORRXrs $xzr, $x[[REG]], 0
-# CHECK-NEXT: $x12 = ADDXri $sp, 48, 0
 # CHECK-NEXT: STRHHroW $w12, $x9, $w30, 1, 1
 # CHECK-NEXT: $lr = ORRXri $xzr, 1
 name:            main
 tracksRegLiveness: true
 body:             |
   bb.0:
+  liveins: $lr
     $sp = frame-setup SUBXri $sp, 16, 0
     renamable $x9 = ADRP target-flags(aarch64-page) @bar
     $x9 = ORRXri $xzr, 1
     $w12 = ORRWri $wzr, 1
     $w30 = ORRWri $wzr, 1
     $lr = ORRXri $xzr, 1
-
+  bb.1:
+  liveins: $lr
     $x20, $x19 = LDPXi $sp, 10
     $w12 = ORRWri $wzr, 1
     $w12 = ORRWri $wzr, 1
@@ -66,8 +66,8 @@ body:             |
     $x12 = ADDXri $sp, 48, 0;
     STRHHroW $w12, $x9, $w30, 1, 1
     $lr = ORRXri $xzr, 1
-    $w3 = ORRWri $wzr, 1993
-
+  bb.2:
+  liveins: $lr
     $x20, $x19 = LDPXi $sp, 10
     $w12 = ORRWri $wzr, 1
     $w12 = ORRWri $wzr, 1
@@ -79,9 +79,8 @@ body:             |
     $x12 = ADDXri $sp, 48, 0;
     STRHHroW $w12, $x9, $w30, 1, 1
     $lr = ORRXri $xzr, 1
-
-    $w4 = ORRWri $wzr, 1994
-
+  bb.3:
+  liveins: $lr
     $x20, $x19 = LDPXi $sp, 10
     $w12 = ORRWri $wzr, 1
     $w12 = ORRWri $wzr, 1
@@ -93,8 +92,9 @@ body:             |
     $x12 = ADDXri $sp, 48, 0;
     STRHHroW $w12, $x9, $w30, 1, 1
     $lr = ORRXri $xzr, 1
-
     $sp = ADDXri $sp, 16, 0
+  bb.4:
+  liveins: $lr
     RET undef $lr
 
 ...