Improve stack-clash implementation on x86

[serge-sans-paille]

Summary

• test both 32 and 64 bit version
• probe the tail in dynamic-alloca
• generate more concise code

[jonpa]

This LGTM, but I would prefer if someone else working on X86 also gave an approval.

[serge-sans-paille]

Rework and improve the patch:

• more documentation
• more assert
• slight update of the way static alloca are probed, leads to slightly smaller code
• completely rework the way dynamic alloca are probed, which leads to smaller and cleaner assembly. As a side effect it also probes the tail.

[serge-sans-paille]

@craig.topper up :-)

[craig.topper]

llvm/lib/Target/X86/X86FrameLowering.cpp

715

This can't be MOV64rr for 32-bit. Can we use TargetOpcode::COPY?

llvm/lib/Target/X86/X86ISelLowering.cpp

31580

use Register instead of unsigned.

31587

This shouldn't be MOV64rr if the stack pointer is 32 bits. Can we just use TargetOpcode::COPY here?

[serge-sans-paille]

Take remarks into account

[craig.topper]

Can we have tests for 32-bit mode too? Which I hope would have caught those previous issues.

llvm/lib/Target/X86/X86FrameLowering.cpp

663

Register

[serge-sans-paille]

@craig.topper test cases for 32bit target added!

[craig.topper]

LGTM

llvm/lib/Target/X86/X86FrameLowering.cpp

void X86FrameLowering::emitSPUpdate(MachineBasicBlock &MBB,
   const X86Subtarget &STI = MF.getSubtarget<X86Subtarget>();
   const X86TargetLowering &TLI = *STI.getTargetLowering();
   const bool EmitInlineStackProbe = TLI.hasInlineStackProbe(MF);
 
   // It's ok to not take into account large chunks when probing, as the
   // allocation is split in smaller chunks anyway.
   if (EmitInlineStackProbe && !InEpilogue) {
 
-    // stack probing may involve looping, and control flow generations is
-    // disallowed at this point. Rely to later processing through
-    // `inlineStackProbe`.
+    // Delegate stack probing to the `inlineStackProbe` mechanism to avoid
+    // complications.
     MachineInstr *Stub = emitStackProbeInlineStub(MF, MBB, MBBI, DL, true);
 
     // Encode the static offset as a metadata attached to the stub.
     LLVMContext &Context = MF.getFunction().getContext();
     MachineInstrBuilder(MF, Stub).addMetadata(
         MDTuple::get(Context, {ConstantAsMetadata::get(ConstantInt::get(
                                   IntegerType::get(Context, 64), Offset))}));
     return;
MachineInstr *MI = BuildMI(MBB, MBBI, DL, TII.get(Opc), StackPtr)
                          .setMIFlag(MachineInstr::FrameSetup);
   MI->getOperand(3).setIsDead(); // The EFLAGS implicit def is dead.
 }
 
 void X86FrameLowering::emitStackProbeInlineGenericLoop(
     MachineFunction &MF, MachineBasicBlock &MBB,
     MachineBasicBlock::iterator MBBI, const DebugLoc &DL,
     uint64_t Offset) const {
+  assert(Offset && "null offset");
 
   const X86Subtarget &STI = MF.getSubtarget<X86Subtarget>();
   const X86TargetLowering &TLI = *STI.getTargetLowering();
   const unsigned MovMIOpc = Is64Bit ? X86::MOV64mi32 : X86::MOV32mi;
   const uint64_t StackProbeSize = TLI.getStackProbeSize(MF);
 
   // Synthesize a loop
   NumFrameLoopProbe++;
   const BasicBlock *LLVM_BB = MBB.getBasicBlock();
 
   MachineBasicBlock *testMBB = MF.CreateMachineBasicBlock(LLVM_BB);
   MachineBasicBlock *tailMBB = MF.CreateMachineBasicBlock(LLVM_BB);
 
   MachineFunction::iterator MBBIter = ++MBB.getIterator();
   MF.insert(MBBIter, testMBB);
   MF.insert(MBBIter, tailMBB);
 
   unsigned FinalStackPtr = Uses64BitFramePtr ? X86::R11 : X86::R11D;
   BuildMI(MBB, MBBI, DL, TII.get(X86::MOV64rr), FinalStackPtr)
       .addReg(StackPtr)
       .setMIFlag(MachineInstr::FrameSetup);
 
   // save loop bound
   {
     const unsigned Opc = getSUBriOpcode(Uses64BitFramePtr, Offset);
     BuildMI(MBB, MBBI, DL, TII.get(Opc), FinalStackPtr)
void X86FrameLowering::emitStackProbeInlineGenericLoop(
   // touch the page
   addRegOffset(BuildMI(testMBB, DL, TII.get(MovMIOpc))
                    .setMIFlag(MachineInstr::FrameSetup),
                StackPtr, false, 0)
       .addImm(0)
       .setMIFlag(MachineInstr::FrameSetup);
 
   // cmp with stack pointer bound
-  BuildMI(testMBB, DL, TII.get(IsLP64 ? X86::CMP64rr : X86::CMP32rr))
+  BuildMI(testMBB, DL, TII.get(Uses64BitFramePtr ? X86::CMP64rr : X86::CMP32rr))
       .addReg(StackPtr)
       .addReg(FinalStackPtr)
       .setMIFlag(MachineInstr::FrameSetup);
 
   // jump
   BuildMI(testMBB, DL, TII.get(X86::JCC_1))
       .addMBB(testMBB)
-      .addImm(X86::COND_NE)
+      .addImm(X86::COND_L)
       .setMIFlag(MachineInstr::FrameSetup);
   testMBB->addSuccessor(testMBB);
   testMBB->addSuccessor(tailMBB);
   testMBB->addLiveIn(FinalStackPtr);
 
-  // allocate a block and touch it
-
+  // BB management
   tailMBB->splice(tailMBB->end(), &MBB, MBBI, MBB.end());
   tailMBB->transferSuccessorsAndUpdatePHIs(&MBB);
   MBB.addSuccessor(testMBB);
 
+  // handle tail
   if (Offset % StackProbeSize) {
-    const unsigned Opc = getSUBriOpcode(Uses64BitFramePtr, Offset);
-    BuildMI(*tailMBB, tailMBB->begin(), DL, TII.get(Opc), StackPtr)
-        .addReg(StackPtr)
-        .addImm(Offset % StackProbeSize)
+    BuildMI(*tailMBB, tailMBB->begin(), DL, TII.get(X86::MOV64rr), StackPtr)
+        .addReg(FinalStackPtr)
         .setMIFlag(MachineInstr::FrameSetup);
   }
 }
 
 void X86FrameLowering::emitStackProbeInlineWindowsCoreCLR64(
     MachineFunction &MF, MachineBasicBlock &MBB,
     MachineBasicBlock::iterator MBBI, const DebugLoc &DL, bool InProlog) const {
   const X86Subtarget &STI = MF.getSubtarget<X86Subtarget>();

llvm/lib/Target/X86/X86ISelLowering.cpp

X86TargetLowering::EmitLoweredSelect(MachineInstr &MI,
   createPHIsForCMOVsInSinkBB(MIItBegin, MIItEnd, ThisMBB, FalseMBB, SinkMBB);
 
   // Now remove the CMOV(s).
   ThisMBB->erase(MIItBegin, MIItEnd);
 
   return SinkMBB;
 }
 
+static unsigned getSUBriOpcode(bool IsLP64, int64_t Imm) {
+  if (IsLP64) {
+    if (isInt<8>(Imm))
+      return X86::SUB64ri8;
+    return X86::SUB64ri32;
+  } else {
+    if (isInt<8>(Imm))
+      return X86::SUB32ri8;
+    return X86::SUB32ri;
+  }
+}
+
 MachineBasicBlock *
 X86TargetLowering::EmitLoweredProbedAlloca(MachineInstr &MI,
-                                           MachineBasicBlock *BB) const {
-  MachineFunction *MF = BB->getParent();
+                                           MachineBasicBlock *MBB) const {
+  MachineFunction *MF = MBB->getParent();
   const TargetInstrInfo *TII = Subtarget.getInstrInfo();
   const X86FrameLowering &TFI = *Subtarget.getFrameLowering();
   DebugLoc DL = MI.getDebugLoc();
-  const BasicBlock *LLVM_BB = BB->getBasicBlock();
+  const BasicBlock *LLVM_BB = MBB->getBasicBlock();
 
   const unsigned ProbeSize = getStackProbeSize(*MF);
 
   MachineRegisterInfo &MRI = MF->getRegInfo();
   MachineBasicBlock *testMBB = MF->CreateMachineBasicBlock(LLVM_BB);
   MachineBasicBlock *tailMBB = MF->CreateMachineBasicBlock(LLVM_BB);
   MachineBasicBlock *blockMBB = MF->CreateMachineBasicBlock(LLVM_BB);
 
-  MachineFunction::iterator MBBIter = ++BB->getIterator();
+  MachineFunction::iterator MBBIter = ++MBB->getIterator();
   MF->insert(MBBIter, testMBB);
   MF->insert(MBBIter, blockMBB);
   MF->insert(MBBIter, tailMBB);
 
   unsigned sizeVReg = MI.getOperand(1).getReg();
 
-  const TargetRegisterClass *SizeRegClass = MRI.getRegClass(sizeVReg);
+  unsigned physSPReg = TFI.Uses64BitFramePtr ? X86::RSP : X86::ESP;
 
-  unsigned tmpSizeVReg = MRI.createVirtualRegister(SizeRegClass);
-  unsigned tmpSizeVReg2 = MRI.createVirtualRegister(SizeRegClass);
+  unsigned TmpStackPtr = MRI.createVirtualRegister(
+      TFI.Uses64BitFramePtr ? &X86::GR64RegClass : &X86::GR32RegClass);
+  unsigned FinalStackPtr = MRI.createVirtualRegister(
+      TFI.Uses64BitFramePtr ? &X86::GR64RegClass : &X86::GR32RegClass);
 
-  unsigned physSPReg = TFI.Uses64BitFramePtr ? X86::RSP : X86::ESP;
+  BuildMI(*MBB, {MI}, DL, TII->get(X86::MOV64rr), TmpStackPtr)
+      .addReg(physSPReg);
+  {
+    const unsigned Opc = TFI.Uses64BitFramePtr ? X86::SUB64rr : X86::SUB32rr;
+    BuildMI(*MBB, {MI}, DL, TII->get(Opc), FinalStackPtr)
+        .addReg(TmpStackPtr)
+        .addReg(sizeVReg);
+  }
 
   // test rsp size
-  BuildMI(testMBB, DL, TII->get(X86::PHI), tmpSizeVReg)
-      .addReg(sizeVReg)
-      .addMBB(BB)
-      .addReg(tmpSizeVReg2)
-      .addMBB(blockMBB);
 
   BuildMI(testMBB, DL,
-          TII->get(TFI.Uses64BitFramePtr ? X86::CMP64ri32 : X86::CMP32ri))
-      .addReg(tmpSizeVReg)
-      .addImm(ProbeSize);
+          TII->get(TFI.Uses64BitFramePtr ? X86::CMP64rr : X86::CMP32rr))
+      .addReg(physSPReg)
+      .addReg(FinalStackPtr);
 
   BuildMI(testMBB, DL, TII->get(X86::JCC_1))
       .addMBB(tailMBB)
       .addImm(X86::COND_L);
   testMBB->addSuccessor(blockMBB);
   testMBB->addSuccessor(tailMBB);
 
   // allocate a block and touch it
 
   BuildMI(blockMBB, DL,
-          TII->get(TFI.Uses64BitFramePtr ? X86::SUB64ri32 : X86::SUB32ri),
-          tmpSizeVReg2)
-      .addReg(tmpSizeVReg)
-      .addImm(ProbeSize);
-
-  BuildMI(blockMBB, DL,
-          TII->get(TFI.Uses64BitFramePtr ? X86::SUB64ri32 : X86::SUB32ri),
-          physSPReg)
+          TII->get(getSUBriOpcode(TFI.Uses64BitFramePtr, ProbeSize)), physSPReg)
       .addReg(physSPReg)
       .addImm(ProbeSize);
 
   const unsigned MovMIOpc =
       TFI.Uses64BitFramePtr ? X86::MOV64mi32 : X86::MOV32mi;
   addRegOffset(BuildMI(blockMBB, DL, TII->get(MovMIOpc)), physSPReg, false, 0)
       .addImm(0);
 
   BuildMI(blockMBB, DL, TII->get(X86::JMP_1)).addMBB(testMBB);
   blockMBB->addSuccessor(testMBB);
 
-  // allocate the tail and continue
-  BuildMI(tailMBB, DL,
-          TII->get(TFI.Uses64BitFramePtr ? X86::SUB64rr : X86::SUB32rr),
-          physSPReg)
-      .addReg(physSPReg)
-      .addReg(tmpSizeVReg);
+  // Replace original instruction by the expected stack ptr
   BuildMI(tailMBB, DL, TII->get(TargetOpcode::COPY), MI.getOperand(0).getReg())
-      .addReg(physSPReg);
+      .addReg(FinalStackPtr);
 
-  tailMBB->splice(tailMBB->end(), BB,
-                  std::next(MachineBasicBlock::iterator(MI)), BB->end());
-  tailMBB->transferSuccessorsAndUpdatePHIs(BB);
-  BB->addSuccessor(testMBB);
+  tailMBB->splice(tailMBB->end(), MBB,
+                  std::next(MachineBasicBlock::iterator(MI)), MBB->end());
+  tailMBB->transferSuccessorsAndUpdatePHIs(MBB);
+  MBB->addSuccessor(testMBB);
 
   // Delete the original pseudo instruction.
   MI.eraseFromParent();
 
   // And we're done.
   return tailMBB;
 }
 

llvm/test/CodeGen/X86/stack-clash-dynamic-alloca.ll

 ; RUN: llc < %s | FileCheck %s
 
 
 target datalayout = "e-m:e-i64:64-f80:128-n8:16:32:64-S128"
 target triple = "x86_64-unknown-linux-gnu"
 
 define i32 @foo(i32 %n) local_unnamed_addr #0 {
 
 ; CHECK-LABEL: foo:
 ; CHECK:       # %bb.0:
 ; CHECK-NEXT:   pushq %rbp
 ; CHECK-NEXT:   .cfi_def_cfa_offset 16
 ; CHECK-NEXT:   .cfi_offset %rbp, -16
 ; CHECK-NEXT:   movq  %rsp, %rbp
 ; CHECK-NEXT:   .cfi_def_cfa_register %rbp
 ; CHECK-NEXT:   movl  %edi, %eax
-; CHECK-NEXT:   leaq  15(,%rax,4), %rax
-; CHECK-NEXT:   andq  $-16, %rax
-; CHECK-NEXT:   cmpq  $4096, %rax # imm = 0x1000
+; CHECK-NEXT:   leaq  15(,%rax,4), %rcx
+; CHECK-NEXT:   andq  $-16, %rcx
+; CHECK-NEXT:   movq  %rsp, %rax
+; CHECK-NEXT:   subq  %rcx, %rax
+; CHECK-NEXT:   cmpq  %rax, %rsp
 ; CHECK-NEXT:   jl  .LBB0_3
 ; CHECK-NEXT:  .LBB0_2: # =>This Inner Loop Header: Depth=1
-; CHECK-NEXT:   subq  $4096, %rax # imm = 0x1000
 ; CHECK-NEXT:   subq  $4096, %rsp # imm = 0x1000
 ; CHECK-NEXT:   movq  $0, (%rsp)
-; CHECK-NEXT:   cmpq  $4096, %rax # imm = 0x1000
+; CHECK-NEXT:   cmpq  %rax, %rsp
 ; CHECK-NEXT:   jge .LBB0_2
 ; CHECK-NEXT:  .LBB0_3:
-; CHECK-NEXT:   subq  %rax, %rsp
-; CHECK-NEXT:   movq  %rsp, %rax
+; CHECK-NEXT:   movq  %rax, %rsp
 ; CHECK-NEXT:   movl  $1, 4792(%rax)
 ; CHECK-NEXT:   movl  (%rax), %eax
 ; CHECK-NEXT:   movq  %rbp, %rsp
 ; CHECK-NEXT:   popq  %rbp
 ; CHECK-NEXT:  .cfi_def_cfa %rsp, 8
 ; CHECK-NEXT:   retq
 
   %a = alloca i32, i32 %n, align 16
   %b = getelementptr inbounds i32, i32* %a, i64 1198
   store volatile i32 1, i32* %b
   %c = load volatile i32, i32* %a
   ret i32 %c
 }
 
 attributes #0 =  {"probe-stack"="inline-asm"}

llvm/test/CodeGen/X86/stack-clash-large.ll

 ; RUN: llc < %s | FileCheck %s
 
 
 target datalayout = "e-m:e-i64:64-f80:128-n8:16:32:64-S128"
 target triple = "x86_64-unknown-linux-gnu"
 
 define i32 @foo() local_unnamed_addr #0 {
 
 ; CHECK-LABEL: foo:
 ; CHECK:        # %bb.0:
 ; CHECK-NEXT: movq  %rsp, %r11
 ; CHECK-NEXT: subq  $69632, %r11 # imm = 0x11000
 ; CHECK-NEXT:   .LBB0_1:
 ; CHECK-NEXT: subq  $4096, %rsp # imm = 0x1000
 ; CHECK-NEXT: movq  $0, (%rsp)
 ; CHECK-NEXT: cmpq  %r11, %rsp
-; CHECK-NEXT: jne .LBB0_1
+; CHECK-NEXT: jl  .LBB0_1
 ; CHECK-NEXT:# %bb.2:
-; CHECK-NEXT: subq  $2248, %rsp # imm = 0x8C8
+; CHECK-NEXT: movq    %r11, %rsp
 ; CHECK-NEXT: .cfi_def_cfa_offset 71888
 ; CHECK-NEXT: movl  $1, 264(%rsp)
 ; CHECK-NEXT: movl  $1, 28664(%rsp)
 ; CHECK-NEXT: movl  -128(%rsp), %eax
 ; CHECK-NEXT: addq  $71880, %rsp # imm = 0x118C8
 ; CHECK-NEXT: .cfi_def_cfa_offset 8
 ; CHECK-NEXT: retq