[AArch64] Extend storeRegToStackSlot to spill SVE registers.

[sdesmalen]

Summary

This patch allows the register allocator to spill SVE registers to the stack.

[efriedma]

llvm/lib/Target/AArch64/AArch64InstrInfo.cpp

2960

I'd prefer to refactor the code so it doesn't depend on what getSpillSize returns for SVE types.

[sdesmalen]

llvm/lib/Target/AArch64/AArch64InstrInfo.cpp

2960

The SVE case is actually outside the switch statement and is therefore independent of what getSpillSize returns. The switch statement seems to be there purely for having a quicker look-up as opposed to having a long if-elseif-elseif-[..]-elseif chain. Do you want me to change that?

[efriedma]

LGTM

llvm/lib/Target/AArch64/AArch64InstrInfo.cpp

2960

Oh, I didn't read the code in the switch that carefully. I guess it only triggers for specific register classes, so it doesn't matter what getSpillSize() returns as long as it doesn't crash.

I guess that's okay.

3048

We'll eventually need to handle ZPR2 etc.? I guess we don't need that immediately.

[sdesmalen]

Thanks for the review!

llvm/lib/Target/AArch64/AArch64InstrInfo.cpp

3048

Correct! We can implement this with Pseudos that expand to individual LDR_ZXI/STR_ZXI instructions.

llvm/lib/Target/AArch64/AArch64InstrInfo.cpp

void AArch64InstrInfo::storeRegToStackSlot(
   case 64:
     if (AArch64::QQQQRegClass.hasSubClassEq(RC)) {
       assert(Subtarget.hasNEON() && "Unexpected register store without NEON");
       Opc = AArch64::ST1Fourv2d;
       Offset = false;
     }
     break;
   }
+  unsigned StackID = TargetStackID::Default;
+  if (AArch64::PPRRegClass.hasSubClassEq(RC)) {
+    assert(Subtarget.hasSVE() && "Unexpected register store without SVE");
+    Opc = AArch64::STR_PXI;
+    StackID = TargetStackID::SVEVector;
+  } else if (AArch64::ZPRRegClass.hasSubClassEq(RC)) {
+    assert(Subtarget.hasSVE() && "Unexpected register store without SVE");
+    Opc = AArch64::STR_ZXI;
+    StackID = TargetStackID::SVEVector;
+  }
   assert(Opc && "Unknown register class");
+  MFI.setStackID(FI, StackID);
 
   const MachineInstrBuilder MI = BuildMI(MBB, MBBI, DebugLoc(), get(Opc))
                                      .addReg(SrcReg, getKillRegState(isKill))
                                      .addFrameIndex(FI);
 
   if (Offset)
     MI.addImm(0);
   MI.addMemOperand(MMO);
void AArch64InstrInfo::loadRegFromStackSlot(
   MachineFrameInfo &MFI = MF.getFrameInfo();
   unsigned Align = MFI.getObjectAlignment(FI);
   MachinePointerInfo PtrInfo = MachinePointerInfo::getFixedStack(MF, FI);
   MachineMemOperand *MMO = MF.getMachineMemOperand(
       PtrInfo, MachineMemOperand::MOLoad, MFI.getObjectSize(FI), Align);
 
   unsigned Opc = 0;
   bool Offset = true;
   switch (TRI->getSpillSize(*RC)) {
   case 1:
     if (AArch64::FPR8RegClass.hasSubClassEq(RC))
       Opc = AArch64::LDRBui;
     break;
   case 2:
     if (AArch64::FPR16RegClass.hasSubClassEq(RC))
       Opc = AArch64::LDRHui;
     break;
void AArch64InstrInfo::loadRegFromStackSlot(
   case 64:
     if (AArch64::QQQQRegClass.hasSubClassEq(RC)) {
       assert(Subtarget.hasNEON() && "Unexpected register load without NEON");
       Opc = AArch64::LD1Fourv2d;
       Offset = false;
     }
     break;
   }
+
+  unsigned StackID = TargetStackID::Default;
+  if (AArch64::PPRRegClass.hasSubClassEq(RC)) {
+    assert(Subtarget.hasSVE() && "Unexpected register load without SVE");
+    Opc = AArch64::LDR_PXI;
+    StackID = TargetStackID::SVEVector;
+  } else if (AArch64::ZPRRegClass.hasSubClassEq(RC)) {
+    assert(Subtarget.hasSVE() && "Unexpected register load without SVE");
+    Opc = AArch64::LDR_ZXI;
+    StackID = TargetStackID::SVEVector;
+  }
   assert(Opc && "Unknown register class");
+  MFI.setStackID(FI, StackID);
 
   const MachineInstrBuilder MI = BuildMI(MBB, MBBI, DebugLoc(), get(Opc))
                                      .addReg(DestReg, getDefRegState(true))
                                      .addFrameIndex(FI);
   if (Offset)
     MI.addImm(0);
   MI.addMemOperand(MMO);
 }

llvm/lib/Target/AArch64/AArch64MachineFunctionInfo.h

 #define LLVM_LIB_TARGET_AARCH64_AARCH64MACHINEFUNCTIONINFO_H
 
 #include "llvm/ADT/ArrayRef.h"
 #include "llvm/ADT/Optional.h"
 #include "llvm/ADT/SmallPtrSet.h"
 #include "llvm/ADT/SmallVector.h"
 #include "llvm/CodeGen/CallingConvLower.h"
 #include "llvm/CodeGen/MachineFunction.h"
+#include "llvm/CodeGen/TargetFrameLowering.h"
 #include "llvm/IR/Function.h"
 #include "llvm/MC/MCLinkerOptimizationHint.h"
 #include <cassert>
 
 namespace llvm {
 
 class MachineInstr;
 
if (!HasCalleeSavedStackSize || ValidateCalleeSavedStackSize) {
       assert(MFI.isCalleeSavedInfoValid() && "CalleeSavedInfo not calculated");
       if (MFI.getCalleeSavedInfo().empty())
         return 0;
 
       int64_t MinOffset = std::numeric_limits<int64_t>::max();
       int64_t MaxOffset = std::numeric_limits<int64_t>::min();
       for (const auto &Info : MFI.getCalleeSavedInfo()) {
         int FrameIdx = Info.getFrameIdx();
+        if (MFI.getStackID(FrameIdx) != TargetStackID::Default)
+          continue;
         int64_t Offset = MFI.getObjectOffset(FrameIdx);
         int64_t ObjSize = MFI.getObjectSize(FrameIdx);
         MinOffset = std::min<int64_t>(Offset, MinOffset);
         MaxOffset = std::max<int64_t>(Offset + ObjSize, MaxOffset);
       }
 
       unsigned Size = alignTo(MaxOffset - MinOffset, 16);
       assert((!HasCalleeSavedStackSize || getCalleeSavedStackSize() == Size) &&

llvm/test/CodeGen/AArch64/spillfill-sve.mir

+# RUN: llc -mtriple=aarch64-linux-gnu -run-pass=greedy %s -o - | FileCheck %s
+# RUN: llc -mtriple=aarch64-linux-gnu -start-before=greedy -stop-after=aarch64-expand-pseudo %s -o - | FileCheck %s --check-prefix=EXPAND
+--- |
+  ; ModuleID = '<stdin>'
+  source_filename = "<stdin>"
+  target datalayout = "e-m:e-i8:8:32-i16:16:32-i64:64-i128:128-n32:64-S128"
+  target triple = "aarch64--linux-gnu"
+
+  define aarch64_sve_vector_pcs void @spills_fills_stack_id_ppr() #0 { entry: unreachable }
+  define aarch64_sve_vector_pcs void @spills_fills_stack_id_zpr() #0 { entry: unreachable }
+
+  attributes #0 = { nounwind "target-features"="+sve" }
+
+...
+---
+name: spills_fills_stack_id_ppr
+tracksRegLiveness: true
+registers:
+  - { id: 0, class: ppr }
+stack:
+liveins:
+  - { reg: '$p0', virtual-reg: '%0' }
+body:             |
+  bb.0.entry:
+    liveins: $p0
+
+    ; CHECK-LABEL: name: spills_fills_stack_id_ppr
+    ; CHECK: stack:
+    ; CHECK:      - { id: 0, name: '', type: spill-slot, offset: 0, size: 2, alignment: 2
+    ; CHECK-NEXT:     stack-id: sve-vec, callee-saved-register: ''
+
+    ; EXPAND-LABEL: name: spills_fills_stack_id_ppr
+    ; EXPAND: STR_PXI $p0, $sp, 7
+    ; EXPAND: $p0 = LDR_PXI $sp, 7
+
+    %0:ppr = COPY $p0
+
+    $p0 = IMPLICIT_DEF
+    $p1 = IMPLICIT_DEF
+    $p2 = IMPLICIT_DEF
+    $p3 = IMPLICIT_DEF
+    $p4 = IMPLICIT_DEF
+    $p5 = IMPLICIT_DEF
+    $p6 = IMPLICIT_DEF
+    $p7 = IMPLICIT_DEF
+    $p8 = IMPLICIT_DEF
+    $p9 = IMPLICIT_DEF
+    $p10 = IMPLICIT_DEF
+    $p11 = IMPLICIT_DEF
+    $p12 = IMPLICIT_DEF
+    $p13 = IMPLICIT_DEF
+    $p14 = IMPLICIT_DEF
+    $p15 = IMPLICIT_DEF
+
+    $p0 = COPY %0
+    RET_ReallyLR
+...
+---
+name: spills_fills_stack_id_zpr
+tracksRegLiveness: true
+registers:
+  - { id: 0, class: zpr }
+stack:
+liveins:
+  - { reg: '$z0', virtual-reg: '%0' }
+body:             |
+  bb.0.entry:
+    liveins: $z0
+
+    ; CHECK-LABEL: name: spills_fills_stack_id_zpr
+    ; CHECK: stack:
+    ; CHECK:      - { id: 0, name: '', type: spill-slot, offset: 0, size: 16, alignment: 16
+    ; CHECK-NEXT:     stack-id: sve-vec, callee-saved-register: ''
+
+    ; EXPAND-LABEL: name: spills_fills_stack_id_zpr
+    ; EXPAND: STR_ZXI $z0, $sp, 0
+    ; EXPAND: $z0 = LDR_ZXI $sp, 0
+
+    %0:zpr = COPY $z0
+
+    $z0_z1_z2_z3     = IMPLICIT_DEF
+    $z4_z5_z6_z7     = IMPLICIT_DEF
+    $z8_z9_z10_z11   = IMPLICIT_DEF
+    $z12_z13_z14_z15 = IMPLICIT_DEF
+    $z16_z17_z18_z19 = IMPLICIT_DEF
+    $z20_z21_z22_z23 = IMPLICIT_DEF
+    $z24_z25_z26_z27 = IMPLICIT_DEF
+    $z28_z29_z30_z31 = IMPLICIT_DEF
+
+    $z0 = COPY %0
+    RET_ReallyLR
+...