[SystemZ][z/OS] Implement llvm.returnaddress for XPLINK (#89440)

The implementation follows the ELF implementation.

Author: redstar

llvm/lib/Target/SystemZ/SystemZFrameLowering.cpp

@@ -56,14 +56,17 @@ static const TargetFrameLowering::SpillSlot XPLINKSpillOffsetTable[] = {
 
 SystemZFrameLowering::SystemZFrameLowering(StackDirection D, Align StackAl,
                                            int LAO, Align TransAl,
-                                           bool StackReal)
-    : TargetFrameLowering(D, StackAl, LAO, TransAl, StackReal) {}
+                                           bool StackReal, unsigned PointerSize)
+    : TargetFrameLowering(D, StackAl, LAO, TransAl, StackReal),
+      PointerSize(PointerSize) {}
 
 std::unique_ptr<SystemZFrameLowering>
 SystemZFrameLowering::create(const SystemZSubtarget &STI) {
+  unsigned PtrSz =
+      STI.getTargetLowering()->getTargetMachine().getPointerSize(0);
   if (STI.isTargetXPLINK64())
-    return std::make_unique<SystemZXPLINKFrameLowering>();
-  return std::make_unique<SystemZELFFrameLowering>();
+    return std::make_unique<SystemZXPLINKFrameLowering>(PtrSz);
+  return std::make_unique<SystemZELFFrameLowering>(PtrSz);
 }
 
 namespace {
@@ -272,9 +275,9 @@ void SystemZELFFrameLowering::determineCalleeSaves(MachineFunction &MF,
   }
 }
 
-SystemZELFFrameLowering::SystemZELFFrameLowering()
+SystemZELFFrameLowering::SystemZELFFrameLowering(unsigned PointerSize)
     : SystemZFrameLowering(TargetFrameLowering::StackGrowsDown, Align(8), 0,
-                           Align(8), /* StackRealignable */ false),
+                           Align(8), /* StackRealignable */ false, PointerSize),
       RegSpillOffsets(0) {
 
   // Due to the SystemZ ABI, the DWARF CFA (Canonical Frame Address) is not
@@ -884,9 +887,10 @@ bool SystemZELFFrameLowering::usePackedStack(MachineFunction &MF) const {
   return HasPackedStackAttr && CallConv;
 }
 
-SystemZXPLINKFrameLowering::SystemZXPLINKFrameLowering()
+SystemZXPLINKFrameLowering::SystemZXPLINKFrameLowering(unsigned PointerSize)
     : SystemZFrameLowering(TargetFrameLowering::StackGrowsDown, Align(32), 0,
-                           Align(32), /* StackRealignable */ false),
+                           Align(32), /* StackRealignable */ false,
+                           PointerSize),
       RegSpillOffsets(-1) {
 
   // Create a mapping from register number to save slot offset.

llvm/lib/Target/SystemZ/SystemZFrameLowering.h

@@ -22,7 +22,7 @@ class SystemZSubtarget;
 class SystemZFrameLowering : public TargetFrameLowering {
 public:
   SystemZFrameLowering(StackDirection D, Align StackAl, int LAO, Align TransAl,
-                       bool StackReal);
+                       bool StackReal, unsigned PointerSize);
 
   static std::unique_ptr<SystemZFrameLowering>
   create(const SystemZSubtarget &STI);
@@ -45,15 +45,24 @@ class SystemZFrameLowering : public TargetFrameLowering {
   // Return the offset of the backchain.
   virtual unsigned getBackchainOffset(MachineFunction &MF) const = 0;
 
+  // Return the offset of the return address.
+  virtual int getReturnAddressOffset(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;
+
+  // Return the size of a pointer (in bytes).
+  unsigned getPointerSize() const { return PointerSize; }
+
+private:
+  unsigned PointerSize;
 };
 
 class SystemZELFFrameLowering : public SystemZFrameLowering {
   IndexedMap<unsigned> RegSpillOffsets;
 
 public:
-  SystemZELFFrameLowering();
+  SystemZELFFrameLowering(unsigned PointerSize);
 
   // Override TargetFrameLowering.
   bool
@@ -97,6 +106,11 @@ class SystemZELFFrameLowering : public SystemZFrameLowering {
     return usePackedStack(MF) ? SystemZMC::ELFCallFrameSize - 8 : 0;
   }
 
+  // Return the offset of the return address.
+  int getReturnAddressOffset(MachineFunction &MF) const override {
+    return (usePackedStack(MF) ? -2 : 14) * getPointerSize();
+  }
+
   // Get or create the frame index of where the old frame pointer is stored.
   int getOrCreateFramePointerSaveIndex(MachineFunction &MF) const override;
 };
@@ -105,7 +119,7 @@ class SystemZXPLINKFrameLowering : public SystemZFrameLowering {
   IndexedMap<unsigned> RegSpillOffsets;
 
 public:
-  SystemZXPLINKFrameLowering();
+  SystemZXPLINKFrameLowering(unsigned PointerSize);
 
   bool
   assignCalleeSavedSpillSlots(MachineFunction &MF,
@@ -146,6 +160,11 @@ class SystemZXPLINKFrameLowering : public SystemZFrameLowering {
     return 0;
   }
 
+  // Return the offset of the return address.
+  int getReturnAddressOffset(MachineFunction &MF) const override {
+    return 3 * getPointerSize();
+  }
+
   // Get or create the frame index of where the old frame pointer is stored.
   int getOrCreateFramePointerSaveIndex(MachineFunction &MF) const override;
 };

llvm/lib/Target/SystemZ/SystemZISelLowering.cpp

@@ -3816,17 +3816,19 @@ SDValue SystemZTargetLowering::lowerRETURNADDR(SDValue Op,
       report_fatal_error("Unsupported stack frame traversal count");
 
     SDValue FrameAddr = lowerFRAMEADDR(Op, DAG);
-    auto *TFL = Subtarget.getFrameLowering<SystemZELFFrameLowering>();
-    int Offset = (TFL->usePackedStack(MF) ? -2 : 14) *
-                 getTargetMachine().getPointerSize(0);
+    const auto *TFL = Subtarget.getFrameLowering<SystemZFrameLowering>();
+    int Offset = TFL->getReturnAddressOffset(MF);
     SDValue Ptr = DAG.getNode(ISD::ADD, DL, PtrVT, FrameAddr,
                               DAG.getConstant(Offset, DL, PtrVT));
     return DAG.getLoad(PtrVT, DL, DAG.getEntryNode(), Ptr,
                        MachinePointerInfo());
   }
 
-  // Return R14D, which has the return address. Mark it an implicit live-in.
-  Register LinkReg = MF.addLiveIn(SystemZ::R14D, &SystemZ::GR64BitRegClass);
+  // Return R14D (Elf) / R7D (XPLINK), which has the return address. Mark it an
+  // implicit live-in.
+  SystemZCallingConventionRegisters *CCR = Subtarget.getSpecialRegisters();
+  Register LinkReg = MF.addLiveIn(CCR->getReturnFunctionAddressRegister(),
+                                  &SystemZ::GR64BitRegClass);
   return DAG.getCopyFromReg(DAG.getEntryNode(), DL, LinkReg, PtrVT);
 }
 

llvm/test/CodeGen/SystemZ/zos-ret-addr.ll

@@ -0,0 +1,41 @@
+; NOTE: Assertions have been autogenerated by utils/update_llc_test_checks.py UTC_ARGS: --version 4
+; RUN: llc < %s -mtriple=s390x-ibm-zos | FileCheck %s
+
+; The current function's return address is in the link register.
+define ptr @rt0() norecurse nounwind readnone {
+; CHECK-LABEL: rt0:
+; CHECK:         lgr 3, 7
+; CHECK-NEXT:    b 2(7)
+entry:
+  %0 = tail call ptr @llvm.returnaddress(i32 0)
+  ret ptr %0
+}
+
+; Check the caller's return address.
+define ptr @rtcaller() nounwind "backchain" {
+; CHECK-LABEL: rtcaller:
+; CHECK:         stmg 4, 7, 2048(4)
+; CHECK-NEXT:    lg 1, 2048(4)
+; CHECK-NEXT:    lg 3, 24(1)
+; CHECK-NEXT:    lmg 4, 7, 2048(4)
+; CHECK-NEXT:    b 2(7)
+entry:
+  %0 = tail call ptr @llvm.returnaddress(i32 1)
+  ret ptr %0
+}
+
+; Check the caller's caller's return address.
+define ptr @rtcallercaller() nounwind "backchain" {
+; CHECK-LABEL: rtcallercaller:
+; CHECK:         stmg 4, 7, 2048(4)
+; CHECK-NEXT:    lg 1, 2048(4)
+; CHECK-NEXT:    lg 1, 0(1)
+; CHECK-NEXT:    lg 3, 24(1)
+; CHECK-NEXT:    lmg 4, 7, 2048(4)
+; CHECK-NEXT:    b 2(7)
+entry:
+  %0 = tail call ptr @llvm.returnaddress(i32 2)
+  ret ptr %0
+}
+
+declare ptr @llvm.returnaddress(i32) nounwind readnone