[AArch64][SVE] Fix PCS for functions taking/returning scalable types.

[sdesmalen]

Summary

The default calling convention needs to save/restore the SVE callee
saves according to the SVE PCS when the function takes or returns
scalable types, even when the aarch64_sve_vector_pcs CC is not
specified for the function.

[paulwalker-arm]

llvm/lib/Target/AArch64/AArch64RegisterInfo.cpp

51–56

Could just do return any_of(...

81

Since the SVE calling convention is really an extension of the "normal" calling convention I think it makes sense for this logic to be at the bottom just before the return. So the flow is:

if (function_specifies_calling_convention_x)
  return calling_convention_x_save_list
if (function_specifies_calling_convention_y)
  return calling_convention_y_save_list
....
if (function_has_sve_args_or_return)
 return sve_save_list

return default_save_list

I guess you could leave the CallingConv::AArch64_SVE_VectorCall code as is given it fits this mould, but since this isn't supposed to be a user visible attribute I wonder if we can remove it at a later date?

[sdesmalen]

• s/needsSVECallingConvention/hasSVEArgsorReturn/
• Moved case to bottom of if/else chain.

[paulwalker-arm]

llvm/lib/Target/AArch64/AArch64RegisterInfo.cpp

43

s/hasSVEArgsorReturn/hasSVEArgsOrReturn/ ?

llvm/lib/Target/AArch64/AArch64RegisterInfo.cpp

 #define GET_REGINFO_TARGET_DESC
 #include "AArch64GenRegisterInfo.inc"
 
 AArch64RegisterInfo::AArch64RegisterInfo(const Triple &TT)
     : AArch64GenRegisterInfo(AArch64::LR), TT(TT) {
   AArch64_MC::initLLVMToCVRegMapping(this);
 }
 
+static bool needsSVECallingConvention(const MachineFunction *MF) {
+  const Function &F = MF->getFunction();
+  if (F.getCallingConv() == CallingConv::AArch64_SVE_VectorCall)
+    return true;
+
+  if (isa<ScalableVectorType>(F.getReturnType()))
+    return true;
+
+  if (any_of(F.args(), [](const Argument &Arg) {
+        return isa<ScalableVectorType>(Arg.getType());
+      }))
+    return true;
+
+  return false;
+}
+
 const MCPhysReg *
 AArch64RegisterInfo::getCalleeSavedRegs(const MachineFunction *MF) const {
   assert(MF && "Invalid MachineFunction pointer.");
 
   if (MF->getFunction().getCallingConv() == CallingConv::GHC)
     // GHC set of callee saved regs is empty as all those regs are
     // used for passing STG regs around
     return CSR_AArch64_NoRegs_SaveList;
   if (MF->getFunction().getCallingConv() == CallingConv::AnyReg)
     return CSR_AArch64_AllRegs_SaveList;
 
   // Darwin has its own CSR_AArch64_AAPCS_SaveList, which means most CSR save
   // lists depending on that will need to have their Darwin variant as well.
   if (MF->getSubtarget<AArch64Subtarget>().isTargetDarwin())
     return getDarwinCalleeSavedRegs(MF);
 
   if (MF->getFunction().getCallingConv() == CallingConv::CFGuard_Check)
     return CSR_Win_AArch64_CFGuard_Check_SaveList;
   if (MF->getSubtarget<AArch64Subtarget>().isTargetWindows())
     return CSR_Win_AArch64_AAPCS_SaveList;
   if (MF->getFunction().getCallingConv() == CallingConv::AArch64_VectorCall)
     return CSR_AArch64_AAVPCS_SaveList;
-  if (MF->getFunction().getCallingConv() == CallingConv::AArch64_SVE_VectorCall)
+  if (needsSVECallingConvention(MF))
     return CSR_AArch64_SVE_AAPCS_SaveList;
   if (MF->getSubtarget<AArch64Subtarget>().getTargetLowering()
           ->supportSwiftError() &&
       MF->getFunction().getAttributes().hasAttrSomewhere(
           Attribute::SwiftError))
     return CSR_AArch64_AAPCS_SwiftError_SaveList;
   if (MF->getFunction().getCallingConv() == CallingConv::PreserveMost)
     return CSR_AArch64_RT_MostRegs_SaveList;

llvm/test/CodeGen/AArch64/sve-calling-convention.ll

 ; RUN: llc -mtriple=aarch64-linux-gnu -mattr=+sve -stop-after=finalize-isel < %s 2>%t | FileCheck %s
+; RUN: llc -mtriple=aarch64-linux-gnu -mattr=+sve -stop-after=prologepilog < %s 2>%t | FileCheck %s --check-prefix=CHECKCSR
 ; RUN: FileCheck --check-prefix=WARN --allow-empty %s <%t
 
 ; If this check fails please read test/CodeGen/AArch64/README for instructions on how to resolve it.
 ; WARN-NOT: warning
 
 ; CHECK-LABEL: name: nosve_signature
 define i32 @nosve_signature() nounwind {
   ret i32 42

 ; CHECK-NEXT: BL @sve_signature_pred, csr_aarch64_sve_aapcs
 ; CHECK: [[RES:%[0-9]+]]:ppr = COPY $p0
 ; CHECK: $p0 = COPY [[RES]]
 ; CHECK: RET_ReallyLR implicit $p0
 define <vscale x 4 x i1> @sve_signature_pred_caller(<vscale x 4 x i1> %arg1, <vscale x 4 x i1> %arg2) nounwind {
   %res = call <vscale x 4 x i1> @sve_signature_pred(<vscale x 4 x i1> %arg2, <vscale x 4 x i1> %arg1)
   ret <vscale x 4 x i1> %res
 }
+
+; Test that functions returning or taking SVE arguments use the correct
+; callee-saved set when using the default C calling convention (as opposed
+; to aarch64_sve_vector_pcs)
+
+; CHECKCSR-LABEL: name: sve_signature_vec_ret_callee
+; CHECKCSR: callee-saved-register: '$z8'
+; CHECKCSR: callee-saved-register: '$p4'
+; CHECKCSR: RET_ReallyLR
+define <vscale x 4 x i32> @sve_signature_vec_ret_callee() nounwind {
+  call void asm sideeffect "nop", "~{z8},~{p4}"()
+  ret <vscale x 4 x i32> zeroinitializer
+}
+
+; CHECKCSR-LABEL: name: sve_signature_vec_arg_callee
+; CHECKCSR: callee-saved-register: '$z8'
+; CHECKCSR: callee-saved-register: '$p4'
+; CHECKCSR: RET_ReallyLR
+define void @sve_signature_vec_arg_callee(<vscale x 4 x i32> %v) nounwind {
+  call void asm sideeffect "nop", "~{z8},~{p4}"()
+  ret void
+}