[SystemZ] Improve handling of backchain offset

[jonpa]

Summary

• New function SDValue getBackchainAddress() used by lowerDYNAMIC_STACKALLOC() and lowerSTACKRESTORE() to properly handle the backchain offset also with packed-stack.

• Make a function getBackchainOffset() for the computation of the backchain and use in three places (NFC).

A new test @f6 in backchain.ll, which is same as @f5 except for using packed-stack. I replaced the offsets as needed and hopefully the test should be correct.

[uweigand]

This all looks good to me, the test case also looks fine (and the offset match what GCC is doing).

llvm/lib/Target/SystemZ/SystemZFrameLowering.h

 //===-- SystemZFrameLowering.h - Frame lowering for SystemZ -----*- C++ -*-===//
 //
 // Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
 // See https://llvm.org/LICENSE.txt for license information.
 // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
 //
 //===----------------------------------------------------------------------===//
 
 #ifndef LLVM_LIB_TARGET_SYSTEMZ_SYSTEMZFRAMELOWERING_H
 #define LLVM_LIB_TARGET_SYSTEMZ_SYSTEMZFRAMELOWERING_H
 
+#include "MCTargetDesc/SystemZMCTargetDesc.h"
 #include "llvm/ADT/IndexedMap.h"
 #include "llvm/CodeGen/TargetFrameLowering.h"
 #include "llvm/Support/TypeSize.h"
 
 namespace llvm {
 class SystemZTargetMachine;
 class SystemZSubtarget;
 
public:
   // ABI-defined save slot.  Return 0 if no slot is defined for Reg.  Adjust
   // the offset in case MF has packed-stack.
   unsigned getRegSpillOffset(MachineFunction &MF, Register Reg) const;
 
   // Get or create the frame index of where the old frame pointer is stored.
   int getOrCreateFramePointerSaveIndex(MachineFunction &MF) const;
 
   bool usePackedStack(MachineFunction &MF) const;
+
+  // Return the offset of the backchain.
+  unsigned getBackchainOffset(MachineFunction &MF) const {
+    // The back chain is stored topmost with packed-stack.
+    return usePackedStack(MF) ? SystemZMC::CallFrameSize - 8 : 0;
+  }
 };
 } // end namespace llvm
 
 #endif

llvm/lib/Target/SystemZ/SystemZFrameLowering.cpp

else {
       bool StoreBackchain = MF.getFunction().hasFnAttribute("backchain");
       // If we need backchain, save current stack pointer.  R1 is free at
       // this point.
       if (StoreBackchain)
         BuildMI(MBB, MBBI, DL, ZII->get(SystemZ::LGR))
           .addReg(SystemZ::R1D, RegState::Define).addReg(SystemZ::R15D);
       emitIncrement(MBB, MBBI, DL, SystemZ::R15D, Delta, ZII);
       buildCFAOffs(MBB, MBBI, DL, SPOffsetFromCFA + Delta, ZII);
-      if (StoreBackchain) {
-        // The back chain is stored topmost with packed-stack.
-        int Offset = usePackedStack(MF) ? SystemZMC::CallFrameSize - 8 : 0;
+      if (StoreBackchain)
         BuildMI(MBB, MBBI, DL, ZII->get(SystemZ::STG))
           .addReg(SystemZ::R1D, RegState::Kill).addReg(SystemZ::R15D)
-          .addImm(Offset).addReg(0);
-      }
+          .addImm(getBackchainOffset(MF)).addReg(0);
     }
     SPOffsetFromCFA += Delta;
   }
 
   if (HasFP) {
     // Copy the base of the frame to R11.
     BuildMI(MBB, MBBI, DL, ZII->get(SystemZ::LGR), SystemZ::R11D)
       .addReg(SystemZ::R15D);
if (NumFullBlocks < 3) {
     MBB = DoneMBB;
     MBBI = DoneMBB->begin();
     buildDefCFAReg(*MBB, MBBI, DL, SystemZ::R15D, ZII);
   }
 
   if (Residual)
     allocateAndProbe(*MBB, MBBI, Residual, true/*EmitCFI*/);
 
-  if (StoreBackchain) {
-    // The back chain is stored topmost with packed-stack.
-    int Offset = usePackedStack(MF) ? SystemZMC::CallFrameSize - 8 : 0;
+  if (StoreBackchain)
     BuildMI(*MBB, MBBI, DL, ZII->get(SystemZ::STG))
       .addReg(SystemZ::R1D, RegState::Kill).addReg(SystemZ::R15D)
-      .addImm(Offset).addReg(0);
-  }
+      .addImm(getBackchainOffset(MF)).addReg(0);
 
   StackAllocMI->eraseFromParent();
   if (DoneMBB != nullptr) {
     // Compute the live-in lists for the new blocks.
     recomputeLiveIns(*DoneMBB);
     recomputeLiveIns(*LoopMBB);
   }
 }

 }
 
 int SystemZFrameLowering::
 getOrCreateFramePointerSaveIndex(MachineFunction &MF) const {
   SystemZMachineFunctionInfo *ZFI = MF.getInfo<SystemZMachineFunctionInfo>();
   int FI = ZFI->getFramePointerSaveIndex();
   if (!FI) {
     MachineFrameInfo &MFFrame = MF.getFrameInfo();
-    // The back chain is stored topmost with packed-stack.
-    int Offset = usePackedStack(MF) ? -8 : -SystemZMC::CallFrameSize;
+    int Offset = getBackchainOffset(MF) - SystemZMC::CallFrameSize;
     FI = MFFrame.CreateFixedObject(8, Offset, false);
     ZFI->setFramePointerSaveIndex(FI);
   }
   return FI;
 }
 
 bool SystemZFrameLowering::usePackedStack(MachineFunction &MF) const {
   bool HasPackedStackAttr = MF.getFunction().hasFnAttribute("packed-stack");
   bool BackChain = MF.getFunction().hasFnAttribute("backchain");
   bool SoftFloat = MF.getSubtarget<SystemZSubtarget>().hasSoftFloat();
   if (HasPackedStackAttr && BackChain && !SoftFloat)
     report_fatal_error("packed-stack + backchain + hard-float is unsupported.");
   bool CallConv = MF.getFunction().getCallingConv() != CallingConv::GHC;
   return HasPackedStackAttr && CallConv;
 }

llvm/lib/Target/SystemZ/SystemZISelLowering.h

MachineBasicBlock *emitTransactionBegin(MachineInstr &MI,
                                           MachineBasicBlock *MBB,
                                           unsigned Opcode, bool NoFloat) const;
   MachineBasicBlock *emitLoadAndTestCmp0(MachineInstr &MI,
                                          MachineBasicBlock *MBB,
                                          unsigned Opcode) const;
   MachineBasicBlock *emitProbedAlloca(MachineInstr &MI,
                                       MachineBasicBlock *MBB) const;
 
+  SDValue getBackchainAddress(SDValue SP, SelectionDAG &DAG) const;
+
   MachineMemOperand::Flags
   getTargetMMOFlags(const Instruction &I) const override;
   const TargetRegisterClass *getRepRegClassFor(MVT VT) const override;
 };
 
 struct SystemZVectorConstantInfo {
 private:
   APInt IntBits;             // The 128 bits as an integer.

llvm/lib/Target/SystemZ/SystemZISelLowering.cpp

lowerDYNAMIC_STACKALLOC(SDValue Op, SelectionDAG &DAG) const {
   SDValue NeededSpace = Size;
 
   // Get a reference to the stack pointer.
   SDValue OldSP = DAG.getCopyFromReg(Chain, DL, SPReg, MVT::i64);
 
   // If we need a backchain, save it now.
   SDValue Backchain;
   if (StoreBackchain)
-    Backchain = DAG.getLoad(MVT::i64, DL, Chain, OldSP, MachinePointerInfo());
+    Backchain = DAG.getLoad(MVT::i64, DL, Chain, getBackchainAddress(OldSP, DAG),
+                            MachinePointerInfo());
 
   // Add extra space for alignment if needed.
   if (ExtraAlignSpace)
     NeededSpace = DAG.getNode(ISD::ADD, DL, MVT::i64, NeededSpace,
                               DAG.getConstant(ExtraAlignSpace, DL, MVT::i64));
 
   // Get the new stack pointer value.
   SDValue NewSP;
Result =
       DAG.getNode(ISD::ADD, DL, MVT::i64, Result,
                   DAG.getConstant(ExtraAlignSpace, DL, MVT::i64));
     Result =
       DAG.getNode(ISD::AND, DL, MVT::i64, Result,
                   DAG.getConstant(~(RequiredAlign - 1), DL, MVT::i64));
   }
 
   if (StoreBackchain)
-    Chain = DAG.getStore(Chain, DL, Backchain, NewSP, MachinePointerInfo());
+    Chain = DAG.getStore(Chain, DL, Backchain, getBackchainAddress(NewSP, DAG),
+                         MachinePointerInfo());
 
   SDValue Ops[2] = { Result, Chain };
   return DAG.getMergeValues(Ops, DL);
 }
 
 SDValue SystemZTargetLowering::lowerGET_DYNAMIC_AREA_OFFSET(
     SDValue Op, SelectionDAG &DAG) const {
   SDLoc DL(Op);
SDValue SystemZTargetLowering::lowerSTACKRESTORE(SDValue Op,
 
   SDValue Chain = Op.getOperand(0);
   SDValue NewSP = Op.getOperand(1);
   SDValue Backchain;
   SDLoc DL(Op);
 
   if (StoreBackchain) {
     SDValue OldSP = DAG.getCopyFromReg(Chain, DL, SystemZ::R15D, MVT::i64);
-    Backchain = DAG.getLoad(MVT::i64, DL, Chain, OldSP, MachinePointerInfo());
+    Backchain = DAG.getLoad(MVT::i64, DL, Chain, getBackchainAddress(OldSP, DAG),
+                            MachinePointerInfo());
   }
 
   Chain = DAG.getCopyToReg(Chain, DL, SystemZ::R15D, NewSP);
 
   if (StoreBackchain)
-    Chain = DAG.getStore(Chain, DL, Backchain, NewSP, MachinePointerInfo());
+    Chain = DAG.getStore(Chain, DL, Backchain, getBackchainAddress(NewSP, DAG),
+                         MachinePointerInfo());
 
   return Chain;
 }
 
 SDValue SystemZTargetLowering::lowerPREFETCH(SDValue Op,
                                              SelectionDAG &DAG) const {
   bool IsData = cast<ConstantSDNode>(Op.getOperand(4))->getZExtValue();
   if (!IsData)
MachineBasicBlock *SystemZTargetLowering::emitProbedAlloca(
   MBB = DoneMBB;
   BuildMI(*MBB, MBB->begin(), DL, TII->get(TargetOpcode::COPY), DstReg)
     .addReg(SystemZ::R15D);
 
   MI.eraseFromParent();
   return DoneMBB;
 }
 
+SDValue SystemZTargetLowering::
+getBackchainAddress(SDValue SP, SelectionDAG &DAG) const {
+  MachineFunction &MF = DAG.getMachineFunction();
+  auto *TFL =
+      static_cast<const SystemZFrameLowering *>(Subtarget.getFrameLowering());
+  SDLoc DL(SP);
+  return DAG.getNode(ISD::ADD, DL, MVT::i64, SP,
+                     DAG.getIntPtrConstant(TFL->getBackchainOffset(MF), DL));
+}
+
 MachineBasicBlock *SystemZTargetLowering::EmitInstrWithCustomInserter(
     MachineInstr &MI, MachineBasicBlock *MBB) const {
   switch (MI.getOpcode()) {
   case SystemZ::Select32:
   case SystemZ::Select64:
   case SystemZ::SelectF32:
   case SystemZ::SelectF64:
   case SystemZ::SelectF128:

llvm/test/CodeGen/SystemZ/backchain.ll

; CHECK: br %r14
   %src = call i8 *@llvm.stacksave()
   %array1 = alloca i8, i32 %count1
   store volatile i8 0, i8 *%array1
   call void @llvm.stackrestore(i8 *%src)
   %array2 = alloca i8, i32 %count2
   store volatile i8 0, i8 *%array2
   ret void
 }
+
+; same, but with the kernel backchain
+define void @f6(i32 %count1, i32 %count2) #0 {
+; CHECK-LABEL: f6:
+; CHECK: stmg %r11, %r15, 112(%r15)
+; CHECK: lgr %r1, %r15
+; CHECK: aghi %r15, -48
+; CHECK: stg %r1, 152(%r15)
+; CHECK: lgr %r11, %r15
+; CHECK-DAG: lgr [[SAVESP:%r[0-9]+]], %r15
+; CHECK-DAG: lg [[BC:%r[0-9]+]], 152(%r15)
+; CHECK-DAG: lgr [[NEWSP:%r[0-9]+]], %r15
+; CHECK-DAG: lgr %r15, [[NEWSP]]
+; CHECK-DAG: stg [[BC]], 152([[NEWSP]])
+; CHECK-DAG: lg [[BC2:%r[0-9]+]], 152(%r15)
+; CHECK-DAG: lgr %r15, [[SAVESP]]
+; CHECK-DAG: stg [[BC2]], 152([[SAVESP]])
+; CHECK-DAG: lg [[BC3:%r[0-9]+]], 152(%r15)
+; CHECK-DAG: lgr [[NEWSP2:%r[0-9]+]], %r15
+; CHECK-DAG: lgr %r15, [[NEWSP2]]
+; CHECK-DAG: stg [[BC3]], 152([[NEWSP2]])
+; CHECK: lmg %r11, %r15, 160(%r11)
+; CHECK: br %r14
+  %src = call i8 *@llvm.stacksave()
+  %array1 = alloca i8, i32 %count1
+  store volatile i8 0, i8 *%array1
+  call void @llvm.stackrestore(i8 *%src)
+  %array2 = alloca i8, i32 %count2
+  store volatile i8 0, i8 *%array2
+  ret void
+}
+
+attributes #0 = { "backchain" "packed-stack" "use-soft-float"="true" }