[X86] Preserve more liveness information in emitStackProbeInline

[thegameg]

Summary

This commit fixes two issues with the liveness information after the call:

1. The code always spills RCX and RDX if InProlog == true, which results in a use of undefined phys reg.
2. FinalReg, JoinReg, RoundedReg, SizeReg are not added as live-ins to the basic blocks that use them, therefore they are seen undefined.

PR38376

[MatzeB]

Thanks for fixing machine verifier errors, LGTM.

PR38376

You can put https://llvm.org/PR38376 into the commit message so it becomes a clickable link in some terminal emulators or e-mail clients.

lib/Target/X86/X86FrameLowering.cpp

619–621

Can you explicitely state the assumption that no earlier prolog instruction changes RCX/RDX so that checking the block live-ins is enough?

[rnk]

Thank you!

lib/Target/X86/X86FrameLowering.cpp

const unsigned SizeReg = InProlog ? (unsigned)X86::RAX
                                     : MRI.createVirtualRegister(RegClass),
                  ProbeReg = InProlog ? (unsigned)X86::RCX
                                      : MRI.createVirtualRegister(RegClass);
 
   // SP-relative offsets where we can save RCX and RDX.
   int64_t RCXShadowSlot = 0;
   int64_t RDXShadowSlot = 0;
 
-  // If inlining in the prolog, save RCX and RDX.     
-  // Future optimization: don't save or restore if not live in.
+  // If inlining in the prolog, save RCX and RDX.
   if (InProlog) {
     // Compute the offsets. We need to account for things already
     // pushed onto the stack at this point: return address, frame
     // pointer (if used), and callee saves.
     X86MachineFunctionInfo *X86FI = MF.getInfo<X86MachineFunctionInfo>();
     const int64_t CalleeSaveSize = X86FI->getCalleeSavedFrameSize();
     const bool HasFP = hasFP(MF);
-    RCXShadowSlot = 8 + CalleeSaveSize + (HasFP ? 8 : 0);
-    RDXShadowSlot = RCXShadowSlot + 8;
-    // Emit the saves.
+
+    // Check if we need to spill RCX and/or RDX.
+    const bool IsRCXLiveIn = MBB.isLiveIn(X86::RCX);
+    const bool IsRDXLiveIn = MBB.isLiveIn(X86::RDX);
+    int64_t InitSlot = 8 + CalleeSaveSize + (HasFP ? 8 : 0);
+    // Assign the initial slot to both registers, then change RDX's slot if both
+    // need to be spilled.
+    if (IsRCXLiveIn)
+      RCXShadowSlot = InitSlot;
+    if (IsRDXLiveIn)
+      RDXShadowSlot = InitSlot;
+    if (IsRDXLiveIn && IsRCXLiveIn)
+      RDXShadowSlot += 8;
+    // Emit the saves if needed.
+    if (IsRCXLiveIn)
       addRegOffset(BuildMI(&MBB, DL, TII.get(X86::MOV64mr)), X86::RSP, false,
                    RCXShadowSlot)
           .addReg(X86::RCX);
+    if (IsRDXLiveIn)
       addRegOffset(BuildMI(&MBB, DL, TII.get(X86::MOV64mr)), X86::RSP, false,
                    RDXShadowSlot)
           .addReg(X86::RDX);
   } else {
     // Not in the prolog. Copy RAX to a virtual reg.
     BuildMI(&MBB, DL, TII.get(X86::MOV64rr), SizeReg).addReg(X86::RAX);
   }
 
   // Add code to MBB to check for overflow and set the new target stack pointer
   // to zero if so.
   BuildMI(&MBB, DL, TII.get(X86::XOR64rr), ZeroReg)
BuildMI(&MBB, DL, TII.get(X86::MOV64rm), LimitReg)
       .addReg(0)
       .addImm(ThreadEnvironmentStackLimit)
       .addReg(X86::GS);
   BuildMI(&MBB, DL, TII.get(X86::CMP64rr)).addReg(FinalReg).addReg(LimitReg);
   // Jump if the desired stack pointer is at or above the stack limit.
   BuildMI(&MBB, DL, TII.get(X86::JAE_1)).addMBB(ContinueMBB);
 
   // Add code to roundMBB to round the final stack pointer to a page boundary.
+  RoundMBB->addLiveIn(FinalReg);
   BuildMI(RoundMBB, DL, TII.get(X86::AND64ri32), RoundedReg)
       .addReg(FinalReg)
       .addImm(PageMask);
   BuildMI(RoundMBB, DL, TII.get(X86::JMP_1)).addMBB(LoopMBB);
 
   // LimitReg now holds the current stack limit, RoundedReg page-rounded
   // final RSP value. Add code to loopMBB to decrement LimitReg page-by-page
   // and probe until we reach RoundedReg.
   if (!InProlog) {
     BuildMI(LoopMBB, DL, TII.get(X86::PHI), JoinReg)
         .addReg(LimitReg)
         .addMBB(RoundMBB)
         .addReg(ProbeReg)
         .addMBB(LoopMBB);
   }
 
+  LoopMBB->addLiveIn(JoinReg);
   addRegOffset(BuildMI(LoopMBB, DL, TII.get(X86::LEA64r), ProbeReg), JoinReg,
                false, -PageSize);
 
   // Probe by storing a byte onto the stack.
   BuildMI(LoopMBB, DL, TII.get(X86::MOV8mi))
       .addReg(ProbeReg)
       .addImm(1)
       .addReg(0)
       .addImm(0)
       .addReg(0)
       .addImm(0);
+
+  LoopMBB->addLiveIn(RoundedReg);
   BuildMI(LoopMBB, DL, TII.get(X86::CMP64rr))
       .addReg(RoundedReg)
       .addReg(ProbeReg);
   BuildMI(LoopMBB, DL, TII.get(X86::JNE_1)).addMBB(LoopMBB);
 
   MachineBasicBlock::iterator ContinueMBBI = ContinueMBB->getFirstNonPHI();
 
   // If in prolog, restore RDX and RCX.
   if (InProlog) {
-    addRegOffset(BuildMI(*ContinueMBB, ContinueMBBI, DL, TII.get(X86::MOV64rm),
-                         X86::RCX),
+    if (RCXShadowSlot) // It means we spilled RCX in the prologue.
+      addRegOffset(BuildMI(*ContinueMBB, ContinueMBBI, DL,
+                           TII.get(X86::MOV64rm), X86::RCX),
                    X86::RSP, false, RCXShadowSlot);
-    addRegOffset(BuildMI(*ContinueMBB, ContinueMBBI, DL, TII.get(X86::MOV64rm),
-                         X86::RDX),
+    if (RDXShadowSlot) // It means we spilled RDX in the prologue.
+      addRegOffset(BuildMI(*ContinueMBB, ContinueMBBI, DL,
+                           TII.get(X86::MOV64rm), X86::RDX),
                    X86::RSP, false, RDXShadowSlot);
   }
 
   // Now that the probing is done, add code to continueMBB to update
   // the stack pointer for real.
+  ContinueMBB->addLiveIn(SizeReg);
   BuildMI(*ContinueMBB, ContinueMBBI, DL, TII.get(X86::SUB64rr), X86::RSP)
       .addReg(X86::RSP)
       .addReg(SizeReg);
 
   // Add the control flow edges we need.
   MBB.addSuccessor(ContinueMBB);
   MBB.addSuccessor(RoundMBB);
   RoundMBB->addSuccessor(LoopMBB);
if (InProlog) {
     for (MachineInstr &MI : *LoopMBB) {
       MI.setFlag(MachineInstr::FrameSetup);
     }
     for (MachineBasicBlock::iterator CMBBI = ContinueMBB->begin();
          CMBBI != ContinueMBBI; ++CMBBI) {
       CMBBI->setFlag(MachineInstr::FrameSetup);
     }
   }
-
-  // Possible TODO: physreg liveness for InProlog case.
 }
 
 void X86FrameLowering::emitStackProbeCall(MachineFunction &MF,
                                           MachineBasicBlock &MBB,
                                           MachineBasicBlock::iterator MBBI,
                                           const DebugLoc &DL,
                                           bool InProlog) const {
   bool IsLargeCodeModel = MF.getTarget().getCodeModel() == CodeModel::Large;

test/CodeGen/X86/win_coreclr_chkstk.ll

 ; RUN: llc < %s -mtriple=x86_64-pc-win32-coreclr | FileCheck %s -check-prefix=WIN_X64
 ; RUN: llc < %s -mtriple=x86_64-pc-linux         | FileCheck %s -check-prefix=LINUX
 
 ; By default, windows CoreCLR requires an inline prologue stack expansion check
 ; if more than 4096 bytes are allocated on the stack.
 
 ; Prolog stack allocation >= 4096 bytes will require the probe sequence
 define i32 @main4k() nounwind {
 entry:
 ; WIN_X64-LABEL:main4k:
 ; WIN_X64: # %bb.0:
 ; WIN_X64:      movl    $4096, %eax
-; WIN_X64:      movq    %rcx, 8(%rsp)
-; WIN_X64:  movq  %rdx, 16(%rsp)
 ; WIN_X64:  xorq  %rcx, %rcx
 ; WIN_X64:  movq  %rsp, %rdx
 ; WIN_X64:  subq  %rax, %rdx
 ; WIN_X64:  cmovbq  %rcx, %rdx
 ; WIN_X64:  movq  %gs:16, %rcx
 ; WIN_X64:  cmpq  %rcx, %rdx
 ; WIN_X64:  jae .LBB0_3
 ; WIN_X64:# %bb.1:
 ; WIN_X64:  andq  $-4096, %rdx
 ; WIN_X64:.LBB0_2:
 ; WIN_X64:  leaq  -4096(%rcx), %rcx
 ; WIN_X64:  movb  $0, (%rcx)
 ; WIN_X64:  cmpq  %rcx, %rdx
 ; WIN_X64:  jne .LBB0_2
 ; WIN_X64:.LBB0_3:
-; WIN_X64:  movq  8(%rsp), %rcx
-; WIN_X64:  movq  16(%rsp), %rdx
 ; WIN_X64:  subq  %rax, %rsp
 ; WIN_X64:  xorl  %eax, %eax
 ; WIN_X64:  addq  $4096, %rsp
 ; WIN_X64:  retq
 ; LINUX-LABEL:main4k:
 ; LINUX-NOT:    movq    %gs:16, %rcx
 ; LINUX:  retq
   %a = alloca [4096 x i8]
   ret i32 0
 }
 
 ; Prolog stack allocation >= 4096 bytes will require the probe sequence
 ; Case with frame pointer
 define i32 @main4k_frame() nounwind "no-frame-pointer-elim"="true" {
 entry:
 ; WIN_X64-LABEL:main4k_frame:
-; WIN_X64:      movq    %rcx,   16(%rsp)
 ; WIN_X64:      movq    %gs:16, %rcx
 ; LINUX-LABEL:main4k_frame:
 ; LINUX-NOT:    movq    %gs:16, %rcx
 ; LINUX:  retq
   %a = alloca [4096 x i8]
   ret i32 0
 }
 
 ; Prolog stack allocation >= 4096 bytes will require the probe sequence
 ; Case with INT args
 define i32 @main4k_intargs(i32 %x, i32 %y) nounwind {
 entry:
-; WIN_X64:      movq    %rcx,   8(%rsp)
 ; WIN_X64:      movq    %gs:16, %rcx
 ; LINUX-NOT:    movq    %gs:16, %rcx
 ; LINUX:  retq
   %a = alloca [4096 x i8]
   %t = add i32 %x, %y
   ret i32 %t
 }
 
 ; Prolog stack allocation >= 4096 bytes will require the probe sequence
 ; Case with FP regs
 define i32 @main4k_fpargs(double %x, double %y) nounwind {
 entry:
-; WIN_X64:      movq    %rcx,   8(%rsp)
 ; WIN_X64:      movq    %gs:16, %rcx
 ; LINUX-NOT:    movq    %gs:16, %rcx
 ; LINUX:  retq
   %a = alloca [4096 x i8]
   ret i32 0
 }
 
 ; Prolog stack allocation >= 4096 bytes will require the probe sequence

test/CodeGen/X86/win_coreclr_chkstk_liveins.mir

+# RUN: llc -verify-machineinstrs %s -run-pass prologepilog -mtriple=x86_64-pc-win32-coreclr -o - | FileCheck %s
+...
+---
+name:            main4k
+# CHECK-LABEL: name: main4k
+
+alignment:       4
+tracksRegLiveness: true
+frameInfo:
+  maxAlignment:    8
+stack:
+  - { id: 0, size: 4096, alignment: 1, stack-id: 0 }
+body:             |
+  bb.0.entry:
+    $eax = IMPLICIT_DEF
+    RET 0, killed $eax
+
+  ; CHECK: bb.1.entry:
+  ; CHECK:   liveins: $rdx
+  ; CHECK: bb.2.entry:
+  ; CHECK:   liveins: $rcx, $rdx
+  ; CHECK: bb.3.entry:
+  ; CHECK:   liveins: $rax
+...