[SystemZ][z/OS] Implement llvm.frameaddr for XPLINK (#89284)

The implementation follows the ELF implementation.

4 files changed, 105 insertions, 14 deletions

diff --git a/llvm/lib/Target/SystemZ/SystemZFrameLowering.cpp b/llvm/lib/Target/SystemZ/SystemZFrameLowering.cpp
index 50ecd6e..6ec46bd 100644
--- a/llvm/lib/Target/SystemZ/SystemZFrameLowering.cpp
+++ b/llvm/lib/Target/SystemZ/SystemZFrameLowering.cpp
@@ -896,6 +896,19 @@ SystemZXPLINKFrameLowering::SystemZXPLINKFrameLowering()
     RegSpillOffsets[Entry.Reg] = Entry.Offset;
 }
 
+int SystemZXPLINKFrameLowering::getOrCreateFramePointerSaveIndex(
+    MachineFunction &MF) const {
+  SystemZMachineFunctionInfo *ZFI = MF.getInfo<SystemZMachineFunctionInfo>();
+  int FI = ZFI->getFramePointerSaveIndex();
+  if (!FI) {
+    MachineFrameInfo &MFFrame = MF.getFrameInfo();
+    FI = MFFrame.CreateFixedObject(8, 0, false);
+    MFFrame.setStackID(FI, TargetStackID::NoAlloc);
+    ZFI->setFramePointerSaveIndex(FI);
+  }
+  return FI;
+}
+
 // Checks if the function is a potential candidate for being a XPLeaf routine.
 static bool isXPLeafCandidate(const MachineFunction &MF) {
   const MachineFrameInfo &MFFrame = MF.getFrameInfo();
@@ -991,6 +1004,9 @@ bool SystemZXPLINKFrameLowering::assignCalleeSavedSpillSlots(
   Register HighGPR = 0;
   int HighOffset = -1;
 
+  // Query index of the saved frame pointer.
+  int FPSI = MFI->getFramePointerSaveIndex();
+
   for (auto &CS : CSI) {
     Register Reg = CS.getReg();
     int Offset = RegSpillOffsets[Reg];
@@ -1013,7 +1029,10 @@ bool SystemZXPLINKFrameLowering::assignCalleeSavedSpillSlots(
         // the bottom of the stack and are not truly part of the "normal" stack
         // frame. Mark the frame index as NoAlloc to indicate it as such.
         unsigned RegSize = 8;
-        int FrameIdx = MFFrame.CreateFixedSpillStackObject(RegSize, Offset);
+        int FrameIdx =
+            (FPSI && Offset == 0)
+                ? FPSI
+                : MFFrame.CreateFixedSpillStackObject(RegSize, Offset);
         CS.setFrameIdx(FrameIdx);
         MFFrame.setStackID(FrameIdx, TargetStackID::NoAlloc);
       }
@@ -1467,15 +1486,16 @@ void SystemZXPLINKFrameLowering::determineFrameLayout(
   StackSize += Regs->getCallFrameSize();
   MFFrame.setStackSize(StackSize);
 
-  // We now know the stack size. Create the fixed spill stack objects for the
-  // register save area now. This has no impact on the stack frame layout, as
-  // this is already computed. However, it makes sure that all callee saved
-  // registers have a valid frame index assigned.
-  const unsigned RegSize = MF.getDataLayout().getPointerSize();
-  for (auto &CS : MFFrame.getCalleeSavedInfo()) {
-    int Offset = RegSpillOffsets[CS.getReg()];
-    if (Offset >= 0)
-      CS.setFrameIdx(
-          MFFrame.CreateFixedSpillStackObject(RegSize, Offset - StackSize));
+  // We now know the stack size. Update the stack objects for the register save
+  // area now. This has no impact on the stack frame layout, as this is already
+  // computed. However, it makes sure that all callee saved registers have a
+  // valid offset assigned.
+  for (int FrameIdx = MFFrame.getObjectIndexBegin(); FrameIdx != 0;
+       ++FrameIdx) {
+    if (MFFrame.getStackID(FrameIdx) == TargetStackID::NoAlloc) {
+      int64_t SPOffset = MFFrame.getObjectOffset(FrameIdx);
+      SPOffset -= StackSize;
+      MFFrame.setObjectOffset(FrameIdx, SPOffset);
+    }
   }
 }
diff --git a/llvm/lib/Target/SystemZ/SystemZFrameLowering.h b/llvm/lib/Target/SystemZ/SystemZFrameLowering.h
index 03ce888..4613181 100644
--- a/llvm/lib/Target/SystemZ/SystemZFrameLowering.h
+++ b/llvm/lib/Target/SystemZ/SystemZFrameLowering.h
@@ -41,6 +41,12 @@ public:
   }
 
   bool hasReservedCallFrame(const MachineFunction &MF) const override;
+
+  // Return the offset of the backchain.
+  virtual unsigned getBackchainOffset(MachineFunction &MF) const = 0;
+
+  // Get or create the frame index of where the old frame pointer is stored.
+  virtual int getOrCreateFramePointerSaveIndex(MachineFunction &MF) const = 0;
 };
 
 class SystemZELFFrameLowering : public SystemZFrameLowering {
@@ -86,13 +92,13 @@ public:
   bool usePackedStack(MachineFunction &MF) const;
 
   // Return the offset of the backchain.
-  unsigned getBackchainOffset(MachineFunction &MF) const {
+  unsigned getBackchainOffset(MachineFunction &MF) const override {
     // The back chain is stored topmost with packed-stack.
     return usePackedStack(MF) ? SystemZMC::ELFCallFrameSize - 8 : 0;
   }
 
   // Get or create the frame index of where the old frame pointer is stored.
-  int getOrCreateFramePointerSaveIndex(MachineFunction &MF) const;
+  int getOrCreateFramePointerSaveIndex(MachineFunction &MF) const override;
 };
 
 class SystemZXPLINKFrameLowering : public SystemZFrameLowering {
@@ -133,6 +139,15 @@ public:
                                            RegScavenger *RS) const override;
 
   void determineFrameLayout(MachineFunction &MF) const;
+
+  // Return the offset of the backchain.
+  unsigned getBackchainOffset(MachineFunction &MF) const override {
+    // The back chain is always the first element of the frame.
+    return 0;
+  }
+
+  // Get or create the frame index of where the old frame pointer is stored.
+  int getOrCreateFramePointerSaveIndex(MachineFunction &MF) const override;
 };
 } // end namespace llvm
 
diff --git a/llvm/lib/Target/SystemZ/SystemZISelLowering.cpp b/llvm/lib/Target/SystemZ/SystemZISelLowering.cpp
index 5c2579f..b3f93e3 100644
--- a/llvm/lib/Target/SystemZ/SystemZISelLowering.cpp
+++ b/llvm/lib/Target/SystemZ/SystemZISelLowering.cpp
@@ -3765,7 +3765,7 @@ SDValue SystemZTargetLowering::lowerConstantPool(ConstantPoolSDNode *CP,
 
 SDValue SystemZTargetLowering::lowerFRAMEADDR(SDValue Op,
                                               SelectionDAG &DAG) const {
-  auto *TFL = Subtarget.getFrameLowering<SystemZELFFrameLowering>();
+  auto *TFL = Subtarget.getFrameLowering<SystemZFrameLowering>();
   MachineFunction &MF = DAG.getMachineFunction();
   MachineFrameInfo &MFI = MF.getFrameInfo();
   MFI.setFrameAddressIsTaken(true);
diff --git a/llvm/test/CodeGen/SystemZ/zos-frameaddr.ll b/llvm/test/CodeGen/SystemZ/zos-frameaddr.ll
new file mode 100644
index 0000000..5977317
--- /dev/null
+++ b/llvm/test/CodeGen/SystemZ/zos-frameaddr.ll
@@ -0,0 +1,56 @@
+; NOTE: Assertions have been autogenerated by utils/update_llc_test_checks.py UTC_ARGS: --version 4
+; RUN: llc < %s -mtriple=s390x-ibm-zos | FileCheck --check-prefix=CHECK %s
+
+; The current function's frame address is the address of
+; the optional back chain slot.
+define ptr @fp0() nounwind {
+; CHECK-LABEL: fp0:
+; CHECK:         la 3, 2048(4)
+; CHECK-NEXT:    b 2(7)
+entry:
+  %0 = tail call ptr @llvm.frameaddress(i32 0)
+  ret ptr %0
+}
+
+; Check that the frame address is correct in a presence
+; of a stack frame.
+define ptr @fp0f() nounwind {
+; CHECK-LABEL: fp0f:
+; CHECK:         stmg 6, 7, 1904(4)
+; CHECK-NEXT:    aghi 4, -160
+; CHECK-NEXT:    la 3, 2048(4)
+; CHECK-NEXT:    lg 7, 2072(4)
+; CHECK-NEXT:    aghi 4, 160
+; CHECK-NEXT:    b 2(7)
+entry:
+  %0 = alloca i64, align 8
+  %1 = tail call ptr @llvm.frameaddress(i32 0)
+  ret ptr %1
+}
+
+; Check the caller's frame address.
+define ptr @fpcaller() nounwind "backchain" {
+; CHECK-LABEL: fpcaller:
+; CHECK:         stmg 4, 7, 2048(4)
+; CHECK-NEXT:    lg 3, 2048(4)
+; CHECK-NEXT:    lmg 4, 7, 2048(4)
+; CHECK-NEXT:    b 2(7)
+entry:
+  %0 = tail call ptr @llvm.frameaddress(i32 1)
+  ret ptr %0
+}
+
+; Check the caller's frame address.
+define ptr @fpcallercaller() nounwind "backchain" {
+; CHECK-LABEL: fpcallercaller:
+; CHECK:         stmg 4, 7, 2048(4)
+; CHECK-NEXT:    lg 1, 2048(4)
+; CHECK-NEXT:    lg 3, 0(1)
+; CHECK-NEXT:    lmg 4, 7, 2048(4)
+; CHECK-NEXT:    b 2(7)
+entry:
+  %0 = tail call ptr @llvm.frameaddress(i32 2)
+  ret ptr %0
+}
+
+declare ptr @llvm.frameaddress(i32) nounwind readnone