LivePhysRegs: Fix addLiveOutsNoPristines() for return blocks past PEI

[MatzeB]

Summary

It turns out addLiveOutsNoPristines missed to add non-pristine callee saved registers in the return block. This probably went unnoticed, because we usually start walking backwards from the block end and at the time we actually queried/looked at the information we were before the epilogue code and it didn't matter anymore. This of course was not true for shrink-wrapping situations leading to bugs like the example in D32156.

In this patch:

• addLiveOutsNoPristines() needs to add callee saved registers that are actually saved and restored somewhere to the set (they are not pristine).
• Cleanup/rewrite the code for addLiveOuts()/addLiveOutsNoPristines().

This fixes the problem from D32156.

[MatzeB]

Note that I'll probably rework the test from D32156 for this if we agree on this patch.

[timshen]

lib/CodeGen/LivePhysRegs.cpp

194 ↗

(On Diff #96492)

If MBB is empty but not a return block, is it the intended behavior here?

[MatzeB]

lib/CodeGen/LivePhysRegs.cpp

194 ↗

(On Diff #96492)

Callee save register only need to be maintained when we actually leave the function.
Having succ_empty() but not a return block only happens in instance where we are not leaving the function: unreachable statements or calls to noreturn functions such as exit or abort are typical examples we don't need to restore CSRs for that.

194 ↗

(On Diff #96492)

I should have said "not returning from the function" instead of "not leaving the function".

[timshen]

LGTM. Let's wait for @qcolombet.

[MatzeB]

ping

[qcolombet]

LGTM

[timshen]

Hi Matthias,

Do you want to use the test case in D32156?

llvm/trunk/lib/CodeGen/LivePhysRegs.cpp

for (const auto &LI : MBB.liveins()) {
     for (; S.isValid(); ++S) {
       unsigned SI = S.getSubRegIndex();
       if ((LI.LaneMask & TRI->getSubRegIndexLaneMask(SI)).any())
         addReg(S.getSubReg());
     }
   }
 }
 
-/// Add pristine registers to the given \p LiveRegs. This function removes
-/// actually saved callee save registers when \p InPrologueEpilogue is false.
-static void addPristines(LivePhysRegs &LiveRegs, const MachineFunction &MF,
-                         const MachineFrameInfo &MFI,
-                         const TargetRegisterInfo &TRI) {
+/// Adds all callee saved registers to \p LiveRegs.
+static void addCalleeSavedRegs(LivePhysRegs &LiveRegs,
+                               const MachineFunction &MF) {
   const MachineRegisterInfo &MRI = MF.getRegInfo();
-  for (const MCPhysReg *CSR = MRI.getCalleeSavedRegs(); CSR && *CSR;
-       ++CSR)
+  for (const MCPhysReg *CSR = MRI.getCalleeSavedRegs(); CSR && *CSR; ++CSR)
     LiveRegs.addReg(*CSR);
+}
+
+/// Adds pristine registers to the given \p LiveRegs. Pristine registers are
+/// callee saved registers that are unused in the function.
+static void addPristines(LivePhysRegs &LiveRegs, const MachineFunction &MF) {
+  const MachineFrameInfo &MFI = MF.getFrameInfo();
+  if (!MFI.isCalleeSavedInfoValid())
+    return;
+  /// Add all callee saved regs, then remove the ones that are saved+restored.
+  addCalleeSavedRegs(LiveRegs, MF);
+  /// Remove the ones that are not saved/restored; they are pristine.
   for (const CalleeSavedInfo &Info : MFI.getCalleeSavedInfo())
     LiveRegs.removeReg(Info.getReg());
 }
 
 void LivePhysRegs::addLiveOutsNoPristines(const MachineBasicBlock &MBB) {
+  if (!MBB.succ_empty()) {
     // To get the live-outs we simply merge the live-ins of all successors.
     for (const MachineBasicBlock *Succ : MBB.successors())
       addBlockLiveIns(*Succ);
-}
-
-void LivePhysRegs::addLiveOuts(const MachineBasicBlock &MBB) {
+  } else if (MBB.isReturnBlock()) {
+    // For the return block: Add all callee saved registers that are saved and
+    // restored (somewhere); This does not include callee saved registers that
+    // are unused and hence not saved and restored; they are called pristine.
     const MachineFunction &MF = *MBB.getParent();
     const MachineFrameInfo &MFI = MF.getFrameInfo();
     if (MFI.isCalleeSavedInfoValid()) {
-    if (MBB.isReturnBlock()) {
-      // The return block has no successors whose live-ins we could merge
-      // below. So instead we add the callee saved registers manually.
-      const MachineRegisterInfo &MRI = MF.getRegInfo();
-      for (const MCPhysReg *I = MRI.getCalleeSavedRegs(); *I; ++I)
-        addReg(*I);
-    } else {
-      addPristines(*this, MF, MFI, *TRI);
-    }
-  }
+      for (const CalleeSavedInfo &Info : MFI.getCalleeSavedInfo())
+        addReg(Info.getReg());
+    }
 
+void LivePhysRegs::addLiveOuts(const MachineBasicBlock &MBB) {
+  if (!MBB.succ_empty()) {
     addLiveOutsNoPristines(MBB);
+    const MachineFunction &MF = *MBB.getParent();
+    addPristines(*this, MF);
+  } else if (MBB.isReturnBlock()) {
+    // For the return block: Add all callee saved registers.
+    const MachineFunction &MF = *MBB.getParent();
+    const MachineFrameInfo &MFI = MF.getFrameInfo();
+    if (MFI.isCalleeSavedInfoValid())
+      addCalleeSavedRegs(*this, MF);
+  }
 }
 
 void LivePhysRegs::addLiveIns(const MachineBasicBlock &MBB) {
   const MachineFunction &MF = *MBB.getParent();
-  const MachineFrameInfo &MFI = MF.getFrameInfo();
-  if (MFI.isCalleeSavedInfoValid())
-    addPristines(*this, MF, MFI, *TRI);
+  addPristines(*this, MF);
   addBlockLiveIns(MBB);
 }
 
 void llvm::computeLiveIns(LivePhysRegs &LiveRegs,
                           const MachineRegisterInfo &MRI,
                           MachineBasicBlock &MBB) {
   const TargetRegisterInfo &TRI = *MRI.getTargetRegisterInfo();
   assert(MBB.livein_empty());