Windows AArch64 frame lowering fails to align large stack frames

[dougallj]

It appears that if __chkstk is called, and the frame requires an alignment greater than 16, the and operation to align sp is missing. This causes a crash in x265 (autovectorization leads to writing to (diff | 16) - 16), which was reported as HandBrake/HandBrake#3692

clang arguments: -target arm64-pc-windows -O2

int external_function(char *p);

void basic_example() {
    char data[4096] __attribute__((aligned(32)));
    external_function(data + 0x10);
}

basic_example:
    // no "and" instruction in prolog to ensure 32-byte alignment
    str     x28, [sp, #-32]!
    stp     x29, x30, [sp, #8]
    add     x29, sp, #8
    mov     x15, #256
    bl      __chkstk
    sub     sp, sp, x15, lsl #4

    // assumes 32-bit alignment, and transforms "add" to "orr"
    mov     x8, sp
    orr     x0, x8, #0x10
    bl      external_function

    add     sp, sp, #1, lsl #12
    ldp     x29, x30, [sp, #8]
    ldr     x28, [sp], #32
    ret

(from Compiler Explorer)

[llvmbot]

@llvm/issue-subscribers-backend-aarch64

[vfdff]

it is interesting, it works fine if doesn't compile with -target arm64-pc-windows

https://godbolt.org/z/zd9nMqejq

[dougallj]

Yeah, it also works fine if the frame is small enough to not require the __chkstk call. (https://godbolt.org/z/YnvfMcGrv)

I think the problem is in AArch64FrameLowering::emitPrologue, where the if (windowsRequiresStackProbe(...)) block builds its own sub instruction then sets NumBytes = 0:

llvm-project/llvm/lib/Target/AArch64/AArch64FrameLowering.cpp

Lines 1708 to 1719 in 0593ce5

	BuildMI(MBB, MBBI, DL, TII->get(AArch64::SUBXrx64), AArch64::SP)
	.addReg(AArch64::SP, RegState::Kill)
	.addReg(AArch64::X15, RegState::Kill)
	.addImm(AArch64_AM::getArithExtendImm(AArch64_AM::UXTX, 4))
	.setMIFlags(MachineInstr::FrameSetup);
	if (NeedsWinCFI) {
	HasWinCFI = true;
	BuildMI(MBB, MBBI, DL, TII->get(AArch64::SEH_StackAlloc))
	.addImm(NumBytes)
	.setMIFlag(MachineInstr::FrameSetup);
	}
	NumBytes = 0;

And, later, realignment is only performed if NumBytes isn't zero:

llvm-project/llvm/lib/Target/AArch64/AArch64FrameLowering.cpp

Lines 1756 to 1760 in 0593ce5

	// Allocate space for the rest of the frame.
	if (NumBytes) {
	// Alignment is required for the parent frame, not the funclet
	const bool NeedsRealignment =
	!IsFunclet && RegInfo->hasStackRealignment(MF);

[efriedma-quic]

Note that in general, since "realignment" is actually allocating stack space, we need to call __chkstk on the allocated space.

We also need to make sure the unwind opcodes correctly reflect the realignment (probably by ending the unwind prologue before we perform the realignment).

[dougallj]

Could this also get the "bug" label? (And probably "miscompilation"?)

@mstorsjo - might you be interested in looking at this?

[llvmbot]

@llvm/issue-subscribers-bug

[mstorsjo]

Note that in general, since "realignment" is actually allocating stack space, we need to call __chkstk on the allocated space.

Not necessarily, as far as I know - we only need to use __chkstk to touch/allocate the stack if we're decrementing the stack by more than one page - any normal stack allocation smaller than that works transparently.

So if we do the large stack allocation with __chkstk, that should touch and allocate everything covered so far - if we then manually decrement the stack pointer by the alignment amount, that's fine, because that can only touch at most one page below the current allocation (assuming alignment smaller than the page size).

[efriedma-quic]

The general rule is just that you have to touch the pages in order. So you don't have to call __chkstk specifically, but you need to ensure each page gets touched before accessing the next page. If you're statically allocating a small amount of stack, avoiding the call probably makes sense; a store is a lot cheaper than a call.

That said, getting this right is much harder if you say that "small" allocations can skip touching the stack. In particular, you need to ensure that it isn't possible to have multiple "small" allocations that add up to more than one page without touching the stack. (If I recall correctly, there was a bug related to this the last time someone was looking at stack probing.)

[mstorsjo]

The general rule is just that you have to touch the pages in order. So you don't have to call __chkstk specifically, but you need to ensure each page gets touched before accessing the next page.

Yup, indeed.

That said, getting this right is much harder if you say that "small" allocations can skip touching the stack. In particular, you need to ensure that it isn't possible to have multiple "small" allocations that add up to more than one page without touching the stack. (If I recall correctly, there was a bug related to this the last time someone was looking at stack probing.)

Right, that's true, if there's other things following and we don't explicitly touch the stack after realigning the stack, then yes - there's opportunity for that to break.