[CFI] cfi directive insertion for callee-save-registers in CFIInstrInserter pass

[wmi]

Summary

https://reviews.llvm.org/D42848 handled CFA related cfi directives but didn't handle callee-save-registers (csr) related cfi directives like .cfi_offset and .cfi_restore. This patch adds that 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.

It fixes the cfi information for the simple case below. The case is shrinkwrap optimized by llvm:

--------------- 1.cc ----------------
void goo();
  
long foo(bool cond, long *a) {
  if (__builtin_expect(cond, 1)) {
    goo();
    return a[0] + a[1] + a[2];
  }
  return 0;
}
------------------------------------

clang_without_patch -O2 -fno-omit-frame-pointer -S 1.cc
	.cfi_startproc
# %bb.0:                                # %entry
	testb	%dil, %dil
	je	.LBB0_1
# %bb.2:                                # %if.then
	pushq	%rbp
	.cfi_def_cfa_offset 16
	.cfi_offset %rbp, -16
	movq	%rsp, %rbp
	.cfi_def_cfa_register %rbp
	pushq	%rbx
	pushq	%rax
	.cfi_offset %rbx, -24
	movq	%rsi, %rbx
	callq	_Z3goov
	movq	8(%rbx), %rax
	addq	(%rbx), %rax
	addq	16(%rbx), %rax
	addq	$8, %rsp
	popq	%rbx
	popq	%rbp
	.cfi_def_cfa %rsp, 8
	retq
.LBB0_1:
	xorl	%eax, %eax          
	retq
.Lfunc_end0:
	.size	_Z3foobPl, .Lfunc_end0-_Z3foobPl
	.cfi_endproc

clang_with_patch -O2 -fno-omit-frame-pointer -S 1.cc
	.cfi_startproc
# %bb.0:                                # %entry
	testb	%dil, %dil
	je	.LBB0_1
# %bb.2:                                # %if.then
	pushq	%rbp
	.cfi_def_cfa_offset 16
	.cfi_offset %rbp, -16
	movq	%rsp, %rbp
	.cfi_def_cfa_register %rbp
	pushq	%rbx
	pushq	%rax
	.cfi_offset %rbx, -24
	movq	%rsi, %rbx
	callq	_Z3goov
	movq	8(%rbx), %rax
	addq	(%rbx), %rax
	addq	16(%rbx), %rax
	addq	$8, %rsp
	popq	%rbx
	popq	%rbp
	.cfi_def_cfa %rsp, 8
	retq
.LBB0_1:
	.cfi_restore %rbx          <==== .cfi_restore inserted.
	.cfi_restore %rbp          <==== .cfi_restore inserted.
	xorl	%eax, %eax           <==== without .cfi_restore above, here libunwind will erroneously think the values of %rbx and %rbp in last stack frame are saved in current stack frame.                                                           
	retq
.Lfunc_end0:
	.size	_Z3foobPl, .Lfunc_end0-_Z3foobPl
	.cfi_endproc

[thegameg]

Thanks for working on this, I think in general it looks good.

Wasn't there a verifier that was running with EXPENSIVE_CHECKS for the previously supported directives? Can we add the registers as well?

More comments inline.

llvm/lib/CodeGen/CFIInstrInserter.cpp

81

You can use BitVectors as sets of registers, with the supported bitwise operations to provide similar set operations as std::set.

llvm/test/CodeGen/X86/cfi-inserter-callee-save-register.mir

7

Does this test actually need the IR?

[wmi]

In D74303#1870079, @thegameg wrote:

Thanks for working on this, I think in general it looks good.

Wasn't there a verifier that was running with EXPENSIVE_CHECKS for the previously supported directives? Can we add the registers as well?

More comments inline.

Thanks for the review and sorry for the late reply.

And thanks for pointing out the verifier. That is a helpful suggestion! I added the check for CSR related CFI, and it helped me to discover a problem I missed before: Previously we skip inserting .cfi_restore in epilogue in prologue/epilogue generation pass, and only insert .cfi_restore in CFIInstrInserter pass. If the epilogue block is not a return block, skipping inserting .cfi_restore in it will lead to the CSR related CFI verification error -- the return block after the epilogue block will see different OutgoingCSRSaved set from its predecessors.

The solution is: In prologue/epilogue generation pass, we need to insert .cfi_restore for epilogue if it is not a return block or other type of exit block. In this way, each block will see consistent OutgoingCSRSaved information from each of its predecessors. At the same time, for the common case that epilogue block is a return block, we still don't need to insert unnecessary .cfi_restore.

To make it clearer, here is a testcase to show it. In the assembly, %bb.1 is an epilogue bb without return inside of it. If no .cfi_restore is inserted in %bb.1, %LBB0_2 will see different OutgoingCSRSaved information from its two predecessors: %bb.0 (no reg saved) and %bb.1 (rbp saved). So we need to insert .cfi_restore in %bb.1.

------------- 1.cc --------------
void goo();

long foo(bool cond) {
  if (__builtin_expect(cond, 1)) {
    goo();
    return 0;
  }
  return 0;
}
-----------------------------

clang -O2 -fno-omit-frame-pointer  -S 1.cc
  .cfi_startproc
# %bb.0:                                # %entry
  testb %dil, %dil
  je  .LBB0_2
# %bb.1:                                # %if.then
  pushq %rbp
  .cfi_def_cfa_offset 16
  .cfi_offset %rbp, -16
  movq  %rsp, %rbp
  .cfi_def_cfa_register %rbp
  callq _Z3goov
  popq  %rbp                     
  .cfi_def_cfa %rsp, 8
.LBB0_2:                                # %return
  xorl  %eax, %eax
  retq
.Lfunc_end0:
  .size _Z3foob, .Lfunc_end0-_Z3foob
  .cfi_endproc

llvm/lib/CodeGen/CFIInstrInserter.cpp

81

Thanks. Done.

llvm/test/CodeGen/X86/cfi-inserter-callee-save-register.mir

7

The function define is needed, but its body can be further simplified.

[wmi]

Address @thegameg's comment.

[wmi]

Friendly ping.

[thegameg]

Sorry for the delay.

That's right, llvm didn't insert any CFI in epilogues before this pass was added, so backends can start emitting them once this pass is ready to deal with them.

I added some more inline comments, this seems pretty good!

llvm/lib/CodeGen/CFIInstrInserter.cpp

186

Can a BitVector be used here too?

273

If CSRSaved is a BitVector, this can be just MBBInfo.OutgoingCSRSaved |= CSRSaved.

348

for (int Reg : SetDifference.set_bits())

404

set_bits() here too.

409–410

set_bits() here too.

llvm/lib/Target/X86/X86FrameLowering.cpp

509

Can we just not call this function at all on no-return blocks?

llvm/test/CodeGen/X86/cfi-epilogue-with-return.mir

10

I would also remove all the unnecessary IR here if possible.

56

Maybe you could use renamable $rbx = IMPLICIT_DEF for all the movs, for a cleaner test here and in the rest of the tests?

llvm/test/CodeGen/X86/cfi-inserter-callee-save-register.mir

7

IIRC, you can even remove the function define by removing the whole IR part. For example: llvm/test/CodeGen/X86/tailcall-pseudo.mir

[wmi]

llvm/lib/Target/X86/X86FrameLowering.cpp

509

Good point. Done.

llvm/test/CodeGen/X86/cfi-epilogue-with-return.mir

10

I remove most of them, but "define _Z3foob" is still needed. In addition, I need to mark the frame-pointer = all attribute to test the cfi related with %rbp is properly set.

56

Another good point. Done.

[wmi]

Address @thegameg's comment.

[thegameg]

LGTM, this looks great, thanks @wmi.

llvm/test/CodeGen/X86/cfi-epilogue-with-return.mir

30

You can also remove the branch probabilities in all the tests, e.g. in this test: (0x7fef9fcb) and (0x00106035).

[wmi]

Thanks for the review and helpful suggestions!

[wmi]

The original committed patch was reverted because a test failure in libc++ was reported.

The original patch only inserts .cfi_restore. A missing piece in the original patch is for some cases, it needs to generate .cfi_offset/.cfi_register as well. I attach a new patch which fixes the problem.

Here is a case just used to show where is the problem. In reality the test below won't expose the problem because we don't generate .cfi_restore in epilogue if it is a return block. Suppose the final machine code has the same layout as the c code below.

int foo(bool cond) {
  // prologue: .cfi_offset $rbp 16
  if (cond) {
    b = 5;
    // epilogue: .cfi_restore $rbp
    return 2;
  } 
  // currently, we don't insert anything here, but we need to insert .cfi_offset $rbp 16 because at the place of 'a = 3', $rbp has not been restored to its value in previous frame yet. 
  a = 3;
  // epilogue: .cfi_restore $rbp
  return 3;
}

[wmi]

Fix an unwinding problem exposed in libc++ test.

[wmi]

reopen it for the review of the patch to be recommitted.

[wmi]

Friendly ping. Because the fix for the patch contains some substantial change so I want it to be reviewed before I recommit the patch. Thanks in advance!

[wmi]

Ping again.

[thegameg]

llvm/lib/CodeGen/CFIInstrInserter.cpp

92

Can we rename this to something like CSRSavedLocation?

94

0/$noreg should be considered an invalid register here, no?

103

Isn't this a per-basicblock property? I actually have a hard time coming up with a test case where this would break, so maybe this is sufficient?

216

Can we also have tests for cfi_register and cfi_rel_offset? My understanding is that we currently don't generate them.

253

it is not used here

254

Does it mean that we assert if we have:

.cfi_offset $rbx, -10
.cfi_offset $rbx, -24

I would think that this is valid and we should just overwrite it. What do you think? If we do want to support it, please add a test as well.

llvm/test/CodeGen/X86/cfi-inserter-callee-save-register-2.mir

74

Can we add the CFI_INSTRUCTION restores for consistency here too?

[wmi]

llvm/lib/CodeGen/CFIInstrInserter.cpp

92

That is better. Done.

94

Register number used in CFI instruction is got from MCRegisterInfo::getDwarfRegNum, which is different from the physical register number. For $rax, CFI.getRegister() returns 0.

103

My assumption is every function will only save a CSR in one place and I havn't seen case outside of the assumption. If every basicblock has its own CSR location array, the cost will be much higher, so I choose to hold the assumption until I see a counterexample.

216

I saw PPC backend has code to generate .cfi_register. I add tests for them.

253

Fixed.

254

Ah, the assertion is not written in a correct way. The original intention was to be consistent with the assumption above that every function will only save the CSR in one place. If one CSR is saved in multiple different locs, the assertion could catch it.

I change the assertion and add a test for it.

llvm/test/CodeGen/X86/cfi-inserter-callee-save-register-2.mir

74

Done.

[wmi]

Address thegameg's comment.

[wmi]

Fix some minor issue in the last version.

[thegameg]

Thank you, this LGTM!

llvm/lib/CodeGen/CFIInstrInserter.cpp

94

Maybe llvm::Optional<T> is cleaner in this case? (I don't feel strongly about this, though)

103

Fair enough. I think it's a safe assumption to make, since for now we only emit the prologue more than once if we have ehfunclet-entry blocks (llvm/test/CodeGen/X86/win64-eh-empty-block-2.mir).

[wmi]

llvm/lib/CodeGen/CFIInstrInserter.cpp

94

Good suggestion. Optional<T> looks cleaner.

[wmi]

Address thegameg@'s comment.

llvm/lib/CodeGen/CFIInstrInserter.cpp

 /// their successors. Then it checks if blocks have correct CFA calculation rule
 /// set and inserts additional CFI instruction at their beginnings if they
 /// don't. CFI instructions are inserted if basic blocks have incorrect offset
 /// or register set by previous blocks, as a result of a non-linear layout of
 /// blocks in a function.
 //===----------------------------------------------------------------------===//
 
 #include "llvm/ADT/DepthFirstIterator.h"
+#include "llvm/ADT/SetOperations.h"
 #include "llvm/CodeGen/MachineFunctionPass.h"
 #include "llvm/CodeGen/MachineInstrBuilder.h"
 #include "llvm/CodeGen/MachineModuleInfo.h"
 #include "llvm/CodeGen/Passes.h"
 #include "llvm/CodeGen/TargetFrameLowering.h"
 #include "llvm/CodeGen/TargetInstrInfo.h"
 #include "llvm/CodeGen/TargetSubtargetInfo.h"
 #include "llvm/InitializePasses.h"
struct MBBCFAInfo {
     /// Value of cfa offset valid at basic block entry.
     int IncomingCFAOffset = -1;
     /// Value of cfa offset valid at basic block exit.
     int OutgoingCFAOffset = -1;
     /// Value of cfa register valid at basic block entry.
     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(unsigned R, int O) : Reg(R), Offset(O) {}
+    unsigned Reg = INVALID_REG;
+    int Offset = INVALID_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);
   /// Calculate cfa offset and register values valid at basic block exit by
   /// checking the block for CFI instructions. Block's incoming CFA info remains
   /// the same.
   void calculateOutgoingCFAInfo(MBBCFAInfo &MBBInfo);
   /// Update in/out cfa offset and register values for successors of the basic
   /// block.
   void updateSuccCFAInfo(MBBCFAInfo &MBBInfo);
 
   /// Check if incoming CFA information of a basic block matches outgoing CFA
   /// information of the previous block. If it doesn't, insert CFI instruction
   /// at the beginning of the block that corrects the CFA calculation rule for
   /// that block.
   bool insertCFIInstrs(MachineFunction &MF);
   /// Return the cfa offset value that should be set at the beginning of a MBB
   /// if needed. The negated value is needed when creating CFI instructions that
   /// set absolute offset.
   int getCorrectCFAOffset(MachineBasicBlock *MBB) {
     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
   /// outgoing offset and register of the MBB.
   unsigned verify(MachineFunction &MF);
 };
 }  // namespace
 
 char CFIInstrInserter::ID = 0;
 INITIALIZE_PASS(CFIInstrInserter, "cfi-instr-inserter",
                 "Check CFA info and insert CFI instructions if needed", false,
                 false)
 FunctionPass *llvm::createCFIInstrInserter() { return new CFIInstrInserter(); }
 
 void CFIInstrInserter::calculateCFAInfo(MachineFunction &MF) {
   // Initial CFA offset value i.e. the one valid at the beginning of the
   // function.
   int InitialOffset =
       MF.getSubtarget().getFrameLowering()->getInitialCFAOffset(MF);
   // Initial CFA register value i.e. the one valid at the beginning of the
   // function.
   unsigned InitialRegister =
       MF.getSubtarget().getFrameLowering()->getInitialCFARegister(MF);
+  const TargetRegisterInfo &TRI = *MF.getSubtarget().getRegisterInfo();
+  unsigned NumRegs = TRI.getNumRegs();
 
   // Initialize MBBMap.
   for (MachineBasicBlock &MBB : MF) {
     MBBCFAInfo MBBInfo;
     MBBInfo.MBB = &MBB;
     MBBInfo.IncomingCFAOffset = InitialOffset;
     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
   // i.e. that it has initial cfa offset and register values as incoming CFA
   // information.
   for (MachineBasicBlock &MBB : MF) {
     if (MBBVector[MBB.getNumber()].Processed) continue;
     updateSuccCFAInfo(MBBVector[MBB.getNumber()]);
   }
 }
 
 void CFIInstrInserter::calculateOutgoingCFAInfo(MBBCFAInfo &MBBInfo) {
   // Outgoing cfa offset set by the block.
   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()) {
+      unsigned CSRReg = INVALID_REG;
+      int CSROffset = INVALID_OFFSET;
       unsigned CFIIndex = MI.getOperand(0).getCFIIndex();
       const MCCFIInstruction &CFI = Instrs[CFIIndex];
       switch (CFI.getOperation()) {
       case MCCFIInstruction::OpDefCfaRegister:
         SetRegister = CFI.getRegister();
         break;
       case MCCFIInstruction::OpDefCfaOffset:
         SetOffset = CFI.getOffset();
         break;
       case MCCFIInstruction::OpAdjustCfaOffset:
         SetOffset += CFI.getOffset();
         break;
       case MCCFIInstruction::OpDefCfa:
         SetRegister = CFI.getRegister();
         SetOffset = CFI.getOffset();
         break;
+      case MCCFIInstruction::OpOffset:
+        CSRReg = INVALID_REG;
+        CSROffset = CFI.getOffset();
+        break;
+      case MCCFIInstruction::OpRegister:
+        CSRReg = CFI.getRegister2();
+        CSROffset = INVALID_OFFSET;
+        break;
+      case MCCFIInstruction::OpRelOffset:
+        CSRReg = INVALID_REG;
+        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
         report_fatal_error(
             "Support for cfi_remember_state not implemented! Value of CFA "
             "may be incorrect!\n");
 #endif
         break;
       case MCCFIInstruction::OpRestoreState:
         // TODO: Add support for handling cfi_restore_state.
 #ifndef NDEBUG
         report_fatal_error(
             "Support for cfi_restore_state not implemented! Value of CFA may "
             "be incorrect!\n");
 #endif
         break;
       // Other CFI directives do not affect CFA value.
+      case MCCFIInstruction::OpUndefined:
       case MCCFIInstruction::OpSameValue:
-      case MCCFIInstruction::OpOffset:
-      case MCCFIInstruction::OpRelOffset:
       case MCCFIInstruction::OpEscape:
-      case MCCFIInstruction::OpRestore:
-      case MCCFIInstruction::OpUndefined:
-      case MCCFIInstruction::OpRegister:
       case MCCFIInstruction::OpWindowSave:
       case MCCFIInstruction::OpNegateRAState:
       case MCCFIInstruction::OpGnuArgsSize:
         break;
       }
+      if ((CSRReg != INVALID_REG || CSROffset != INVALID_OFFSET)) {
+        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());
+      }
     }
   }
 
   MBBInfo.Processed = true;
 
   // 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) {
   SmallVector<MachineBasicBlock *, 4> Stack;
   Stack.push_back(MBBInfo.MBB);
 
   do {
     MachineBasicBlock *Current = Stack.pop_back_val();
     MBBCFAInfo &CurrentInfo = MBBVector[Current->getNumber()];
     if (CurrentInfo.Processed)
       continue;
 
     calculateOutgoingCFAInfo(CurrentInfo);
     for (auto *Succ : CurrentInfo.MBB->successors()) {
       MBBCFAInfo &SuccInfo = MBBVector[Succ->getNumber()];
       if (!SuccInfo.Processed) {
         SuccInfo.IncomingCFAOffset = CurrentInfo.OutgoingCFAOffset;
         SuccInfo.IncomingCFARegister = CurrentInfo.OutgoingCFARegister;
+        SuccInfo.IncomingCSRSaved = CurrentInfo.OutgoingCSRSaved;
         Stack.push_back(Succ);
       }
     }
   } while (!Stack.empty());
 }
 
 bool CFIInstrInserter::insertCFIInstrs(MachineFunction &MF) {
   const MBBCFAInfo *PrevMBBInfo = &MBBVector[MF.front().getNumber()];
if (PrevMBBInfo->OutgoingCFAOffset != MBBInfo.IncomingCFAOffset) {
                MBBInfo.IncomingCFARegister) {
       unsigned CFIIndex =
           MF.addFrameInst(MCCFIInstruction::createDefCfaRegister(
               nullptr, MBBInfo.IncomingCFARegister));
       BuildMI(*MBBInfo.MBB, MBBI, DL, TII->get(TargetOpcode::CFI_INSTRUCTION))
           .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 == INVALID_REG && RO.Offset != INVALID_OFFSET) {
+        CFIIndex = MF.addFrameInst(
+            MCCFIInstruction::createOffset(nullptr, Reg, RO.Offset));
+      } else if (RO.Reg != INVALID_REG && RO.Offset == INVALID_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()
          << " in " << Pred.MBB->getParent()->getName()
          << " outgoing CFA Reg:" << Pred.OutgoingCFARegister << "\n";
   errs() << "Pred: " << Pred.MBB->getName() << " #" << Pred.MBB->getNumber()
          << " in " << Pred.MBB->getParent()->getName()
          << " outgoing CFA Offset:" << Pred.OutgoingCFAOffset << "\n";
   errs() << "Succ: " << Succ.MBB->getName() << " #" << Succ.MBB->getNumber()
          << " incoming CFA Reg:" << Succ.IncomingCFARegister << "\n";
   errs() << "Succ: " << Succ.MBB->getName() << " #" << Succ.MBB->getNumber()
          << " 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)) {
     const MBBCFAInfo &CurrMBBInfo = MBBVector[CurrMBB->getNumber()];
     for (MachineBasicBlock *Succ : CurrMBB->successors()) {
       const MBBCFAInfo &SuccMBBInfo = MBBVector[Succ->getNumber()];
       // Check that incoming offset and register values of successors match the
       // outgoing offset and register values of CurrMBB
       if (SuccMBBInfo.IncomingCFAOffset != CurrMBBInfo.OutgoingCFAOffset ||
           SuccMBBInfo.IncomingCFARegister != CurrMBBInfo.OutgoingCFARegister) {
         // Inconsistent offsets/registers are ok for 'noreturn' blocks because
         // 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++;
       }
     }
   }
   return ErrorNum;
 }

llvm/lib/Target/X86/X86FrameLowering.h

void emitStackProbe(MachineFunction &MF, MachineBasicBlock &MBB,
                       bool InProlog) const;
 
   /// Replace a StackProbe inline-stub with the actual probe code inline.
   void inlineStackProbe(MachineFunction &MF,
                         MachineBasicBlock &PrologMBB) const override;
 
   void emitCalleeSavedFrameMoves(MachineBasicBlock &MBB,
                                  MachineBasicBlock::iterator MBBI,
-                                 const DebugLoc &DL) const;
+                                 const DebugLoc &DL, bool IsPrologue) const;
 
   /// emitProlog/emitEpilog - These methods insert prolog and epilog code into
   /// the function.
   void emitPrologue(MachineFunction &MF, MachineBasicBlock &MBB) const override;
   void emitEpilogue(MachineFunction &MF, MachineBasicBlock &MBB) const override;
 
   void adjustForSegmentedStacks(MachineFunction &MF,
                                 MachineBasicBlock &PrologueMBB) const override;

llvm/lib/Target/X86/X86FrameLowering.cpp

void X86FrameLowering::BuildCFI(MachineBasicBlock &MBB,
   MachineFunction &MF = *MBB.getParent();
   unsigned CFIIndex = MF.addFrameInst(CFIInst);
   BuildMI(MBB, MBBI, DL, TII.get(TargetOpcode::CFI_INSTRUCTION))
       .addCFIIndex(CFIIndex);
 }
 
 void X86FrameLowering::emitCalleeSavedFrameMoves(
     MachineBasicBlock &MBB, MachineBasicBlock::iterator MBBI,
-    const DebugLoc &DL) const {
+    const DebugLoc &DL, bool IsPrologue) const {
   MachineFunction &MF = *MBB.getParent();
   MachineFrameInfo &MFI = MF.getFrameInfo();
   MachineModuleInfo &MMI = MF.getMMI();
   const MCRegisterInfo *MRI = MMI.getContext().getRegisterInfo();
 
   // Add callee saved registers to move list.
   const std::vector<CalleeSavedInfo> &CSI = MFI.getCalleeSavedInfo();
   if (CSI.empty()) return;
 
   // Calculate offsets.
   for (std::vector<CalleeSavedInfo>::const_iterator
          I = CSI.begin(), E = CSI.end(); I != E; ++I) {
     int64_t Offset = MFI.getObjectOffset(I->getFrameIdx());
     unsigned Reg = I->getReg();
-
     unsigned DwarfReg = MRI->getDwarfRegNum(Reg, true);
+
+    if (IsPrologue) {
       BuildCFI(MBB, MBBI, DL,
                MCCFIInstruction::createOffset(nullptr, DwarfReg, Offset));
+    } else {
+      BuildCFI(MBB, MBBI, DL,
+               MCCFIInstruction::createRestore(nullptr, DwarfReg));
+    }
   }
 }
 
 void X86FrameLowering::emitStackProbe(MachineFunction &MF,
                                       MachineBasicBlock &MBB,
                                       MachineBasicBlock::iterator MBBI,
                                       const DebugLoc &DL, bool InProlog) const {
   const X86Subtarget &STI = MF.getSubtarget<X86Subtarget>();
if (((!HasFP && NumBytes) || PushedRegs) && NeedsDwarfCFI) {
     if (!HasFP && NumBytes) {
       // Define the current CFA rule to use the provided offset.
       assert(StackSize);
       BuildCFI(MBB, MBBI, DL, MCCFIInstruction::createDefCfaOffset(
                                   nullptr, -StackSize + stackGrowth));
     }
 
     // Emit DWARF info specifying the offsets of the callee-saved registers.
-    emitCalleeSavedFrameMoves(MBB, MBBI, DL);
+    emitCalleeSavedFrameMoves(MBB, MBBI, DL, true);
   }
 
   // X86 Interrupt handling function cannot assume anything about the direction
   // flag (DF in EFLAGS register). Clear this flag by creating "cld" instruction
   // in each prologue of interrupt handler function.
   //
   // FIXME: Create "cld" instruction only in these cases:
   // 1. The interrupt handling function uses any of the "rep" instructions.
if (IsFunclet) {
     // realigned.
     if (TRI->needsStackRealignment(MF) && !IsWin64Prologue)
       NumBytes = alignTo(FrameSize, MaxAlign);
   } else {
     NumBytes = StackSize - CSSize;
   }
   uint64_t SEHStackAllocAmt = NumBytes;
 
+  // AfterPop is the position to insert .cfi_restore.
+  MachineBasicBlock::iterator AfterPop = MBBI;
   if (HasFP) {
     // Pop EBP.
     BuildMI(MBB, MBBI, DL, TII.get(Is64Bit ? X86::POP64r : X86::POP32r),
             MachineFramePtr)
         .setMIFlag(MachineInstr::FrameDestroy);
     if (NeedsDwarfCFI) {
       unsigned DwarfStackPtr =
           TRI->getDwarfRegNum(Is64Bit ? X86::RSP : X86::ESP, true);
       BuildCFI(MBB, MBBI, DL, MCCFIInstruction::createDefCfa(
                                   nullptr, DwarfStackPtr, -SlotSize));
+      if (!MBB.succ_empty() && !MBB.isReturnBlock()) {
+        unsigned DwarfFramePtr = TRI->getDwarfRegNum(MachineFramePtr, true);
+        BuildCFI(MBB, AfterPop, DL,
+                 MCCFIInstruction::createRestore(nullptr, DwarfFramePtr));
+        --MBBI;
+        --AfterPop;
+      }
       --MBBI;
     }
   }
 
   MachineBasicBlock::iterator FirstCSPop = MBBI;
   // Skip the callee-saved pop instructions.
   while (MBBI != MBB.begin()) {
     MachineBasicBlock::iterator PI = std::prev(MBBI);
while (MBBI != MBB.end()) {
       if (Opc == X86::POP32r || Opc == X86::POP64r) {
         Offset += SlotSize;
         BuildCFI(MBB, MBBI, DL,
                  MCCFIInstruction::createDefCfaOffset(nullptr, Offset));
       }
     }
   }
 
+  // Emit DWARF info specifying the restores of the callee-saved registers.
+  // For epilogue with return inside or being other block without successor,
+  // no need to generate .cfi_restore for callee-saved registers.
+  if (NeedsDwarfCFI && !MBB.succ_empty() && !MBB.isReturnBlock()) {
+    emitCalleeSavedFrameMoves(MBB, AfterPop, DL, false);
+  }
+
   if (Terminator == MBB.end() || !isTailCallOpcode(Terminator->getOpcode())) {
     // Add the return addr area delta back since we are not tail calling.
     int Offset = -1 * X86FI->getTCReturnAddrDelta();
     assert(Offset >= 0 && "TCDelta should never be positive");
     if (Offset) {
       // Check for possible merge with preceding ADD instruction.
       Offset += mergeSPUpdates(MBB, Terminator, true);
       emitSPUpdate(MBB, Terminator, DL, Offset, /*InEpilogue=*/true);

llvm/test/CodeGen/X86/cfi-epilogue-with-return.mir

+# RUN: llc -o - %s -mtriple=x86_64-- -run-pass=prologepilog 2>&1 | FileCheck %s
+--- |
+  define i64 @_Z3foob(i1 zeroext %cond) #0 {
+    ret i64 0
+  }
+  attributes #0 = {"frame-pointer"="all"}
+...
+---
+# If the epilogue bb.1 is a return block, no .cfi_restore is
+# needed in it.
+# CHECK:    bb.1:
+# CHECK-NOT:  CFI_INSTRUCTION restore
+# CHECK:      RET 0
+# CHECK:    bb.2:
+# CHECK:      RET 0
+name:            _Z3foob
+alignment:       16
+tracksRegLiveness: true
+liveins:
+  - { reg: '$edi' }
+frameInfo:
+  maxAlignment:    1
+  hasCalls:        true
+  savePoint:       '%bb.1'
+  restorePoint:    '%bb.1'
+machineFunctionInfo: {}
+body:             |
+  bb.0:
+    liveins: $edi
+  
+    TEST8rr renamable $dil, renamable $dil, implicit-def $eflags, implicit killed $edi
+    JCC_1 %bb.2, 4, implicit killed $eflags
+    JMP_1 %bb.1
+  
+  bb.1:
+    renamable $rbx = IMPLICIT_DEF
+    renamable $r14 = IMPLICIT_DEF
+    renamable $r15 = IMPLICIT_DEF
+    renamable $r12 = IMPLICIT_DEF
+    renamable $r13 = IMPLICIT_DEF
+    dead $eax = MOV32r0 implicit-def dead $eflags, implicit-def $rax
+    RET 0, killed $rax
+  
+  bb.2:
+    dead $eax = MOV32r0 implicit-def dead $eflags, implicit-def $rax
+    RET 0, killed $rax
+
+...

llvm/test/CodeGen/X86/cfi-epilogue-without-return.mir

+# RUN: llc -o - %s -mtriple=x86_64-- -run-pass=prologepilog 2>&1 | FileCheck %s
+--- |
+  declare dso_local void @_Z3goov()
+  define i64 @_Z3foob(i1 zeroext %cond) #0 {
+    ret i64 0
+  }
+  attributes #0 = {"frame-pointer"="all"}
+...
+---
+# If the epilogue bb.1.if.then is not a return block, .cfi_restore is
+# needed in it, otherwise bb.2.return will see different outgoing CFI
+# information from its predecessors.
+# CHECK:    bb.1:
+# CHECK:      CFI_INSTRUCTION restore $rbx
+# CHECK-NEXT: CFI_INSTRUCTION restore $r12
+# CHECK-NEXT: CFI_INSTRUCTION restore $r13
+# CHECK-NEXT: CFI_INSTRUCTION restore $r14
+# CHECK-NEXT: CFI_INSTRUCTION restore $r15
+# CHECK-NEXT: CFI_INSTRUCTION restore $rbp
+# CHECK-NOT:  RET 0
+# CHECK:    bb.2:
+# CHECK:      RET 0
+name:            _Z3foob
+alignment:       16
+tracksRegLiveness: true
+liveins:
+  - { reg: '$edi' }
+frameInfo:
+  maxAlignment:    1
+  hasCalls:        true
+  savePoint:       '%bb.1'
+  restorePoint:    '%bb.1'
+machineFunctionInfo: {}
+body:             |
+  bb.0:
+    liveins: $edi
+  
+    TEST8rr renamable $dil, renamable $dil, implicit-def $eflags, implicit killed $edi
+    JCC_1 %bb.2, 4, implicit killed $eflags
+    JMP_1 %bb.1
+  
+  bb.1:
+    renamable $rbx = IMPLICIT_DEF
+    renamable $r14 = IMPLICIT_DEF
+    renamable $r15 = IMPLICIT_DEF
+    renamable $r12 = IMPLICIT_DEF
+    renamable $r13 = IMPLICIT_DEF
+  
+  bb.2:
+    dead $eax = MOV32r0 implicit-def dead $eflags, implicit-def $rax
+    RET 0, killed $rax
+
+...

llvm/test/CodeGen/X86/cfi-inserter-callee-save-register-2.mir

+# RUN: llc -o - %s -mtriple=x86_64-- -verify-cfiinstrs \
+# RUN:     -run-pass=cfi-instr-inserter 2>&1 | FileCheck %s
+# Test that CFI inserter inserts .cfi_offset/.cfi_register/.cfi_rel_offset
+# properly for callee saved registers.
+--- |
+  define void @foo() {
+    ret void
+  }
+...
+---
+# CHECK:      bb.3:
+# CHECK:      CFI_INSTRUCTION offset $rbp, -16
+# CHECK-NEXT: CFI_INSTRUCTION offset $r12, -24
+# CHECK-NEXT: CFI_INSTRUCTION register $r13, $rcx
+# CHECK-NEXT: CFI_INSTRUCTION offset $r14, -40
+name:            foo
+alignment:       16
+tracksRegLiveness: true
+liveins:
+  - { reg: '$edi' }
+frameInfo:
+  stackSize:       16
+  offsetAdjustment: -16
+  maxAlignment:    16
+  hasCalls:        true
+  maxCallFrameSize: 0
+  cvBytesOfCalleeSavedRegisters: 8
+fixedStack:
+  - { id: 0, type: spill-slot, offset: -40, size: 8, alignment: 8, callee-saved-register: '$r14' }
+  - { id: 1, type: spill-slot, offset: -24, size: 8, alignment: 8, callee-saved-register: '$r12' }
+  - { id: 2, type: spill-slot, offset: -16, size: 8, alignment: 16 }
+  - { id: 3, type: spill-slot, offset: -16, size: 8, alignment: 16 }
+  - { id: 4, type: spill-slot, offset: -24, size: 8, alignment: 8 }
+machineFunctionInfo: {}
+body:             |
+  bb.0:
+    successors: %bb.2(0x40000000), %bb.1(0x40000000)
+    liveins: $edi, $r12, $r13, $r14
+  
+    frame-setup PUSH64r killed $rbp, implicit-def $rsp, implicit $rsp
+    CFI_INSTRUCTION def_cfa_offset 16
+    CFI_INSTRUCTION offset $rbp, -16
+    $rbp = frame-setup MOV64rr $rsp
+    CFI_INSTRUCTION def_cfa_register $rbp
+    frame-setup PUSH64r killed $r12, implicit-def $rsp, implicit $rsp
+    $rcx = frame-setup COPY $r13
+    frame-setup PUSH64r killed $r14, implicit-def $rsp, implicit $rsp
+    CFI_INSTRUCTION offset $r12, -24
+    CFI_INSTRUCTION register $r13, $rcx
+    CFI_INSTRUCTION rel_offset $r14, -24
+    TEST8rr renamable $dil, renamable $dil, implicit-def $eflags, implicit killed $edi
+    JCC_1 %bb.2, 4, implicit killed $eflags
+    JMP_1 %bb.1
+  
+  bb.1:
+    successors: %bb.3(0x80000000)
+  
+    renamable $r12 = IMPLICIT_DEF
+    renamable $r13 = IMPLICIT_DEF
+    renamable $r14 = IMPLICIT_DEF
+    JMP_1 %bb.3
+  
+  bb.2:
+    liveins: $rcx
+    dead $eax = MOV32r0 implicit-def dead $eflags, implicit-def $rax
+    $r12 = frame-destroy POP64r implicit-def $rsp, implicit $rsp
+    $r13 = frame-destroy COPY $rcx
+    $r14 = frame-destroy POP64r implicit-def $rsp, implicit $rsp
+    $rbp = frame-destroy POP64r implicit-def $rsp, implicit $rsp
+    CFI_INSTRUCTION restore $rbp
+    CFI_INSTRUCTION restore $r12
+    CFI_INSTRUCTION restore $r13
+    CFI_INSTRUCTION restore $r14
+    CFI_INSTRUCTION def_cfa $rsp, 8
+    RET 0, killed $rax
+  
+  bb.3:
+    successors: %bb.4(0x80000000)
+  
+    renamable $rdi = IMPLICIT_DEF
+    renamable $rsi = IMPLICIT_DEF
+  
+  bb.4:
+    liveins: $rcx
+    dead $eax = MOV32r0 implicit-def dead $eflags, implicit-def $rax
+    $r12 = frame-destroy POP64r implicit-def $rsp, implicit $rsp
+    $r13 = frame-destroy COPY $rcx
+    $r14 = frame-destroy POP64r implicit-def $rsp, implicit $rsp
+    $rbp = frame-destroy POP64r implicit-def $rsp, implicit $rsp
+    CFI_INSTRUCTION restore $rbp
+    CFI_INSTRUCTION restore $r12
+    CFI_INSTRUCTION restore $r13
+    CFI_INSTRUCTION restore $r14
+    CFI_INSTRUCTION def_cfa $rsp, 8
+    RET 0, killed $rax
+
+...

llvm/test/CodeGen/X86/cfi-inserter-callee-save-register.mir

+# RUN: llc -o - %s -mtriple=x86_64-- -verify-cfiinstrs \
+# RUN:     -run-pass=cfi-instr-inserter 2>&1 | FileCheck %s
+# Test that CFI inserter inserts .cfi_restore properly for
+# callee saved registers.
+--- |
+  define void @foo() {
+    ret void
+  }
+...
+---
+# CHECK:      bb.3:
+# CHECK:      CFI_INSTRUCTION restore $rbx
+# CHECK-NEXT: CFI_INSTRUCTION restore $rbp
+name:            foo
+body:             |
+  bb.0:
+    TEST8rr renamable $dil, renamable $dil, implicit-def $eflags, implicit killed $edi
+    JCC_1 %bb.2, 5, implicit killed $eflags
+
+  bb.1:
+    JMP_1 %bb.3
+
+  bb.2:
+    CFI_INSTRUCTION def_cfa_offset 16
+    CFI_INSTRUCTION offset $rbp, -16
+    CFI_INSTRUCTION def_cfa_register $rbp
+    CFI_INSTRUCTION offset $rbx, -24
+    CFI_INSTRUCTION def_cfa $rsp, 8
+    RET 0, $rax
+
+  bb.3:
+    RET 0, $rax
+
+...

llvm/test/CodeGen/X86/cfi-inserter-verify-inconsistent-csr.mir

+# RUN: not --crash llc -o - %s -mtriple=x86_64-- -verify-cfiinstrs \
+# RUN:     -run-pass=cfi-instr-inserter 2>&1 | FileCheck %s
+# Test that CFI verifier finds inconsistent csr saved set between bb.end and
+# one of its precedessors.
+--- |
+  define void @inconsistentCSR() {
+  entry:
+    br label %then
+  then:
+    br label %end
+  end:
+    ret void
+  }
+...
+---
+# CHECK: *** Inconsistent CSR Saved between pred and succ in function inconsistentCSR ***
+# CHECK: LLVM ERROR: Found 1 in/out CFI information errors.
+name: inconsistentCSR
+body: |
+  bb.0.entry:
+    JCC_1 %bb.2, 5, implicit undef $eflags
+
+  bb.1.then:
+    CFI_INSTRUCTION offset $rbp, -16
+
+  bb.2.end:
+    RET 0
+...

llvm/test/CodeGen/X86/cfi-inserter-verify-inconsistent-loc.mir

+# REQUIRES: asserts
+# RUN: not --crash llc -o - %s -mtriple=x86_64-- \
+# RUN:     -run-pass=cfi-instr-inserter 2>&1 | FileCheck %s
+# Test that CSR being saved in multiple locations can be caught by
+# an assertion.
+--- |
+  define void @inconsistentlocs() {
+    bb.end:
+      ret void
+  }
+...
+---
+# CHECK: Different saved locations for the same CSR
+# CHECK-NEXT: UNREACHABLE executed 
+name: inconsistentlocs
+body: |
+  bb.0:
+    CFI_INSTRUCTION offset $r12, -8
+    JCC_1 %bb.2, 5, implicit undef $eflags
+
+  bb.1:
+    CFI_INSTRUCTION offset $r12, -16
+
+  bb.2.bb.end:
+    RET 0
+...