[ShrinkWrap] Use underlying object to rule out stack access.

Allow shrink-wrapping past memory accesses that only access globals or
function arguments. This patch uses getUnderlyingObject to try to
identify the accessed object by a given memory operand. If it is a
global or an argument, it does not access the stack of the current
function and should not block shrink wrapping.

Note that the caller's stack may get accessed when passing an argument
via the stack, but not the stack of the current function.

This addresses part of the TODO from D63152.

Reviewed By: thegameg

Differential Revision: https://reviews.llvm.org/D149668

Author: fhahn

llvm/lib/CodeGen/ShrinkWrap.cpp

@@ -53,6 +53,7 @@
 #include "llvm/ADT/SmallVector.h"
 #include "llvm/ADT/Statistic.h"
 #include "llvm/Analysis/CFG.h"
+#include "llvm/Analysis/ValueTracking.h"
 #include "llvm/CodeGen/MachineBasicBlock.h"
 #include "llvm/CodeGen/MachineBlockFrequencyInfo.h"
 #include "llvm/CodeGen/MachineDominators.h"
@@ -259,13 +260,30 @@ INITIALIZE_PASS_END(ShrinkWrap, DEBUG_TYPE, "Shrink Wrap Pass", false, false)
 
 bool ShrinkWrap::useOrDefCSROrFI(const MachineInstr &MI,
                                  RegScavenger *RS) const {
+  /// Check if \p Op is known to access an address not on the function's stack .
+  /// At the moment, accesses where the underlying object is a global or a
+  /// function argument are considered non-stack accesses. Note that the
+  /// caller's stack may get accessed when passing an argument via the stack,
+  /// but not the stack of the current function.
+  ///
+  auto IsKnownNonStackPtr = [](MachineMemOperand *Op) {
+    if (Op->getValue()) {
+      const Value *UO = getUnderlyingObject(Op->getValue());
+      if (!UO)
+        return false;
+      if (auto *Arg = dyn_cast<Argument>(UO))
+        return !Arg->hasPassPointeeByValueCopyAttr();
+      return isa<GlobalValue>(UO);
+    }
+    return false;
+  };
   // This prevents premature stack popping when occurs a indirect stack
-  // access. It is overly aggressive for the moment.
-  // TODO: - Obvious non-stack loads and store, such as global values,
-  //         are known to not access the stack.
+  // access.  It is overly aggressive for the moment.
+  // TODO:
   //       - Further, data dependency and alias analysis can validate
   //         that load and stores never derive from the stack pointer.
-  if (MI.mayLoadOrStore())
+  if (MI.mayLoadOrStore() && (MI.isCall() || MI.hasUnmodeledSideEffects() ||
+                              !all_of(MI.memoperands(), IsKnownNonStackPtr)))
     return true;
 
   if (MI.getOpcode() == FrameSetupOpcode ||

llvm/test/CodeGen/AArch64/addsub.ll

@@ -651,9 +651,6 @@ declare dso_local i32 @crng_reseed(...) local_unnamed_addr
 define dso_local i32 @_extract_crng_crng() {
 ; CHECK-LABEL: _extract_crng_crng:
 ; CHECK:       // %bb.0: // %entry
-; CHECK-NEXT:    str x30, [sp, #-16]! // 8-byte Folded Spill
-; CHECK-NEXT:    .cfi_def_cfa_offset 16
-; CHECK-NEXT:    .cfi_offset w30, -16
 ; CHECK-NEXT:    adrp x8, _extract_crng_crng
 ; CHECK-NEXT:    add x8, x8, :lo12:_extract_crng_crng
 ; CHECK-NEXT:    tbnz x8, #63, .LBB36_2
@@ -665,15 +662,20 @@ define dso_local i32 @_extract_crng_crng() {
 ; CHECK-NEXT:    cmn x8, #1272
 ; CHECK-NEXT:    b.pl .LBB36_3
 ; CHECK-NEXT:  .LBB36_2: // %if.then
+; CHECK-NEXT:    str x30, [sp, #-16]! // 8-byte Folded Spill
+; CHECK-NEXT:    .cfi_def_cfa_offset 16
+; CHECK-NEXT:    .cfi_offset w30, -16
 ; CHECK-NEXT:    adrp x8, primary_crng
 ; CHECK-NEXT:    adrp x9, input_pool
 ; CHECK-NEXT:    add x9, x9, :lo12:input_pool
 ; CHECK-NEXT:    ldr w8, [x8, :lo12:primary_crng]
 ; CHECK-NEXT:    cmp w8, #0
 ; CHECK-NEXT:    csel x0, xzr, x9, eq
 ; CHECK-NEXT:    bl crng_reseed
-; CHECK-NEXT:  .LBB36_3: // %if.end
 ; CHECK-NEXT:    ldr x30, [sp], #16 // 8-byte Folded Reload
+; CHECK-NEXT:  .LBB36_3: // %if.end
+; CHECK-NEXT:    .cfi_def_cfa wsp, 0
+; CHECK-NEXT:    .cfi_same_value w30
 ; CHECK-NEXT:    ret
 entry:
   br i1 icmp slt (ptr @_extract_crng_crng, ptr null), label %if.then, label %lor.lhs.false

llvm/test/CodeGen/AArch64/arm64-shrink-wrapping.ll

@@ -1004,25 +1004,28 @@ end:
 define i32 @stack_realign(i32 %a, i32 %b, ptr %ptr1, ptr %ptr2) {
 ; ENABLE-LABEL: stack_realign:
 ; ENABLE:       ; %bb.0:
+; ENABLE-NEXT:    lsl w8, w0, w1
+; ENABLE-NEXT:    lsl w9, w1, w0
+; ENABLE-NEXT:    cmp w0, w1
+; ENABLE-NEXT:    b.ge LBB13_2
+; ENABLE-NEXT:  ; %bb.1: ; %true
 ; ENABLE-NEXT:    stp x29, x30, [sp, #-16]! ; 16-byte Folded Spill
 ; ENABLE-NEXT:    .cfi_def_cfa_offset 16
 ; ENABLE-NEXT:    mov x29, sp
 ; ENABLE-NEXT:    .cfi_def_cfa w29, 16
 ; ENABLE-NEXT:    .cfi_offset w30, -8
 ; ENABLE-NEXT:    .cfi_offset w29, -16
-; ENABLE-NEXT:    sub x9, sp, #16
-; ENABLE-NEXT:    and sp, x9, #0xffffffffffffffe0
-; ENABLE-NEXT:    lsl w8, w0, w1
-; ENABLE-NEXT:    lsl w9, w1, w0
-; ENABLE-NEXT:    cmp w0, w1
-; ENABLE-NEXT:    b.ge LBB13_2
-; ENABLE-NEXT:  ; %bb.1: ; %true
+; ENABLE-NEXT:    sub x1, sp, #16
+; ENABLE-NEXT:    and sp, x1, #0xffffffffffffffe0
 ; ENABLE-NEXT:    str w0, [sp]
+; ENABLE-NEXT:    mov sp, x29
+; ENABLE-NEXT:    ldp x29, x30, [sp], #16 ; 16-byte Folded Reload
 ; ENABLE-NEXT:  LBB13_2: ; %false
+; ENABLE-NEXT:    .cfi_def_cfa wsp, 0
+; ENABLE-NEXT:    .cfi_same_value w30
+; ENABLE-NEXT:    .cfi_same_value w29
 ; ENABLE-NEXT:    str w8, [x2]
 ; ENABLE-NEXT:    str w9, [x3]
-; ENABLE-NEXT:    mov sp, x29
-; ENABLE-NEXT:    ldp x29, x30, [sp], #16 ; 16-byte Folded Reload
 ; ENABLE-NEXT:    ret
 ;
 ; DISABLE-LABEL: stack_realign:

llvm/test/CodeGen/AArch64/combine-comparisons-by-cse.ll

@@ -658,28 +658,28 @@ return:                                           ; preds = %if.end, %land.lhs.t
 define i32 @fcmpri(i32 %argc, ptr nocapture readonly %argv) #0 {
 ; CHECK-LABEL: fcmpri:
 ; CHECK:       // %bb.0: // %entry
-; CHECK-NEXT:    str d8, [sp, #-32]! // 8-byte Folded Spill
-; CHECK-NEXT:    .cfi_def_cfa_offset 32
-; CHECK-NEXT:    stp x30, x19, [sp, #16] // 16-byte Folded Spill
-; CHECK-NEXT:    .cfi_offset w19, -8
-; CHECK-NEXT:    .cfi_offset w30, -16
-; CHECK-NEXT:    .cfi_offset b8, -32
 ; CHECK-NEXT:    cmp w0, #2
 ; CHECK-NEXT:    b.lt .LBB9_3
 ; CHECK-NEXT:  // %bb.1: // %land.lhs.true
 ; CHECK-NEXT:    ldr x8, [x1, #8]
 ; CHECK-NEXT:    cbz x8, .LBB9_3
 ; CHECK-NEXT:  // %bb.2:
-; CHECK-NEXT:    mov w0, #3
-; CHECK-NEXT:    b .LBB9_4
+; CHECK-NEXT:    mov w0, #3 // =0x3
+; CHECK-NEXT:    ret
 ; CHECK-NEXT:  .LBB9_3: // %if.end
-; CHECK-NEXT:    mov w0, #1
+; CHECK-NEXT:    str d8, [sp, #-32]! // 8-byte Folded Spill
+; CHECK-NEXT:    .cfi_def_cfa_offset 32
+; CHECK-NEXT:    stp x30, x19, [sp, #16] // 16-byte Folded Spill
+; CHECK-NEXT:    .cfi_offset w19, -8
+; CHECK-NEXT:    .cfi_offset w30, -16
+; CHECK-NEXT:    .cfi_offset b8, -32
+; CHECK-NEXT:    mov w0, #1 // =0x1
 ; CHECK-NEXT:    bl zoo
 ; CHECK-NEXT:    mov w19, w0
-; CHECK-NEXT:    mov w0, #-1
+; CHECK-NEXT:    mov w0, #-1 // =0xffffffff
 ; CHECK-NEXT:    bl yoo
 ; CHECK-NEXT:    cmp w19, #0
-; CHECK-NEXT:    mov w1, #2
+; CHECK-NEXT:    mov w1, #2 // =0x2
 ; CHECK-NEXT:    cinc w0, w19, gt
 ; CHECK-NEXT:    fmov d8, d0
 ; CHECK-NEXT:    bl xoo
@@ -689,9 +689,8 @@ define i32 @fcmpri(i32 %argc, ptr nocapture readonly %argv) #0 {
 ; CHECK-NEXT:    fadd d0, d8, d0
 ; CHECK-NEXT:    fcsel d0, d8, d0, gt
 ; CHECK-NEXT:    bl woo
-; CHECK-NEXT:    mov w0, #4
-; CHECK-NEXT:  .LBB9_4: // %return
 ; CHECK-NEXT:    ldp x30, x19, [sp, #16] // 16-byte Folded Reload
+; CHECK-NEXT:    mov w0, #4 // =0x4
 ; CHECK-NEXT:    ldr d8, [sp], #32 // 8-byte Folded Reload
 ; CHECK-NEXT:    .cfi_def_cfa_offset 0
 ; CHECK-NEXT:    .cfi_restore w19

llvm/test/CodeGen/AArch64/csr-split.ll

@@ -9,23 +9,17 @@
 define dso_local signext i32 @test1(ptr %b) local_unnamed_addr uwtable  {
 ; CHECK-LABEL: test1:
 ; CHECK:       // %bb.0: // %entry
-; CHECK-NEXT:    stp x30, x19, [sp, #-16]! // 16-byte Folded Spill
-; CHECK-NEXT:    .cfi_def_cfa_offset 16
-; CHECK-NEXT:    .cfi_offset w19, -8
-; CHECK-NEXT:    .cfi_offset w30, -16
-; CHECK-NEXT:    .cfi_remember_state
 ; CHECK-NEXT:    adrp x8, a
 ; CHECK-NEXT:    ldrsw x8, [x8, :lo12:a]
 ; CHECK-NEXT:    cmp x8, x0
 ; CHECK-NEXT:    b.eq .LBB0_2
 ; CHECK-NEXT:  // %bb.1: // %if.end
-; CHECK-NEXT:    ldp x30, x19, [sp], #16 // 16-byte Folded Reload
-; CHECK-NEXT:    .cfi_def_cfa_offset 0
-; CHECK-NEXT:    .cfi_restore w19
-; CHECK-NEXT:    .cfi_restore w30
 ; CHECK-NEXT:    ret
 ; CHECK-NEXT:  .LBB0_2: // %if.then
-; CHECK-NEXT:    .cfi_restore_state
+; CHECK-NEXT:    stp x30, x19, [sp, #-16]! // 16-byte Folded Spill
+; CHECK-NEXT:    .cfi_def_cfa_offset 16
+; CHECK-NEXT:    .cfi_offset w19, -8
+; CHECK-NEXT:    .cfi_offset w30, -16
 ; CHECK-NEXT:    mov x19, x0
 ; CHECK-NEXT:    bl callVoid
 ; CHECK-NEXT:    mov x0, x19
@@ -37,31 +31,22 @@ define dso_local signext i32 @test1(ptr %b) local_unnamed_addr uwtable  {
 ;
 ; CHECK-APPLE-LABEL: test1:
 ; CHECK-APPLE:       ; %bb.0: ; %entry
-; CHECK-APPLE-NEXT:    stp x20, x19, [sp, #-32]! ; 16-byte Folded Spill
-; CHECK-APPLE-NEXT:    .cfi_def_cfa_offset 32
-; CHECK-APPLE-NEXT:    stp x29, x30, [sp, #16] ; 16-byte Folded Spill
-; CHECK-APPLE-NEXT:    .cfi_offset w30, -8
-; CHECK-APPLE-NEXT:    .cfi_offset w29, -16
-; CHECK-APPLE-NEXT:    .cfi_offset w19, -24
-; CHECK-APPLE-NEXT:    .cfi_offset w20, -32
-; CHECK-APPLE-NEXT:    .cfi_remember_state
 ; CHECK-APPLE-NEXT:  Lloh0:
 ; CHECK-APPLE-NEXT:    adrp x8, _a@PAGE
 ; CHECK-APPLE-NEXT:  Lloh1:
 ; CHECK-APPLE-NEXT:    ldrsw x8, [x8, _a@PAGEOFF]
 ; CHECK-APPLE-NEXT:    cmp x8, x0
 ; CHECK-APPLE-NEXT:    b.eq LBB0_2
 ; CHECK-APPLE-NEXT:  ; %bb.1: ; %if.end
-; CHECK-APPLE-NEXT:    ldp x29, x30, [sp, #16] ; 16-byte Folded Reload
-; CHECK-APPLE-NEXT:    ldp x20, x19, [sp], #32 ; 16-byte Folded Reload
-; CHECK-APPLE-NEXT:    .cfi_def_cfa_offset 0
-; CHECK-APPLE-NEXT:    .cfi_restore w30
-; CHECK-APPLE-NEXT:    .cfi_restore w29
-; CHECK-APPLE-NEXT:    .cfi_restore w19
-; CHECK-APPLE-NEXT:    .cfi_restore w20
 ; CHECK-APPLE-NEXT:    ret
 ; CHECK-APPLE-NEXT:  LBB0_2: ; %if.then
-; CHECK-APPLE-NEXT:    .cfi_restore_state
+; CHECK-APPLE-NEXT:    stp x20, x19, [sp, #-32]! ; 16-byte Folded Spill
+; CHECK-APPLE-NEXT:    .cfi_def_cfa_offset 32
+; CHECK-APPLE-NEXT:    stp x29, x30, [sp, #16] ; 16-byte Folded Spill
+; CHECK-APPLE-NEXT:    .cfi_offset w30, -8
+; CHECK-APPLE-NEXT:    .cfi_offset w29, -16
+; CHECK-APPLE-NEXT:    .cfi_offset w19, -24
+; CHECK-APPLE-NEXT:    .cfi_offset w20, -32
 ; CHECK-APPLE-NEXT:    mov x19, x0
 ; CHECK-APPLE-NEXT:    bl _callVoid
 ; CHECK-APPLE-NEXT:    ldp x29, x30, [sp, #16] ; 16-byte Folded Reload

llvm/test/CodeGen/AArch64/shrink-wrap-byval-inalloca-preallocated.ll

@@ -12,13 +12,6 @@ declare void @fn(ptr, ptr)
 define void @test_regular_pointers(ptr %a, ptr %b) {
 ; CHECK-LABEL: test_regular_pointers:
 ; CHECK:       ; %bb.0: ; %entry
-; CHECK-NEXT:    stp x20, x19, [sp, #-32]! ; 16-byte Folded Spill
-; CHECK-NEXT:    .cfi_def_cfa_offset 32
-; CHECK-NEXT:    stp x29, x30, [sp, #16] ; 16-byte Folded Spill
-; CHECK-NEXT:    .cfi_offset w30, -8
-; CHECK-NEXT:    .cfi_offset w29, -16
-; CHECK-NEXT:    .cfi_offset w19, -24
-; CHECK-NEXT:    .cfi_offset w20, -32
 ; CHECK-NEXT:    ldr d0, [x0]
 ; CHECK-NEXT:    mov x8, #1 ; =0x1
 ; CHECK-NEXT:    ldr d1, [x1, #8]
@@ -30,12 +23,24 @@ define void @test_regular_pointers(ptr %a, ptr %b) {
 ; CHECK-NEXT:    b.mi LBB0_2
 ; CHECK-NEXT:    b.gt LBB0_2
 ; CHECK-NEXT:  ; %bb.1: ; %then
+; CHECK-NEXT:    stp x20, x19, [sp, #-32]! ; 16-byte Folded Spill
+; CHECK-NEXT:    .cfi_def_cfa_offset 32
+; CHECK-NEXT:    stp x29, x30, [sp, #16] ; 16-byte Folded Spill
+; CHECK-NEXT:    .cfi_offset w30, -8
+; CHECK-NEXT:    .cfi_offset w29, -16
+; CHECK-NEXT:    .cfi_offset w19, -24
+; CHECK-NEXT:    .cfi_offset w20, -32
 ; CHECK-NEXT:    mov x19, x1
 ; CHECK-NEXT:    bl _fn
-; CHECK-NEXT:    str xzr, [x19]
-; CHECK-NEXT:  LBB0_2: ; %exit
 ; CHECK-NEXT:    ldp x29, x30, [sp, #16] ; 16-byte Folded Reload
+; CHECK-NEXT:    str xzr, [x19]
 ; CHECK-NEXT:    ldp x20, x19, [sp], #32 ; 16-byte Folded Reload
+; CHECK-NEXT:  LBB0_2: ; %exit
+; CHECK-NEXT:    .cfi_def_cfa wsp, 0
+; CHECK-NEXT:    .cfi_same_value w30
+; CHECK-NEXT:    .cfi_same_value w29
+; CHECK-NEXT:    .cfi_same_value w19
+; CHECK-NEXT:    .cfi_same_value w20
 ; CHECK-NEXT:    ret
 entry:
   %l.a = load double, ptr %a, align 8

llvm/test/CodeGen/AArch64/taildup-cfi.ll

@@ -32,7 +32,7 @@ if.then:                                          ; preds = %entry
   store i32 0, ptr @f, align 4, !tbaa !2
   br label %if.end
 
-; DARWIN:           Merging into block
+; DARWIN-NOT:           Merging into block
 ; LINUX:    	      Merging into block
 
 if.end:                                           ; preds = %entry.if.end_crit_edge, %if.then