[llvm] Support multiple save/restore points in mir

llvm/include/llvm/CodeGen/MIRYamlMapping.h

@@ -631,6 +631,79 @@ LLVM_YAML_IS_SEQUENCE_VECTOR(llvm::yaml::CalledGlobal)
 namespace llvm {
 namespace yaml {
 
+// Struct representing one save/restore point in the
+// 'savePoint' / 'restorePoint' list. One point consists of machine basic block
+// name and list of saved/restored in this basic block registers. There are
+// two forms of Save/Restore point representation:
+// 1. Without explicit register enumeration:
+//      savePoint:       '%bb.n'
+//      restorePoint:    '%bb.n'
+// in this case we assume that all CalleeSavedRegisters
+// are splilled/restored in these points
+// 2. With explicit register:
+//  savePoint:
+//    - point:           '%bb.1'
+//      registers:
+//        - '$rbx'
+//        - '$r12'
+//        ...
+//  restorePoint:
+//    - point:           '%bb.1'
+//      registers:
+//        - '$rbx'
+//        - '$r12'
+// If this representation form is used and no register is saved/restored in the
+// selected BB, the empty list of register should be specified ( i.e. registers:
+// [])
+struct SaveRestorePointEntry {
+  StringValue Point;
+  std::vector<StringValue> Registers;
+
+  bool operator==(const SaveRestorePointEntry &Other) const {
+    return Point == Other.Point && Registers == Other.Registers;
+  }
+};
+
+using SaveRestorePoints =
+    std::variant<std::vector<SaveRestorePointEntry>, StringValue>;
+
+template <> struct PolymorphicTraits<SaveRestorePoints> {
+
+  static NodeKind getKind(const SaveRestorePoints &SRPoints) {
+    if (std::holds_alternative<std::vector<SaveRestorePointEntry>>(SRPoints))
+      return NodeKind::Sequence;
+    if (std::holds_alternative<StringValue>(SRPoints))
+      return NodeKind::Scalar;
+    llvm_unreachable("Unsupported NodeKind of SaveRestorePoints");
+  }
+
+  static SaveRestorePointEntry &getAsMap(SaveRestorePoints &SRPoints) {
+    llvm_unreachable("SaveRestorePoints can't be represented as Map");
+  }
+
+  static std::vector<SaveRestorePointEntry> &
+  getAsSequence(SaveRestorePoints &SRPoints) {
+    if (!std::holds_alternative<std::vector<SaveRestorePointEntry>>(SRPoints))
+      SRPoints = std::vector<SaveRestorePointEntry>();
+
+    return std::get<std::vector<SaveRestorePointEntry>>(SRPoints);
+  }
+
+  static StringValue &getAsScalar(SaveRestorePoints &SRPoints) {
+    if (!std::holds_alternative<StringValue>(SRPoints))
+      SRPoints = StringValue();
+
+    return std::get<StringValue>(SRPoints);
+  }
+};
+
+template <> struct MappingTraits<SaveRestorePointEntry> {
+  static void mapping(IO &YamlIO, SaveRestorePointEntry &Entry) {
+    YamlIO.mapRequired("point", Entry.Point);
+    YamlIO.mapRequired("registers", Entry.Registers);
+  }
+};
+
 template <> struct MappingTraits<MachineJumpTable> {
   static void mapping(IO &YamlIO, MachineJumpTable &JT) {
     YamlIO.mapRequired("kind", JT.Kind);
@@ -639,6 +712,14 @@ template <> struct MappingTraits<MachineJumpTable> {
   }
 };
 
+} // namespace yaml
+} // namespace llvm
+
+LLVM_YAML_IS_SEQUENCE_VECTOR(llvm::yaml::SaveRestorePointEntry)
+
+namespace llvm {
+namespace yaml {
+
 /// Serializable representation of MachineFrameInfo.
 ///
 /// Doesn't serialize attributes like 'StackAlignment', 'IsStackRealignable' and
@@ -666,8 +747,8 @@ struct MachineFrameInfo {
   bool HasTailCall = false;
   bool IsCalleeSavedInfoValid = false;
   unsigned LocalFrameSize = 0;
-  StringValue SavePoint;
-  StringValue RestorePoint;
+  SaveRestorePoints SavePoints;
+  SaveRestorePoints RestorePoints;
 
   bool operator==(const MachineFrameInfo &Other) const {
     return IsFrameAddressTaken == Other.IsFrameAddressTaken &&
@@ -688,7 +769,8 @@ struct MachineFrameInfo {
            HasMustTailInVarArgFunc == Other.HasMustTailInVarArgFunc &&
            HasTailCall == Other.HasTailCall &&
            LocalFrameSize == Other.LocalFrameSize &&
-           SavePoint == Other.SavePoint && RestorePoint == Other.RestorePoint &&
+           SavePoints == Other.SavePoints &&
+           RestorePoints == Other.RestorePoints &&
            IsCalleeSavedInfoValid == Other.IsCalleeSavedInfoValid;
   }
 };
@@ -720,10 +802,14 @@ template <> struct MappingTraits<MachineFrameInfo> {
     YamlIO.mapOptional("isCalleeSavedInfoValid", MFI.IsCalleeSavedInfoValid,
                        false);
     YamlIO.mapOptional("localFrameSize", MFI.LocalFrameSize, (unsigned)0);
-    YamlIO.mapOptional("savePoint", MFI.SavePoint,
-                       StringValue()); // Don't print it out when it's empty.
-    YamlIO.mapOptional("restorePoint", MFI.RestorePoint,
-                       StringValue()); // Don't print it out when it's empty.
+    YamlIO.mapOptional(
+        "savePoint", MFI.SavePoints,
+        SaveRestorePoints(
+            StringValue())); // Don't print it out when it's empty.
+    YamlIO.mapOptional(
+        "restorePoint", MFI.RestorePoints,
+        SaveRestorePoints(
+            StringValue())); // Don't print it out when it's empty.
   }
 };
 

llvm/include/llvm/CodeGen/MachineFrameInfo.h

@@ -58,9 +58,9 @@ class CalleeSavedInfo {
   explicit CalleeSavedInfo(unsigned R, int FI = 0) : Reg(R), FrameIdx(FI) {}
 
   // Accessors.
-  Register getReg()                        const { return Reg; }
-  int getFrameIdx()                        const { return FrameIdx; }
-  unsigned getDstReg()                     const { return DstReg; }
+  Register getReg() const { return Reg; }
+  int getFrameIdx() const { return FrameIdx; }
+  unsigned getDstReg() const { return DstReg; }
   void setFrameIdx(int FI) {
     FrameIdx = FI;
     SpilledToReg = false;
@@ -74,6 +74,36 @@ class CalleeSavedInfo {
   bool isSpilledToReg()                    const { return SpilledToReg; }
 };
 
+class SaveRestorePoints {
+public:
+  using PointsMap = DenseMap<MachineBasicBlock *, std::vector<CalleeSavedInfo>>;
+
+private:
+  PointsMap Map;
+
+public:
+  const PointsMap &get() const { return Map; }
+
+  const std::vector<CalleeSavedInfo> getCSInfo(MachineBasicBlock *MBB) const {
+    return Map.lookup(MBB);
+  }
+
+  void set(PointsMap &&CSI) { Map = std::move(CSI); }
+
+  MachineBasicBlock *findAny(const CalleeSavedInfo &Match) const {
+    for (auto [BB, CSIV] : Map) {
+      for (auto &CSI : CSIV) {
+        if (CSI.getReg() == Match.getReg())
+          return BB;
+      }
+    }
+    return nullptr;
+  }
+
+  void clear() { Map.clear(); }
+  bool empty() const { return Map.empty(); }
+};
+
 /// 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
@@ -331,9 +361,16 @@ class MachineFrameInfo {
   bool HasTailCall = false;
 
   /// Not null, if shrink-wrapping found a better place for the prologue.
-  MachineBasicBlock *Save = nullptr;
+  MachineBasicBlock *Prolog = nullptr;
   /// Not null, if shrink-wrapping found a better place for the epilogue.
-  MachineBasicBlock *Restore = nullptr;
+  MachineBasicBlock *Epilog = nullptr;
+
+  /// Not empty, if shrink-wrapping found a better place for saving callee
+  /// saves.
+  SaveRestorePoints SavePoints;
+  /// Not empty, if shrink-wrapping found a better place for restoring callee
+  /// saves.
+  SaveRestorePoints RestorePoints;
 
   /// Size of the UnsafeStack Frame
   uint64_t UnsafeStackSize = 0;
@@ -809,6 +846,28 @@ class MachineFrameInfo {
   /// \copydoc getCalleeSavedInfo()
   std::vector<CalleeSavedInfo> &getCalleeSavedInfo() { return CSInfo; }
 
+  /// Returns callee saved info vector for provided save point in
+  /// the current function.
+  const std::vector<CalleeSavedInfo>
+  getSaveCSInfo(MachineBasicBlock *MBB) const {
+    return SavePoints.getCSInfo(MBB);
+  }
+
+  /// Returns callee saved info vector for provided restore point
+  /// in the current function.
+  const std::vector<CalleeSavedInfo>
+  getRestoreCSInfo(MachineBasicBlock *MBB) const {
+    return RestorePoints.getCSInfo(MBB);
+  }
+
+  MachineBasicBlock *findSpilledIn(const CalleeSavedInfo &CSI) const {
+    return SavePoints.findAny(CSI);
+  }
+
+  MachineBasicBlock *findRestoredIn(const CalleeSavedInfo &CSI) const {
+    return RestorePoints.findAny(CSI);
+  }
+
   /// Used by prolog/epilog inserter to set the function's callee saved
   /// information.
   void setCalleeSavedInfo(std::vector<CalleeSavedInfo> CSI) {
@@ -820,10 +879,38 @@ class MachineFrameInfo {
 
   void setCalleeSavedInfoValid(bool v) { CSIValid = v; }
 
-  MachineBasicBlock *getSavePoint() const { return Save; }
-  void setSavePoint(MachineBasicBlock *NewSave) { Save = NewSave; }
-  MachineBasicBlock *getRestorePoint() const { return Restore; }
-  void setRestorePoint(MachineBasicBlock *NewRestore) { Restore = NewRestore; }
+  const SaveRestorePoints::PointsMap &getRestorePoints() const {
+    return RestorePoints.get();
+  }
+
+  const SaveRestorePoints::PointsMap &getSavePoints() const {
+    return SavePoints.get();
+  }
+
+  void setSavePoints(SaveRestorePoints::PointsMap NewSavePoints) {
+    SavePoints.set(std::move(NewSavePoints));
+  }
+
+  void setRestorePoints(SaveRestorePoints::PointsMap NewRestorePoints) {
+    RestorePoints.set(std::move(NewRestorePoints));
+  }
+
+  static const SaveRestorePoints::PointsMap constructSaveRestorePoints(
+      const SaveRestorePoints::PointsMap &SRPoints,
+      const DenseMap<MachineBasicBlock *, MachineBasicBlock *> &BBMap) {
+    SaveRestorePoints::PointsMap Pts{};
+    for (auto &Src : SRPoints)
+      Pts.insert({BBMap.find(Src.first)->second, Src.second});
+    return Pts;
+  }
+
+  MachineBasicBlock *getProlog() const { return Prolog; }
+  void setProlog(MachineBasicBlock *BB) { Prolog = BB; }
+  MachineBasicBlock *getEpilog() const { return Epilog; }
+  void setEpilog(MachineBasicBlock *BB) { Epilog = BB; }
+
+  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

@@ -124,6 +124,10 @@ class MIRParserImpl {
   bool initializeFrameInfo(PerFunctionMIParsingState &PFS,
                            const yaml::MachineFunction &YamlMF);
 
+  bool initializeSaveRestorePoints(PerFunctionMIParsingState &PFS,
+                                   const yaml::SaveRestorePoints &YamlSRPoints,
+                                   bool IsSavePoints);
+
   bool initializeCallSiteInfo(PerFunctionMIParsingState &PFS,
                               const yaml::MachineFunction &YamlMF);
 
@@ -851,18 +855,12 @@ bool MIRParserImpl::initializeFrameInfo(PerFunctionMIParsingState &PFS,
   MFI.setHasTailCall(YamlMFI.HasTailCall);
   MFI.setCalleeSavedInfoValid(YamlMFI.IsCalleeSavedInfoValid);
   MFI.setLocalFrameSize(YamlMFI.LocalFrameSize);
-  if (!YamlMFI.SavePoint.Value.empty()) {
-    MachineBasicBlock *MBB = nullptr;
-    if (parseMBBReference(PFS, MBB, YamlMFI.SavePoint))
-      return true;
-    MFI.setSavePoint(MBB);
-  }
-  if (!YamlMFI.RestorePoint.Value.empty()) {
-    MachineBasicBlock *MBB = nullptr;
-    if (parseMBBReference(PFS, MBB, YamlMFI.RestorePoint))
-      return true;
-    MFI.setRestorePoint(MBB);
-  }
+  if (initializeSaveRestorePoints(PFS, YamlMFI.SavePoints,
+                                  /*IsSavePoints=*/true))
+    return true;
+  if (initializeSaveRestorePoints(PFS, YamlMFI.RestorePoints,
+                                  /*IsSavePoints=*/false))
+    return true;
 
   std::vector<CalleeSavedInfo> CSIInfo;
   // Initialize the fixed frame objects.
@@ -1077,6 +1075,52 @@ bool MIRParserImpl::initializeConstantPool(PerFunctionMIParsingState &PFS,
   return false;
 }
 
+// Return true if basic block was incorrectly specified in MIR
+bool MIRParserImpl::initializeSaveRestorePoints(
+    PerFunctionMIParsingState &PFS, const yaml::SaveRestorePoints &YamlSRPoints,
+    bool IsSavePoints) {
+  SMDiagnostic Error;
+  MachineBasicBlock *MBB = nullptr;
+  llvm::SaveRestorePoints::PointsMap SRPoints;
+  MachineFunction &MF = PFS.MF;
+  MachineFrameInfo &MFI = MF.getFrameInfo();
+
+  if (std::holds_alternative<std::vector<yaml::SaveRestorePointEntry>>(
+          YamlSRPoints)) {
+    const auto &VectorRepr =
+        std::get<std::vector<yaml::SaveRestorePointEntry>>(YamlSRPoints);
+    if (VectorRepr.empty())
+      return false;
+    std::vector<CalleeSavedInfo> Registers;
+    for (const auto &Entry : VectorRepr) {
+      const auto &MBBSource = Entry.Point;
+      if (parseMBBReference(PFS, MBB, MBBSource.Value))
+        return true;
+      Registers.clear();
+      for (auto &RegStr : Entry.Registers) {
+        Register Reg;
+        if (parseNamedRegisterReference(PFS, Reg, RegStr.Value, Error))
+          return error(Error, RegStr.SourceRange);
+        Registers.push_back(CalleeSavedInfo(Reg));
+      }
+      SRPoints.insert(std::make_pair(MBB, Registers));
+    }
+  } else {
+    yaml::StringValue StringRepr = std::get<yaml::StringValue>(YamlSRPoints);
+    if (StringRepr.Value.empty())
+      return false;
+    if (parseMBBReference(PFS, MBB, StringRepr))
+      return true;
+    SRPoints.insert(std::make_pair(MBB, MFI.getCalleeSavedInfo()));
+  }
+
+  if (IsSavePoints)
+    MFI.setSavePoints(SRPoints);
+  else
+    MFI.setRestorePoints(SRPoints);
+  return false;
+}
+
 bool MIRParserImpl::initializeJumpTableInfo(PerFunctionMIParsingState &PFS,
     const yaml::MachineJumpTable &YamlJTI) {
   MachineJumpTableInfo *JTI = PFS.MF.getOrCreateJumpTableInfo(YamlJTI.Kind);