Refactor recomputeLiveIns to operate on whole CFG (#79498)

Currently, the way that recomputeLiveIns works is that it will recompute the livein registers for that MachineBasicBlock but it matters what order you call recomputeLiveIn which can result in incorrect register allocations down the line. This PR fixes that by simply recomputing the liveins for the entire CFG until convergence is achieved. This makes it harder to introduce subtle bugs which alter liveness.

Author: oskarwirga

llvm/include/llvm/CodeGen/LivePhysRegs.h

@@ -31,6 +31,7 @@
 
 #include "llvm/ADT/SparseSet.h"
 #include "llvm/CodeGen/MachineBasicBlock.h"
+#include "llvm/CodeGen/MachineFunction.h"
 #include "llvm/CodeGen/TargetRegisterInfo.h"
 #include "llvm/MC/MCRegister.h"
 #include "llvm/MC/MCRegisterInfo.h"
@@ -193,11 +194,31 @@ void addLiveIns(MachineBasicBlock &MBB, const LivePhysRegs &LiveRegs);
 void computeAndAddLiveIns(LivePhysRegs &LiveRegs,
  MachineBasicBlock &MBB);
 
-/// Convenience function for recomputing live-in's for \p MBB.
-static inline void recomputeLiveIns(MachineBasicBlock &MBB) {
+/// Function to update the live-in's for a basic block and return whether any
+/// changes were made.
+static inline bool updateBlockLiveInfo(MachineBasicBlock &MBB) {
  LivePhysRegs LPR;
+ auto oldLiveIns = MBB.getLiveIns();
+
  MBB.clearLiveIns();
  computeAndAddLiveIns(LPR, MBB);
+ MBB.sortUniqueLiveIns();
+
+ auto newLiveIns = MBB.getLiveIns();
+ return oldLiveIns != newLiveIns;
+}
+
+/// Convenience function for recomputing live-in's for the entire CFG until
+/// convergence is reached.
+static inline void recomputeLiveIns(MachineFunction &MF) {
+ bool anyChanged;
+ do {
+ anyChanged = false;
+ for (auto MFI = MF.rbegin(), MFE = MF.rend(); MFI != MFE; ++MFI) {
+ MachineBasicBlock &MBB = *MFI;
+ anyChanged |= updateBlockLiveInfo(MBB);
+ }
+ } while (anyChanged);
 }
 
 } // end namespace llvm

llvm/include/llvm/CodeGen/MachineBasicBlock.h

@@ -111,6 +111,10 @@ class MachineBasicBlock
 
  RegisterMaskPair(MCPhysReg PhysReg, LaneBitmask LaneMask)
  : PhysReg(PhysReg), LaneMask(LaneMask) {}
+
+ bool operator==(const RegisterMaskPair &other) const {
+ return PhysReg == other.PhysReg && LaneMask == other.LaneMask;
+ }
  };
 
 private:
@@ -473,6 +477,8 @@ class MachineBasicBlock
  /// Remove entry from the livein set and return iterator to the next.
  livein_iterator removeLiveIn(livein_iterator I);
 
+ std::vector<RegisterMaskPair> getLiveIns() const { return LiveIns; }
+
  class liveout_iterator {
  public:
  using iterator_category = std::input_iterator_tag;

llvm/lib/CodeGen/BranchFolding.cpp

@@ -2048,8 +2048,7 @@ bool BranchFolder::HoistCommonCodeInSuccs(MachineBasicBlock *MBB) {
  FBB->erase(FBB->begin(), FIB);
 
  if (UpdateLiveIns) {
- recomputeLiveIns(*TBB);
- recomputeLiveIns(*FBB);
+ recomputeLiveIns(*MBB->getParent());
  }
 
  ++NumHoist;

llvm/lib/Target/AArch64/AArch64FrameLowering.cpp

@@ -4339,8 +4339,7 @@ AArch64FrameLowering::inlineStackProbeLoopExactMultiple(
  ExitMBB->transferSuccessorsAndUpdatePHIs(&MBB);
  MBB.addSuccessor(LoopMBB);
  // Update liveins.
- recomputeLiveIns(*LoopMBB);
- recomputeLiveIns(*ExitMBB);
+ recomputeLiveIns(MF);
 
  return ExitMBB->begin();
 }

llvm/lib/Target/AArch64/AArch64InstrInfo.cpp

@@ -9597,9 +9597,7 @@ AArch64InstrInfo::probedStackAlloc(MachineBasicBlock::iterator MBBI,
 
  // Update liveins.
  if (MF.getRegInfo().reservedRegsFrozen()) {
- recomputeLiveIns(*LoopTestMBB);
- recomputeLiveIns(*LoopBodyMBB);
- recomputeLiveIns(*ExitMBB);
+ recomputeLiveIns(MF);
  }
 
  return ExitMBB->begin();

llvm/lib/Target/ARM/ARMLowOverheadLoops.cpp

@@ -1806,12 +1806,7 @@ void ARMLowOverheadLoops::Expand(LowOverheadLoop &LoLoop) {
  PostOrderLoopTraversal DFS(LoLoop.ML, *MLI);
  DFS.ProcessLoop();
  const SmallVectorImpl<MachineBasicBlock*> &PostOrder = DFS.getOrder();
- for (auto *MBB : PostOrder) {
- recomputeLiveIns(*MBB);
- // FIXME: For some reason, the live-in print order is non-deterministic for
- // our tests and I can't out why... So just sort them.
- MBB->sortUniqueLiveIns();
- }
+ recomputeLiveIns(*LoLoop.MF);
 
  for (auto *MBB : reverse(PostOrder))
  recomputeLivenessFlags(*MBB);

llvm/lib/Target/PowerPC/PPCExpandAtomicPseudoInsts.cpp

@@ -208,8 +208,7 @@ bool PPCExpandAtomicPseudo::expandAtomicRMW128(
  .addMBB(LoopMBB);
  CurrentMBB->addSuccessor(LoopMBB);
  CurrentMBB->addSuccessor(ExitMBB);
- recomputeLiveIns(*LoopMBB);
- recomputeLiveIns(*ExitMBB);
+ recomputeLiveIns(*MF);
  NMBBI = MBB.end();
  MI.eraseFromParent();
  return true;
@@ -286,9 +285,7 @@ bool PPCExpandAtomicPseudo::expandAtomicCmpSwap128(
  CurrentMBB->addSuccessor(LoopCmpMBB);
  CurrentMBB->addSuccessor(ExitMBB);
 
- recomputeLiveIns(*LoopCmpMBB);
- recomputeLiveIns(*CmpSuccMBB);
- recomputeLiveIns(*ExitMBB);
+ recomputeLiveIns(*MF);
  NMBBI = MBB.end();
  MI.eraseFromParent();
  return true;

llvm/lib/Target/PowerPC/PPCFrameLowering.cpp

@@ -1441,8 +1441,7 @@ void PPCFrameLowering::inlineStackProbe(MachineFunction &MF,
  ProbeLoopBodyMBB->addSuccessor(ProbeLoopBodyMBB);
  }
  // Update liveins.
- recomputeLiveIns(*ProbeLoopBodyMBB);
- recomputeLiveIns(*ProbeExitMBB);
+ recomputeLiveIns(MF);
  return ProbeExitMBB;
  };
  // For case HasBP && MaxAlign > 1, we have to realign the SP by performing
@@ -1534,8 +1533,7 @@ void PPCFrameLowering::inlineStackProbe(MachineFunction &MF,
  buildDefCFAReg(*ExitMBB, ExitMBB->begin(), SPReg);
  }
  // Update liveins.
- recomputeLiveIns(*LoopMBB);
- recomputeLiveIns(*ExitMBB);
+ recomputeLiveIns(MF);
  }
  }
  ++NumPrologProbed;

llvm/lib/Target/SystemZ/SystemZFrameLowering.cpp

@@ -840,8 +840,7 @@ void SystemZELFFrameLowering::inlineStackProbe(
  StackAllocMI->eraseFromParent();
  if (DoneMBB != nullptr) {
  // Compute the live-in lists for the new blocks.
- recomputeLiveIns(*DoneMBB);
- recomputeLiveIns(*LoopMBB);
+ recomputeLiveIns(MF);
  }
 }
 
@@ -1439,8 +1438,7 @@ void SystemZXPLINKFrameLowering::inlineStackProbe(
  StackAllocMI->eraseFromParent();
 
  // Compute the live-in lists for the new blocks.
- recomputeLiveIns(*NextMBB);
- recomputeLiveIns(*StackExtMBB);
+ recomputeLiveIns(MF);
 }
 
 bool SystemZXPLINKFrameLowering::hasFP(const MachineFunction &MF) const {

llvm/lib/Target/X86/X86FrameLowering.cpp

@@ -885,8 +885,7 @@ void X86FrameLowering::emitStackProbeInlineGenericLoop(
  }
 
  // Update Live In information
- recomputeLiveIns(*testMBB);
- recomputeLiveIns(*tailMBB);
+ recomputeLiveIns(MF);
 }
 
 void X86FrameLowering::emitStackProbeInlineWindowsCoreCLR64(
@@ -1378,10 +1377,7 @@ void X86FrameLowering::BuildStackAlignAND(MachineBasicBlock &MBB,
  footMBB->addSuccessor(&MBB);
  }
 
- recomputeLiveIns(*headMBB);
- recomputeLiveIns(*bodyMBB);
- recomputeLiveIns(*footMBB);
- recomputeLiveIns(MBB);
+ recomputeLiveIns(MF);
  }
  } else {
  MachineInstr *MI = BuildMI(MBB, MBBI, DL, TII.get(AndOp), Reg)