[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

10 files changed, 182 insertions, 9 deletions

diff --git a/include/llvm/CodeGen/TargetFrameLowering.h b/include/llvm/CodeGen/TargetFrameLowering.h
index 72edb27964c..2100a64c8f5 100644
--- a/include/llvm/CodeGen/TargetFrameLowering.h
+++ b/include/llvm/CodeGen/TargetFrameLowering.h
@@ -282,6 +282,11 @@ public:
     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 @@ public:
   /// 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 f1b237d83e8..ca20b111d30 100644
--- a/lib/CodeGen/LiveDebugValues.cpp
+++ b/lib/CodeGen/LiveDebugValues.cpp
@@ -1439,8 +1439,7 @@ bool LiveDebugValues::runOnMachineFunction(MachineFunction &MF) {
   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 757ff0e4495..5a79ac44dcf 100644
--- a/lib/CodeGen/RegUsageInfoCollector.cpp
+++ b/lib/CodeGen/RegUsageInfoCollector.cpp
@@ -56,7 +56,7 @@ public:
 
   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 @@ computeCalleeSavedRegs(BitVector &SavedRegs, MachineFunction &MF) {
 
   // 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 9eeacc2584c..bc59be890c9 100644
--- a/lib/CodeGen/TargetFrameLoweringImpl.cpp
+++ b/lib/CodeGen/TargetFrameLoweringImpl.cpp
@@ -60,6 +60,19 @@ bool TargetFrameLowering::needsFrameIndexResolution(
   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 68e1e6a3022..042d8fdcc51 100644
--- a/lib/Target/AArch64/AArch64FrameLowering.cpp
+++ b/lib/Target/AArch64/AArch64FrameLowering.cpp
@@ -1588,7 +1588,8 @@ static StackOffset getFPOffset(const MachineFunction &MF, int ObjectOffset) {
   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 @@ StackOffset AArch64FrameLowering::resolveFrameOffsetReference(
   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 @@ void AArch64FrameLowering::determineCalleeSaves(MachineFunction &MF,
                << 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 0009fb7b552..32661860934 100644
--- a/lib/Target/AArch64/AArch64MachineFunctionInfo.h
+++ b/lib/Target/AArch64/AArch64MachineFunctionInfo.h
@@ -55,6 +55,7 @@ class AArch64FunctionInfo final : public MachineFunctionInfo {
 
   /// 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 @@ public:
   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 01ae93086dc..5428bd6c94b 100644
--- a/lib/Target/ARM/ARMFrameLowering.cpp
+++ b/lib/Target/ARM/ARMFrameLowering.cpp
@@ -2128,10 +2128,16 @@ void ARMFrameLowering::determineCalleeSaves(MachineFunction &MF,
     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 6d8aee59794..0462b01af70 100644
--- a/lib/Target/ARM/ARMFrameLowering.h
+++ b/lib/Target/ARM/ARMFrameLowering.h
@@ -53,6 +53,8 @@ public:
   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 00000000000..231de4e1896
--- /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 ba748b989c4..3abde776904 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