[MIR] Support save/restore points with independent sets of registers (#119358)

This patch adds the MIR parsing and serialization support for save and
restore points with subsets of callee saved registers. That is, it
syntactically allows a function to contain two or more distinct
sub-regions in which distinct subsets of registers are spilled/filled as
callee save. This is useful if e.g. one of the CSRs isn't modified in
one of the sub-regions, but is in the other(s).

Support for actually using this capability in code generation is still
forthcoming. This patch is the next logical step for multiple
save/restore points support.

All points are now stored in DenseMap from MBB to vector of
CalleeSavedInfo.

Shrink-Wrap points split Part 4.
RFC:
https://discourse.llvm.org/t/shrink-wrap-save-restore-points-splitting/83581

Part 1: #117862 (landed)
Part 2: #119355 (landed)
Part 3: #119357 (landed)
Part 5: #119359 (likely to be
further split)

Author: enoskova-sc

llvm/include/llvm/CodeGen/MIRYamlMapping.h

@@ -634,19 +634,36 @@ LLVM_YAML_IS_SEQUENCE_VECTOR(llvm::yaml::CalledGlobal)
 namespace llvm {
 namespace yaml {
 
-// Struct representing one save/restore point in the 'savePoint'/'restorePoint'
-// list
+// Struct representing one save/restore point in the 'savePoint' /
+// 'restorePoint' list. One point consists of machine basic block name and list
+// of registers saved/restored in this basic block. In MIR it looks like:
+//  savePoint:
+//    - point:           '%bb.1'
+//      registers:
+//        - '$rbx'
+//        - '$r12'
+//        ...
+//  restorePoint:
+//    - point:           '%bb.1'
+//      registers:
+//        - '$rbx'
+//        - '$r12'
+// If no register is saved/restored in the selected BB,
+// field 'registers' is not specified.
 struct SaveRestorePointEntry {
   StringValue Point;
+  std::vector<StringValue> Registers;
 
   bool operator==(const SaveRestorePointEntry &Other) const {
-    return Point == Other.Point;
+    return Point == Other.Point && Registers == Other.Registers;
   }
 };
 
 template <> struct MappingTraits<SaveRestorePointEntry> {
   static void mapping(IO &YamlIO, SaveRestorePointEntry &Entry) {
     YamlIO.mapRequired("point", Entry.Point);
+    YamlIO.mapOptional("registers", Entry.Registers,
+                       std::vector<StringValue>());
   }
 };
 

llvm/include/llvm/CodeGen/MachineFrameInfo.h

@@ -76,6 +76,9 @@ class CalleeSavedInfo {
   bool isSpilledToReg()                    const { return SpilledToReg; }
 };
 
+using SaveRestorePoints =
+    DenseMap<MachineBasicBlock *, std::vector<CalleeSavedInfo>>;
+
 /// The MachineFrameInfo class represents an abstract stack frame until
 /// prolog/epilog code is inserted.  This class is key to allowing stack frame
 /// representation optimizations, such as frame pointer elimination.  It also
@@ -333,9 +336,9 @@ class MachineFrameInfo {
   bool HasTailCall = false;
 
   /// Not empty, if shrink-wrapping found a better place for the prologue.
-  SmallVector<MachineBasicBlock *, 4> SavePoints;
+  SaveRestorePoints SavePoints;
   /// Not empty, if shrink-wrapping found a better place for the epilogue.
-  SmallVector<MachineBasicBlock *, 4> RestorePoints;
+  SaveRestorePoints RestorePoints;
 
   /// Size of the UnsafeStack Frame
   uint64_t UnsafeStackSize = 0;
@@ -825,17 +828,21 @@ class MachineFrameInfo {
 
   void setCalleeSavedInfoValid(bool v) { CSIValid = v; }
 
-  ArrayRef<MachineBasicBlock *> getSavePoints() const { return SavePoints; }
-  void setSavePoints(ArrayRef<MachineBasicBlock *> NewSavePoints) {
-    SavePoints.assign(NewSavePoints.begin(), NewSavePoints.end());
-  }
-  ArrayRef<MachineBasicBlock *> getRestorePoints() const {
-    return RestorePoints;
+  const SaveRestorePoints &getRestorePoints() const { return RestorePoints; }
+
+  const SaveRestorePoints &getSavePoints() const { return SavePoints; }
+
+  void setSavePoints(SaveRestorePoints NewSavePoints) {
+    SavePoints = std::move(NewSavePoints);
   }
-  void setRestorePoints(ArrayRef<MachineBasicBlock *> NewRestorePoints) {
-    RestorePoints.assign(NewRestorePoints.begin(), NewRestorePoints.end());
+
+  void setRestorePoints(SaveRestorePoints NewRestorePoints) {
+    RestorePoints = std::move(NewRestorePoints);
   }
 
+  void clearSavePoints() { SavePoints.clear(); }
+  void clearRestorePoints() { RestorePoints.clear(); }
+
   uint64_t getUnsafeStackSize() const { return UnsafeStackSize; }
   void setUnsafeStackSize(uint64_t Size) { UnsafeStackSize = Size; }
 

llvm/lib/CodeGen/MIRParser/MIRParser.cpp

@@ -127,7 +127,7 @@ class MIRParserImpl {
   bool initializeSaveRestorePoints(
       PerFunctionMIParsingState &PFS,
       const std::vector<yaml::SaveRestorePointEntry> &YamlSRPoints,
-      SmallVectorImpl<MachineBasicBlock *> &SaveRestorePoints);
+      llvm::SaveRestorePoints &SaveRestorePoints);
 
   bool initializeCallSiteInfo(PerFunctionMIParsingState &PFS,
                               const yaml::MachineFunction &YamlMF);
@@ -872,11 +872,11 @@ bool MIRParserImpl::initializeFrameInfo(PerFunctionMIParsingState &PFS,
   MFI.setHasTailCall(YamlMFI.HasTailCall);
   MFI.setCalleeSavedInfoValid(YamlMFI.IsCalleeSavedInfoValid);
   MFI.setLocalFrameSize(YamlMFI.LocalFrameSize);
-  SmallVector<MachineBasicBlock *, 4> SavePoints;
+  llvm::SaveRestorePoints SavePoints;
   if (initializeSaveRestorePoints(PFS, YamlMFI.SavePoints, SavePoints))
     return true;
   MFI.setSavePoints(SavePoints);
-  SmallVector<MachineBasicBlock *, 4> RestorePoints;
+  llvm::SaveRestorePoints RestorePoints;
   if (initializeSaveRestorePoints(PFS, YamlMFI.RestorePoints, RestorePoints))
     return true;
   MFI.setRestorePoints(RestorePoints);
@@ -1098,14 +1098,22 @@ bool MIRParserImpl::initializeConstantPool(PerFunctionMIParsingState &PFS,
 bool MIRParserImpl::initializeSaveRestorePoints(
     PerFunctionMIParsingState &PFS,
     const std::vector<yaml::SaveRestorePointEntry> &YamlSRPoints,
-    SmallVectorImpl<MachineBasicBlock *> &SaveRestorePoints) {
+    llvm::SaveRestorePoints &SaveRestorePoints) {
+  SMDiagnostic Error;
   MachineBasicBlock *MBB = nullptr;
   for (const yaml::SaveRestorePointEntry &Entry : YamlSRPoints) {
     if (parseMBBReference(PFS, MBB, Entry.Point.Value))
       return true;
-    SaveRestorePoints.push_back(MBB);
-  }
 
+    std::vector<CalleeSavedInfo> Registers;
+    for (auto &RegStr : Entry.Registers) {
+      Register Reg;
+      if (parseNamedRegisterReference(PFS, Reg, RegStr.Value, Error))
+        return error(Error, RegStr.SourceRange);
+      Registers.push_back(CalleeSavedInfo(Reg));
+    }
+    SaveRestorePoints.try_emplace(MBB, std::move(Registers));
+  }
   return false;
 }
 

llvm/lib/CodeGen/MIRPrinter.cpp

@@ -149,11 +149,13 @@ static void convertMCP(yaml::MachineFunction &MF,
 static void convertMJTI(ModuleSlotTracker &MST, yaml::MachineJumpTable &YamlJTI,
                         const MachineJumpTableInfo &JTI);
 static void convertMFI(ModuleSlotTracker &MST, yaml::MachineFrameInfo &YamlMFI,
-                       const MachineFrameInfo &MFI);
+                       const MachineFrameInfo &MFI,
+                       const TargetRegisterInfo *TRI);
 static void
 convertSRPoints(ModuleSlotTracker &MST,
                 std::vector<yaml::SaveRestorePointEntry> &YamlSRPoints,
-                ArrayRef<MachineBasicBlock *> SaveRestorePoints);
+                const llvm::SaveRestorePoints &SRPoints,
+                const TargetRegisterInfo *TRI);
 static void convertStackObjects(yaml::MachineFunction &YMF,
                                 const MachineFunction &MF,
                                 ModuleSlotTracker &MST, MFPrintState &State);
@@ -204,7 +206,8 @@ static void printMF(raw_ostream &OS, const MachineModuleInfo &MMI,
   convertMRI(YamlMF, MF, MF.getRegInfo(), MF.getSubtarget().getRegisterInfo());
   MachineModuleSlotTracker &MST = State.MST;
   MST.incorporateFunction(MF.getFunction());
-  convertMFI(MST, YamlMF.FrameInfo, MF.getFrameInfo());
+  convertMFI(MST, YamlMF.FrameInfo, MF.getFrameInfo(),
+             MF.getSubtarget().getRegisterInfo());
   convertStackObjects(YamlMF, MF, MST, State);
   convertEntryValueObjects(YamlMF, MF, MST);
   convertCallSiteObjects(YamlMF, MF, MST);
@@ -339,7 +342,8 @@ static void convertMRI(yaml::MachineFunction &YamlMF, const MachineFunction &MF,
 }
 
 static void convertMFI(ModuleSlotTracker &MST, yaml::MachineFrameInfo &YamlMFI,
-                       const MachineFrameInfo &MFI) {
+                       const MachineFrameInfo &MFI,
+                       const TargetRegisterInfo *TRI) {
   YamlMFI.IsFrameAddressTaken = MFI.isFrameAddressTaken();
   YamlMFI.IsReturnAddressTaken = MFI.isReturnAddressTaken();
   YamlMFI.HasStackMap = MFI.hasStackMap();
@@ -360,9 +364,9 @@ static void convertMFI(ModuleSlotTracker &MST, yaml::MachineFrameInfo &YamlMFI,
   YamlMFI.IsCalleeSavedInfoValid = MFI.isCalleeSavedInfoValid();
   YamlMFI.LocalFrameSize = MFI.getLocalFrameSize();
   if (!MFI.getSavePoints().empty())
-    convertSRPoints(MST, YamlMFI.SavePoints, MFI.getSavePoints());
+    convertSRPoints(MST, YamlMFI.SavePoints, MFI.getSavePoints(), TRI);
   if (!MFI.getRestorePoints().empty())
-    convertSRPoints(MST, YamlMFI.RestorePoints, MFI.getRestorePoints());
+    convertSRPoints(MST, YamlMFI.RestorePoints, MFI.getRestorePoints(), TRI);
 }
 
 static void convertEntryValueObjects(yaml::MachineFunction &YMF,
@@ -619,16 +623,35 @@ static void convertMCP(yaml::MachineFunction &MF,
 static void
 convertSRPoints(ModuleSlotTracker &MST,
                 std::vector<yaml::SaveRestorePointEntry> &YamlSRPoints,
-                ArrayRef<MachineBasicBlock *> SRPoints) {
-  for (const auto &MBB : SRPoints) {
+                const llvm::SaveRestorePoints &SRPoints,
+                const TargetRegisterInfo *TRI) {
+  for (const auto &[MBB, CSInfos] : SRPoints) {
     SmallString<16> Str;
     yaml::SaveRestorePointEntry Entry;
     raw_svector_ostream StrOS(Str);
     StrOS << printMBBReference(*MBB);
     Entry.Point = StrOS.str().str();
     Str.clear();
+    for (const CalleeSavedInfo &Info : CSInfos) {
+      if (Info.getReg()) {
+        StrOS << printReg(Info.getReg(), TRI);
+        Entry.Registers.push_back(StrOS.str().str());
+        Str.clear();
+      }
+    }
+    // Sort here needed for stable output for lit tests
+    std::sort(Entry.Registers.begin(), Entry.Registers.end(),
+              [](const yaml::StringValue &Lhs, const yaml::StringValue &Rhs) {
+                return Lhs.Value < Rhs.Value;
+              });
     YamlSRPoints.push_back(std::move(Entry));
   }
+  // Sort here needed for stable output for lit tests
+  std::sort(YamlSRPoints.begin(), YamlSRPoints.end(),
+            [](const yaml::SaveRestorePointEntry &Lhs,
+               const yaml::SaveRestorePointEntry &Rhs) {
+              return Lhs.Point.Value < Rhs.Point.Value;
+            });
 }
 
 static void convertMJTI(ModuleSlotTracker &MST, yaml::MachineJumpTable &YamlJTI,

llvm/lib/CodeGen/MachineFrameInfo.cpp

@@ -250,14 +250,14 @@ void MachineFrameInfo::print(const MachineFunction &MF, raw_ostream &OS) const{
     OS << "save points:\n";
 
     for (auto &item : SavePoints)
-      OS << printMBBReference(*item) << "\n";
+      OS << printMBBReference(*item.first) << "\n";
   } else
     OS << "save points are empty\n";
 
   if (!RestorePoints.empty()) {
     OS << "restore points:\n";
     for (auto &item : RestorePoints)
-      OS << printMBBReference(*item) << "\n";
+      OS << printMBBReference(*item.first) << "\n";
   } else
     OS << "restore points are empty\n";
 }

llvm/lib/CodeGen/PrologEpilogInserter.cpp

@@ -351,8 +351,8 @@ bool PEIImpl::run(MachineFunction &MF) {
   delete RS;
   SaveBlocks.clear();
   RestoreBlocks.clear();
-  MFI.setSavePoints({});
-  MFI.setRestorePoints({});
+  MFI.clearSavePoints();
+  MFI.clearRestorePoints();
   return true;
 }
 
@@ -431,10 +431,12 @@ void PEIImpl::calculateSaveRestoreBlocks(MachineFunction &MF) {
   if (!MFI.getSavePoints().empty()) {
     assert(MFI.getSavePoints().size() == 1 &&
            "Multiple save points are not yet supported!");
-    SaveBlocks.push_back(MFI.getSavePoints().front());
+    const auto &SavePoint = *MFI.getSavePoints().begin();
+    SaveBlocks.push_back(SavePoint.first);
     assert(MFI.getRestorePoints().size() == 1 &&
            "Multiple restore points are not yet supported!");
-    MachineBasicBlock *RestoreBlock = MFI.getRestorePoints().front();
+    const auto &RestorePoint = *MFI.getRestorePoints().begin();
+    MachineBasicBlock *RestoreBlock = RestorePoint.first;
     // If RestoreBlock does not have any successor and is not a return block
     // then the end point is unreachable and we do not need to insert any
     // epilogue.
@@ -563,8 +565,9 @@ static void updateLiveness(MachineFunction &MF) {
 
   assert(MFI.getSavePoints().size() < 2 &&
          "Multiple save points not yet supported!");
-  MachineBasicBlock *Save =
-      MFI.getSavePoints().empty() ? nullptr : MFI.getSavePoints().front();
+  MachineBasicBlock *Save = MFI.getSavePoints().empty()
+                                ? nullptr
+                                : (*MFI.getSavePoints().begin()).first;
 
   if (!Save)
     Save = Entry;
@@ -577,8 +580,9 @@ static void updateLiveness(MachineFunction &MF) {
 
   assert(MFI.getRestorePoints().size() < 2 &&
          "Multiple restore points not yet supported!");
-  MachineBasicBlock *Restore =
-      MFI.getRestorePoints().empty() ? nullptr : MFI.getRestorePoints().front();
+  MachineBasicBlock *Restore = MFI.getRestorePoints().empty()
+                                   ? nullptr
+                                   : (*MFI.getRestorePoints().begin()).first;
   if (Restore)
     // By construction Restore cannot be visited, otherwise it
     // means there exists a path to Restore that does not go
@@ -687,6 +691,20 @@ void PEIImpl::spillCalleeSavedRegs(MachineFunction &MF) {
     MFI.setCalleeSavedInfoValid(true);
 
     std::vector<CalleeSavedInfo> &CSI = MFI.getCalleeSavedInfo();
+
+    // Fill SavePoints and RestorePoints with CalleeSavedRegisters
+    if (!MFI.getSavePoints().empty()) {
+      SaveRestorePoints SaveRestorePts;
+      for (const auto &SavePoint : MFI.getSavePoints())
+        SaveRestorePts.insert({SavePoint.first, CSI});
+      MFI.setSavePoints(std::move(SaveRestorePts));
+
+      SaveRestorePts.clear();
+      for (const auto &RestorePoint : MFI.getRestorePoints())
+        SaveRestorePts.insert({RestorePoint.first, CSI});
+      MFI.setRestorePoints(std::move(SaveRestorePts));
+    }
+
     if (!CSI.empty()) {
       if (!MFI.hasCalls())
         NumLeafFuncWithSpills++;

llvm/lib/CodeGen/ShrinkWrap.cpp

@@ -967,12 +967,12 @@ bool ShrinkWrapImpl::run(MachineFunction &MF) {
                     << "\nRestore: " << printMBBReference(*Restore) << '\n');
 
   MachineFrameInfo &MFI = MF.getFrameInfo();
-  SmallVector<MachineBasicBlock *, 4> SavePoints;
-  SmallVector<MachineBasicBlock *, 4> RestorePoints;
-  if (Save) {
-    SavePoints.push_back(Save);
-    RestorePoints.push_back(Restore);
-  }
+
+  // List of CalleeSavedInfo for registers will be added during prologepilog
+  // pass
+  SaveRestorePoints SavePoints({{Save, {}}});
+  SaveRestorePoints RestorePoints({{Restore, {}}});
+
   MFI.setSavePoints(SavePoints);
   MFI.setRestorePoints(RestorePoints);
   ++NumCandidates;

llvm/lib/Target/AMDGPU/SILowerSGPRSpills.cpp

@@ -213,10 +213,12 @@ void SILowerSGPRSpills::calculateSaveRestoreBlocks(MachineFunction &MF) {
   if (!MFI.getSavePoints().empty()) {
     assert(MFI.getSavePoints().size() == 1 &&
            "Multiple save points not yet supported!");
-    SaveBlocks.push_back(MFI.getSavePoints().front());
+    const auto &SavePoint = *MFI.getSavePoints().begin();
+    SaveBlocks.push_back(SavePoint.first);
     assert(MFI.getRestorePoints().size() == 1 &&
            "Multiple restore points not yet supported!");
-    MachineBasicBlock *RestoreBlock = MFI.getRestorePoints().front();
+    const auto &RestorePoint = *MFI.getRestorePoints().begin();
+    MachineBasicBlock *RestoreBlock = RestorePoint.first;
     // If RestoreBlock does not have any successor and is not a return block
     // then the end point is unreachable and we do not need to insert any
     // epilogue.

llvm/lib/Target/PowerPC/PPCFrameLowering.cpp

@@ -2081,9 +2081,8 @@ void PPCFrameLowering::processFunctionBeforeFrameFinalized(MachineFunction &MF,
   if (!MFI.getSavePoints().empty() && MFI.hasTailCall()) {
     assert(MFI.getRestorePoints().size() < 2 &&
            "MFI can't contain multiple restore points!");
-    MachineBasicBlock *RestoreBlock = MFI.getRestorePoints().front();
     for (MachineBasicBlock &MBB : MF) {
-      if (MBB.isReturnBlock() && (&MBB) != RestoreBlock)
+      if (MBB.isReturnBlock() && (!MFI.getRestorePoints().contains(&MBB)))
         createTailCallBranchInstr(MBB);
     }
   }

llvm/test/CodeGen/MIR/X86/frame-info-multiple-save-restore-points-parse.mir

@@ -32,10 +32,14 @@ liveins:
 # CHECK: frameInfo:
 # CHECK:      savePoint:
 # CHECK-NEXT:   - point:           '%bb.1'
+# CHECK-NEXT:     registers: []
 # CHECK-NEXT:   - point:           '%bb.2'
+# CHECK-NEXT:     registers: []
 # CHECK:      restorePoint:
 # CHECK-NEXT:   - point:           '%bb.2'
+# CHECK-NEXT:     registers: []
 # CHECK-NEXT:   - point:           '%bb.3'
+# CHECK-NEXT:     registers: []
 # CHECK: stack
 frameInfo:
   maxAlignment:  4