[SVE] Add fatal error when running out of registers for SVE tuple call arguments

[david-arm]

Summary

When passing SVE types as arguments to function calls we can run
out of hardware SVE registers. This is normally fine, since we
switch to an indirect mode where we pass a pointer to a SVE stack
object in a GPR. However, if we switch over part-way through
processing a SVE tuple then part of it will be in registers and
the other part will be on the stack. This is wrong and we'd like
to avoid any silent ABI compatibility issues in future. For now,
I've added a fatal error when this happens until we can get a
proper fix.

[efriedma]

Can you add a testcase to check we properly trigger the error? Otherwise LGTM.

llvm/lib/Target/AArch64/AArch64CallingConvention.cpp

static const MCPhysReg SRegList[] = {AArch64::S0, AArch64::S1, AArch64::S2,
                                      AArch64::S3, AArch64::S4, AArch64::S5,
                                      AArch64::S6, AArch64::S7};
 static const MCPhysReg DRegList[] = {AArch64::D0, AArch64::D1, AArch64::D2,
                                      AArch64::D3, AArch64::D4, AArch64::D5,
                                      AArch64::D6, AArch64::D7};
 static const MCPhysReg QRegList[] = {AArch64::Q0, AArch64::Q1, AArch64::Q2,
                                      AArch64::Q3, AArch64::Q4, AArch64::Q5,
                                      AArch64::Q6, AArch64::Q7};
+static const MCPhysReg ZRegList[] = {AArch64::Z0, AArch64::Z1, AArch64::Z2,
+                                     AArch64::Z3, AArch64::Z4, AArch64::Z5,
+                                     AArch64::Z6, AArch64::Z7};
 
 static bool finishStackBlock(SmallVectorImpl<CCValAssign> &PendingMembers,
                              MVT LocVT, ISD::ArgFlagsTy &ArgFlags,
                              CCState &State, Align SlotAlign) {
   unsigned Size = LocVT.getSizeInBits() / 8;
   const Align StackAlign =
       State.getMachineFunction().getDataLayout().getStackAlignment();
   const Align OrigAlign = ArgFlags.getNonZeroOrigAlign();
static bool CC_AArch64_Custom_Block(unsigned &ValNo, MVT &ValVT, MVT &LocVT,
   else if (LocVT.SimpleTy == MVT::f16)
     RegList = HRegList;
   else if (LocVT.SimpleTy == MVT::f32 || LocVT.is32BitVector())
     RegList = SRegList;
   else if (LocVT.SimpleTy == MVT::f64 || LocVT.is64BitVector())
     RegList = DRegList;
   else if (LocVT.SimpleTy == MVT::f128 || LocVT.is128BitVector())
     RegList = QRegList;
+  else if (LocVT.isScalableVector())
+    RegList = ZRegList;
   else {
     // Not an array we want to split up after all.
     return false;
   }
 
   SmallVectorImpl<CCValAssign> &PendingMembers = State.getPendingLocs();
 
   // Add the argument to the list to be allocated once we know the size of the
for (auto &It : PendingMembers) {
       UseHigh = !UseHigh;
       if (!UseHigh)
         ++RegResult;
     }
     PendingMembers.clear();
     return true;
   }
 
+  if (LocVT.isScalableVector())
+    report_fatal_error(
+        "Passing consecutive scalable vector registers unsupported");
+
   // Mark all regs in the class as unavailable
   for (auto Reg : RegList)
     State.AllocateReg(Reg);
 
   const Align SlotAlign = Subtarget.isTargetDarwin() ? Align(1) : Align(8);
 
   return finishStackBlock(PendingMembers, LocVT, ArgFlags, State, SlotAlign);
 }
 
 // TableGen provides definitions of the calling convention analysis entry
 // points.
 #include "AArch64GenCallingConv.inc"

llvm/lib/Target/AArch64/AArch64ISelLowering.cpp

Value *AArch64TargetLowering::emitStoreConditional(IRBuilder<> &Builder,
   return Builder.CreateCall(Stxr,
                             {Builder.CreateZExtOrBitCast(
                                  Val, Stxr->getFunctionType()->getParamType(0)),
                              Addr});
 }
 
 bool AArch64TargetLowering::functionArgumentNeedsConsecutiveRegisters(
     Type *Ty, CallingConv::ID CallConv, bool isVarArg) const {
-  return Ty->isArrayTy();
+  if (Ty->isArrayTy())
+    return true;
+
+  const TypeSize &TySize = Ty->getPrimitiveSizeInBits();
+  if (TySize.isScalable() && TySize.getKnownMinSize() > 128)
+    return true;
+
+  return false;
 }
 
 bool AArch64TargetLowering::shouldNormalizeToSelectSequence(LLVMContext &,
                                                             EVT) const {
   return false;
 }
 
 static Value *UseTlsOffset(IRBuilder<> &IRB, unsigned Offset) {

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

+; RUN: not --crash llc < %s -mtriple aarch64-linux-gnu -mattr=+sve >/dev/null 2>%t
+; RUN: FileCheck %s < %t
+
+; CHECK: Passing consecutive scalable vector registers unsupported
+
+target datalayout = "e-m:e-i8:8:32-i16:16:32-i64:64-i128:128-n32:64-S128"
+target triple = "aarch64-unknown-linux-gnu"
+
+define float @foo(double* %x0, double* %x1) {
+entry:
+  %0 = call <vscale x 16 x i1> @llvm.aarch64.sve.ptrue.nxv16i1(i32 31)
+  %1 = call <vscale x 2 x i1> @llvm.aarch64.sve.convert.from.svbool.nxv2i1(<vscale x 16 x i1> %0)
+  %2 = call <vscale x 8 x double> @llvm.aarch64.sve.ld4.nxv8f64.nxv2i1(<vscale x 2 x i1> %1, double* %x0)
+  %3 = call <vscale x 8 x double> @llvm.aarch64.sve.ld4.nxv8f64.nxv2i1(<vscale x 2 x i1> %1, double* %x1)
+  %call = call float @callee(float 1.000000e+00, <vscale x 8 x double> %2, <vscale x 8 x double> %3)
+  ret float %call
+}
+
+declare float @callee(float, <vscale x 8 x double>, <vscale x 8 x double>)
+
+declare <vscale x 16 x i1> @llvm.aarch64.sve.ptrue.nxv16i1(i32 immarg)
+declare <vscale x 2 x i1> @llvm.aarch64.sve.convert.from.svbool.nxv2i1(<vscale x 16 x i1>)
+declare <vscale x 8 x double> @llvm.aarch64.sve.ld4.nxv8f64.nxv2i1(<vscale x 2 x i1>, double*)