[AArch64][DebugInfo] Do not recompute CalleeSavedStackSize (Take 2)

Commit message from D66935:

This patch fixes a bug exposed by D65653 where a subsequent invocation
of `determineCalleeSaves` ends up with a different size for the callee
save area, leading to different frame-offsets in debug information.

In the invocation by PEI, `determineCalleeSaves` tries to determine
whether it needs to spill an extra callee-saved register to get an
emergency spill slot. To do this, it calls 'estimateStackSize' and
manually adds the size of the callee-saves to this. PEI then allocates
the spill objects for the callee saves and the remaining frame layout
is calculated accordingly.

A second invocation in LiveDebugValues causes estimateStackSize to return
the size of the stack frame including the callee-saves. Given that the
size of the callee-saves is added to this, these callee-saves are counted
twice, which leads `determineCalleeSaves` to believe the stack has
become big enough to require spilling an extra callee-save as emergency
spillslot. It then updates CalleeSavedStackSize with a larger value.

Since CalleeSavedStackSize is used in the calculation of the frame
offset in getFrameIndexReference, this leads to incorrect offsets for
variables/locals when this information is recalculated after PEI.

This patch fixes the lldb unit tests in `functionalities/thread/concurrent_events/*`

Changes after D66935:

Ensures AArch64FunctionInfo::getCalleeSavedStackSize does not return
the uninitialized CalleeSavedStackSize when running `llc` on a specific
pass where the MIR code has already been expected to have gone through PEI.

Instead, getCalleeSavedStackSize (when passed the MachineFrameInfo) will try
to recalculate the CalleeSavedStackSize from the CalleeSavedInfo. In debug
mode, the compiler will assert the recalculated size equals the cached
size as calculated through a call to determineCalleeSaves.

This fixes two tests:
  test/DebugInfo/AArch64/asan-stack-vars.mir
  test/DebugInfo/AArch64/compiler-gen-bbs-livedebugvalues.mir
that otherwise fail when compiled using msan.

Reviewed By: omjavaid, efriedma

Tags: #llvm

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

git-svn-id: https://llvm.org/svn/llvm-project/llvm/trunk@375425 91177308-0d34-0410-b5e6-96231b3b80d8
diff --git a/include/llvm/CodeGen/TargetFrameLowering.h b/include/llvm/CodeGen/TargetFrameLowering.h
index 72edb27..2100a64 100644
--- a/include/llvm/CodeGen/TargetFrameLowering.h
+++ b/include/llvm/CodeGen/TargetFrameLowering.h
@@ -282,6 +282,11 @@
     return getFrameIndexReference(MF, FI, FrameReg);
   }
 
+  /// Returns the callee-saved registers as computed by determineCalleeSaves
+  /// in the BitVector \p SavedRegs.
+  virtual void getCalleeSaves(const MachineFunction &MF,
+                                  BitVector &SavedRegs) const;
+
   /// This method determines which of the registers reported by
   /// TargetRegisterInfo::getCalleeSavedRegs() should actually get saved.
   /// The default implementation checks populates the \p SavedRegs bitset with
@@ -289,6 +294,9 @@
   /// this function to save additional registers.
   /// This method also sets up the register scavenger ensuring there is a free
   /// register or a frameindex available.
+  /// This method should not be called by any passes outside of PEI, because
+  /// it may change state passed in by \p MF and \p RS. The preferred
+  /// interface outside PEI is getCalleeSaves.
   virtual void determineCalleeSaves(MachineFunction &MF, BitVector &SavedRegs,
                                     RegScavenger *RS = nullptr) const;
 
diff --git a/lib/CodeGen/LiveDebugValues.cpp b/lib/CodeGen/LiveDebugValues.cpp
index f1b237d..ca20b11 100644
--- a/lib/CodeGen/LiveDebugValues.cpp
+++ b/lib/CodeGen/LiveDebugValues.cpp
@@ -1439,8 +1439,7 @@
   TRI = MF.getSubtarget().getRegisterInfo();
   TII = MF.getSubtarget().getInstrInfo();
   TFI = MF.getSubtarget().getFrameLowering();
-  TFI->determineCalleeSaves(MF, CalleeSavedRegs,
-                            std::make_unique<RegScavenger>().get());
+  TFI->getCalleeSaves(MF, CalleeSavedRegs);
   LS.initialize(MF);
 
   bool Changed = ExtendRanges(MF);
diff --git a/lib/CodeGen/RegUsageInfoCollector.cpp b/lib/CodeGen/RegUsageInfoCollector.cpp
index 757ff0e..5a79ac4 100644
--- a/lib/CodeGen/RegUsageInfoCollector.cpp
+++ b/lib/CodeGen/RegUsageInfoCollector.cpp
@@ -56,7 +56,7 @@
 
   bool runOnMachineFunction(MachineFunction &MF) override;
 
-  // Call determineCalleeSaves and then also set the bits for subregs and
+  // Call getCalleeSaves and then also set the bits for subregs and
   // fully saved superregs.
   static void computeCalleeSavedRegs(BitVector &SavedRegs, MachineFunction &MF);
 
@@ -199,7 +199,7 @@
 
   // Target will return the set of registers that it saves/restores as needed.
   SavedRegs.clear();
-  TFI.determineCalleeSaves(MF, SavedRegs);
+  TFI.getCalleeSaves(MF, SavedRegs);
   if (SavedRegs.none())
     return;
 
diff --git a/lib/CodeGen/TargetFrameLoweringImpl.cpp b/lib/CodeGen/TargetFrameLoweringImpl.cpp
index 9eeacc2..bc59be8 100644
--- a/lib/CodeGen/TargetFrameLoweringImpl.cpp
+++ b/lib/CodeGen/TargetFrameLoweringImpl.cpp
@@ -60,6 +60,19 @@
   return MF.getFrameInfo().hasStackObjects();
 }
 
+void TargetFrameLowering::getCalleeSaves(const MachineFunction &MF,
+                                         BitVector &CalleeSaves) const {
+  const TargetRegisterInfo &TRI = *MF.getSubtarget().getRegisterInfo();
+  CalleeSaves.resize(TRI.getNumRegs());
+
+  const MachineFrameInfo &MFI = MF.getFrameInfo();
+  if (!MFI.isCalleeSavedInfoValid())
+    return;
+
+  for (const CalleeSavedInfo &Info : MFI.getCalleeSavedInfo())
+    CalleeSaves.set(Info.getReg());
+}
+
 void TargetFrameLowering::determineCalleeSaves(MachineFunction &MF,
                                                BitVector &SavedRegs,
                                                RegScavenger *RS) const {
diff --git a/lib/Target/AArch64/AArch64FrameLowering.cpp b/lib/Target/AArch64/AArch64FrameLowering.cpp
index 68e1e6a..042d8fd 100644
--- a/lib/Target/AArch64/AArch64FrameLowering.cpp
+++ b/lib/Target/AArch64/AArch64FrameLowering.cpp
@@ -1588,7 +1588,8 @@
   bool IsWin64 =
       Subtarget.isCallingConvWin64(MF.getFunction().getCallingConv());
   unsigned FixedObject = IsWin64 ? alignTo(AFI->getVarArgsGPRSize(), 16) : 0;
-  unsigned FPAdjust = isTargetDarwin(MF) ? 16 : AFI->getCalleeSavedStackSize();
+  unsigned FPAdjust = isTargetDarwin(MF)
+                        ? 16 : AFI->getCalleeSavedStackSize(MF.getFrameInfo());
   return {ObjectOffset + FixedObject + FPAdjust, MVT::i8};
 }
 
@@ -1630,7 +1631,7 @@
   int FPOffset = getFPOffset(MF, ObjectOffset).getBytes();
   int Offset = getStackOffset(MF, ObjectOffset).getBytes();
   bool isCSR =
-      !isFixed && ObjectOffset >= -((int)AFI->getCalleeSavedStackSize());
+      !isFixed && ObjectOffset >= -((int)AFI->getCalleeSavedStackSize(MFI));
 
   const StackOffset &SVEStackSize = getSVEStackSize(MF);
 
@@ -2304,6 +2305,10 @@
                << EstimatedStackSize + AlignedCSStackSize
                << " bytes.\n");
 
+  assert((!MFI.isCalleeSavedInfoValid() ||
+          AFI->getCalleeSavedStackSize() == AlignedCSStackSize) &&
+         "Should not invalidate callee saved info");
+
   // Round up to register pair alignment to avoid additional SP adjustment
   // instructions.
   AFI->setCalleeSavedStackSize(AlignedCSStackSize);
diff --git a/lib/Target/AArch64/AArch64MachineFunctionInfo.h b/lib/Target/AArch64/AArch64MachineFunctionInfo.h
index 0009fb7..3266186 100644
--- a/lib/Target/AArch64/AArch64MachineFunctionInfo.h
+++ b/lib/Target/AArch64/AArch64MachineFunctionInfo.h
@@ -55,6 +55,7 @@
 
   /// Amount of stack frame size used for saving callee-saved registers.
   unsigned CalleeSavedStackSize;
+  bool HasCalleeSavedStackSize = false;
 
   /// Number of TLS accesses using the special (combinable)
   /// _TLS_MODULE_BASE_ symbol.
@@ -167,8 +168,55 @@
   void setLocalStackSize(unsigned Size) { LocalStackSize = Size; }
   unsigned getLocalStackSize() const { return LocalStackSize; }
 
-  void setCalleeSavedStackSize(unsigned Size) { CalleeSavedStackSize = Size; }
-  unsigned getCalleeSavedStackSize() const { return CalleeSavedStackSize; }
+  void setCalleeSavedStackSize(unsigned Size) {
+    CalleeSavedStackSize = Size;
+    HasCalleeSavedStackSize = true;
+  }
+
+  // When CalleeSavedStackSize has not been set (for example when
+  // some MachineIR pass is run in isolation), then recalculate
+  // the CalleeSavedStackSize directly from the CalleeSavedInfo.
+  // Note: This information can only be recalculated after PEI
+  // has assigned offsets to the callee save objects.
+  unsigned getCalleeSavedStackSize(const MachineFrameInfo &MFI) const {
+    bool ValidateCalleeSavedStackSize = false;
+
+#ifndef NDEBUG
+    // Make sure the calculated size derived from the CalleeSavedInfo
+    // equals the cached size that was calculated elsewhere (e.g. in
+    // determineCalleeSaves).
+    ValidateCalleeSavedStackSize = HasCalleeSavedStackSize;
+#endif
+
+    if (!HasCalleeSavedStackSize || ValidateCalleeSavedStackSize) {
+      assert(MFI.isCalleeSavedInfoValid() && "CalleeSavedInfo not calculated");
+      if (MFI.getCalleeSavedInfo().empty())
+        return 0;
+
+      int64_t MinOffset = std::numeric_limits<int64_t>::max();
+      int64_t MaxOffset = std::numeric_limits<int64_t>::min();
+      for (const auto &Info : MFI.getCalleeSavedInfo()) {
+        int FrameIdx = Info.getFrameIdx();
+        int64_t Offset = MFI.getObjectOffset(FrameIdx);
+        int64_t ObjSize = MFI.getObjectSize(FrameIdx);
+        MinOffset = std::min<int64_t>(Offset, MinOffset);
+        MaxOffset = std::max<int64_t>(Offset + ObjSize, MaxOffset);
+      }
+
+      unsigned Size = alignTo(MaxOffset - MinOffset, 16);
+      assert((!HasCalleeSavedStackSize || getCalleeSavedStackSize() == Size) &&
+             "Invalid size calculated for callee saves");
+      return Size;
+    }
+
+    return getCalleeSavedStackSize();
+  }
+
+  unsigned getCalleeSavedStackSize() const {
+    assert(HasCalleeSavedStackSize &&
+           "CalleeSavedStackSize has not been calculated");
+    return CalleeSavedStackSize;
+  }
 
   void incNumLocalDynamicTLSAccesses() { ++NumLocalDynamicTLSAccesses; }
   unsigned getNumLocalDynamicTLSAccesses() const {
diff --git a/lib/Target/ARM/ARMFrameLowering.cpp b/lib/Target/ARM/ARMFrameLowering.cpp
index 01ae930..5428bd6 100644
--- a/lib/Target/ARM/ARMFrameLowering.cpp
+++ b/lib/Target/ARM/ARMFrameLowering.cpp
@@ -2128,10 +2128,16 @@
     AFI->setLRIsSpilledForFarJump(true);
   }
   AFI->setLRIsSpilled(SavedRegs.test(ARM::LR));
+}
+
+void ARMFrameLowering::getCalleeSaves(const MachineFunction &MF,
+                                      BitVector &SavedRegs) const {
+  TargetFrameLowering::getCalleeSaves(MF, SavedRegs);
 
   // If we have the "returned" parameter attribute which guarantees that we
   // return the value which was passed in r0 unmodified (e.g. C++ 'structors),
   // record that fact for IPRA.
+  const ARMFunctionInfo *AFI = MF.getInfo<ARMFunctionInfo>();
   if (AFI->getPreservesR0())
     SavedRegs.set(ARM::R0);
 }
diff --git a/lib/Target/ARM/ARMFrameLowering.h b/lib/Target/ARM/ARMFrameLowering.h
index 6d8aee5..0462b01 100644
--- a/lib/Target/ARM/ARMFrameLowering.h
+++ b/lib/Target/ARM/ARMFrameLowering.h
@@ -53,6 +53,8 @@
   int ResolveFrameIndexReference(const MachineFunction &MF, int FI,
                                  unsigned &FrameReg, int SPAdj) const;
 
+  void getCalleeSaves(const MachineFunction &MF,
+                      BitVector &SavedRegs) const override;
   void determineCalleeSaves(MachineFunction &MF, BitVector &SavedRegs,
                             RegScavenger *RS) const override;
 
diff --git a/test/CodeGen/AArch64/wrong-callee-save-size-after-livedebugvariables.mir b/test/CodeGen/AArch64/wrong-callee-save-size-after-livedebugvariables.mir
new file mode 100644
index 0000000..231de4e
--- /dev/null
+++ b/test/CodeGen/AArch64/wrong-callee-save-size-after-livedebugvariables.mir
@@ -0,0 +1,92 @@
+# RUN: llc -start-before=prologepilog -filetype=obj -o %t %s
+# RUN: llvm-dwarfdump --name=obj1 %t | FileCheck %s --check-prefix=CHECKDWARF1
+# RUN: llvm-dwarfdump --name=obj2 %t | FileCheck %s --check-prefix=CHECKDWARF2
+# RUN: llvm-objdump --disassemble %t | FileCheck %s --check-prefix=CHECKASM
+#
+# Test that the location for obj1 and obj2 in the debug information is
+# the same as the location used by load instructions.
+#
+# CHECKDWARF1: DW_AT_location    (DW_OP_fbreg -1)
+# CHECKDWARF2: DW_AT_location    (DW_OP_fbreg -2)
+# CHECKASM: ldurb   w0, [x29, #-1]
+# CHECKASM: ldurb   w1, [x29, #-2]
+--- |
+  ; ModuleID = 'wrong-callee-save-size-after-livedebugvariables.c'
+  source_filename = "wrong-callee-save-size-after-livedebugvariables.c"
+  target datalayout = "e-m:e-i8:8:32-i16:16:32-i64:64-i128:128-n32:64-S128"
+  target triple = "aarch64-unknown-linux-gnu"
+
+  ; Function Attrs: noinline nounwind optnone
+  define dso_local i8 @foo() #0 !dbg !7 {
+  entry:
+    %obj1 = alloca i8, align 1
+    %obj2 = alloca i8, align 1
+    %obj3 = alloca [238 x i8], align 1
+    ret i8 undef, !dbg !24
+  }
+
+  declare dso_local i8 @bar(i8, i8, i8*) #0
+
+  attributes #0 = { noinline nounwind optnone "frame-pointer"="all" }
+
+  !llvm.dbg.cu = !{!0}
+  !llvm.module.flags = !{!3, !4, !5}
+  !llvm.ident = !{!6}
+
+  !0 = distinct !DICompileUnit(language: DW_LANG_C99, file: !1, producer: "clang version 10.0.0", isOptimized: false, runtimeVersion: 0, emissionKind: FullDebug, enums: !2, nameTableKind: None)
+  !1 = !DIFile(filename: "wrong-callee-save-size-after-livedebugvariables.c", directory: "")
+  !2 = !{}
+  !3 = !{i32 2, !"Dwarf Version", i32 4}
+  !4 = !{i32 2, !"Debug Info Version", i32 3}
+  !5 = !{i32 1, !"wchar_size", i32 4}
+  !6 = !{!"clang version 10.0.0"}
+  !7 = distinct !DISubprogram(name: "foo", scope: !1, file: !1, line: 3, type: !8, scopeLine: 3, flags: DIFlagPrototyped, spFlags: DISPFlagDefinition, unit: !0, retainedNodes: !2)
+  !8 = !DISubroutineType(types: !9)
+  !9 = !{!10}
+  !10 = !DIBasicType(name: "char", size: 8, encoding: DW_ATE_unsigned_char)
+  !11 = !DILocalVariable(name: "obj1", scope: !7, file: !1, line: 4, type: !10)
+  !12 = !DILocation(line: 4, column: 8, scope: !7)
+  !13 = !DILocalVariable(name: "obj2", scope: !7, file: !1, line: 5, type: !10)
+  !14 = !DILocation(line: 5, column: 8, scope: !7)
+  !15 = !DILocalVariable(name: "obj3", scope: !7, file: !1, line: 6, type: !16)
+  !16 = !DICompositeType(tag: DW_TAG_array_type, baseType: !10, size: 1904, elements: !17)
+  !17 = !{!18}
+  !18 = !DISubrange(count: 238)
+  !19 = !DILocation(line: 6, column: 8, scope: !7)
+  !20 = !DILocation(line: 7, column: 14, scope: !7)
+  !21 = !DILocation(line: 7, column: 20, scope: !7)
+  !22 = !DILocation(line: 7, column: 27, scope: !7)
+  !23 = !DILocation(line: 7, column: 10, scope: !7)
+  !24 = !DILocation(line: 7, column: 3, scope: !7)
+
+...
+---
+name:            foo
+tracksRegLiveness: true
+frameInfo:
+  hasCalls:        true
+fixedStack:      []
+stack:
+  - { id: 0, name: obj1, type: default, offset: 0, size: 1, alignment: 1,
+      stack-id: default, callee-saved-register: '', callee-saved-restored: true,
+      local-offset: -1, debug-info-variable: '!11', debug-info-expression: '!DIExpression()',
+      debug-info-location: '!12' }
+  - { id: 1, name: obj2, type: default, offset: 0, size: 1, alignment: 1,
+      stack-id: default, callee-saved-register: '', callee-saved-restored: true,
+      local-offset: -2, debug-info-variable: '!13', debug-info-expression: '!DIExpression()',
+      debug-info-location: '!14' }
+  - { id: 2, name: obj3, type: default, offset: 0, size: 238, alignment: 1,
+      stack-id: default, callee-saved-register: '', callee-saved-restored: true,
+      local-offset: -240, debug-info-variable: '!15', debug-info-expression: '!DIExpression()',
+      debug-info-location: '!19' }
+body:             |
+  bb.1.entry:
+    renamable $x2 = ADDXri %stack.2.obj3, 0, 0
+    renamable $w0 = LDRBBui %stack.0.obj1, 0, debug-location !20 :: (load 1 from %ir.obj1)
+    renamable $w1 = LDRBBui %stack.1.obj2, 0, debug-location !21 :: (load 1 from %ir.obj2)
+    ADJCALLSTACKDOWN 0, 0, implicit-def $sp, implicit $sp, debug-location !23
+    BL @bar, csr_aarch64_aapcs, implicit-def $lr, implicit $sp, implicit killed $w0, implicit killed $w1, implicit killed $x2, implicit-def $w0, debug-location !23
+    ADJCALLSTACKUP 0, 0, implicit-def $sp, implicit $sp, debug-location !23
+    RET_ReallyLR implicit killed $w0, debug-location !24
+
+...
diff --git a/test/DebugInfo/MIR/X86/live-debug-values-reg-copy.mir b/test/DebugInfo/MIR/X86/live-debug-values-reg-copy.mir
index ba748b9..3abde77 100644
--- a/test/DebugInfo/MIR/X86/live-debug-values-reg-copy.mir
+++ b/test/DebugInfo/MIR/X86/live-debug-values-reg-copy.mir
@@ -1,4 +1,4 @@
-# RUN: llc -run-pass=livedebugvalues %s -o - | FileCheck %s
+# RUN: llc -start-before=prologepilog -stop-after=livedebugvalues %s -o - | FileCheck %s
 #
 # This test tests tracking variables value transferring from one register to another.
 # This example is altered additionally in order to test transferring from one float register