[X86] Fix stack probing on x32 (PR41477)

[nikic]

Summary

Fix for https://bugs.llvm.org/show_bug.cgi?id=41477. On the x32 ABI with stack probing a dynamic alloca will result in a WIN_ALLOCA_32 with a 32-bit size. The current implementation tries to copy it into RAX, resulting in a physreg copy error. Fix this by copying to EAX instead.

[efriedma]

llvm/lib/Target/X86/X86WinAllocaExpander.cpp

242–243

Are you sure this subtraction is using the right opcode?

260–261

Also the subtraction here.

llvm/test/CodeGen/X86/pr41477.ll

18

subq is the wrong opcode; should be using subl.

[nikic]

Fix more registers/opcodes.

[nikic]

Thanks @efriedma. I've added test cases that exercise all three codepaths now and fixed the subtract regs/ops.

[efriedma]

LGTM

[nikic]

Committed in rL358807, forgot the Differential Revision tag...

llvm/lib/Target/X86/X86FrameLowering.cpp

if (Is64Bit && MF.getTarget().getCodeModel() == CodeModel::Large) {
     BuildMI(MBB, MBBI, DL, TII.get(X86::MOV64ri), X86::R11)
         .addExternalSymbol(MF.createExternalSymbolName(Symbol));
     CI = BuildMI(MBB, MBBI, DL, TII.get(CallOp)).addReg(X86::R11);
   } else {
     CI = BuildMI(MBB, MBBI, DL, TII.get(CallOp))
         .addExternalSymbol(MF.createExternalSymbolName(Symbol));
   }
 
-  unsigned AX = Is64Bit ? X86::RAX : X86::EAX;
-  unsigned SP = Is64Bit ? X86::RSP : X86::ESP;
+  unsigned AX = Uses64BitFramePtr ? X86::RAX : X86::EAX;
+  unsigned SP = Uses64BitFramePtr ? X86::RSP : X86::ESP;
   CI.addReg(AX, RegState::Implicit)
       .addReg(SP, RegState::Implicit)
       .addReg(AX, RegState::Define | RegState::Implicit)
       .addReg(SP, RegState::Define | RegState::Implicit)
       .addReg(X86::EFLAGS, RegState::Define | RegState::Implicit);
 
   if (STI.isTargetWin64() || !STI.isOSWindows()) {
     // MSVC x32's _chkstk and cygwin/mingw's _alloca adjust %esp themselves.
     // MSVC x64's __chkstk and cygwin/mingw's ___chkstk_ms do not adjust %rsp
     // themselves. They also does not clobber %rax so we can reuse it when
     // adjusting %rsp.
     // All other platforms do not specify a particular ABI for the stack probe
     // function, so we arbitrarily define it to not adjust %esp/%rsp itself.
-    BuildMI(MBB, MBBI, DL, TII.get(getSUBrrOpcode(Is64Bit)), SP)
+    BuildMI(MBB, MBBI, DL, TII.get(getSUBrrOpcode(Uses64BitFramePtr)), SP)
         .addReg(SP)
         .addReg(AX);
   }
 
   if (InProlog) {
     // Apply the frame setup flag to all inserted instrs.
     for (++ExpansionMBBI; ExpansionMBBI != MBBI; ++ExpansionMBBI)
       ExpansionMBBI->setFlag(MachineInstr::FrameSetup);

llvm/lib/Target/X86/X86WinAllocaExpander.cpp

void X86WinAllocaExpander::lower(MachineInstr* MI, Lowering L) {
   MachineBasicBlock::iterator I = *MI;
 
   int64_t Amount = getWinAllocaAmount(MI, MRI);
   if (Amount == 0) {
     MI->eraseFromParent();
     return;
   }
 
+  // These two variables differ on x32, which is a 64-bit target with a
+  // 32-bit alloca.
   bool Is64Bit = STI->is64Bit();
+  bool Is64BitAlloca = MI->getOpcode() == X86::WIN_ALLOCA_64;
   assert(SlotSize == 4 || SlotSize == 8);
-  unsigned RegA = (SlotSize == 8) ? X86::RAX : X86::EAX;
 
   switch (L) {
-  case TouchAndSub:
+  case TouchAndSub: {
     assert(Amount >= SlotSize);
 
     // Use a push to touch the top of the stack.
+    unsigned RegA = Is64Bit ? X86::RAX : X86::EAX;
     BuildMI(*MBB, I, DL, TII->get(Is64Bit ? X86::PUSH64r : X86::PUSH32r))
         .addReg(RegA, RegState::Undef);
     Amount -= SlotSize;
     if (!Amount)
       break;
 
     // Fall through to make any remaining adjustment.
     LLVM_FALLTHROUGH;
+  }
   case Sub:
     assert(Amount > 0);
     if (Amount == SlotSize) {
       // Use push to save size.
+      unsigned RegA = Is64Bit ? X86::RAX : X86::EAX;
       BuildMI(*MBB, I, DL, TII->get(Is64Bit ? X86::PUSH64r : X86::PUSH32r))
           .addReg(RegA, RegState::Undef);
     } else {
       // Sub.
-      BuildMI(*MBB, I, DL, TII->get(getSubOpcode(Is64Bit, Amount)), StackPtr)
+      BuildMI(*MBB, I, DL,
+              TII->get(getSubOpcode(Is64BitAlloca, Amount)), StackPtr)
           .addReg(StackPtr)
           .addImm(Amount);
     }
     break;
   case Probe:
     if (!NoStackArgProbe) {
       // The probe lowering expects the amount in RAX/EAX.
+      unsigned RegA = Is64BitAlloca ? X86::RAX : X86::EAX;
       BuildMI(*MBB, MI, DL, TII->get(TargetOpcode::COPY), RegA)
           .addReg(MI->getOperand(0).getReg());
 
       // Do the probe.
       STI->getFrameLowering()->emitStackProbe(*MBB->getParent(), *MBB, MI, DL,
                                               /*InPrologue=*/false);
     } else {
       // Sub
-      BuildMI(*MBB, I, DL, TII->get(Is64Bit ? X86::SUB64rr : X86::SUB32rr),
-              StackPtr)
+      BuildMI(*MBB, I, DL,
+              TII->get(Is64BitAlloca ? X86::SUB64rr : X86::SUB32rr), StackPtr)
           .addReg(StackPtr)
           .addReg(MI->getOperand(0).getReg());
     }
     break;
   }
 
   unsigned AmountReg = MI->getOperand(0).getReg();
   MI->eraseFromParent();

llvm/test/CodeGen/X86/pr41477.ll

+; NOTE: Assertions have been autogenerated by utils/update_llc_test_checks.py
+; RUN: llc < %s -mtriple=x86_64-unknown-linux-gnux32 -verify-machineinstrs | FileCheck %s
+
+target datalayout = "e-m:e-p:32:32-i64:64-f80:128-n8:16:32:64-S128"
+target triple = "x86_64-unknown-linux-gnux32"
+
+; probe-stack + dynamic size alloca
+define void @test1(i32 %size) nounwind #0 {
+; CHECK-LABEL: test1:
+; CHECK:       # %bb.0: # %start
+; CHECK-NEXT:    pushq %rbp
+; CHECK-NEXT:    movl %esp, %ebp
+; CHECK-NEXT:    # kill: def $edi killed $edi def $rdi
+; CHECK-NEXT:    leal 15(%rdi), %eax
+; CHECK-NEXT:    andl $-16, %eax
+; CHECK-NEXT:    callq __rust_probestack
+; CHECK-NEXT:    subl %eax, %esp
+start:
+  %alloca = alloca i8, i32 %size
+  unreachable
+}
+
+; probe-stack + no-stack-arg-probe + dynamic size alloca
+define void @test2(i32 %size) nounwind #1 {
+; CHECK-LABEL: test2:
+; CHECK:       # %bb.0: # %start
+; CHECK-NEXT:    pushq %rbp
+; CHECK-NEXT:    movl %esp, %ebp
+; CHECK-NEXT:    addl $15, %edi
+; CHECK-NEXT:    andl $-16, %edi
+; CHECK-NEXT:    subl %edi, %esp
+start:
+  %alloca = alloca i8, i32 %size
+  unreachable
+}
+
+; probe-stack + fixed size alloca not in entry block
+define void @test3() nounwind #0 {
+; CHECK-LABEL: test3:
+; CHECK:       # %bb.0: # %start
+; CHECK-NEXT:    pushq %rbp
+; CHECK-NEXT:    movl %esp, %ebp
+; CHECK-NEXT:    pushq %rax
+; CHECK-NEXT:    subl $1992, %esp # imm = 0x7C8
+start:
+  br label %block
+
+block:
+  %alloca = alloca i8, i32 2000
+  unreachable
+}
+
+attributes #0 = { "probe-stack"="__rust_probestack" }
+attributes #1 = { "probe-stack"="__rust_probestack" "no-stack-arg-probe" }