[stack-clash] Fix probing of dynamic alloca

[serge-sans-paille]

Summary

Perform the probing in the *cough* correct direction.

Related to https://github.com/rust-lang/rust/pull/77885#issuecomment-711062924

[cuviper]

I tried the C reproducer from that Rust bug, and it turns out that alloca(0) was already fixed by D88548. But if I change that to a larger size, then it was skipping right over the probe loop. I confirmed that your change here does make it go through the loop.

It would be nice to have tests that actually execute this stuff, rather than just relying on manual reviews of the expected assembly, but I don't know if that's feasible in the current infrastructure.

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

29–30

Shouldn't this sub before mov 0? I think right now, the first iteration is going to clobber the most recent thing on the stack, in this case the saved value from pushq %rbp.

[serge-sans-paille]

Fix issue with clobbered stack memory

[cuviper]

Works for me!

llvm/lib/Target/X86/X86ISelLowering.cpp

X86TargetLowering::EmitLoweredProbedAlloca(MachineInstr &MI,
 
   BuildMI(testMBB, DL,
           TII->get(TFI.Uses64BitFramePtr ? X86::CMP64rr : X86::CMP32rr))
       .addReg(FinalStackPtr)
       .addReg(physSPReg);
 
   BuildMI(testMBB, DL, TII->get(X86::JCC_1))
       .addMBB(tailMBB)
-      .addImm(X86::COND_LE);
+      .addImm(X86::COND_GE);
   testMBB->addSuccessor(blockMBB);
   testMBB->addSuccessor(tailMBB);
 
   // Touch the block then extend it. This is done on the opposite side of
   // static probe where we allocate then touch, to avoid the need of probing the
   // tail of the static alloca. Possible scenarios are:
   //
   //       + ---- <- ------------ <- ------------- <- ------------ +
   //       |                                                       |
   // [free probe] -> [page alloc] -> [alloc probe] -> [tail alloc] + -> [dyn probe] -> [page alloc] -> [dyn probe] -> [tail alloc] +
   //                                                               |                                                               |
   //                                                               + <- ----------- <- ------------ <- ----------- <- ------------ +
   //
   // The property we want to enforce is to never have more than [page alloc] between two probes.
 
-  const unsigned MovMIOpc =
-      TFI.Uses64BitFramePtr ? X86::MOV64mi32 : X86::MOV32mi;
-  addRegOffset(BuildMI(blockMBB, DL, TII->get(MovMIOpc)), physSPReg, false, 0)
+  const unsigned XORMIOpc =
+      TFI.Uses64BitFramePtr ? X86::XOR64mi8 : X86::XOR32mi8;
+  addRegOffset(BuildMI(blockMBB, DL, TII->get(XORMIOpc)), physSPReg, false, 0)
       .addImm(0);
 
   BuildMI(blockMBB, DL,
           TII->get(getSUBriOpcode(TFI.Uses64BitFramePtr, ProbeSize)), physSPReg)
       .addReg(physSPReg)
       .addImm(ProbeSize);
 
 

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

 ; CHECK-X86-64-NEXT:    movq  %rsp, %rbp
 ; CHECK-X86-64-NEXT:    .cfi_def_cfa_register %rbp
 ; CHECK-X86-64-NEXT:    movq    %rsp, %rax
 ; CHECK-X86-64-NEXT:    movl    %edi, %ecx
 ; CHECK-X86-64-NEXT:    leaq 15(,%rcx,4), %rcx
 ; CHECK-X86-64-NEXT:    andq  $-16, %rcx
 ; CHECK-X86-64-NEXT:    subq  %rcx, %rax
 ; CHECK-X86-64-NEXT:    cmpq  %rsp, %rax
-; CHECK-X86-64-NEXT:    jle .LBB0_3
+; CHECK-X86-64-NEXT:    jge .LBB0_3
 ; CHECK-X86-64-NEXT:  .LBB0_2: # =>This Inner Loop Header: Depth=1
-; CHECK-X86-64-NEXT:    movq  $0, (%rsp)
+; CHECK-X86-64-NEXT:    xorq  $0, (%rsp)
 ; CHECK-X86-64-NEXT:    subq  $4096, %rsp # imm = 0x1000
 ; CHECK-X86-64-NEXT:    cmpq  %rsp, %rax
-; CHECK-X86-64-NEXT:    jg  .LBB0_2
+; CHECK-X86-64-NEXT:    jl  .LBB0_2
 ; CHECK-X86-64-NEXT:  .LBB0_3:
 ; CHECK-X86-64-NEXT:    movq  %rax, %rsp
 ; CHECK-X86-64-NEXT:    movl  $1, 4792(%rax)
 ; CHECK-X86-64-NEXT:    movl  (%rax), %eax
 ; CHECK-X86-64-NEXT:    movq  %rbp, %rsp
 ; CHECK-X86-64-NEXT:    popq  %rbp
 ; CHECK-X86-64-NEXT:  .cfi_def_cfa %rsp, 8
 ; CHECK-X86-64-NEXT:   retq
 
 
 ; CHECK-X86-32-LABEL: foo:
 ; CHECK-X86-32:       # %bb.0:
 ; CHECK-X86-32-NEXT:    pushl   %ebp
 ; CHECK-X86-32-NEXT:    .cfi_def_cfa_offset 8
 ; CHECK-X86-32-NEXT:    .cfi_offset %ebp, -8
 ; CHECK-X86-32-NEXT:    movl    %esp, %ebp
 ; CHECK-X86-32-NEXT:    .cfi_def_cfa_register %ebp
 ; CHECK-X86-32-NEXT:    subl    $8, %esp
 ; CHECK-X86-32-NEXT:    movl    8(%ebp), %ecx
 ; CHECK-X86-32-NEXT:    movl    %esp, %eax
 ; CHECK-X86-32-NEXT:    leal    15(,%ecx,4), %ecx
 ; CHECK-X86-32-NEXT:    andl    $-16, %ecx
 ; CHECK-X86-32-NEXT:    subl    %ecx, %eax
 ; CHECK-X86-32-NEXT:    cmpl    %esp, %eax
-; CHECK-X86-32-NEXT:    jle  .LBB0_3
+; CHECK-X86-32-NEXT:    jge  .LBB0_3
 ; CHECK-X86-32-NEXT:  .LBB0_2: # =>This Inner Loop Header: Depth=1
-; CHECK-X86-32-NEXT:    movl    $0, (%esp)
+; CHECK-X86-32-NEXT:    xorl    $0, (%esp)
 ; CHECK-X86-32-NEXT:    subl    $4096, %esp # imm = 0x1000
 ; CHECK-X86-32-NEXT:    cmpl    %esp, %eax
-; CHECK-X86-32-NEXT:    jg .LBB0_2
+; CHECK-X86-32-NEXT:    jl .LBB0_2
 ; CHECK-X86-32-NEXT:  .LBB0_3:
 ; CHECK-X86-32-NEXT:    movl    %eax, %esp
 ; CHECK-X86-32-NEXT:    movl    $1, 4792(%eax)
 ; CHECK-X86-32-NEXT:    movl    (%eax), %eax
 ; CHECK-X86-32-NEXT:    movl    %ebp, %esp
 ; CHECK-X86-32-NEXT:    popl    %ebp
 ; CHECK-X86-32-NEXT:    .cfi_def_cfa %esp, 4
 ; CHECK-X86-32-NEXT:    retl
 

llvm/test/CodeGen/X86/stack-clash-small-alloc-medium-align.ll

 ; CHECK-NEXT: .cfi_offset %rbx, -24
 ; CHECK-NEXT: movl  $1, (%rbx,%rdi,4)
 ; CHECK-NEXT: movl  (%rbx), %ecx
 ; CHECK-NEXT: movq  %rsp, %rax
 ; CHECK-NEXT: leaq  15(,%rcx,4), %rcx
 ; CHECK-NEXT: andq  $-16, %rcx
 ; CHECK-NEXT: subq  %rcx, %rax
 ; CHECK-NEXT: cmpq  %rsp, %rax
-; CHECK-NEXT: jle .LBB3_3
+; CHECK-NEXT: jge .LBB3_3
 ; CHECK-NEXT:.LBB3_2:                                # =>This Inner Loop Header: Depth=1
-; CHECK-NEXT: movq  $0, (%rsp)
+; CHECK-NEXT: xorq  $0, (%rsp)
 ; CHECK-NEXT: subq  $4096, %rsp                     # imm = 0x1000
 ; CHECK-NEXT: cmpq  %rsp, %rax
-; CHECK-NEXT: jg  .LBB3_2
+; CHECK-NEXT: jl  .LBB3_2
 ; CHECK-NEXT:.LBB3_3:
 ; CHECK-NEXT: andq  $-64, %rax
 ; CHECK-NEXT: movq  %rax, %rsp
 ; CHECK-NEXT: movl  (%rax), %eax
 ; CHECK-NEXT: leaq  -8(%rbp), %rsp
 ; CHECK-NEXT: popq  %rbx
 ; CHECK-NEXT: popq  %rbp
 ; CHECK-NEXT: .cfi_def_cfa %rsp, 8