fix stack probe lowering for x86_intrcc

The x86_intrcc calling convention will build two STACKALLOC_W_PROBING machine instructions if the function takes an error code. This is caused by an additional call to emitSPUpdate in llvm/lib/Target/X86/X86FrameLowering.cpp:1650. Previously only the first STACKALLOC_W_PROBING machine instruction was properly handled, the second one was simply ignored. This lead to miscompilations where the stack pointer wasn't properly updated (see https://github.com/rust-lang/rust/issues/109918). This patch fixes this by handling all STACKALLOC_W_PROBING machine instructions.

To be honest I don't quite understand why this didn't lead to more noticeable miscompilations previously.

This is my first time contributing to LLVM.

Reviewed By: pengfei

Differential Revision: https://reviews.llvm.org/D150033

(cherry picked from commit f6154364f65709df234f07ad1fe8799e68d84134)

2 files changed, 30 insertions, 1 deletions

diff --git a/llvm/lib/Target/X86/X86FrameLowering.cpp b/llvm/lib/Target/X86/X86FrameLowering.cpp
index 1606413c382b..df9aaddd5729 100644
--- a/llvm/lib/Target/X86/X86FrameLowering.cpp
+++ b/llvm/lib/Target/X86/X86FrameLowering.cpp
@@ -1557,7 +1557,19 @@ void X86FrameLowering::emitPrologue(MachineFunction &MF,
       Fn.arg_size() == 2) {
     StackSize += 8;
     MFI.setStackSize(StackSize);
-    emitSPUpdate(MBB, MBBI, DL, -8, /*InEpilogue=*/false);
+
+    // Update the stack pointer by pushing a register. This is the instruction
+    // emitted that would be end up being emitted by a call to `emitSPUpdate`.
+    // Hard-coding the update to a push avoids emitting a second
+    // `STACKALLOC_W_PROBING` instruction in the save block: We know that stack
+    // probing isn't needed anyways for an 8-byte update.
+    // Pushing a register leaves us in a similar situation to a regular
+    // function call where we know that the address at (rsp-8) is writeable.
+    // That way we avoid any off-by-ones with stack probing for additional
+    // stack pointer updates later on.
+    BuildMI(MBB, MBBI, DL, TII.get(X86::PUSH64r))
+        .addReg(X86::RAX, RegState::Undef)
+        .setMIFlag(MachineInstr::FrameSetup);
   }
 
   // If this is x86-64 and the Red Zone is not disabled, if we are a leaf
diff --git a/llvm/test/CodeGen/X86/x86-64-intrcc.ll b/llvm/test/CodeGen/X86/x86-64-intrcc.ll
index 1e1e01214c90..443d4c2fa464 100644
--- a/llvm/test/CodeGen/X86/x86-64-intrcc.ll
+++ b/llvm/test/CodeGen/X86/x86-64-intrcc.ll
@@ -174,5 +174,22 @@ entry:
   ret void
 }
 
+define x86_intrcc void @test_stack_allocation(ptr byval(%struct.interrupt_frame) %frame, i64 %err) #1 {
+  ; CHECK-LABEL: test_stack_allocation:
+  ; CHECK: # %bb.0: # %entry
+
+  ;; Ensure that STACKALLOC_W_PROBING isn't emitted.
+  ; CHECK-NOT: # fixed size alloca with probing
+  ;; Ensure that stack space is allocated.
+  ; CHECK: subq    $280, %rsp
+entry:
+  %some_allocation = alloca i64
+  ;; Call a un-inlineable function to ensure the allocation isn't put in the red zone.
+  call void @external_function(ptr %some_allocation)
+  ret void
+}
+
+declare void @external_function(ptr)
 
 attributes #0 = { nounwind "frame-pointer"="all" }
+attributes #1 = { nounwind "probe-stack"="inline-asm" }