Preserve existing registers when adding pristines to LivePhysRegs

[kparzysz]

Summary

If a non-pristine, callee-saved register is already present in LivePhysRegs, calling addPristines would cause it to be removed. Make sure that the existing contents of LivePhysRegs are preserved.

[MatzeB]

When I wrote this I assumed that addLiveOuts()/addLiveIns() is only called on empty sets looks like I didn't add an assert() for that and ifconversion does indeed call it no non-empty sets now. Thanks for tracking this down!

• Can you maintain an optimized code path that avoids the creation of an extra LivePhysRegs instance in the common case where the set is empty?
• We probably want the same fix (or an assert(empty())) in LiveRegUnits as well.

[kparzysz]

• Kept the original code for the case of empty set.
• Added the same fix to LiveRegUnits.

[MatzeB]

LGTM, thanks

include/llvm/CodeGen/LiveRegUnits.h

54

Good catch

include/llvm/CodeGen/LivePhysRegs.h

public:
 
   /// Dumps the currently live registers to the debug output.
   void dump() const;
 
 private:
   /// \brief Adds live-in registers from basic block \p MBB, taking associated
   /// lane masks into consideration.
   void addBlockLiveIns(const MachineBasicBlock &MBB);
+
+  /// Adds pristine registers. Pristine registers are callee saved registers
+  /// that are unused in the function.
+  void addPristines(const MachineFunction &MF);
 };
 
 inline raw_ostream &operator<<(raw_ostream &OS, const LivePhysRegs& LR) {
   LR.print(OS);
   return OS;
 }
 
 /// \brief Computes registers live-in to \p MBB assuming all of its successors

include/llvm/CodeGen/LiveRegUnits.h

void init(const TargetRegisterInfo &TRI) {
     Units.reset();
     Units.resize(TRI.getNumRegUnits());
   }
 
   /// Clears the set.
   void clear() { Units.reset(); }
 
   /// Returns true if the set is empty.
-  bool empty() const { return Units.empty(); }
+  bool empty() const { return Units.none(); }
 
   /// Adds register units covered by physical register \p Reg.
   void addReg(unsigned Reg) {
     for (MCRegUnitIterator Unit(Reg, TRI); Unit.isValid(); ++Unit)
       Units.set(*Unit);
   }
 
   /// \brief Adds register units covered by physical register \p Reg that are
public:
   /// Removes all register units marked in the bitvector \p RegUnits.
   void removeUnits(const BitVector &RegUnits) {
     Units.reset(RegUnits);
   }
   /// Return the internal bitvector representation of the set.
   const BitVector &getBitVector() const {
     return Units;
   }
+
+private:
+  /// Adds pristine registers. Pristine registers are callee saved registers
+  /// that are unused in the function.
+  void addPristines(const MachineFunction &MF);
 };
 
 } // end namespace llvm
 
 #endif // LLVM_CODEGEN_LIVEREGUNITS_H

lib/CodeGen/LivePhysRegs.cpp

 /// 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)
     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) {
+void LivePhysRegs::addPristines(const MachineFunction &MF) {
   const MachineFrameInfo &MFI = MF.getFrameInfo();
   if (!MFI.isCalleeSavedInfoValid())
     return;
+  /// This function will usually be called on an empty object, handle this
+  /// as a special case.
+  if (empty()) {
+    /// Add all callee saved regs, then remove the ones that are saved and
+    /// restored.
+    addCalleeSavedRegs(*this, MF);
+    /// Remove the ones that are not saved/restored; they are pristine.
+    for (const CalleeSavedInfo &Info : MFI.getCalleeSavedInfo())
+      removeReg(Info.getReg());
+    return;
+  }
+  /// If a callee-saved register that is not pristine is already present
+  /// in the set, we should make sure that it stays in it. Precompute the
+  /// set of pristine registers in a separate object.
   /// Add all callee saved regs, then remove the ones that are saved+restored.
-  addCalleeSavedRegs(LiveRegs, MF);
+  LivePhysRegs Pristine(*TRI);
+  addCalleeSavedRegs(Pristine, MF);
   /// Remove the ones that are not saved/restored; they are pristine.
   for (const CalleeSavedInfo &Info : MFI.getCalleeSavedInfo())
-    LiveRegs.removeReg(Info.getReg());
+    Pristine.removeReg(Info.getReg());
+  for (MCPhysReg R : Pristine)
+    addReg(R);
 }
 
 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);
   } 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()) {
       for (const CalleeSavedInfo &Info : MFI.getCalleeSavedInfo())
         if (Info.isRestored())
           addReg(Info.getReg());
     }
   }
 }
 
 void LivePhysRegs::addLiveOuts(const MachineBasicBlock &MBB) {
   const MachineFunction &MF = *MBB.getParent();
   if (!MBB.succ_empty()) {
-    addPristines(*this, MF);
+    addPristines(MF);
     addLiveOutsNoPristines(MBB);
   } else if (MBB.isReturnBlock()) {
     // For the return block: Add all callee saved registers.
     const MachineFrameInfo &MFI = MF.getFrameInfo();
     if (MFI.isCalleeSavedInfoValid())
       addCalleeSavedRegs(*this, MF);
   }
 }
 
 void LivePhysRegs::addLiveIns(const MachineBasicBlock &MBB) {
   const MachineFunction &MF = *MBB.getParent();
-  addPristines(*this, MF);
+  addPristines(MF);
   addBlockLiveIns(MBB);
 }
 
 void llvm::computeLiveIns(LivePhysRegs &LiveRegs,
                           const MachineBasicBlock &MBB) {
   const MachineFunction &MF = *MBB.getParent();
   const MachineRegisterInfo &MRI = MF.getRegInfo();
   const TargetRegisterInfo &TRI = *MRI.getTargetRegisterInfo();

lib/CodeGen/LiveRegUnits.cpp

 /// Adds all callee saved registers to \p LiveUnits.
 static void addCalleeSavedRegs(LiveRegUnits &LiveUnits,
                                const MachineFunction &MF) {
   const MachineRegisterInfo &MRI = MF.getRegInfo();
   for (const MCPhysReg *CSR = MRI.getCalleeSavedRegs(); CSR && *CSR; ++CSR)
     LiveUnits.addReg(*CSR);
 }
 
-/// Adds pristine registers to the given \p LiveUnits. Pristine registers are
-/// callee saved registers that are unused in the function.
-static void addPristines(LiveRegUnits &LiveUnits, const MachineFunction &MF) {
+void LiveRegUnits::addPristines(const MachineFunction &MF) {
   const MachineFrameInfo &MFI = MF.getFrameInfo();
   if (!MFI.isCalleeSavedInfoValid())
     return;
+  /// This function will usually be called on an empty object, handle this
+  /// as a special case.
+  if (empty()) {
+    /// Add all callee saved regs, then remove the ones that are saved and
+    /// restored.
+    addCalleeSavedRegs(*this, MF);
+    /// Remove the ones that are not saved/restored; they are pristine.
+    for (const CalleeSavedInfo &Info : MFI.getCalleeSavedInfo())
+      removeReg(Info.getReg());
+    return;
+  }
+  /// If a callee-saved register that is not pristine is already present
+  /// in the set, we should make sure that it stays in it. Precompute the
+  /// set of pristine registers in a separate object.
   /// Add all callee saved regs, then remove the ones that are saved+restored.
-  addCalleeSavedRegs(LiveUnits, MF);
+  LiveRegUnits Pristine(*TRI);
+  addCalleeSavedRegs(Pristine, MF);
   /// Remove the ones that are not saved/restored; they are pristine.
   for (const CalleeSavedInfo &Info : MFI.getCalleeSavedInfo())
-    LiveUnits.removeReg(Info.getReg());
+    Pristine.removeReg(Info.getReg());
+  addUnits(Pristine.getBitVector());
 }
 
 void LiveRegUnits::addLiveOuts(const MachineBasicBlock &MBB) {
   const MachineFunction &MF = *MBB.getParent();
   if (!MBB.succ_empty()) {
-    addPristines(*this, MF);
+    addPristines(MF);
     // To get the live-outs we simply merge the live-ins of all successors.
     for (const MachineBasicBlock *Succ : MBB.successors())
       addBlockLiveIns(*this, *Succ);
   } else if (MBB.isReturnBlock()) {
     // For the return block: Add all callee saved registers.
     const MachineFrameInfo &MFI = MF.getFrameInfo();
     if (MFI.isCalleeSavedInfoValid())
       addCalleeSavedRegs(*this, MF);
   }
 }
 
 void LiveRegUnits::addLiveIns(const MachineBasicBlock &MBB) {
   const MachineFunction &MF = *MBB.getParent();
-  addPristines(*this, MF);
+  addPristines(MF);
   addBlockLiveIns(*this, MBB);
 }

test/CodeGen/Hexagon/livephysregs-add-pristines.mir

+# RUN: llc -march=hexagon -run-pass if-converter -o - %s -verify-machineinstrs | FileCheck %s
+
+# The register r23 is live on the path bb.0->bb.2->bb.3. Make sure we add
+# an implicit use of r23 to the predicated redefinition:
+# CHECK: %r23 = A2_tfrt %p0, killed %r1, implicit %r23
+
+# LivePhysRegs::addPristines could accidentally remove a callee-saved
+# register, if it determined that it wasn't pristine. Doing that caused
+# r23 in this testcase to be dropped from the Redefs set, and subsequently
+# the necessary implicit use was not added for it.
+
+---
+name: foo
+tracksRegLiveness: true
+fixedStack:
+  - { id: 0, offset: 0, size: 4, alignment: 4, callee-saved-register: '%r23' }
+body: |
+  bb.0:
+    successors: %bb.1, %bb.2
+    liveins: %r0, %r1, %r23
+      %p0 = C2_cmpgti killed %r0, 0
+      J2_jumpf killed %p0, %bb.2, implicit-def %pc
+
+  bb.1:
+    successors: %bb.3
+    liveins: %r1
+      %r23 = A2_tfr killed %r1
+      J2_jump %bb.3, implicit-def %pc
+
+  bb.2:
+    successors: %bb.3
+    liveins: %r1, %r23
+      %r0 = A2_tfr %r1
+
+  bb.3:
+    liveins: %r23
+...