[AArch64] Fix an immediate out of range for large realignments on Windows

[mstorsjo]

Summary

This fixes https://github.com/llvm/llvm-project/issues/63701.

Not sure if there's any more elegant way to do this at this level;
this seems to fix the issue at least.

(Should the MI operand registers be marked with more flags, as killed etc?)

[efriedma]

LGTM

llvm/lib/Target/AArch64/AArch64FrameLowering.cpp

1748

You could stick a kill flag on x16 here, I guess, although it doesn't matter much this late in the pipeline.

Probably we should stick a FrameSetup flag on the add (for both cases).

1753

If we care, for constants up to 1<<23, we can use a shifted immediate. ("add sp, sp, #1, lsl #12" etc.) But probably not worth bothering.

[mstorsjo]

llvm/lib/Target/AArch64/AArch64FrameLowering.cpp

1748

Thanks, will do.

1753

Yeah, but that doesn't work if the immediate contains bit set both below and above the 12 bit boundary; the constants here are like 0x1ff0, so it's still 0x1, lsl #12 combined with 0xff0.

[mstorsjo]

Added some kill/define flags for registers, and missing FrameSetup flags.

Will push later.

[mstorsjo]

llvm/lib/Target/AArch64/AArch64FrameLowering.cpp

1748

SP wasn’t a valid register for ADDXrr here, as shown by the bots with expensive checks enabled.

I tried grepping around to see what instruction name I should use here instead, but didn’t have any luck.

Is it possible to get some dump of what instruction names it uses internally if I assemble the desired instruction from a .s file, or does that not pass through those layers at all?

[efriedma]

llvm/lib/Target/AArch64/AArch64FrameLowering.cpp

1748

$ echo "add x16, sp, x16" | llvm-mc -show-inst -triple=aarch64
        .text
        add     x16, sp, x16                    // <MCInst #1439 ADDXrx64
                                        //  <MCOperand Reg:235>
                                        //  <MCOperand Reg:6>
                                        //  <MCOperand Reg:235>
                                        //  <MCOperand Imm:24>>

This is an MCInst, not a MachineInstr, but should be enough to get the idea. The immediate is an extend (i.e. the instruction is actually "add x16, sp, x16, uxtx #0").

[mstorsjo]

Updated to use ADDXrx64 which works properly for SP, fixing the tests when expensive checks are enabled, and also fixing the actually encoded instruction.

[efriedma]

LGTM

llvm/lib/Target/AArch64/AArch64FrameLowering.cpp

if (NeedsWinCFI) {
       HasWinCFI = true;
       BuildMI(MBB, MBBI, DL, TII->get(AArch64::SEH_StackAlloc))
           .addImm(NumBytes)
           .setMIFlag(MachineInstr::FrameSetup);
     }
     NumBytes = 0;
 
     if (RealignmentPadding > 0) {
+      if (RealignmentPadding >= 4096) {
+        BuildMI(MBB, MBBI, DL, TII->get(AArch64::MOVi64imm))
+            .addReg(AArch64::X16, RegState::Define)
+            .addImm(RealignmentPadding)
+            .setMIFlags(MachineInstr::FrameSetup);
+        BuildMI(MBB, MBBI, DL, TII->get(AArch64::ADDXrx64), AArch64::X15)
+            .addReg(AArch64::SP)
+            .addReg(AArch64::X16, RegState::Kill)
+            .addImm(AArch64_AM::getArithExtendImm(AArch64_AM::UXTX, 0))
+            .setMIFlag(MachineInstr::FrameSetup);
+      } else {
         BuildMI(MBB, MBBI, DL, TII->get(AArch64::ADDXri), AArch64::X15)
             .addReg(AArch64::SP)
             .addImm(RealignmentPadding)
-          .addImm(0);
+            .addImm(0)
+            .setMIFlag(MachineInstr::FrameSetup);
+      }
 
       uint64_t AndMask = ~(MFI.getMaxAlign().value() - 1);
       BuildMI(MBB, MBBI, DL, TII->get(AArch64::ANDXri), AArch64::SP)
           .addReg(AArch64::X15, RegState::Kill)
           .addImm(AArch64_AM::encodeLogicalImmediate(AndMask, 64));
       AFI->setStackRealigned(true);
 
       // No need for SEH instructions here; if we're realigning the stack,

llvm/test/CodeGen/AArch64/win-align-chkstk.ll

 ; CHECK-NEXT: add x29, sp, #8
 ; CHECK-NEXT: .seh_add_fp 8
 ; CHECK-NEXT: .seh_endprologue
 ; CHECK-NEXT: mov x15, #513
 ; CHECK-NEXT: bl __chkstk
 ; CHECK-NEXT: sub sp, sp, x15, lsl #4
 ; CHECK-NEXT: add x15, sp, #16
 ; CHECK-NEXT: and sp, x15, #0xffffffffffffffe0
+
+define dso_local void @func2() {
+entry:
+  %buf = alloca [8192 x i8], align 8192
+  %arraydecay = getelementptr inbounds [8192 x i8], ptr %buf, i64 0, i64 0
+  call void @other(ptr noundef %arraydecay)
+  ret void
+}
+
+; CHECK-LABEL: func2:
+; CHECK-NEXT: .seh_proc func2
+; CHECK-NEXT: // %bb.0:
+; CHECK-NEXT: str x28, [sp, #-32]!
+; CHECK-NEXT: .seh_save_reg_x x28, 32
+; CHECK-NEXT: stp x29, x30, [sp, #8]
+; CHECK-NEXT: .seh_save_fplr 8
+; CHECK-NEXT: add x29, sp, #8
+; CHECK-NEXT: .seh_add_fp 8
+; CHECK-NEXT: .seh_endprologue
+; CHECK-NEXT: mov x15, #1533
+; CHECK-NEXT: bl __chkstk
+; CHECK-NEXT: sub sp, sp, x15, lsl #4
+; CHECK-NEXT: mov x16, #8176
+; CHECK-NEXT: add x15, sp, x16
+; CHECK-NEXT: and sp, x15, #0xffffffffffffe000