[SystemZ][z/OS] Implement detection and handling for XPLink Leaf procedures.

This PR adds support for creating leaf functions when there are no CSRs used, no function calls are made, no stack frame is acquired, and contain no try/catch/throw statements.

Reviewed By: uweigand

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

llvm/lib/Target/SystemZ/SystemZFrameLowering.cpp

@@ -911,6 +911,54 @@ SystemZXPLINKFrameLowering::SystemZXPLINKFrameLowering()
         XPLINKSpillOffsetTable[I].Offset;
 }
 
+// Checks if the function is a potential candidate for being a XPLeaf routine.
+static bool isXPLeafCandidate(const MachineFunction &MF) {
+  const MachineFrameInfo &MFFrame = MF.getFrameInfo();
+  const MachineRegisterInfo &MRI = MF.getRegInfo();
+  const SystemZSubtarget &Subtarget = MF.getSubtarget<SystemZSubtarget>();
+  auto *Regs =
+      static_cast<SystemZXPLINK64Registers *>(Subtarget.getSpecialRegisters());
+
+  // If function calls other functions including alloca, then it is not a XPLeaf
+  // routine.
+  if (MFFrame.hasCalls())
+    return false;
+
+  // If the function has var Sized Objects, then it is not a XPLeaf routine.
+  if (MFFrame.hasVarSizedObjects())
+    return false;
+
+  // If the function adjusts the stack, then it is not a XPLeaf routine.
+  if (MFFrame.adjustsStack())
+    return false;
+
+  // If function modifies the stack pointer register, then it is not a XPLeaf
+  // routine.
+  if (MRI.isPhysRegModified(Regs->getStackPointerRegister()))
+    return false;
+
+  // If function modifies the ADA register, then it is not a XPLeaf routine.
+  if (MRI.isPhysRegModified(Regs->getAddressOfCalleeRegister()))
+    return false;
+
+  // If function modifies the return address register, then it is not a XPLeaf
+  // routine.
+  if (MRI.isPhysRegModified(Regs->getReturnFunctionAddressRegister()))
+    return false;
+
+  // If the backchain pointer should be stored, then it is not a XPLeaf routine.
+  if (MF.getFunction().hasFnAttribute("backchain"))
+    return false;
+
+  // If function acquires its own stack frame, then it is not a XPLeaf routine.
+  // At the time this function is called, only slots for local variables are
+  // allocated, so this is a very rough estimate.
+  if (MFFrame.estimateStackSize(MF) > 0)
+    return false;
+
+  return true;
+}
+
 bool SystemZXPLINKFrameLowering::assignCalleeSavedSpillSlots(
     MachineFunction &MF, const TargetRegisterInfo *TRI,
     std::vector<CalleeSavedInfo> &CSI) const {
@@ -920,6 +968,18 @@ bool SystemZXPLINKFrameLowering::assignCalleeSavedSpillSlots(
   auto &Regs = Subtarget.getSpecialRegisters<SystemZXPLINK64Registers>();
   auto &GRRegClass = SystemZ::GR64BitRegClass;
 
+  // At this point, the result of isXPLeafCandidate() is not accurate because
+  // the size of the save area has not yet been determined. If
+  // isXPLeafCandidate() indicates a potential leaf function, and there are no
+  // callee-save registers, then it is indeed a leaf function, and we can early
+  // exit.
+  // TODO: It is possible for leaf functions to use callee-saved registers.
+  // It can use the 0-2k range between R4 and the caller's stack frame without
+  // acquiring its own stack frame.
+  bool IsLeaf = CSI.empty() && isXPLeafCandidate(MF);
+  if (IsLeaf)
+    return true;
+
   // For non-leaf functions:
   // - the address of callee (entry point) register R6 must be saved
   CSI.push_back(CalleeSavedInfo(Regs.getAddressOfCalleeRegister()));
@@ -1137,16 +1197,16 @@ void SystemZXPLINKFrameLowering::emitPrologue(MachineFunction &MF,
   auto &Regs = Subtarget.getSpecialRegisters<SystemZXPLINK64Registers>();
   MachineFrameInfo &MFFrame = MF.getFrameInfo();
   MachineInstr *StoreInstr = nullptr;
+
+  determineFrameLayout(MF);
+
   bool HasFP = hasFP(MF);
   // Debug location must be unknown since the first debug location is used
   // to determine the end of the prologue.
   DebugLoc DL;
   uint64_t Offset = 0;
 
-  // TODO: Support leaf functions; only add size of save+reserved area when
-  // function is non-leaf.
-  MFFrame.setStackSize(MFFrame.getStackSize() + Regs.getCallFrameSize());
-  uint64_t StackSize = MFFrame.getStackSize();
+  const uint64_t StackSize = MFFrame.getStackSize();
 
   if (ZFI->getSpillGPRRegs().LowGPR) {
     // Skip over the GPR saves.
@@ -1321,3 +1381,32 @@ void SystemZXPLINKFrameLowering::processFunctionBeforeFrameFinalized(
   // Setup stack frame offset
   MFFrame.setOffsetAdjustment(Regs.getStackPointerBias());
 }
+
+// Determines the size of the frame, and creates the deferred spill objects.
+void SystemZXPLINKFrameLowering::determineFrameLayout(
+    MachineFunction &MF) const {
+  MachineFrameInfo &MFFrame = MF.getFrameInfo();
+  const SystemZSubtarget &Subtarget = MF.getSubtarget<SystemZSubtarget>();
+  auto *Regs =
+      static_cast<SystemZXPLINK64Registers *>(Subtarget.getSpecialRegisters());
+
+  uint64_t StackSize = MFFrame.getStackSize();
+  if (StackSize == 0)
+    return;
+
+  // Add the size of the register save area and the reserved area to the size.
+  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));
+  }
+}

llvm/lib/Target/SystemZ/SystemZFrameLowering.h

@@ -134,6 +134,8 @@ public:
 
   void processFunctionBeforeFrameFinalized(MachineFunction &MF,
                                            RegScavenger *RS) const override;
+
+  void determineFrameLayout(MachineFunction &MF) const;
 };
 } // end namespace llvm
 

llvm/test/CodeGen/SystemZ/call-zos-01.ll

@@ -86,7 +86,7 @@ entry:
 }
 
 ; CHECK-LABEL: pass_integrals0:
-; CHECK: ag  2, 2328(4)
+; CHECK: ag  2, 2200(4)
 ; CHECK-NEXT: lgr 3, 2
 define signext i64 @pass_integrals0(i64 signext %arg0, i32 signext %arg1, i16 signext %arg2, i64 signext %arg3) {
 entry:

llvm/test/CodeGen/SystemZ/call-zos-vec.ll

@@ -14,7 +14,7 @@ entry:
 ; CHECK: vaf 1, 1, 27
 ; CHECK: vaf 1, 1, 28
 ; CHECK: vaf 1, 1, 29
-; CHECK: vl  0, 2432(4), 4
+; CHECK: vl  0, 2304(4), 4
 ; CHECK: vaf 1, 1, 30
 ; CHECK: vaf 1, 1, 31
 ; CHECK: vaf 24, 1, 0

llvm/test/CodeGen/SystemZ/zos-prologue-epilog.ll

@@ -328,6 +328,21 @@ define void @large_stack() {
   ret void
 }
 
+; CHECK-LABEL: leaf_func
+; CHECK-NOT: aghi  4,
+; CHECK-NOT: stmg
+; CHECK: agr	1, 2
+; CHECK: msgr	1, 3
+; CHECK: aghik	3, 1, -4
+; CHECK-NOT: aghi  4,
+; CHECK-NOT: lmg
+define i64 @leaf_func0(i64 %a, i64 %b, i64 %c) {
+  %n = add i64 %a, %b
+  %m = mul i64 %n, %c
+  %o = sub i64 %m, 4
+  ret i64 %o
+}
+
 declare i64 @fun(i64 %arg0)
 declare i64 @fun1(i8* %ptr)
 declare i64 @fun2(i64 %n, i64* %arr0, i64* %arr1)

llvm/test/MC/GOFF/ppa1.ll

@@ -7,7 +7,7 @@
 ; CHECK: .short  197
 ; CHECK: .byte   0
 ; CHECK: .byte   241 * Mark Type C'1'
-; CHECK: .long   128 * DSA Size 0x80
+; CHECK: .long   0 * DSA Size 0x0
 ; CHECK: * Entry Flags
 ; CHECK: *   Bit 2: 0 = Does not use alloca
 ; CHECK: @@func_end0:
@@ -15,7 +15,7 @@
 ; CHECK: @@PPA1_void_test_0:                     * PPA1
 ; CHECK:        .byte   2                               * Version
 ; CHECK:        .byte   206                             * LE Signature X'CE'
-; CHECK:        .short  768                             * Saved GPR Mask
+; CHECK:        .short  0                               * Saved GPR Mask
 ; CHECK:        .byte   128                             * PPA1 Flags 1
 ; CHECK:                                        *   Bit 0: 1 = 64-bit DSA
 ; CHECK:        .byte   128                             * PPA1 Flags 2