[AArch64][SME] Add support for Copy/Spill/Fill of strided ZPR2/ZPR4 registers.

[sdesmalen]

Summary

This patch contains a few changes:

• It changes the alignment of the strided/contiguous ZPR2/ZPR4 registers to 128-bits. This is important, because when we spill these registers to the stack, the address doesn't need to be 256/512 bits aligned because we split the single-store/reload pseudo instruction up into multiple STR_ZXI/LDR_ZXI (single vector store/load) instructions, which only require a 128-bit alignment. Additionally, an alignment larger than the stack-alignment is not supported for scalable vectors.

• It adds support for these register classes in storeRegToStackSlot, loadRegFromStackSlot and copyPhysReg.

• It adds tests only for the strided forms. There is no need to also test the contiguous forms, because a register such as z2_z3 or z4_z5_z6_z7 are also part of the regular ZPR2 and ZPR4 register classes, respectively, which are already covered and tested.

[dtemirbulatov]

LGTM.

llvm/lib/Target/AArch64/AArch64InstrInfo.cpp

if (AArch64::ZPRRegClass.contains(DestReg) &&
     assert(Subtarget.hasSVEorSME() && "Unexpected SVE register.");
     BuildMI(MBB, I, DL, get(AArch64::ORR_ZZZ), DestReg)
       .addReg(SrcReg)
       .addReg(SrcReg, getKillRegState(KillSrc));
     return;
   }
 
   // Copy a Z register pair by copying the individual sub-registers.
-  if (AArch64::ZPR2RegClass.contains(DestReg) &&
-      AArch64::ZPR2RegClass.contains(SrcReg)) {
+  if ((AArch64::ZPR2RegClass.contains(DestReg) ||
+       AArch64::ZPR2StridedOrContiguousRegClass.contains(DestReg)) &&
+      (AArch64::ZPR2RegClass.contains(SrcReg) ||
+       AArch64::ZPR2StridedOrContiguousRegClass.contains(SrcReg))) {
     assert(Subtarget.hasSVEorSME() && "Unexpected SVE register.");
     static const unsigned Indices[] = {AArch64::zsub0, AArch64::zsub1};
     copyPhysRegTuple(MBB, I, DL, DestReg, SrcReg, KillSrc, AArch64::ORR_ZZZ,
                      Indices);
     return;
   }
 
   // Copy a Z register triple by copying the individual sub-registers.
   if (AArch64::ZPR3RegClass.contains(DestReg) &&
       AArch64::ZPR3RegClass.contains(SrcReg)) {
     assert(Subtarget.hasSVEorSME() && "Unexpected SVE register.");
     static const unsigned Indices[] = {AArch64::zsub0, AArch64::zsub1,
                                        AArch64::zsub2};
     copyPhysRegTuple(MBB, I, DL, DestReg, SrcReg, KillSrc, AArch64::ORR_ZZZ,
                      Indices);
     return;
   }
 
   // Copy a Z register quad by copying the individual sub-registers.
-  if (AArch64::ZPR4RegClass.contains(DestReg) &&
-      AArch64::ZPR4RegClass.contains(SrcReg)) {
+  if ((AArch64::ZPR4RegClass.contains(DestReg) ||
+       AArch64::ZPR4StridedOrContiguousRegClass.contains(DestReg)) &&
+      (AArch64::ZPR4RegClass.contains(SrcReg) ||
+       AArch64::ZPR4StridedOrContiguousRegClass.contains(SrcReg))) {
     assert(Subtarget.hasSVEorSME() && "Unexpected SVE register.");
     static const unsigned Indices[] = {AArch64::zsub0, AArch64::zsub1,
                                        AArch64::zsub2, AArch64::zsub3};
     copyPhysRegTuple(MBB, I, DL, DestReg, SrcReg, KillSrc, AArch64::ORR_ZZZ,
                      Indices);
     return;
   }
 
case 32:
     if (AArch64::DDDDRegClass.hasSubClassEq(RC)) {
       assert(Subtarget.hasNEON() && "Unexpected register store without NEON");
       Opc = AArch64::ST1Fourv1d;
       Offset = false;
     } else if (AArch64::QQRegClass.hasSubClassEq(RC)) {
       assert(Subtarget.hasNEON() && "Unexpected register store without NEON");
       Opc = AArch64::ST1Twov2d;
       Offset = false;
-    } else if (AArch64::ZPR2RegClass.hasSubClassEq(RC)) {
+    } else if (AArch64::ZPR2RegClass.hasSubClassEq(RC) ||
+               AArch64::ZPR2StridedOrContiguousRegClass.hasSubClassEq(RC)) {
       assert(Subtarget.hasSVE() && "Unexpected register store without SVE");
       Opc = AArch64::STR_ZZXI;
       StackID = TargetStackID::ScalableVector;
     }
     break;
   case 48:
     if (AArch64::QQQRegClass.hasSubClassEq(RC)) {
       assert(Subtarget.hasNEON() && "Unexpected register store without NEON");
       Opc = AArch64::ST1Threev2d;
       Offset = false;
     } else if (AArch64::ZPR3RegClass.hasSubClassEq(RC)) {
       assert(Subtarget.hasSVE() && "Unexpected register store without SVE");
       Opc = AArch64::STR_ZZZXI;
       StackID = TargetStackID::ScalableVector;
     }
     break;
   case 64:
     if (AArch64::QQQQRegClass.hasSubClassEq(RC)) {
       assert(Subtarget.hasNEON() && "Unexpected register store without NEON");
       Opc = AArch64::ST1Fourv2d;
       Offset = false;
-    } else if (AArch64::ZPR4RegClass.hasSubClassEq(RC)) {
+    } else if (AArch64::ZPR4RegClass.hasSubClassEq(RC) ||
+               AArch64::ZPR4StridedOrContiguousRegClass.hasSubClassEq(RC)) {
       assert(Subtarget.hasSVE() && "Unexpected register store without SVE");
       Opc = AArch64::STR_ZZZZXI;
       StackID = TargetStackID::ScalableVector;
     }
     break;
   }
   assert(Opc && "Unknown register class");
   MFI.setStackID(FI, StackID);
case 32:
     if (AArch64::DDDDRegClass.hasSubClassEq(RC)) {
       assert(Subtarget.hasNEON() && "Unexpected register load without NEON");
       Opc = AArch64::LD1Fourv1d;
       Offset = false;
     } else if (AArch64::QQRegClass.hasSubClassEq(RC)) {
       assert(Subtarget.hasNEON() && "Unexpected register load without NEON");
       Opc = AArch64::LD1Twov2d;
       Offset = false;
-    } else if (AArch64::ZPR2RegClass.hasSubClassEq(RC)) {
+    } else if (AArch64::ZPR2RegClass.hasSubClassEq(RC) ||
+               AArch64::ZPR2StridedOrContiguousRegClass.hasSubClassEq(RC)) {
       assert(Subtarget.hasSVE() && "Unexpected register load without SVE");
       Opc = AArch64::LDR_ZZXI;
       StackID = TargetStackID::ScalableVector;
     }
     break;
   case 48:
     if (AArch64::QQQRegClass.hasSubClassEq(RC)) {
       assert(Subtarget.hasNEON() && "Unexpected register load without NEON");
       Opc = AArch64::LD1Threev2d;
       Offset = false;
     } else if (AArch64::ZPR3RegClass.hasSubClassEq(RC)) {
       assert(Subtarget.hasSVE() && "Unexpected register load without SVE");
       Opc = AArch64::LDR_ZZZXI;
       StackID = TargetStackID::ScalableVector;
     }
     break;
   case 64:
     if (AArch64::QQQQRegClass.hasSubClassEq(RC)) {
       assert(Subtarget.hasNEON() && "Unexpected register load without NEON");
       Opc = AArch64::LD1Fourv2d;
       Offset = false;
-    } else if (AArch64::ZPR4RegClass.hasSubClassEq(RC)) {
+    } else if (AArch64::ZPR4RegClass.hasSubClassEq(RC) ||
+               AArch64::ZPR4StridedOrContiguousRegClass.hasSubClassEq(RC)) {
       assert(Subtarget.hasSVE() && "Unexpected register load without SVE");
       Opc = AArch64::LDR_ZZZZXI;
       StackID = TargetStackID::ScalableVector;
     }
     break;
   }
 
   assert(Opc && "Unknown register class");

llvm/lib/Target/AArch64/AArch64RegisterInfo.td

def ZStridedQuadsLo : RegisterTuples<[zsub0, zsub1, zsub2, zsub3], [
   (trunc (rotl ZPR, 8), 4), (trunc (rotl ZPR, 12), 4)
 ]>;
 // Z16_Z20_Z24_Z28, Z17_Z21_Z25_Z29, Z18_Z22_Z26_Z30, Z19_Z23_Z27_Z31
 def ZStridedQuadsHi : RegisterTuples<[zsub0, zsub1, zsub2, zsub3], [
   (trunc (rotl ZPR, 16), 4), (trunc (rotl ZPR, 20), 4),
   (trunc (rotl ZPR, 24), 4), (trunc (rotl ZPR, 28), 4)
 ]>;
 
-def ZPR2Strided : RegisterClass<"AArch64", [untyped], 256,
+def ZPR2Strided : RegisterClass<"AArch64", [untyped], 128,
                                 (add ZStridedPairsLo, ZStridedPairsHi)>  {
   let Size = 256;
 }
-def ZPR4Strided : RegisterClass<"AArch64", [untyped], 512,
+def ZPR4Strided : RegisterClass<"AArch64", [untyped], 128,
                                 (add ZStridedQuadsLo, ZStridedQuadsHi)>  {
   let Size = 512;
 }
 
-def ZPR2StridedOrContiguous : RegisterClass<"AArch64", [untyped], 256,
+def ZPR2StridedOrContiguous : RegisterClass<"AArch64", [untyped], 128,
                                 (add ZStridedPairsLo, ZStridedPairsHi,
                                 (decimate ZSeqPairs, 2))> {
   let Size = 256;
 }
 
 class ZPRVectorListStrided<int ElementWidth, int NumRegs, int Stride>
     : ZPRVectorList<ElementWidth, NumRegs> {
   let Name = "SVEVectorListStrided" # NumRegs # "x" # ElementWidth;
let EncoderMethod = "EncodeZPR2StridedRegisterClass",
   def ZZ_h_strided_and_contiguous
       : RegisterOperand<ZPR2StridedOrContiguous, "printTypedVectorList<0,'h'>">;
   def ZZ_s_strided_and_contiguous
       : RegisterOperand<ZPR2StridedOrContiguous, "printTypedVectorList<0,'s'>">;
   def ZZ_d_strided_and_contiguous
       : RegisterOperand<ZPR2StridedOrContiguous, "printTypedVectorList<0,'d'>">;
 }
 
-def ZPR4StridedOrContiguous : RegisterClass<"AArch64", [untyped], 512,
+def ZPR4StridedOrContiguous : RegisterClass<"AArch64", [untyped], 128,
                                 (add ZStridedQuadsLo, ZStridedQuadsHi,
                                 (decimate ZSeqQuads, 4))>  {
   let Size = 512;
 }
 
 let EncoderMethod = "EncodeZPR4StridedRegisterClass",
     DecoderMethod = "DecodeZPR4StridedRegisterClass" in {
   def ZZZZ_b_strided

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 }
   define aarch64_sve_vector_pcs void @spills_fills_stack_id_zpr2() #0 { entry: unreachable }
+  define aarch64_sve_vector_pcs void @spills_fills_stack_id_zpr2strided() #0 { entry: unreachable }
   define aarch64_sve_vector_pcs void @spills_fills_stack_id_zpr3() #0 { entry: unreachable }
   define aarch64_sve_vector_pcs void @spills_fills_stack_id_zpr4() #0 { entry: unreachable }
+  define aarch64_sve_vector_pcs void @spills_fills_stack_id_zpr4strided() #0 { entry: unreachable }
 
   attributes #0 = { nounwind "target-features"="+sve" }
 
 ...
 ---
 name: spills_fills_stack_id_ppr
 tracksRegLiveness: true
 registers:
bb.0.entry:
     $z20_z21_z22_z23 = IMPLICIT_DEF
     $z24_z25_z26_z27 = IMPLICIT_DEF
     $z28_z29_z30_z31 = IMPLICIT_DEF
 
     $z0_z1 = COPY %0
     RET_ReallyLR
 ...
 ---
+name: spills_fills_stack_id_zpr2strided
+tracksRegLiveness: true
+registers:
+  - { id: 0, class: zpr2strided }
+stack:
+liveins:
+  - { reg: '$z0_z8', virtual-reg: '%0' }
+body:             |
+  bb.0.entry:
+    liveins: $z0_z8
+
+    ; CHECK-LABEL: name: spills_fills_stack_id_zpr2strided
+    ; CHECK: stack:
+    ; CHECK:      - { id: 0, name: '', type: spill-slot, offset: 0, size: 32, alignment: 16
+    ; CHECK-NEXT:     stack-id: scalable-vector
+
+    ; EXPAND-LABEL: name: spills_fills_stack_id_zpr2strided
+    ; EXPAND: STR_ZXI $z0, $sp, 0
+    ; EXPAND: STR_ZXI $z8, $sp, 1
+    ; EXPAND: $z0 = LDR_ZXI $sp, 0
+    ; EXPAND: $z8 = LDR_ZXI $sp, 1
+
+    %0:zpr2strided = COPY $z0_z8
+
+    $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_z8 = COPY %0
+    RET_ReallyLR
+...
+---
 name: spills_fills_stack_id_zpr3
 tracksRegLiveness: true
 registers:
   - { id: 0, class: zpr3 }
 stack:
 liveins:
   - { reg: '$z0_z1_z2', virtual-reg: '%0' }
 body:             |
bb.0.entry:
     $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_z1_z2_z3 = COPY %0
     RET_ReallyLR
 ...
+---
+name: spills_fills_stack_id_zpr4strided
+tracksRegLiveness: true
+registers:
+  - { id: 0, class: zpr4strided }
+stack:
+liveins:
+  - { reg: '$z0_z4_z8_z12', virtual-reg: '%0' }
+body:             |
+  bb.0.entry:
+    liveins: $z0_z4_z8_z12
+
+    ; CHECK-LABEL: name: spills_fills_stack_id_zpr4strided
+    ; CHECK: stack:
+    ; CHECK:      - { id: 0, name: '', type: spill-slot, offset: 0, size: 64, alignment: 16
+    ; CHECK-NEXT:     stack-id: scalable-vector
+
+    ; EXPAND-LABEL: name: spills_fills_stack_id_zpr4strided
+    ; EXPAND: STR_ZXI $z0, $sp, 0
+    ; EXPAND: STR_ZXI $z4, $sp, 1
+    ; EXPAND: STR_ZXI $z8, $sp, 2
+    ; EXPAND: STR_ZXI $z12, $sp, 3
+    ; EXPAND: $z0 = LDR_ZXI $sp, 0
+    ; EXPAND: $z4 = LDR_ZXI $sp, 1
+    ; EXPAND: $z8 = LDR_ZXI $sp, 2
+    ; EXPAND: $z12 = LDR_ZXI $sp, 3
+
+    %0:zpr4strided = COPY $z0_z4_z8_z12
+
+    $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_z4_z8_z12 = COPY %0
+    RET_ReallyLR
+...

llvm/test/CodeGen/AArch64/sve-copy-zprpair.mir

bb.0:
     ; CHECK: $z1 = ORR_ZZZ $z2, $z2
     ; CHECK: RET_ReallyLR
     $z1_z2 = COPY $z0_z1
     $z0_z1 = COPY $z1_z2
     RET_ReallyLR
 
 ...
 ---
+name:            copy_zpr2strided
+alignment:       4
+tracksRegLiveness: true
+liveins:
+  - { reg: '$z0_z1' }
+frameInfo:
+  maxCallFrameSize: 0
+body:             |
+  bb.0:
+    liveins: $z0_z1
+    ; CHECK-LABEL: name: copy_zpr2strided
+    ; CHECK: liveins: $z0_z1
+    ; CHECK: $z8 = ORR_ZZZ $z1, $z1
+    ; CHECK: $z0 = ORR_ZZZ $z0, $z0
+    ; CHECK: $z1 = ORR_ZZZ $z8, $z8
+    ; CHECK: $z0 = ORR_ZZZ $z0, $z0
+    ; CHECK: RET_ReallyLR
+    $z0_z8 = COPY $z0_z1
+    $z0_z1 = COPY $z0_z8
+    RET_ReallyLR
+
+...
+---
 name:            copy_zpr3
 alignment:       4
 tracksRegLiveness: true
 liveins:
   - { reg: '$z0_z1_z2' }
 frameInfo:
   maxCallFrameSize: 0
 body:             |
bb.0:
     ; CHECK: $z2 = ORR_ZZZ $z3, $z3
     ; CHECK: $z3 = ORR_ZZZ $z4, $z4
     ; CHECK: RET_ReallyLR
     $z1_z2_z3_z4 = COPY $z0_z1_z2_z3
     $z0_z1_z2_z3 = COPY $z1_z2_z3_z4
     RET_ReallyLR
 
 ...
+---
+name:            copy_zpr4strided
+alignment:       4
+tracksRegLiveness: true
+liveins:
+  - { reg: '$z0_z1_z2_z3' }
+frameInfo:
+  maxCallFrameSize: 0
+body:             |
+  bb.0:
+    liveins: $z0_z1_z2_z3
+    ; CHECK-LABEL: name: copy_zpr4
+    ; CHECK: liveins: $z0_z1_z2_z3
+    ; CHECK: $z12 = ORR_ZZZ $z3, $z3
+    ; CHECK: $z8 = ORR_ZZZ $z2, $z2
+    ; CHECK: $z4 = ORR_ZZZ $z1, $z1
+    ; CHECK: $z0 = ORR_ZZZ $z0, $z0
+    ; CHECK: $z3 = ORR_ZZZ $z12, $z12
+    ; CHECK: $z2 = ORR_ZZZ $z8, $z8
+    ; CHECK: $z1 = ORR_ZZZ $z4, $z4
+    ; CHECK: $z0 = ORR_ZZZ $z0, $z0
+    ; CHECK: RET_ReallyLR
+    $z0_z4_z8_z12 = COPY $z0_z1_z2_z3
+    $z0_z1_z2_z3 = COPY $z0_z4_z8_z12
+    RET_ReallyLR
+
+...