AArch64: support extended spills in SEH on WoS (llvm#166849)

When lowering code for Windows, we might be using a non-standard calling
convention (e.g. `preserve_most`). In such a case, we might be spilling
registers which are unexpected (i.e. x9-x15). Use the ARM64EC opcodes to
indicate such spills.

This adds support for the handling for these spills but is insufficient
on its own. The encoded results are incorrect due to the expectation
that the pair wise spills are always 16-byte aligned which we currently
do not enforce. Fixing that is beyond the scope of emitting the SEH
directives for the spill.

Author: compnerd

llvm/lib/Target/AArch64/AArch64AsmPrinter.cpp

@@ -3364,6 +3364,22 @@ void AArch64AsmPrinter::emitInstruction(const MachineInstr *MI) {
     TS->emitARM64WinCFIPACSignLR();
     return;
 
+  case AArch64::SEH_SaveAnyRegI:
+    assert(MI->getOperand(1).getImm() <= 1008 &&
+           "SaveAnyRegQP SEH opcode offset must fit into 6 bits");
+    TS->emitARM64WinCFISaveAnyRegI(MI->getOperand(0).getImm(),
+                                   MI->getOperand(1).getImm());
+    return;
+
+  case AArch64::SEH_SaveAnyRegIP:
+    assert(MI->getOperand(1).getImm() - MI->getOperand(0).getImm() == 1 &&
+           "Non-consecutive registers not allowed for save_any_reg");
+    assert(MI->getOperand(2).getImm() <= 1008 &&
+           "SaveAnyRegQP SEH opcode offset must fit into 6 bits");
+    TS->emitARM64WinCFISaveAnyRegIP(MI->getOperand(0).getImm(),
+                                    MI->getOperand(2).getImm());
+    return;
+
   case AArch64::SEH_SaveAnyRegQP:
     assert(MI->getOperand(1).getImm() - MI->getOperand(0).getImm() == 1 &&
            "Non-consecutive registers not allowed for save_any_reg");

llvm/lib/Target/AArch64/AArch64FrameLowering.cpp

@@ -1082,25 +1082,41 @@ AArch64FrameLowering::insertSEH(MachineBasicBlock::iterator MBBI,
   case AArch64::LDPXi: {
     Register Reg0 = MBBI->getOperand(0).getReg();
     Register Reg1 = MBBI->getOperand(1).getReg();
+
+    int SEHReg0 = RegInfo->getSEHRegNum(Reg0);
+    int SEHReg1 = RegInfo->getSEHRegNum(Reg1);
+
     if (Reg0 == AArch64::FP && Reg1 == AArch64::LR)
       MIB = BuildMI(MF, DL, TII.get(AArch64::SEH_SaveFPLR))
                 .addImm(Imm * 8)
                 .setMIFlag(Flag);
-    else
+    else if (SEHReg0 >= 19 && SEHReg1 >= 19)
       MIB = BuildMI(MF, DL, TII.get(AArch64::SEH_SaveRegP))
-                .addImm(RegInfo->getSEHRegNum(Reg0))
-                .addImm(RegInfo->getSEHRegNum(Reg1))
+                .addImm(SEHReg0)
+                .addImm(SEHReg1)
+                .addImm(Imm * 8)
+                .setMIFlag(Flag);
+    else
+      MIB = BuildMI(MF, DL, TII.get(AArch64::SEH_SaveAnyRegIP))
+                .addImm(SEHReg0)
+                .addImm(SEHReg1)
                 .addImm(Imm * 8)
                 .setMIFlag(Flag);
     break;
   }
   case AArch64::STRXui:
   case AArch64::LDRXui: {
     int Reg = RegInfo->getSEHRegNum(MBBI->getOperand(0).getReg());
-    MIB = BuildMI(MF, DL, TII.get(AArch64::SEH_SaveReg))
-              .addImm(Reg)
-              .addImm(Imm * 8)
-              .setMIFlag(Flag);
+    if (Reg >= 19)
+      MIB = BuildMI(MF, DL, TII.get(AArch64::SEH_SaveReg))
+                .addImm(Reg)
+                .addImm(Imm * 8)
+                .setMIFlag(Flag);
+    else
+      MIB = BuildMI(MF, DL, TII.get(AArch64::SEH_SaveAnyRegI))
+                .addImm(Reg)
+                .addImm(Imm * 8)
+                .setMIFlag(Flag);
     break;
   }
   case AArch64::STRDui:

llvm/lib/Target/AArch64/AArch64InstrInfo.cpp

@@ -1217,6 +1217,8 @@ bool AArch64InstrInfo::isSEHInstruction(const MachineInstr &MI) {
     case AArch64::SEH_EpilogStart:
     case AArch64::SEH_EpilogEnd:
     case AArch64::SEH_PACSignLR:
+    case AArch64::SEH_SaveAnyRegI:
+    case AArch64::SEH_SaveAnyRegIP:
     case AArch64::SEH_SaveAnyRegQP:
     case AArch64::SEH_SaveAnyRegQPX:
     case AArch64::SEH_AllocZ:

llvm/lib/Target/AArch64/AArch64InstrInfo.td

@@ -5666,6 +5666,8 @@ let isPseudo = 1 in {
   def SEH_EpilogStart : Pseudo<(outs), (ins), []>, Sched<[]>;
   def SEH_EpilogEnd : Pseudo<(outs), (ins), []>, Sched<[]>;
   def SEH_PACSignLR : Pseudo<(outs), (ins), []>, Sched<[]>;
+  def SEH_SaveAnyRegI : Pseudo<(outs), (ins i32imm:$reg0, i32imm:$offs), []>, Sched<[]>;
+  def SEH_SaveAnyRegIP : Pseudo<(outs), (ins i32imm:$reg0, i32imm:$reg1, i32imm:$offs), []>, Sched<[]>;
   def SEH_SaveAnyRegQP : Pseudo<(outs), (ins i32imm:$reg0, i32imm:$reg1, i32imm:$offs), []>, Sched<[]>;
   def SEH_SaveAnyRegQPX : Pseudo<(outs), (ins i32imm:$reg0, i32imm:$reg1, i32imm:$offs), []>, Sched<[]>;
   def SEH_AllocZ : Pseudo<(outs), (ins i32imm:$offs), []>, Sched<[]>;

llvm/lib/Target/AArch64/AArch64PrologueEpilogue.cpp

@@ -253,6 +253,8 @@ static void fixupSEHOpcode(MachineBasicBlock::iterator MBBI,
   case AArch64::SEH_SaveReg:
   case AArch64::SEH_SaveFRegP:
   case AArch64::SEH_SaveFReg:
+  case AArch64::SEH_SaveAnyRegI:
+  case AArch64::SEH_SaveAnyRegIP:
   case AArch64::SEH_SaveAnyRegQP:
   case AArch64::SEH_SaveAnyRegQPX:
     ImmOpnd = &MBBI->getOperand(ImmIdx);

llvm/test/CodeGen/AArch64/seh-extended-spills.ll

@@ -0,0 +1,34 @@
+; RUN: llc -mtriple aarch64-unknown-windows-msvc -filetype asm -o - %s | FileCheck %s
+
+declare dso_local void @g(ptr noundef)
+define dso_local preserve_mostcc void @f(ptr noundef %p) #0 {
+entry:
+  %p.addr = alloca ptr, align 8
+  store ptr %p, ptr %p.addr, align 8
+  %0 = load ptr, ptr %p.addr, align 8
+  call void @g(ptr noundef %0)
+  ret void
+}
+
+attributes #0 = { nounwind uwtable(sync) }
+
+; CHECK: stp x9, x10, [sp, #[[OFFSET_0:[0-9]+]]]
+; CHECK-NEXT: .seh_save_any_reg_p x9, [[OFFSET_0]]
+; CHECK: stp x11, x12, [sp, #[[OFFSET_1:[0-9]+]]]
+; CHECK-NEXT: .seh_save_any_reg_p x11, [[OFFSET_1]]
+; CHECK: stp x13, x14, [sp, #[[OFFSET_2:[0-9]+]]]
+; CHECK-NEXT: .seh_save_any_reg_p x13, [[OFFSET_2]]
+; CHECK: str x15, [sp, #[[OFFSET_3:[0-9]+]]]
+; CHECK-NEXT: .seh_save_any_reg x15, [[OFFSET_3]]
+; CHECK: .seh_endprologue
+
+; CHECK: .seh_startepilogue
+; CHECK: ldr x15, [sp, #[[OFFSET_3]]]
+; CHECK-NEXT: .seh_save_any_reg x15, [[OFFSET_3]]
+; CHECK: ldp x13, x14, [sp, #[[OFFSET_2]]]
+; CHECK-NEXT: .seh_save_any_reg_p x13, [[OFFSET_2]]
+; CHECK: ldp x11, x12, [sp, #[[OFFSET_1]]]
+; CHECK-NEXT: .seh_save_any_reg_p x11, [[OFFSET_1]]
+; CHECK: ldp x9, x10, [sp, #[[OFFSET_0]]]
+; CHECK-NEXT: .seh_save_any_reg_p x9, [[OFFSET_0]]
+; CHECK: .seh_endepilogue