Recommit "Generate Callee Saved Register (CSR) related cfi directives
like .cfi_restore"

Insert .cfi_offset/.cfi_register when IncomingCSRSaved of current block
is larger than OutgoingCSRSaved of its previous block.

Original commit message:
https://reviews.llvm.org/D42848 only handled CFA related cfi directives but
didn't handle CSR related cfi. The patch adds the CSR part. Basically it reuses
the framework created in D42848. For each basicblock, the patch tracks which
CSR set have been saved at its CFG predecessors's exits, and compare the CSR
set with the set at its previous basicblock's exit (The previous block is the
block laid before the current block). If the saved CSR set at its previous
basicblock's exit is larger, .cfi_restore will be inserted.

The patch also generates proper .cfi_restore in epilogue to make sure the
saved CSR set is consistent for the incoming edges of each block.

Differential Revision: https://reviews.llvm.org/D74303
diff --git a/llvm/lib/CodeGen/CFIInstrInserter.cpp b/llvm/lib/CodeGen/CFIInstrInserter.cpp
index ef548c8..0269936 100644
--- a/llvm/lib/CodeGen/CFIInstrInserter.cpp
+++ b/llvm/lib/CodeGen/CFIInstrInserter.cpp
@@ -18,6 +18,8 @@
 //===----------------------------------------------------------------------===//
 
 #include "llvm/ADT/DepthFirstIterator.h"
+#include "llvm/ADT/Optional.h"
+#include "llvm/ADT/SetOperations.h"
 #include "llvm/CodeGen/MachineFunctionPass.h"
 #include "llvm/CodeGen/MachineInstrBuilder.h"
 #include "llvm/CodeGen/MachineModuleInfo.h"
@@ -76,15 +78,32 @@
     unsigned IncomingCFARegister = 0;
     /// Value of cfa register valid at basic block exit.
     unsigned OutgoingCFARegister = 0;
+    /// Set of callee saved registers saved at basic block entry.
+    BitVector IncomingCSRSaved;
+    /// Set of callee saved registers saved at basic block exit.
+    BitVector OutgoingCSRSaved;
     /// If in/out cfa offset and register values for this block have already
     /// been set or not.
     bool Processed = false;
   };
 
+#define INVALID_REG UINT_MAX
+#define INVALID_OFFSET INT_MAX
+  /// contains the location where CSR register is saved.
+  struct CSRSavedLocation {
+    CSRSavedLocation(Optional<unsigned> R, Optional<int> O)
+        : Reg(R), Offset(O) {}
+    Optional<unsigned> Reg;
+    Optional<int> Offset;
+  };
+
   /// Contains cfa offset and register values valid at entry and exit of basic
   /// blocks.
   std::vector<MBBCFAInfo> MBBVector;
 
+  /// Map the callee save registers to the locations where they are saved.
+  SmallDenseMap<unsigned, CSRSavedLocation, 16> CSRLocMap;
+
   /// Calculate cfa offset and register values valid at entry and exit for all
   /// basic blocks in a function.
   void calculateCFAInfo(MachineFunction &MF);
@@ -108,7 +127,8 @@
     return -MBBVector[MBB->getNumber()].IncomingCFAOffset;
   }
 
-  void report(const MBBCFAInfo &Pred, const MBBCFAInfo &Succ);
+  void reportCFAError(const MBBCFAInfo &Pred, const MBBCFAInfo &Succ);
+  void reportCSRError(const MBBCFAInfo &Pred, const MBBCFAInfo &Succ);
   /// Go through each MBB in a function and check that outgoing offset and
   /// register of its predecessors match incoming offset and register of that
   /// MBB, as well as that incoming offset and register of its successors match
@@ -132,6 +152,8 @@
   // function.
   unsigned InitialRegister =
       MF.getSubtarget().getFrameLowering()->getInitialCFARegister(MF);
+  const TargetRegisterInfo &TRI = *MF.getSubtarget().getRegisterInfo();
+  unsigned NumRegs = TRI.getNumRegs();
 
   // Initialize MBBMap.
   for (MachineBasicBlock &MBB : MF) {
@@ -141,8 +163,11 @@
     MBBInfo.OutgoingCFAOffset = InitialOffset;
     MBBInfo.IncomingCFARegister = InitialRegister;
     MBBInfo.OutgoingCFARegister = InitialRegister;
+    MBBInfo.IncomingCSRSaved.resize(NumRegs);
+    MBBInfo.OutgoingCSRSaved.resize(NumRegs);
     MBBVector[MBB.getNumber()] = MBBInfo;
   }
+  CSRLocMap.clear();
 
   // Set in/out cfa info for all blocks in the function. This traversal is based
   // on the assumption that the first block in the function is the entry block
@@ -159,12 +184,17 @@
   int SetOffset = MBBInfo.IncomingCFAOffset;
   // Outgoing cfa register set by the block.
   unsigned SetRegister = MBBInfo.IncomingCFARegister;
-  const std::vector<MCCFIInstruction> &Instrs =
-      MBBInfo.MBB->getParent()->getFrameInstructions();
+  MachineFunction *MF = MBBInfo.MBB->getParent();
+  const std::vector<MCCFIInstruction> &Instrs = MF->getFrameInstructions();
+  const TargetRegisterInfo &TRI = *MF->getSubtarget().getRegisterInfo();
+  unsigned NumRegs = TRI.getNumRegs();
+  BitVector CSRSaved(NumRegs), CSRRestored(NumRegs);
 
   // Determine cfa offset and register set by the block.
   for (MachineInstr &MI : *MBBInfo.MBB) {
     if (MI.isCFIInstruction()) {
+      Optional<unsigned> CSRReg;
+      Optional<int> CSROffset;
       unsigned CFIIndex = MI.getOperand(0).getCFIIndex();
       const MCCFIInstruction &CFI = Instrs[CFIIndex];
       switch (CFI.getOperation()) {
@@ -181,6 +211,18 @@
         SetRegister = CFI.getRegister();
         SetOffset = CFI.getOffset();
         break;
+      case MCCFIInstruction::OpOffset:
+        CSROffset = CFI.getOffset();
+        break;
+      case MCCFIInstruction::OpRegister:
+        CSRReg = CFI.getRegister2();
+        break;
+      case MCCFIInstruction::OpRelOffset:
+        CSROffset = CFI.getOffset() - SetOffset;
+        break;
+      case MCCFIInstruction::OpRestore:
+        CSRRestored.set(CFI.getRegister());
+        break;
       case MCCFIInstruction::OpRememberState:
         // TODO: Add support for handling cfi_remember_state.
 #ifndef NDEBUG
@@ -198,18 +240,24 @@
 #endif
         break;
       // Other CFI directives do not affect CFA value.
-      case MCCFIInstruction::OpSameValue:
-      case MCCFIInstruction::OpOffset:
-      case MCCFIInstruction::OpRelOffset:
-      case MCCFIInstruction::OpEscape:
-      case MCCFIInstruction::OpRestore:
       case MCCFIInstruction::OpUndefined:
-      case MCCFIInstruction::OpRegister:
+      case MCCFIInstruction::OpSameValue:
+      case MCCFIInstruction::OpEscape:
       case MCCFIInstruction::OpWindowSave:
       case MCCFIInstruction::OpNegateRAState:
       case MCCFIInstruction::OpGnuArgsSize:
         break;
       }
+      if (CSRReg || CSROffset) {
+        auto It = CSRLocMap.find(CFI.getRegister());
+        if (It == CSRLocMap.end()) {
+          CSRLocMap.insert(
+              {CFI.getRegister(), CSRSavedLocation(CSRReg, CSROffset)});
+        } else if (It->second.Reg != CSRReg || It->second.Offset != CSROffset) {
+          llvm_unreachable("Different saved locations for the same CSR");
+        }
+        CSRSaved.set(CFI.getRegister());
+      }
     }
   }
 
@@ -218,6 +266,11 @@
   // Update outgoing CFA info.
   MBBInfo.OutgoingCFAOffset = SetOffset;
   MBBInfo.OutgoingCFARegister = SetRegister;
+
+  // Update outgoing CSR info.
+  MBBInfo.OutgoingCSRSaved = MBBInfo.IncomingCSRSaved;
+  MBBInfo.OutgoingCSRSaved |= CSRSaved;
+  MBBInfo.OutgoingCSRSaved.reset(CSRRestored);
 }
 
 void CFIInstrInserter::updateSuccCFAInfo(MBBCFAInfo &MBBInfo) {
@@ -236,6 +289,7 @@
       if (!SuccInfo.Processed) {
         SuccInfo.IncomingCFAOffset = CurrentInfo.OutgoingCFAOffset;
         SuccInfo.IncomingCFARegister = CurrentInfo.OutgoingCFARegister;
+        SuccInfo.IncomingCSRSaved = CurrentInfo.OutgoingCSRSaved;
         Stack.push_back(Succ);
       }
     }
@@ -287,12 +341,45 @@
           .addCFIIndex(CFIIndex);
       InsertedCFIInstr = true;
     }
+
+    BitVector SetDifference = PrevMBBInfo->OutgoingCSRSaved;
+    SetDifference.reset(MBBInfo.IncomingCSRSaved);
+    for (int Reg : SetDifference.set_bits()) {
+      unsigned CFIIndex =
+          MF.addFrameInst(MCCFIInstruction::createRestore(nullptr, Reg));
+      BuildMI(*MBBInfo.MBB, MBBI, DL, TII->get(TargetOpcode::CFI_INSTRUCTION))
+          .addCFIIndex(CFIIndex);
+      InsertedCFIInstr = true;
+    }
+
+    SetDifference = MBBInfo.IncomingCSRSaved;
+    SetDifference.reset(PrevMBBInfo->OutgoingCSRSaved);
+    for (int Reg : SetDifference.set_bits()) {
+      auto it = CSRLocMap.find(Reg);
+      assert(it != CSRLocMap.end() && "Reg should have an entry in CSRLocMap");
+      unsigned CFIIndex;
+      CSRSavedLocation RO = it->second;
+      if (!RO.Reg && RO.Offset) {
+        CFIIndex = MF.addFrameInst(
+            MCCFIInstruction::createOffset(nullptr, Reg, *RO.Offset));
+      } else if (RO.Reg && !RO.Offset) {
+        CFIIndex = MF.addFrameInst(
+            MCCFIInstruction::createRegister(nullptr, Reg, *RO.Reg));
+      } else {
+        llvm_unreachable("RO.Reg and RO.Offset cannot both be valid/invalid");
+      }
+      BuildMI(*MBBInfo.MBB, MBBI, DL, TII->get(TargetOpcode::CFI_INSTRUCTION))
+          .addCFIIndex(CFIIndex);
+      InsertedCFIInstr = true;
+    }
+
     PrevMBBInfo = &MBBInfo;
   }
   return InsertedCFIInstr;
 }
 
-void CFIInstrInserter::report(const MBBCFAInfo &Pred, const MBBCFAInfo &Succ) {
+void CFIInstrInserter::reportCFAError(const MBBCFAInfo &Pred,
+                                      const MBBCFAInfo &Succ) {
   errs() << "*** Inconsistent CFA register and/or offset between pred and succ "
             "***\n";
   errs() << "Pred: " << Pred.MBB->getName() << " #" << Pred.MBB->getNumber()
@@ -307,6 +394,22 @@
          << " incoming CFA Offset:" << Succ.IncomingCFAOffset << "\n";
 }
 
+void CFIInstrInserter::reportCSRError(const MBBCFAInfo &Pred,
+                                      const MBBCFAInfo &Succ) {
+  errs() << "*** Inconsistent CSR Saved between pred and succ in function "
+         << Pred.MBB->getParent()->getName() << " ***\n";
+  errs() << "Pred: " << Pred.MBB->getName() << " #" << Pred.MBB->getNumber()
+         << " outgoing CSR Saved: ";
+  for (int Reg : Pred.OutgoingCSRSaved.set_bits())
+    errs() << Reg << " ";
+  errs() << "\n";
+  errs() << "Succ: " << Succ.MBB->getName() << " #" << Succ.MBB->getNumber()
+         << " incoming CSR Saved: ";
+  for (int Reg : Succ.IncomingCSRSaved.set_bits())
+    errs() << Reg << " ";
+  errs() << "\n";
+}
+
 unsigned CFIInstrInserter::verify(MachineFunction &MF) {
   unsigned ErrorNum = 0;
   for (auto *CurrMBB : depth_first(&MF)) {
@@ -321,7 +424,13 @@
         // we don't generate epilogues inside such blocks.
         if (SuccMBBInfo.MBB->succ_empty() && !SuccMBBInfo.MBB->isReturnBlock())
           continue;
-        report(CurrMBBInfo, SuccMBBInfo);
+        reportCFAError(CurrMBBInfo, SuccMBBInfo);
+        ErrorNum++;
+      }
+      // Check that IncomingCSRSaved of every successor matches the
+      // OutgoingCSRSaved of CurrMBB
+      if (SuccMBBInfo.IncomingCSRSaved != CurrMBBInfo.OutgoingCSRSaved) {
+        reportCSRError(CurrMBBInfo, SuccMBBInfo);
         ErrorNum++;
       }
     }