[llvm] [LLVM][AArch64]Use load/store with consecutive registers in SME2 or S… (PR #77665)

[Sander de Smalen via llvm-commits]

================
@@ -3287,8 +3351,39 @@ void AArch64FrameLowering::determineCalleeSaves(MachineFunction &MF,
           !RegInfo->isReservedReg(MF, PairedReg))
         ExtraCSSpill = PairedReg;
     }
+
+    // Save PReg in FunctionInfo to build PTRUE instruction later. The PTRUE is
+    // being used in the function to save and restore the pair of ZReg
+    AArch64FunctionInfo *AFI = MF.getInfo<AArch64FunctionInfo>();
+    if (Subtarget.hasSVE2p1() || Subtarget.hasSME2()) {
+      if (AArch64::PPRRegClass.contains(Reg) &&
+          (Reg >= AArch64::P8 && Reg <= AArch64::P15) && SavedRegs.test(Reg) &&
+          AFI->getPredicateRegForFillSpill() == 0)
+        AFI->setPredicateRegForFillSpill((Reg - AArch64::P0) + AArch64::PN0);
+
+      // Check if there is a pair of ZRegs, so it can select P8 to create PTRUE,
+      // in case there is no PRege being saved(above)
+      HasPairZReg =
+          HasPairZReg || (AArch64::ZPRRegClass.contains(Reg, CSRegs[i ^ 1]) &&
+                          SavedRegs.test(CSRegs[i ^ 1]));
----------------
sdesmalen-arm wrote:

nit:
```suggestion
      HasPairZReg |= (AArch64::ZPRRegClass.contains(Reg, CSRegs[i ^ 1]) && SavedRegs.test(CSRegs[i ^ 1]));
```

Also just as a suggestion, personally I find the logic a bit easier to follow if you move the searching for a free PN register to after this loop, so it would only has to check for SVE2p1||SME2 once, i.e.

```
  for (unsigned i = 0; CSRegs[i]; ++i) {

    ....

    // Check if there is a pair of ZRegs, so it can select P8 to create PTRUE,
    // in case there is no PRege being saved(above)
    HasPairZReg |= AArch64::ZPRRegClass.contains(Reg, CSRegs[i ^ 1]) &&
                   SavedRegs.test(CSRegs[i ^ 1]);
  }

  // Find a suitable predicate register for the multi-vector spill/fill
  // instructions.
  if (HasPairZReg && (Subtarget.hasSVE2p1() || Subtarget.hasSME2())) {
    AArch64FunctionInfo *AFI = MF.getInfo<AArch64FunctionInfo>();
    for (unsigned Reg = AArch64::P8; Reg <= AArch64::P15; ++Reg) {
      if (SavedRegs.test(Reg)) {
        AFI->setPredicateRegForFillSpill(Reg - AArch64::P0 + AArch64::PN0);
        break;
      }
    }

    // If no free callee-save has been found assign one.
    if (!AFI->getPredicateRegForFillSpill() &&
        MF.getFunction().getCallingConv() ==
            CallingConv::AArch64_SVE_VectorCall) {
      SavedRegs.set(AArch64::P8);
      AFI->setPredicateRegForFillSpill(AArch64::PN8);
    }

    assert(!RegInfo->isReservedReg(MF, AFI->getPredicateRegForFillSpill()) &&
           "Predicate cannot be a reserved register");
  }
```

https://github.com/llvm/llvm-project/pull/77665