[llvm] [AArch64] Reorder callee-saved spills to maximise load/store register groupings (PR #229349)

[Kieran B via llvm-commits]

================
@@ -2863,102 +2863,94 @@ static void orderZPRCalleeSavesForGroups(MachineFunction &MF,
   if (ZPRSaves.size() < 2)
     return;
 
-  llvm::sort(ZPRSaves, [](const auto &A, const auto &B) {
-    return A.getReg() < B.getReg();
-  });
+  int ZPRByteOffset = AFI->getZPRCalleeSavedStackSize();
+  if (!AFI->hasSplitSVEObjects())
+    ZPRByteOffset += AFI->getPPRCalleeSavedStackSize();
+
+  const int Scale = RegInfo->getSpillSize(AArch64::ZPRRegClass);
+  const int LowestOffset =
+      (ZPRByteOffset / Scale) - static_cast<int>(ZPRSaves.size());
+
+  int NumToAlign = alignTo(LowestOffset, 4) - LowestOffset;
+  if (NumToAlign > int(ZPRSaves.size()))
+    return;
+
+  sort(ZPRSaves,
+       [](const auto &A, const auto &B) { return A.getReg() < B.getReg(); });
 
   SmallVector<std::array<CalleeSavedInfo, 4>> Quads;
   SmallVector<std::array<CalleeSavedInfo, 2>> Pairs;
   SmallVector<CalleeSavedInfo> Singles;
+
+  int ZPRsNeededForAlignment = NumToAlign;
   for (size_t i = 0; i < ZPRSaves.size();) {
+    int ZPRsRemaining = int(ZPRSaves.size()) - i;
     if (i + 3 < ZPRSaves.size() &&
         (ZPRSaves[i].getReg() - AArch64::Z0) % 4 == 0 &&
-        (ZPRSaves[i].getReg() + 3 == ZPRSaves[i + 3].getReg())) {
+        ZPRSaves[i].getReg() + 3 == ZPRSaves[i + 3].getReg() &&
+        ZPRsRemaining - 4 >= ZPRsNeededForAlignment) {
       Quads.push_back(
           {ZPRSaves[i], ZPRSaves[i + 1], ZPRSaves[i + 2], ZPRSaves[i + 3]});
       i += 4;
-    } else if (i + 1 < ZPRSaves.size() &&
-               (ZPRSaves[i].getReg() - AArch64::Z0) % 2 == 0 &&
-               ZPRSaves[i].getReg() + 1 == ZPRSaves[i + 1].getReg()) {
-      Pairs.push_back({ZPRSaves[i], ZPRSaves[i + 1]});
-      i += 2;
-    } else {
-      Singles.push_back(ZPRSaves[i++]);
+      continue;
     }
-  }
 
-  if (Quads.empty() && Pairs.empty())
-    return;
-
-  int ZPRByteOffset = AFI->getZPRCalleeSavedStackSize();
-  if (!AFI->hasSplitSVEObjects())
-    ZPRByteOffset += AFI->getPPRCalleeSavedStackSize();
+    bool MustSplitPairForSingle =
+        ZPRsNeededForAlignment == 1 && Singles.empty() && ZPRsRemaining == 2;
 
-  const int Scale = RegInfo->getSpillSize(AArch64::ZPRRegClass);
-  const int LowestOffset =
-      (ZPRByteOffset / Scale) - static_cast<int>(ZPRSaves.size());
+    if (i + 1 < ZPRSaves.size() &&
+        (ZPRSaves[i].getReg() - AArch64::Z0) % 2 == 0 &&
+        ZPRSaves[i].getReg() + 1 == ZPRSaves[i + 1].getReg() &&
+        !MustSplitPairForSingle) {
+      Pairs.push_back({ZPRSaves[i], ZPRSaves[i + 1]});
 
-  auto TakeSingle = [&]() -> CalleeSavedInfo {
-    if (!Singles.empty())
-      return Singles.pop_back_val();
+      if (ZPRsNeededForAlignment >= 2) {
+        ZPRsNeededForAlignment -= 2;
+      } else if (ZPRsNeededForAlignment == 1 && !Singles.empty()) {
+        // In this case, we have already found a single to use for alignment but
+        // we prefer to use this pair instead to align the offset
+        ZPRsNeededForAlignment = 0;
+      }
 
-    if (!Pairs.empty()) {
-      auto [a, b] = Pairs.pop_back_val();
-      Singles.push_back(a);
-      return b;
+      i += 2;
+      continue;
     }
 
-    if (!Quads.empty()) {
-      auto [a, b, c, d] = Quads.pop_back_val();
-      Pairs.push_back({a, b});
-      Singles.push_back(c);
-      return d;
-    }
+    Singles.push_back(ZPRSaves[i++]);
+    if (ZPRsNeededForAlignment > 0)
+      ZPRsNeededForAlignment--;
+  }
 
-    llvm_unreachable("Register groups exhausted before output was complete");
-  };
+  if (Quads.empty() && Pairs.empty())
+    return;
 
   int NumUnassigned = ZPRSaves.size();
-  auto Emit = [&](const CalleeSavedInfo &CS) {
+  auto Emit = [&](ArrayRef<CalleeSavedInfo> Saves) {
     // During pair/quad formation, we check for the register groups in reverse
     // order. This is to ensures that stack indexes remain consistent with the
     // instruction.
-    CSI[ZPRPositions[--NumUnassigned]] = CS;
+    for (const CalleeSavedInfo &CS : Saves)
+      CSI[ZPRPositions[--NumUnassigned]] = CS;
   };
-  //  With an odd offset, the Emit(single) branch will be reached first,
-  //  aligning to an even offset. Then the Emit(pair) branch will be reached,
-  //  aligning to an offset divisible by 4. The algorithm minimally decomposes
-  //  quads/pairs when needed.
-  int Offset = LowestOffset;
-  while (NumUnassigned != 0) {
-    if (Offset % 4 == 0 && !Quads.empty()) {
-      for (const auto &ZReg : Quads.pop_back_val())
-        Emit(ZReg);
-
-      Offset += 4;
-      continue;
-    }
 
-    if (Offset % 2 == 0) {
-      // Two singles can advance an even offset to the quad-aligned
-      // offset, preserve the quad in this case
-      if (Pairs.empty() && Singles.size() < 2 && !Quads.empty()) {
-        const auto &[a, b, c, d] = Quads.pop_back_val();
-        Pairs.push_back({a, b});
-        Pairs.push_back({c, d});
-      }
+  // Try to align the offset to mul 4 (to allow for the most quads/pairs).
+  if (NumToAlign & 0b01)
+    Emit(Singles.pop_back_val());
+  if (NumToAlign & 0b10) {
+    if (!Pairs.empty())
+      Emit(Pairs.pop_back_val());
+    else
+      Emit({Singles.pop_back_val(), Singles.pop_back_val()});
+  }
 
-      if (!Pairs.empty()) {
-        for (const auto &ZReg : Pairs.pop_back_val())
-          Emit(ZReg);
-        Offset += 2;
-        continue;
-      }
-    }
+  for (const auto &Quad : Quads)
----------------
kieroxide wrote:

We could reverse all Quads, Pairs and singles here to keep register spills from high to low (somewhat). But doesn't change anything but style.

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