[COFF, ARM64] Fix localaddress to handle stack realignment and variable size objects

Summary: This fixes using the correct stack registers for SEH when stack realignment is needed or when variable size objects are present.

Reviewers: rnk, efriedma, ssijaric, TomTan

Reviewed By: rnk, efriedma

Subscribers: javed.absar, kristof.beyls, llvm-commits

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

llvm-svn: 352923

Author: Mandeep Singh Grang

llvm/include/llvm/CodeGen/MachineFunction.h

@@ -328,7 +328,6 @@ class MachineFunction {
   bool CallsUnwindInit = false;
   bool HasEHScopes = false;
   bool HasEHFunclets = false;
-  bool HasLocalEscape = false;
 
   /// List of C++ TypeInfo used.
   std::vector<const GlobalValue *> TypeInfos;
@@ -811,9 +810,6 @@ class MachineFunction {
   bool hasEHFunclets() const { return HasEHFunclets; }
   void setHasEHFunclets(bool V) { HasEHFunclets = V; }
 
-  bool hasLocalEscape() const { return HasLocalEscape; }
-  void setHasLocalEscape(bool V) { HasLocalEscape = V; }
-
   /// Find or create an LandingPadInfo for the specified MachineBasicBlock.
   LandingPadInfo &getOrCreateLandingPadInfo(MachineBasicBlock *LandingPad);
 

llvm/include/llvm/CodeGen/TargetFrameLowering.h

@@ -261,6 +261,17 @@ class TargetFrameLowering {
     return getFrameIndexReference(MF, FI, FrameReg);
   }
 
+  /// getNonLocalFrameIndexReference - This method returns the offset used to
+  /// reference a frame index location. The offset can be from either FP/BP/SP
+  /// based on which base register is returned by llvm.localaddress.
+  virtual int getNonLocalFrameIndexReference(const MachineFunction &MF,
+                                       int FI) const {
+    // By default, dispatch to getFrameIndexReference. Interested targets can
+    // override this.
+    unsigned FrameReg;
+    return getFrameIndexReference(MF, FI, FrameReg);
+  }
+
   /// 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

llvm/lib/CodeGen/AsmPrinter/WinException.cpp

@@ -937,11 +937,7 @@ void WinException::emitEHRegistrationOffsetLabel(const WinEHFuncInfo &FuncInfo,
   int FI = FuncInfo.EHRegNodeFrameIndex;
   if (FI != INT_MAX) {
     const TargetFrameLowering *TFI = Asm->MF->getSubtarget().getFrameLowering();
-    unsigned UnusedReg;
-    // FIXME: getFrameIndexReference needs to match the behavior of
-    // AArch64RegisterInfo::hasBasePointer in which one of the scenarios where
-    // SP is used is if frame size >= 256.
-    Offset = TFI->getFrameIndexReference(*Asm->MF, FI, UnusedReg);
+    Offset = TFI->getNonLocalFrameIndexReference(*Asm->MF, FI);
   }
 
   MCContext &Ctx = Asm->OutContext;

llvm/lib/CodeGen/SelectionDAG/SelectionDAGBuilder.cpp

@@ -6220,8 +6220,6 @@ SelectionDAGBuilder::visitIntrinsicCall(const CallInst &I, unsigned Intrinsic) {
           .addFrameIndex(FI);
     }
 
-    MF.setHasLocalEscape(true);
-
     return nullptr;
   }
 

llvm/lib/Target/AArch64/AArch64FrameLowering.cpp

@@ -227,10 +227,6 @@ bool AArch64FrameLowering::hasFP(const MachineFunction &MF) const {
       MFI.getMaxCallFrameSize() > DefaultSafeSPDisplacement)
     return true;
 
-  // Win64 SEH requires frame pointer if funclets are present.
-  if (MF.hasLocalEscape())
-    return true;
-
   return false;
 }
 
@@ -1469,19 +1465,44 @@ int AArch64FrameLowering::getFrameIndexReference(const MachineFunction &MF,
   return resolveFrameIndexReference(MF, FI, FrameReg);
 }
 
+int AArch64FrameLowering::getNonLocalFrameIndexReference(
+  const MachineFunction &MF, int FI) const {
+  return getSEHFrameIndexOffset(MF, FI);
+}
+
+static int getFPOffset(const MachineFunction &MF, int FI) {
+  const auto &MFI = MF.getFrameInfo();
+  const auto *AFI = MF.getInfo<AArch64FunctionInfo>();
+  const auto &Subtarget = MF.getSubtarget<AArch64Subtarget>();
+  bool IsWin64 =
+      Subtarget.isCallingConvWin64(MF.getFunction().getCallingConv());
+  unsigned FixedObject = IsWin64 ? alignTo(AFI->getVarArgsGPRSize(), 16) : 0;
+  return MFI.getObjectOffset(FI) + FixedObject + 16;
+}
+
+static int getStackOffset(const MachineFunction &MF, int FI) {
+  const auto &MFI = MF.getFrameInfo();
+  return MFI.getObjectOffset(FI) + MFI.getStackSize();
+}
+
+int AArch64FrameLowering::getSEHFrameIndexOffset(const MachineFunction &MF,
+                                                 int FI) const {
+  const auto *RegInfo = static_cast<const AArch64RegisterInfo *>(
+      MF.getSubtarget().getRegisterInfo());
+  return RegInfo->getLocalAddressRegister(MF) == AArch64::FP ?
+         getFPOffset(MF, FI) : getStackOffset(MF, FI);
+}
+
 int AArch64FrameLowering::resolveFrameIndexReference(const MachineFunction &MF,
                                                      int FI, unsigned &FrameReg,
                                                      bool PreferFP) const {
-  const MachineFrameInfo &MFI = MF.getFrameInfo();
-  const AArch64RegisterInfo *RegInfo = static_cast<const AArch64RegisterInfo *>(
+  const auto &MFI = MF.getFrameInfo();
+  const auto *RegInfo = static_cast<const AArch64RegisterInfo *>(
       MF.getSubtarget().getRegisterInfo());
-  const AArch64FunctionInfo *AFI = MF.getInfo<AArch64FunctionInfo>();
-  const AArch64Subtarget &Subtarget = MF.getSubtarget<AArch64Subtarget>();
-  bool IsWin64 =
-      Subtarget.isCallingConvWin64(MF.getFunction().getCallingConv());
-  unsigned FixedObject = IsWin64 ? alignTo(AFI->getVarArgsGPRSize(), 16) : 0;
-  int FPOffset = MFI.getObjectOffset(FI) + FixedObject + 16;
-  int Offset = MFI.getObjectOffset(FI) + MFI.getStackSize();
+  const auto *AFI = MF.getInfo<AArch64FunctionInfo>();
+  const auto &Subtarget = MF.getSubtarget<AArch64Subtarget>();
+  int FPOffset = getFPOffset(MF, FI);
+  int Offset = getStackOffset(MF, FI);
   bool isFixed = MFI.isFixedObjectIndex(FI);
   bool isCSR = !isFixed && MFI.getObjectOffset(FI) >=
                                -((int)AFI->getCalleeSavedStackSize());

llvm/lib/Target/AArch64/AArch64FrameLowering.h

@@ -78,6 +78,9 @@ class AArch64FrameLowering : public TargetFrameLowering {
   int getFrameIndexReferencePreferSP(const MachineFunction &MF, int FI,
                                      unsigned &FrameReg,
                                      bool IgnoreSPUpdates) const override;
+  int getNonLocalFrameIndexReference(const MachineFunction &MF,
+                               int FI) const override;
+  int getSEHFrameIndexOffset(const MachineFunction &MF, int FI) const;
 
 private:
   bool shouldCombineCSRLocalStackBump(MachineFunction &MF,

llvm/lib/Target/AArch64/AArch64ISelLowering.cpp

@@ -2744,15 +2744,9 @@ SDValue AArch64TargetLowering::LowerINTRINSIC_WO_CHAIN(SDValue Op,
                        Op.getOperand(1), Op.getOperand(2));
 
   case Intrinsic::localaddress: {
-    // Returns one of the stack, base, or frame pointer registers, depending on
-    // which is used to reference local variables.
-    MachineFunction &MF = DAG.getMachineFunction();
-    const AArch64RegisterInfo *RegInfo = Subtarget->getRegisterInfo();
-    unsigned Reg;
-    if (RegInfo->hasBasePointer(MF))
-      Reg = RegInfo->getBaseRegister();
-    else // This function handles the SP or FP case.
-      Reg = RegInfo->getFrameRegister(MF);
+    const auto &MF = DAG.getMachineFunction();
+    const auto *RegInfo = Subtarget->getRegisterInfo();
+    unsigned Reg = RegInfo->getLocalAddressRegister(MF);
     return DAG.getCopyFromReg(DAG.getEntryNode(), dl, Reg,
                               Op.getSimpleValueType());
   }

llvm/lib/Target/AArch64/AArch64RegisterInfo.cpp

@@ -463,15 +463,16 @@ void AArch64RegisterInfo::eliminateFrameIndex(MachineBasicBlock::iterator II,
     return;
   }
 
-  // Modify MI as necessary to handle as much of 'Offset' as possible
-  Offset = TFI->resolveFrameIndexReference(MF, FrameIndex, FrameReg);
-
   if (MI.getOpcode() == TargetOpcode::LOCAL_ESCAPE) {
     MachineOperand &FI = MI.getOperand(FIOperandNum);
+    Offset = TFI->getNonLocalFrameIndexReference(MF, FrameIndex);
     FI.ChangeToImmediate(Offset);
     return;
   }
 
+  // Modify MI as necessary to handle as much of 'Offset' as possible
+  Offset = TFI->getFrameIndexReference(MF, FrameIndex, FrameReg);
+
   if (rewriteAArch64FrameIndex(MI, FIOperandNum, FrameReg, Offset, TII))
     return;
 
@@ -525,3 +526,13 @@ unsigned AArch64RegisterInfo::getRegPressureLimit(const TargetRegisterClass *RC,
     return 16;
   }
 }
+
+unsigned AArch64RegisterInfo::getLocalAddressRegister(
+  const MachineFunction &MF) const {
+  const auto &MFI = MF.getFrameInfo();
+  if (!MF.hasEHFunclets() && !MFI.hasVarSizedObjects())
+    return AArch64::SP;
+  else if (needsStackRealignment(MF))
+    return getBaseRegister();
+  return getFrameRegister(MF);
+}

llvm/lib/Target/AArch64/AArch64RegisterInfo.h

@@ -121,6 +121,8 @@ class AArch64RegisterInfo final : public AArch64GenRegisterInfo {
   bool trackLivenessAfterRegAlloc(const MachineFunction&) const override {
     return true;
   }
+
+  unsigned getLocalAddressRegister(const MachineFunction &MF) const;
 };
 
 } // end namespace llvm