[SystemZ][z/OS] Save (and restore) R3 to avoid clobbering parameter when call stack frame extension is invoked

[Everybody0523]

Summary

When the stack frame extension routine is used, the contents of r3 is overwritten. However, if r3 is live in the prologue (ie. one of the function's parameters resides in r3), it needs to be saved. We save r3 in r0 if r0 is available (ie. r0 is not used as temporary storage for r4), and in the corresponding stack slot for the third parameter otherwise.

[uweigand]

llvm/lib/Target/SystemZ/SystemZFrameLowering.cpp

1354

%r0 cannot be used as base register. Is this really supposed to store to the low core? (On Linux, that might be unsafe as it could be clobbered if we're e.g. interrupted by an async signal handler - not sure if there are similar concerns on z/OS.) If yes, you need to use addReg(0) here. If no, I guess some other changes are required.

[Everybody0523]

Don't use r0 as base register...

[Everybody0523]

llvm/lib/Target/SystemZ/SystemZFrameLowering.cpp

1354

Right, this was indeed a mistake. I've amended it to use the incoming value of r4 (ie. r4 before the stack update) for the store, and then to use r3 as the base register (after a load from r0 to r3) in the load.

[uweigand]

LGTM.

[Everybody0523]

Fix build failure with expensive checks.

llvm/lib/Target/SystemZ/SystemZFrameLowering.cpp

void SystemZXPLINKFrameLowering::inlineStackProbe(
   for (MachineInstr &MI : PrologMBB)
     if (MI.getOpcode() == SystemZ::XPLINK_STACKALLOC) {
       StackAllocMI = &MI;
       break;
     }
   if (StackAllocMI == nullptr)
     return;
 
+  bool NeedSaveSP = hasFP(MF);
+  bool NeedSaveArg = PrologMBB.isLiveIn(SystemZ::R3D);
+  const int64_t SaveSlotR3 = 2192;
+
   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);
-
+  if (NeedSaveArg) {
+    if (!NeedSaveSP) {
+      // LGR r0,r3
+      BuildMI(MBB, StackAllocMI, DL, ZII->get(SystemZ::LGR))
+          .addReg(SystemZ::R0D, RegState::Define)
+          .addReg(SystemZ::R3D);
+    } else {
+      // In this case, the incoming value of r4 is saved in r0 so the
+      // latter register is unavailable. Store r3 in its corresponding
+      // slot in the parameter list instead. Do this at the start of
+      // the prolog before r4 is manipulated by anything else.
+      // STG r3, 2192(r4)
+      BuildMI(MBB, MBB.begin(), DL, ZII->get(SystemZ::STG))
+          .addReg(SystemZ::R3D)
+          .addReg(SystemZ::R4D)
+          .addImm(SaveSlotR3)
+          .addReg(0);
+    }
+  }
   // 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);
+  if (NeedSaveArg) {
+    if (!NeedSaveSP) {
+      // LGR r3, r0
+      BuildMI(*NextMBB, StackAllocMI, DL, ZII->get(SystemZ::LGR))
+          .addReg(SystemZ::R3D, RegState::Define)
+          .addReg(SystemZ::R0D, RegState::Kill);
+    } else {
+      // In this case, the incoming value of r4 is saved in r0 so the
+      // latter register is unavailable. We stored r3 in its corresponding
+      // slot in the parameter list instead and we now restore it from there.
+      // LGR r3, r0
+      BuildMI(*NextMBB, StackAllocMI, DL, ZII->get(SystemZ::LGR))
+          .addReg(SystemZ::R3D)
+          .addReg(SystemZ::R0D);
+      // LG r3, 2192(r3)
+      BuildMI(*NextMBB, StackAllocMI, DL, ZII->get(SystemZ::LG))
+          .addReg(SystemZ::R3D, RegState::Define)
+          .addReg(SystemZ::R3D)
+          .addImm(SaveSlotR3)
+          .addReg(0);
+    }
+  }
 
   // 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.

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

 ; 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() {
+define void @large_stack0() {
   %arr = alloca [131072 x i64], align 8
   call i64 (ptr) @fun1(ptr %arr)
   ret void
 }
 
+; CHECK-LABEL: large_stack1
+; CHECK64: agfi  4, -1048768
+; CHECK64: lgr 0, 3
+; CHECK64: llgt  3, 1208
+; CHECK64: cg  4, 64(3)
+; CHECK64: jhe @BB7_2
+; CHECK64: %bb.1:
+; CHECK64: lg  3, 72(3)
+; CHECK64: basr  3, 3
+; CHECK64: bcr 0, 7
+; CHECK64: @BB7_2:
+; CHECK64: stmg  6, 7, 2064(4)
+; CHECK64: lgr 3, 0
+define void @large_stack1(i64 %n1, i64 %n2, i64 %n3) {
+  %arr = alloca [131072 x i64], align 8
+  call i64 (ptr, i64, i64, i64) @fun3(ptr %arr,
+            i64 %n1, i64 %n2, i64 %n3)
+  ret void
+}
+
+
+; CHECK-LABEL: large_stack2
+; CHECK64: lgr 0, 4
+; CHECK64: stg 3, 2192(4)
+; CHECK64: agfi  4, -1048768
+; CHECK64: llgt  3, 1208
+; CHECK64: cg  4, 64(3)
+; CHECK64: jhe @BB8_2
+; CHECK64: %bb.1:
+; CHECK64: lg  3, 72(3)
+; CHECK64: basr  3, 3
+; CHECK64: bcr 0, 7
+; CHECK64: @BB8_2:
+; CHECK64: lgr 3, 0
+; CHECK64: lg  3, 2192(3)
+; CHECK64: stmg  4, 11, 2048(4)
+; CHECK64: lgr 8, 4
+define void @large_stack2(i64 %n1, i64 %n2, i64 %n3) {
+  %arr0 = alloca [131072 x i64], align 8
+  %arr1 = alloca i64, i64 %n1, align 8
+  call i64 (ptr, ptr, i64, i64, i64) @fun4(ptr %arr0,
+            ptr %arr1, i64 %n1, i64 %n2, i64 %n3)
+  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(ptr %ptr)
 declare i64 @fun2(i64 %n, ptr %arr0, ptr %arr1)
+declare i64 @fun3(ptr %ptr, i64 %n1, i64 %n2, i64 %n3)
+declare i64 @fun4(ptr %ptr0, ptr %ptr1, i64 %n1, i64 %n2, i64 %n3)