[SystemZ] [z/OS] Add support for generating huge (1 MiB) stack frames in XPLINK64

[Everybody0523]

Summary

This patch extends support for generating huge stack frames on 64-bit XPLINK by implementing the ABI-mandated call to the stack extension routine.

[uweigand]

LGTM, thanks!

llvm/lib/Target/SystemZ/SystemZAsmPrinter.cpp

void SystemZAsmPrinter::emitInstruction(const MachineInstr *MI) {
 
   case SystemZ::CallBASR_XPLINK64:
     EmitToStreamer(*OutStreamer, MCInstBuilder(SystemZ::BASR)
                                      .addReg(SystemZ::R7D)
                                      .addReg(MI->getOperand(0).getReg()));
     emitCallInformation(CallType::BASR76);
     return;
 
+  case SystemZ::CallBASR_STACKEXT:
+    EmitToStreamer(*OutStreamer, MCInstBuilder(SystemZ::BASR)
+                                     .addReg(SystemZ::R3D)
+                                     .addReg(MI->getOperand(0).getReg()));
+    emitCallInformation(CallType::BASR33);
+    return;
+
   case SystemZ::CallBRASL:
     LoweredMI = MCInstBuilder(SystemZ::BRASL)
       .addReg(SystemZ::R14D)
       .addExpr(Lower.getExpr(MI->getOperand(0), MCSymbolRefExpr::VK_PLT));
     break;
 
   case SystemZ::CallBASR:
     LoweredMI = MCInstBuilder(SystemZ::BASR)

llvm/lib/Target/SystemZ/SystemZFrameLowering.h

restoreCalleeSavedRegisters(MachineBasicBlock &MBB,
                               MachineBasicBlock::iterator MBBII,
                               MutableArrayRef<CalleeSavedInfo> CSI,
                               const TargetRegisterInfo *TRI) const override;
 
   void emitEpilogue(MachineFunction &MF, MachineBasicBlock &MBB) const override;
 
   void emitPrologue(MachineFunction &MF, MachineBasicBlock &MBB) const override;
 
+  void inlineStackProbe(MachineFunction &MF,
+                        MachineBasicBlock &PrologMBB) const override;
+
   bool hasFP(const MachineFunction &MF) const override;
 
   void processFunctionBeforeFrameFinalized(MachineFunction &MF,
                                            RegScavenger *RS) const override;
 };
 } // end namespace llvm
 
 #endif

llvm/lib/Target/SystemZ/SystemZFrameLowering.cpp

void SystemZXPLINKFrameLowering::emitPrologue(MachineFunction &MF,
   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();
 
-  // FIXME: Implement support for large stack sizes, when the stack extension
-  // routine needs to be called.
-  if (StackSize > 1024 * 1024) {
-    llvm_unreachable("Huge Stack Frame not yet supported on z/OS");
-  }
-
   if (ZFI->getSpillGPRRegs().LowGPR) {
     // Skip over the GPR saves.
     if ((MBBI != MBB.end()) && ((MBBI->getOpcode() == SystemZ::STMG))) {
       const int Operand = 3;
       // Now we can set the offset for the operation, since now the Stack
       // has been finalized.
       Offset = Regs.getStackPointerBias() + MBBI->getOperand(Operand).getImm();
       // Maximum displacement for STMG instruction.
if (StoreInstr && HasFP) {
           .addReg(SystemZ::R0D, RegState::Kill)
           .addReg(SystemZ::R4D)
           .addImm(Offset)
           .addReg(0);
     }
 
     emitIncrement(MBB, InsertPt, DL, Regs.getStackPointerRegister(), Delta,
                   ZII);
+
+    // If the requested stack size is larger than the guard page, then we need
+    // to check if we need to call the stack extender. This requires adding a
+    // conditional branch, but splitting the prologue block is not possible at
+    // this point since it would invalidate the SaveBlocks / RestoreBlocks sets
+    // of PEI in the single block function case. Build a pseudo to be handled
+    // later by inlineStackProbe().
+    const uint64_t GuardPageSize = 1024 * 1024;
+    if (StackSize > GuardPageSize) {
+      assert(StoreInstr && "Wrong insertion point");
+      BuildMI(MBB, InsertPt, DL, ZII->get(SystemZ::XPLINK_STACKALLOC));
+    }
   }
 
   if (HasFP) {
     // Copy the base of the frame to Frame Pointer Register.
     BuildMI(MBB, MBBI, DL, ZII->get(SystemZ::LGR),
             Regs.getFramePointerRegister())
         .addReg(Regs.getStackPointerRegister());
 
if (StackSize) {
     unsigned SPReg = Regs.getStackPointerRegister();
     if (ZFI->getRestoreGPRRegs().LowGPR != SPReg) {
       DebugLoc DL = MBBI->getDebugLoc();
       emitIncrement(MBB, MBBI, DL, SPReg, StackSize, ZII);
     }
   }
 }
 
+// Emit a compare of the stack pointer against the stack floor, and a call to
+// the LE stack extender if needed.
+void SystemZXPLINKFrameLowering::inlineStackProbe(
+    MachineFunction &MF, MachineBasicBlock &PrologMBB) const {
+  auto *ZII =
+      static_cast<const SystemZInstrInfo *>(MF.getSubtarget().getInstrInfo());
+
+  MachineInstr *StackAllocMI = nullptr;
+  for (MachineInstr &MI : PrologMBB)
+    if (MI.getOpcode() == SystemZ::XPLINK_STACKALLOC) {
+      StackAllocMI = &MI;
+      break;
+    }
+  if (StackAllocMI == nullptr)
+    return;
+
+  MachineBasicBlock &MBB = PrologMBB;
+  const DebugLoc DL = StackAllocMI->getDebugLoc();
+
+  // The 2nd half of block MBB after split.
+  MachineBasicBlock *NextMBB;
+
+  // Add new basic block for the call to the stack overflow function.
+  MachineBasicBlock *StackExtMBB =
+      MF.CreateMachineBasicBlock(MBB.getBasicBlock());
+  MF.push_back(StackExtMBB);
+
+  // LG r3,72(,r3)
+  BuildMI(StackExtMBB, DL, ZII->get(SystemZ::LG), SystemZ::R3D)
+      .addReg(SystemZ::R3D)
+      .addImm(72)
+      .addReg(0);
+  // BASR r3,r3
+  BuildMI(StackExtMBB, DL, ZII->get(SystemZ::CallBASR_STACKEXT))
+      .addReg(SystemZ::R3D);
+
+  // LLGT r3,1208
+  BuildMI(MBB, StackAllocMI, DL, ZII->get(SystemZ::LLGT), SystemZ::R3D)
+      .addReg(0)
+      .addImm(1208)
+      .addReg(0);
+  // CG r4,64(,r3)
+  BuildMI(MBB, StackAllocMI, DL, ZII->get(SystemZ::CG))
+      .addReg(SystemZ::R4D)
+      .addReg(SystemZ::R3D)
+      .addImm(64)
+      .addReg(0);
+  // JLL b'0100',F'37'
+  BuildMI(MBB, StackAllocMI, DL, ZII->get(SystemZ::BRC))
+      .addImm(SystemZ::CCMASK_ICMP)
+      .addImm(SystemZ::CCMASK_CMP_LT)
+      .addMBB(StackExtMBB);
+
+  NextMBB = SystemZ::splitBlockBefore(StackAllocMI, &MBB);
+  MBB.addSuccessor(NextMBB);
+  MBB.addSuccessor(StackExtMBB);
+
+  // Add jump back from stack extension BB.
+  BuildMI(StackExtMBB, DL, ZII->get(SystemZ::J)).addMBB(NextMBB);
+  StackExtMBB->addSuccessor(NextMBB);
+
+  StackAllocMI->eraseFromParent();
+
+  // Compute the live-in lists for the new blocks.
+  recomputeLiveIns(*NextMBB);
+  recomputeLiveIns(*StackExtMBB);
+}
+
 bool SystemZXPLINKFrameLowering::hasFP(const MachineFunction &MF) const {
   return (MF.getFrameInfo().hasVarSizedObjects());
 }
 
 void SystemZXPLINKFrameLowering::processFunctionBeforeFrameFinalized(
     MachineFunction &MF, RegScavenger *RS) const {
   MachineFrameInfo &MFFrame = MF.getFrameInfo();
   const SystemZSubtarget &Subtarget = MF.getSubtarget<SystemZSubtarget>();
   auto &Regs = Subtarget.getSpecialRegisters<SystemZXPLINK64Registers>();
 
   // Setup stack frame offset
   MFFrame.setOffsetAdjustment(Regs.getStackPointerBias());
 }

llvm/lib/Target/SystemZ/SystemZInstrInfo.td

let Defs = [R15D, CC], Uses = [R15D], hasNoSchedulingInfo = 1,
   def PROBED_ALLOCA : Pseudo<(outs GR64:$dst),
                              (ins GR64:$oldSP, GR64:$space),
            [(set GR64:$dst, (z_probed_alloca GR64:$oldSP, GR64:$space))]>;
 
 let Defs = [R1D, R15D, CC], Uses = [R15D], hasNoSchedulingInfo = 1,
     hasSideEffects = 1 in
   def PROBED_STACKALLOC : Pseudo<(outs), (ins i64imm:$stacksize), []>;
 
+let Defs = [R3D, CC], Uses = [R3D, R4D], hasNoSchedulingInfo = 1,
+    hasSideEffects = 1 in
+  def XPLINK_STACKALLOC : Pseudo<(outs), (ins), []>;
+
 //===----------------------------------------------------------------------===//
 // Branch instructions
 //===----------------------------------------------------------------------===//
 
 // Conditional branches.
 let isBranch = 1, isTerminator = 1, Uses = [CC] in {
   // It's easier for LLVM to handle these branches in their raw BRC/BRCL form
   // with the condition-code mask being the first operand.  It seems friendlier

 // z/OS XPLINK
 let Predicates = [IsTargetXPLINK64] in {
   let isCall = 1, Defs = [R7D, CC], Uses = [FPC] in {
     def CallBRASL_XPLINK64 : Alias<8, (outs), (ins pcrel32:$I2, variable_ops),
                           [(z_call pcrel32:$I2)]>;
     def CallBASR_XPLINK64  : Alias<4, (outs), (ins ADDR64:$R2, variable_ops),
                           [(z_call ADDR64:$R2)]>;
   }
+
+  let isCall = 1, Defs = [R3D, CC], Uses = [FPC] in {
+    def CallBASR_STACKEXT  : Alias<4, (outs), (ins ADDR64:$R2), []>;
+  }
 }
 
 // Regular calls.
 // z/Linux ELF
 let Predicates = [IsTargetELF] in {
   let isCall = 1, Defs = [R14D, CC], Uses = [FPC] in {
     def CallBRASL : Alias<6, (outs), (ins pcrel32:$I2, variable_ops),
                           [(z_call pcrel32:$I2)]>;

llvm/lib/Target/SystemZ/SystemZScheduleZ13.td

 
 //===----------------------------------------------------------------------===//
 // Call and return instructions
 //===----------------------------------------------------------------------===//
 
 // Call
 def : InstRW<[WLat1, VBU, FXa2, GroupAlone], (instregex "(Call)?BRAS$")>;
 def : InstRW<[WLat1, FXa2, FXb, GroupAlone], (instregex "(Call)?BRASL(_XPLINK64)?$")>;
-def : InstRW<[WLat1, FXa2, FXb, GroupAlone], (instregex "(Call)?BAS(R)?(_XPLINK64)?$")>;
+def : InstRW<[WLat1, FXa2, FXb, GroupAlone], (instregex "(Call)?BAS(R)?(_XPLINK64|_STACKEXT)?$")>;
 def : InstRW<[WLat1, FXa2, FXb, GroupAlone], (instregex "TLS_(G|L)DCALL$")>;
 
 // Return
 def : InstRW<[WLat1, FXb, EndGroup], (instregex "Return(_XPLINK)?$")>;
 def : InstRW<[WLat1, FXb, NormalGr], (instregex "CondReturn(_XPLINK)?$")>;
 
 //===----------------------------------------------------------------------===//
 // Move instructions

llvm/lib/Target/SystemZ/SystemZScheduleZ14.td

 
 //===----------------------------------------------------------------------===//
 // Call and return instructions
 //===----------------------------------------------------------------------===//
 
 // Call
 def : InstRW<[WLat1, VBU, FXa2, GroupAlone], (instregex "(Call)?BRAS$")>;
 def : InstRW<[WLat1, FXa2, FXb, GroupAlone], (instregex "(Call)?BRASL(_XPLINK64)?$")>;
-def : InstRW<[WLat1, FXa2, FXb, GroupAlone], (instregex "(Call)?BAS(R)?(_XPLINK64)?$")>;
+def : InstRW<[WLat1, FXa2, FXb, GroupAlone], (instregex "(Call)?BAS(R)?(_XPLINK64|_STACKEXT)?$")>;
 def : InstRW<[WLat1, FXa2, FXb, GroupAlone], (instregex "TLS_(G|L)DCALL$")>;
 
 // Return
 def : InstRW<[WLat1, FXb, EndGroup], (instregex "Return(_XPLINK)?$")>;
 def : InstRW<[WLat1, FXb, NormalGr], (instregex "CondReturn(_XPLINK)?$")>;
 
 //===----------------------------------------------------------------------===//
 // Move instructions

llvm/lib/Target/SystemZ/SystemZScheduleZ15.td

 
 //===----------------------------------------------------------------------===//
 // Call and return instructions
 //===----------------------------------------------------------------------===//
 
 // Call
 def : InstRW<[WLat1, VBU, FXa2, GroupAlone], (instregex "(Call)?BRAS$")>;
 def : InstRW<[WLat1, FXa2, FXb, GroupAlone], (instregex "(Call)?BRASL(_XPLINK64)?$")>;
-def : InstRW<[WLat1, FXa2, FXb, GroupAlone], (instregex "(Call)?BAS(R)?(_XPLINK64)?$")>;
+def : InstRW<[WLat1, FXa2, FXb, GroupAlone], (instregex "(Call)?BAS(R)?(_XPLINK64|_STACKEXT)?$")>;
 def : InstRW<[WLat1, FXa2, FXb, GroupAlone], (instregex "TLS_(G|L)DCALL$")>;
 
 // Return
 def : InstRW<[WLat1, FXb, EndGroup], (instregex "Return(_XPLINK)?$")>;
 def : InstRW<[WLat1, FXb, NormalGr], (instregex "CondReturn(_XPLINK)?$")>;
 
 //===----------------------------------------------------------------------===//
 // Move instructions

llvm/lib/Target/SystemZ/SystemZScheduleZ196.td

 
 //===----------------------------------------------------------------------===//
 // Call and return instructions
 //===----------------------------------------------------------------------===//
 
 // Call
 def : InstRW<[WLat1, LSU, FXU2, GroupAlone], (instregex "(Call)?BRAS$")>;
 def : InstRW<[WLat1, LSU, FXU2, GroupAlone], (instregex "(Call)?BRASL(_XPLINK64)?$")>;
-def : InstRW<[WLat1, LSU, FXU2, GroupAlone], (instregex "(Call)?BAS(R)?(_XPLINK64)?$")>;
+def : InstRW<[WLat1, LSU, FXU2, GroupAlone], (instregex "(Call)?BAS(R)?(_XPLINK64|_STACKEXT)?$")>;
 def : InstRW<[WLat1, LSU, FXU2, GroupAlone], (instregex "TLS_(G|L)DCALL$")>;
 
 // Return
 def : InstRW<[WLat1, LSU, EndGroup], (instregex "Return(_XPLINK)?$")>;
 def : InstRW<[WLat1, LSU, EndGroup], (instregex "CondReturn(_XPLINK)?$")>;
 
 //===----------------------------------------------------------------------===//
 // Move instructions

llvm/lib/Target/SystemZ/SystemZScheduleZEC12.td

 
 //===----------------------------------------------------------------------===//
 // Call and return instructions
 //===----------------------------------------------------------------------===//
 
 // Call
 def : InstRW<[WLat1, FXU2, VBU, GroupAlone], (instregex "(Call)?BRAS$")>;
 def : InstRW<[WLat1, FXU2, LSU, GroupAlone], (instregex "(Call)?BRASL(_XPLINK64)?$")>;
-def : InstRW<[WLat1, FXU2, LSU, GroupAlone], (instregex "(Call)?BAS(R)?(_XPLINK64)?$")>;
+def : InstRW<[WLat1, FXU2, LSU, GroupAlone], (instregex "(Call)?BAS(R)?(_XPLINK64|_STACKEXT)?$")>;
 def : InstRW<[WLat1, FXU2, LSU, GroupAlone], (instregex "TLS_(G|L)DCALL$")>;
 
 // Return
 def : InstRW<[WLat1, LSU, EndGroup], (instregex "Return(_XPLINK)?$")>;
 def : InstRW<[WLat1, LSU, NormalGr], (instregex "CondReturn(_XPLINK)?$")>;
 
 //===----------------------------------------------------------------------===//
 // Move instructions

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

 define i64 @func5(i64 %n) {
   %vla = alloca i64, i64 %n, align 8
   %arr = alloca [130000 x i64], align 8
   %ptr = bitcast [130000 x i64]* %arr to i64*
   %call = call i64 @fun2(i64 %n, i64* nonnull %vla, i64* %ptr)
   ret i64 %call
 }
 
+; CHECK-LABEL: large_stack
+; CHECK64: agfi  4, -1048768
+; CHECK64-NEXT: llgt  3, 1208
+; CHECK64-NEXT: cg  4, 64(3)
+; CHECK64-NEXT: jhe
+; CHECK64: * %bb.1:
+; CHECK64: lg  3, 72(3)
+; CHECK64: basr  3, 3
+; CHECK64: stmg  6, 7, 2064(4)
+define void @large_stack() {
+  %arr = alloca [131072 x i64], align 8
+  %ptr = bitcast [131072 x i64]* %arr to i8*
+  call i64 (i8*) @fun1(i8* %ptr)
+  ret void
+}
+
 declare i64 @fun(i64 %arg0)
 declare i64 @fun1(i8* %ptr)
 declare i64 @fun2(i64 %n, i64* %arr0, i64* %arr1)