[llvm] Support multiple save/restore points in mir

llvm/include/llvm/CodeGen/MIRYamlMapping.h

@@ -631,6 +631,57 @@ LLVM_YAML_IS_SEQUENCE_VECTOR(llvm::yaml::CalledGlobal)
 namespace llvm {
 namespace yaml {
 
+// Struct representing one save/restore point in the 'savePoint'/'restorePoint'
+// list
+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 +690,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 +725,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 +747,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 +780,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

@@ -27,6 +27,21 @@ class MachineBasicBlock;
 class BitVector;
 class AllocaInst;
 
+using SaveRestorePoints = DenseMap<MachineBasicBlock *, std::vector<Register>>;
+
+class CalleeSavedInfoPerBB {
+  DenseMap<MachineBasicBlock *, std::vector<CalleeSavedInfo>> Map;
+
+public:
+  std::vector<CalleeSavedInfo> get(MachineBasicBlock *MBB) const {
+    return Map.lookup(MBB);
+  }
+
+  void set(DenseMap<MachineBasicBlock *, std::vector<CalleeSavedInfo>> CSI) {
+    Map = std::move(CSI);
+  }
+};
+
 /// The CalleeSavedInfo class tracks the information need to locate where a
 /// callee saved register is in the current frame.
 /// Callee saved reg can also be saved to a different register rather than
@@ -37,6 +52,8 @@ class CalleeSavedInfo {
     int FrameIdx;
     unsigned DstReg;
   };
+  std::vector<MachineBasicBlock *> SpilledIn;
+  std::vector<MachineBasicBlock *> RestoredIn;
   /// Flag indicating whether the register is actually restored in the epilog.
   /// In most cases, if a register is saved, it is also restored. There are
   /// some situations, though, when this is not the case. For example, the
@@ -58,9 +75,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;
@@ -72,6 +89,16 @@ class CalleeSavedInfo {
   bool isRestored()                        const { return Restored; }
   void setRestored(bool R)                       { Restored = R; }
   bool isSpilledToReg()                    const { return SpilledToReg; }
+  ArrayRef<MachineBasicBlock *> spilledIn() const { return SpilledIn; }
+  ArrayRef<MachineBasicBlock *> restoredIn() const { return RestoredIn; }
+  void addSpilledIn(MachineBasicBlock *MBB) { SpilledIn.push_back(MBB); }
+  void addRestoredIn(MachineBasicBlock *MBB) { RestoredIn.push_back(MBB); }
+  void setSpilledIn(std::vector<MachineBasicBlock *> BBV) {
+    SpilledIn = std::move(BBV);
+  }
+  void setRestoredIn(std::vector<MachineBasicBlock *> BBV) {
+    RestoredIn = std::move(BBV);
+  }
 };
 
 /// The MachineFrameInfo class represents an abstract stack frame until
@@ -295,6 +322,10 @@ class MachineFrameInfo {
   /// Has CSInfo been set yet?
   bool CSIValid = false;
 
+  CalleeSavedInfoPerBB CSInfoPerSave;
+
+  CalleeSavedInfoPerBB CSInfoPerRestore;
+
   /// References to frame indices which are mapped
   /// into the local frame allocation block. <FrameIdx, LocalOffset>
   SmallVector<std::pair<int, int64_t>, 32> LocalFrameObjects;
@@ -331,9 +362,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,21 +847,99 @@ class MachineFrameInfo {
   /// \copydoc getCalleeSavedInfo()
   std::vector<CalleeSavedInfo> &getCalleeSavedInfo() { return CSInfo; }
 
+  /// Returns callee saved info vector for provided save point in
+  /// the current function.
+  std::vector<CalleeSavedInfo> getCSInfoPerSave(MachineBasicBlock *MBB) const {
+    return CSInfoPerSave.get(MBB);
+  }
+
+  /// Returns callee saved info vector for provided restore point
+  /// in the current function.
+  std::vector<CalleeSavedInfo>
+  getCSInfoPerRestore(MachineBasicBlock *MBB) const {
+    return CSInfoPerRestore.get(MBB);
+  }
+
   /// Used by prolog/epilog inserter to set the function's callee saved
   /// information.
   void setCalleeSavedInfo(std::vector<CalleeSavedInfo> CSI) {
     CSInfo = std::move(CSI);
   }
 
+  /// Used by prolog/epilog inserter to set the function's callee saved
+  /// information for particular save point.
+  void setCSInfoPerSave(
+      DenseMap<MachineBasicBlock *, std::vector<CalleeSavedInfo>> CSI) {
+    CSInfoPerSave.set(CSI);
+  }
+
+  /// Used by prolog/epilog inserter to set the function's callee saved
+  /// information for particular restore point.
+  void setCSInfoPerRestore(
+      DenseMap<MachineBasicBlock *, std::vector<CalleeSavedInfo>> CSI) {
+    CSInfoPerRestore.set(CSI);
+  }
+
   /// Has the callee saved info been calculated yet?
   bool isCalleeSavedInfoValid() const { return CSIValid; }
 
   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 &getRestorePoints() const { return RestorePoints; }
+
+  const SaveRestorePoints &getSavePoints() const { return SavePoints; }
+
+  std::pair<MachineBasicBlock *, std::vector<Register>>
+  getRestorePoint(MachineBasicBlock *MBB) const {
+    if (auto It = RestorePoints.find(MBB); It != RestorePoints.end())
+      return *It;
+
+    std::vector<Register> Regs = {};
+    return std::make_pair(nullptr, Regs);
+  }
+
+  std::pair<MachineBasicBlock *, std::vector<Register>>
+  getSavePoint(MachineBasicBlock *MBB) const {
+    if (auto It = SavePoints.find(MBB); It != SavePoints.end())
+      return *It;
+
+    std::vector<Register> Regs = {};
+    return std::make_pair(nullptr, Regs);
+  }
+
+  void setSavePoints(SaveRestorePoints NewSavePoints) {
+    SavePoints = std::move(NewSavePoints);
+  }
+
+  void setRestorePoints(SaveRestorePoints NewRestorePoints) {
+    RestorePoints = std::move(NewRestorePoints);
+  }
+
+  void setSavePoint(MachineBasicBlock *MBB, const std::vector<Register> &Regs) {
+    SavePoints[MBB] = Regs;
+  }
+
+  static const SaveRestorePoints constructSaveRestorePoints(
+      const SaveRestorePoints &SRP,
+      const DenseMap<MachineBasicBlock *, MachineBasicBlock *> &BBMap) {
+    SaveRestorePoints Pts{};
+    for (auto &Src : SRP) {
+      Pts.insert(std::make_pair(BBMap.find(Src.first)->second, Src.second));
+    }
+    return Pts;
+  }
+
+  void setRestorePoint(MachineBasicBlock *MBB, std::vector<Register> &Regs) {
+    RestorePoints[MBB] = Regs;
+  }
+
+  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,
+                                  true /*IsSavePoints*/))
+    return true;
+  if (initializeSaveRestorePoints(PFS, YamlMFI.RestorePoints,
+                                  false /*IsSavePoints*/))
+    return true;
 
   std::vector<CalleeSavedInfo> CSIInfo;
   // Initialize the fixed frame objects.
@@ -1077,6 +1075,54 @@ 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 SRPoints;
+  std::vector<Register> Registers;
+
+  if (std::holds_alternative<std::vector<yaml::SaveRestorePointEntry>>(
+          YamlSRPoints)) {
+    const auto &VectorRepr =
+        std::get<std::vector<yaml::SaveRestorePointEntry>>(YamlSRPoints);
+    if (VectorRepr.empty())
+      return false;
+    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(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, Registers));
+  }
+
+  MachineFunction &MF = PFS.MF;
+  MachineFrameInfo &MFI = MF.getFrameInfo();
+
+  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);