fix stack probe lowering for x86_intrcc

[Freax13]

Summary

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.

[pengfei]

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

I don't know much about STACKALLOC_W_PROBING, just curious what's special to x86_intrcc? My real concern is if it's still well defined behavior if the function takes an error code.

llvm/lib/Target/X86/X86FrameLowering.cpp

658

Use for (MachineInstr &MI : PrologMBB) { to iterate rather than loop find_if.

llvm/test/CodeGen/X86/x86-64-intrcc.ll

181

Maybe better to use ;; for comments. The same below.

187–188

This can be in one line. There's no strict column limitation for test case.

[Freax13]

In D150033#4324280, @pengfei wrote:

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

I don't know much about STACKALLOC_W_PROBING, just curious what's special to x86_intrcc? My real concern is if it's still well defined behavior if the function takes an error code.

x86_intrcc is used for exception/interrupt handlers on x86. Before pushing anything onto the stack the cpu will align the stack to 16-bytes. For exceptions without an error code it will push 5 8-byte values onto the stack and for exceptions with an error code it will push 6 8-byte values. This means that for exceptions without an error code the stack will be misaligned by 8 bytes and for exceptions with an error code it will be correctly 16-byte aligned.

Usually the x86-64 ABI requires that the stack will be 16-byte aligned before a function call meaning that it will by misaligned by 8 bytes after the call. The function will then push an additional value to the stack to align the stack to 16 bytes. IIUC for exceptions without an error code no special handling is needed because the stack is 8-byte misaligned just like it would be for a normal call, whereas it's off by 8 for an exception with an error code which is why the stack pointer is updated to bring it back to that misaligned state.

I also don't know much about STACKALLOC_W_PROBING but from I can gather it's used to allocate space on the stack while also probing it to ensure that that space isn't part of a guard page. The probing is done by writing to the memory to see if it will trigger a fault. IIUC in the context of x86_intrcc the 8 byte stack pointer change should never cause any stack probing to be emitted because it's smaller than the stack probe size (4096 bytes). Even if probing were to occur it should probe the memory at the opposite end of the allocation nowhere near the interrupt stack frame pushed by the cpu, so corruption of the interrupt stack frame shouldn't happen. I'm not sure if updating the stack pointer by 8 breaks any assumptions made for stack probing (previously the called function knew that (%rsp) contained the return address i.e. is used and writeable memory, this is no longer the case if we update the stack pointer by 8 when the function is first entered).

llvm/lib/Target/X86/X86FrameLowering.cpp

658

Hmm, this doesn't work and I think my whole approach is flawed: Initially I was afraid to use a for loop because I was afraid that the iterator would get invalidated by the call to emitStackProbeInline and eraseFromParent. The invalidation caused eraseFromParent isn't a huge problem and could probably be bypassed by using make_early_inc_range. The real problem is that emitStackProbeInline will split the basic block when it actually inserts stack probe if it chooses to insert a loop. Crucially this means that if the first STACKALLOC_W_PROBING machine instruction causes one of these stack probe loops to be created all other STACKALLOC_W_PROBINGs will be moved into a new basic block after the loop. This means that inlineStackProbe will no longer find them and we're back to where we started with STACKALLOC_W_PROBING not being inlined.
I'm not yet sure how to fix this, I'd welcome suggestions.

[pengfei]

IIUC in the context of x86_intrcc the 8 byte stack pointer change should never cause any stack probing to be emitted because it's smaller than the stack probe size (4096 bytes)

I have the same conclusion but for a different reason. It is not because 8 bytes < 4096. We still have chance to enter a guard page if the current RSP is near the end a a page, e.g., 4095. The reason is no matter what the size we want for stack realignment, the rest space we used to align is always in the same page unless the alignment requirement > 4096 bytes, which I don't think existing in reality.
Given that, I think should prevent to build the first STACKALLOC_W_PROBING rather than patch it later.

x86_intrcc is used for exception/interrupt handlers on x86... The probing is done by writing to the memory to see if it will trigger a fault.

IIUC, the exception/interrupt handlers are under ring 0 mode, do we really have a chance to trigger fault? Maybe we shouldn't do probe-stack with x86_intrcc?

llvm/lib/Target/X86/X86FrameLowering.cpp

658

I think this can be solved by recording MI in a loop first and then emitting stack probe later.

[pengfei]

What's more, faulty during exception/interrupt seems UB because we may have disabled interrupt before entering the handlers.

[Freax13]

In D150033#4324769, @pengfei wrote:

IIUC in the context of x86_intrcc the 8 byte stack pointer change should never cause any stack probing to be emitted because it's smaller than the stack probe size (4096 bytes)

I have the same conclusion but for a different reason. It is not because 8 bytes < 4096. We still have chance to enter a guard page if the current RSP is near the end a a page, e.g., 4095. The reason is no matter what the size we want for stack realignment, the rest space we used to align is always in the same page unless the alignment requirement > 4096 bytes, which I don't think existing in reality.
Given that, I think should prevent to build the first STACKALLOC_W_PROBING rather than patch it later.

So it turns out that we currently emit push rax for the stack update. This instruction obviously writes to the page and would trigger faults if it landed on a guard page. We can probably get away with just hard-coding the stack update to push rax and avoid any bad interactions with stack probing that way.

x86_intrcc is used for exception/interrupt handlers on x86... The probing is done by writing to the memory to see if it will trigger a fault.

IIUC, the exception/interrupt handlers are under ring 0 mode, do we really have a chance to trigger fault? Maybe we shouldn't do probe-stack with x86_intrcc?

Exceptions handlers typically run in ring 0, but this is not a constraint by the architecture, it's also possible to have exception handlers in rings 1, 2 and 3. In general, faults happen like normal in ring 0 (with few exceptions). One relevant exception to this is that the write-bit is ignored in ring 0 if the WriteProtect bit is not set in CR0, meaning that ring 0 can write to read-only pages (but not unmapped pages). Kernels that want to disable write-protection in ring 0 and also want to use stack probing should use unmapped pages as guard pages. In practice the WriteProtect bit is often enabled e.g. by some (all?) UEFI implementations and the Linux kernel, so stack probing should work just fine.

In D150033#4324774, @pengfei wrote:

What's more, faulty during exception/interrupt seems UB because we may have disabled interrupt before entering the handlers.

Exceptions/faults can't be disabled.

[Freax13]

Instead of handling multiple STACKALLOC_W_PROBING machine instructions, we now just hard-code the stack update for x86_intrcc with an error code to push rax.

[pengfei]

This approach look great.

In D150033#4324861, @Freax13 wrote:

In D150033#4324769, @pengfei wrote:

IIUC in the context of x86_intrcc the 8 byte stack pointer change should never cause any stack probing to be emitted because it's smaller than the stack probe size (4096 bytes)

I have the same conclusion but for a different reason. It is not because 8 bytes < 4096. We still have chance to enter a guard page if the current RSP is near the end a a page, e.g., 4095. The reason is no matter what the size we want for stack realignment, the rest space we used to align is always in the same page unless the alignment requirement > 4096 bytes, which I don't think existing in reality.
Given that, I think should prevent to build the first STACKALLOC_W_PROBING rather than patch it later.

So it turns out that we currently emit push rax for the stack update. This instruction obviously writes to the page and would trigger faults if it landed on a guard page. We can probably get away with just hard-coding the stack update to push rax and avoid any bad interactions with stack probing that way.

You missed my point. push rax won't trigger a page fault if it's used for stack alignment, because the current value of RSP must be 0xXXXXXXX8, the instruction cannot cross the page boundary.

x86_intrcc is used for exception/interrupt handlers on x86... The probing is done by writing to the memory to see if it will trigger a fault.

IIUC, the exception/interrupt handlers are under ring 0 mode, do we really have a chance to trigger fault? Maybe we shouldn't do probe-stack with x86_intrcc?

Exceptions handlers typically run in ring 0, but this is not a constraint by the architecture, it's also possible to have exception handlers in rings 1, 2 and 3. In general, faults happen like normal in ring 0 (with few exceptions). One relevant exception to this is that the write-bit is ignored in ring 0 if the WriteProtect bit is not set in CR0, meaning that ring 0 can write to read-only pages (but not unmapped pages). Kernels that want to disable write-protection in ring 0 and also want to use stack probing should use unmapped pages as guard pages. In practice the WriteProtect bit is often enabled e.g. by some (all?) UEFI implementations and the Linux kernel, so stack probing should work just fine.

In D150033#4324774, @pengfei wrote:

What's more, faulty during exception/interrupt seems UB because we may have disabled interrupt before entering the handlers.

Exceptions/faults can't be disabled.

Thanks! I didn't touch kernel things for a long time. Only leaving an impression we should be cautious to avoid page fault in kernel mode. With a qucik review of IST (Interrupt Stack Table) and NMI (non-maskable interrupts), I think you are correct, we still have chance to handle these cases.

[Freax13]

In D150033#4324920, @pengfei wrote:

This approach look great.

In D150033#4324861, @Freax13 wrote:

In D150033#4324769, @pengfei wrote:

IIUC in the context of x86_intrcc the 8 byte stack pointer change should never cause any stack probing to be emitted because it's smaller than the stack probe size (4096 bytes)

I have the same conclusion but for a different reason. It is not because 8 bytes < 4096. We still have chance to enter a guard page if the current RSP is near the end a a page, e.g., 4095. The reason is no matter what the size we want for stack realignment, the rest space we used to align is always in the same page unless the alignment requirement > 4096 bytes, which I don't think existing in reality.
Given that, I think should prevent to build the first STACKALLOC_W_PROBING rather than patch it later.

So it turns out that we currently emit push rax for the stack update. This instruction obviously writes to the page and would trigger faults if it landed on a guard page. We can probably get away with just hard-coding the stack update to push rax and avoid any bad interactions with stack probing that way.

You missed my point. push rax won't trigger a page fault if it's used for stack alignment, because the current value of RSP must be 0xXXXXXXX8, the instruction cannot cross the page boundary.

RSP won't be 0xXXXXXXX8, it will be 16-byte aligned. When handling an exception with an error code the CPU aligns the stack to 16 bytes and then pushes 6 8-bytes values, which will result in the stack still being 16-byte aligned.

[pengfei]

In D150033#4324928, @Freax13 wrote:

In D150033#4324920, @pengfei wrote:

This approach look great.

In D150033#4324861, @Freax13 wrote:

In D150033#4324769, @pengfei wrote:

IIUC in the context of x86_intrcc the 8 byte stack pointer change should never cause any stack probing to be emitted because it's smaller than the stack probe size (4096 bytes)

I have the same conclusion but for a different reason. It is not because 8 bytes < 4096. We still have chance to enter a guard page if the current RSP is near the end a a page, e.g., 4095. The reason is no matter what the size we want for stack realignment, the rest space we used to align is always in the same page unless the alignment requirement > 4096 bytes, which I don't think existing in reality.
Given that, I think should prevent to build the first STACKALLOC_W_PROBING rather than patch it later.

So it turns out that we currently emit push rax for the stack update. This instruction obviously writes to the page and would trigger faults if it landed on a guard page. We can probably get away with just hard-coding the stack update to push rax and avoid any bad interactions with stack probing that way.

You missed my point. push rax won't trigger a page fault if it's used for stack alignment, because the current value of RSP must be 0xXXXXXXX8, the instruction cannot cross the page boundary.

RSP won't be 0xXXXXXXX8, it will be 16-byte aligned. When handling an exception with an error code the CPU aligns the stack to 16 bytes and then pushes 6 8-bytes values, which will result in the stack still being 16-byte aligned.

Thanks! I got the point now. This is used to make RSP be 0xXXXXXXX8 again to match with normal functions. I was confused it with stack realignment. So yeah, there might be a chance to trigger faults. It looks to me we have confidence to use push in prologue without prob, so I'm inclined the reason is they are small than 4096.

[Freax13]

Thanks for the review!

This is my first time contributing to LLVM, so I'm not sure about the workflow. Do we need to do anything else to get this merged?

[pengfei]

In D150033#4327315, @Freax13 wrote:

Thanks for the review!

This is my first time contributing to LLVM, so I'm not sure about the workflow. Do we need to do anything else to get this merged?

I can commit it for you. Please provide your preferred name and email address.

[Freax13]

In D150033#4328421, @pengfei wrote:

In D150033#4327315, @Freax13 wrote:

Thanks for the review!

This is my first time contributing to LLVM, so I'm not sure about the workflow. Do we need to do anything else to get this merged?

I can commit it for you. Please provide your preferred name and email address.

Perfect, thanks! My preferred name would be "Tom Dohrmann" and my preferred email would be "[e-mail]".

llvm/lib/Target/X86/X86FrameLowering.cpp

 }
 
 void X86FrameLowering::inlineStackProbe(MachineFunction &MF,
                                         MachineBasicBlock &PrologMBB) const {
   auto Where = llvm::find_if(PrologMBB, [](MachineInstr &MI) {
     return MI.getOpcode() == X86::STACKALLOC_W_PROBING;
   });
   if (Where != PrologMBB.end()) {
     DebugLoc DL = PrologMBB.findDebugLoc(Where);
     emitStackProbeInline(MF, PrologMBB, Where, DL, true);
     Where->eraseFromParent();
   }
 }
 
 void X86FrameLowering::emitStackProbeInline(MachineFunction &MF,
                                             MachineBasicBlock &MBB,
                                             MachineBasicBlock::iterator MBBI,
void X86FrameLowering::emitPrologue(MachineFunction &MF,
 
   // Re-align the stack on 64-bit if the x86-interrupt calling convention is
   // used and an error code was pushed, since the x86-64 ABI requires a 16-byte
   // stack alignment.
   if (Fn.getCallingConv() == CallingConv::X86_INTR && Is64Bit &&
       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
   // function, and use up to 128 bytes of stack space, don't have a frame
   // pointer, calls, or dynamic alloca then we do not need to adjust the
   // stack pointer (we fit in the Red Zone). We also check that we don't
   // push and pop from the stack.
   if (has128ByteRedZone(MF) && !TRI->hasStackRealignment(MF) &&

llvm/test/CodeGen/X86/x86-64-intrcc.ll

define x86_intrcc void @test_copy_elide(ptr byval(%struct.interrupt_frame) %frame, i64 %err) #0 {
   ; CHECK: movq %[[R1]], sink_address(%rip)
 entry:
   %err.addr = alloca i64, align 4
   store i64 %err, ptr %err.addr, align 4
   store volatile ptr %err.addr, ptr @sink_address
   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" }