[AArch64] Fix argument passing in reserved registers for preserve_nonecc (#96259)

These registers include:
- X19, used by LLVM as the base pointer
- X15 on Windows, where it is used for stack allocation. It can still be
used on Linux/Darwin.
- Adjust FrameLowering scratch register code to not assume X9 is
available if the calling convention is preserve_nonecc. The code will
then pick an unused register as scratch, and allow X9 to continue being
used for argument passing.

Author: antangelo

clang/include/clang/Basic/AttrDocs.td

@@ -5672,9 +5672,10 @@ may be changed in the future.
   be used to pass function arguments. Floating-point registers (XMMs/YMMs) still
   follow the C calling convention.
 - On AArch64, only LR and FP are preserved by the callee.
-  Registers X19-X28, X0-X7, and X9-X15 are used to pass function arguments.
-  X8, X16-X18, SIMD and floating-point registers follow the AAPCS calling
-  convention.
+  Registers X20-X28, X0-X7, and X9-X14 are used to pass function arguments.
+  X8, X16-X19, SIMD and floating-point registers follow the AAPCS calling
+  convention. X15 is not available for argument passing on Windows, but is
+  used to pass arguments on other platforms.
   }];
 }
 

llvm/lib/Target/AArch64/AArch64CallingConvention.td

@@ -500,23 +500,31 @@ def CC_AArch64_Preserve_None : CallingConv<[
     // - X8, used for sret
     // - X16/X17, used by the linker as IP0/IP1
     // - X18, the platform register
+    // - X19, the base pointer
     // - X29, the frame pointer
     // - X30, the link register
     // General registers are not preserved with the exception of
     // FP, LR, and X18
     // Non-volatile registers are used first, so functions may call
     // normal functions without saving and reloading arguments.
-    CCIfType<[i32], CCAssignToReg<[W19, W20, W21, W22, W23,
+    // X9 is assigned last as it is used in FrameLowering as the first
+    // choice for a scratch register.
+    CCIfType<[i32], CCAssignToReg<[W20, W21, W22, W23,
                                    W24, W25, W26, W27, W28,
                                    W0, W1, W2, W3, W4, W5,
-                                   W6, W7, W9, W10, W11,
-                                   W12, W13, W14, W15]>>,
-    CCIfType<[i64], CCAssignToReg<[X19, X20, X21, X22, X23,
+                                   W6, W7, W10, W11,
+                                   W12, W13, W14, W9]>>,
+    CCIfType<[i64], CCAssignToReg<[X20, X21, X22, X23,
                                    X24, X25, X26, X27, X28,
                                    X0, X1, X2, X3, X4, X5,
-                                   X6, X7, X9, X10, X11,
-                                   X12, X13, X14, X15]>>,
-
+                                   X6, X7, X10, X11,
+                                   X12, X13, X14, X9]>>,
+
+    // Windows uses X15 for stack allocation
+    CCIf<"!State.getMachineFunction().getSubtarget<AArch64Subtarget>().isTargetWindows()",
+        CCIfType<[i32], CCAssignToReg<[W15]>>>,
+    CCIf<"!State.getMachineFunction().getSubtarget<AArch64Subtarget>().isTargetWindows()",
+        CCIfType<[i64], CCAssignToReg<[X15]>>>,
     CCDelegateTo<CC_AArch64_AAPCS>
 ]>;
 

llvm/lib/Target/AArch64/AArch64FrameLowering.cpp

@@ -1036,7 +1036,10 @@ static Register findScratchNonCalleeSaveRegister(MachineBasicBlock *MBB) {
   MachineFunction *MF = MBB->getParent();
 
   // If MBB is an entry block, use X9 as the scratch register
-  if (&MF->front() == MBB)
+  // preserve_none functions may be using X9 to pass arguments,
+  // so prefer to pick an available register below.
+  if (&MF->front() == MBB &&
+      MF->getFunction().getCallingConv() != CallingConv::PreserveNone)
     return AArch64::X9;
 
   const AArch64Subtarget &Subtarget = MF->getSubtarget<AArch64Subtarget>();