3 files changed, 106 insertions, 59 deletions

diff --git a/llvm/lib/CodeGen/CMakeLists.txt b/llvm/lib/CodeGen/CMakeLists.txt
index 398ea88363b..fb30218a838 100644
--- a/llvm/lib/CodeGen/CMakeLists.txt
+++ b/llvm/lib/CodeGen/CMakeLists.txt
@@ -46,6 +46,7 @@ add_llvm_library(LLVMCodeGen
   LiveRangeCalc.cpp
   LiveRangeEdit.cpp
   LiveRegMatrix.cpp
+  LiveRegUnits.cpp
   LiveStackAnalysis.cpp
   LiveVariables.cpp
   LLVMTargetMachine.cpp
diff --git a/llvm/lib/CodeGen/LiveRegUnits.cpp b/llvm/lib/CodeGen/LiveRegUnits.cpp
new file mode 100644
index 00000000000..14da799a63f
--- /dev/null
+++ b/llvm/lib/CodeGen/LiveRegUnits.cpp
@@ -0,0 +1,97 @@
+//===--- LiveRegUnits.cpp - Register Unit Set -----------------------------===//
+//
+//                     The LLVM Compiler Infrastructure
+//
+// This file is distributed under the University of Illinois Open Source
+// License. See LICENSE.TXT for details.
+//
+//===----------------------------------------------------------------------===//
+//
+/// \file This file imlements the LiveRegUnits set.
+//
+//===----------------------------------------------------------------------===//
+
+#include "llvm/CodeGen/LiveRegUnits.h"
+#include "llvm/CodeGen/MachineFrameInfo.h"
+#include "llvm/CodeGen/MachineFunction.h"
+#include "llvm/CodeGen/MachineInstrBundle.h"
+using namespace llvm;
+
+void LiveRegUnits::removeRegsNotPreserved(const uint32_t *RegMask) {
+  for (unsigned U = 0, E = TRI->getNumRegUnits(); U != E; ++U) {
+    for (MCRegUnitRootIterator RootReg(U, TRI); RootReg.isValid(); ++RootReg) {
+      if (MachineOperand::clobbersPhysReg(RegMask, *RootReg))
+        Units.reset(U);
+    }
+  }
+}
+
+void LiveRegUnits::stepBackward(const MachineInstr &MI) {
+  // Remove defined registers and regmask kills from the set.
+  for (ConstMIBundleOperands O(MI); O.isValid(); ++O) {
+    if (O->isReg()) {
+      if (!O->isDef())
+        continue;
+      unsigned Reg = O->getReg();
+      if (!TargetRegisterInfo::isPhysicalRegister(Reg))
+        continue;
+      removeReg(Reg);
+    } else if (O->isRegMask())
+      removeRegsNotPreserved(O->getRegMask());
+  }
+
+  // Add uses to the set.
+  for (ConstMIBundleOperands O(MI); O.isValid(); ++O) {
+    if (!O->isReg() || !O->readsReg())
+      continue;
+    unsigned Reg = O->getReg();
+    if (!TargetRegisterInfo::isPhysicalRegister(Reg))
+      continue;
+    addReg(Reg);
+  }
+}
+
+/// Add live-in registers of basic block \p MBB to \p LiveUnits.
+static void addLiveIns(LiveRegUnits &LiveUnits, const MachineBasicBlock &MBB) {
+  for (const auto &LI : MBB.liveins())
+    LiveUnits.addRegMasked(LI.PhysReg, LI.LaneMask);
+}
+
+static void addLiveOuts(LiveRegUnits &LiveUnits, const MachineBasicBlock &MBB) {
+  // To get the live-outs we simply merge the live-ins of all successors.
+  for (const MachineBasicBlock *Succ : MBB.successors())
+    addLiveIns(LiveUnits, *Succ);
+}
+
+/// Add pristine registers to the given \p LiveUnits. This function removes
+/// actually saved callee save registers when \p InPrologueEpilogue is false.
+static void removeSavedRegs(LiveRegUnits &LiveUnits, const MachineFunction &MF,
+                            const MachineFrameInfo &MFI,
+                            const TargetRegisterInfo &TRI) {
+  for (const CalleeSavedInfo &Info : MFI.getCalleeSavedInfo())
+    LiveUnits.removeReg(Info.getReg());
+}
+
+void LiveRegUnits::addLiveOuts(const MachineBasicBlock &MBB) {
+  const MachineFunction &MF = *MBB.getParent();
+  const MachineFrameInfo &MFI = MF.getFrameInfo();
+  if (MFI.isCalleeSavedInfoValid()) {
+    for (const MCPhysReg *I = TRI->getCalleeSavedRegs(&MF); *I; ++I)
+      addReg(*I);
+    if (!MBB.isReturnBlock())
+      removeSavedRegs(*this, MF, MFI, *TRI);
+  }
+  ::addLiveOuts(*this, MBB);
+}
+
+void LiveRegUnits::addLiveIns(const MachineBasicBlock &MBB) {
+  const MachineFunction &MF = *MBB.getParent();
+  const MachineFrameInfo &MFI = MF.getFrameInfo();
+  if (MFI.isCalleeSavedInfoValid()) {
+    for (const MCPhysReg *I = TRI->getCalleeSavedRegs(&MF); *I; ++I)
+      addReg(*I);
+    if (&MBB != &MF.front())
+      removeSavedRegs(*this, MF, MFI, *TRI);
+  }
+  ::addLiveIns(*this, MBB);
+}
diff --git a/llvm/lib/CodeGen/RegisterScavenging.cpp b/llvm/lib/CodeGen/RegisterScavenging.cpp
index fdf741fd58f..23e07581fc8 100644
--- a/llvm/lib/CodeGen/RegisterScavenging.cpp
+++ b/llvm/lib/CodeGen/RegisterScavenging.cpp
@@ -32,11 +32,7 @@ using namespace llvm;
 #define DEBUG_TYPE "reg-scavenging"
 
 void RegScavenger::setRegUsed(unsigned Reg, LaneBitmask LaneMask) {
-  for (MCRegUnitMaskIterator RUI(Reg, TRI); RUI.isValid(); ++RUI) {
-    LaneBitmask UnitMask = (*RUI).second;
-    if (UnitMask.none() || (LaneMask & UnitMask).any())
-      RegUnitsAvailable.reset((*RUI).first);
-  }
+  LiveUnits.addRegMasked(Reg, LaneMask);
 }
 
 void RegScavenger::init(MachineBasicBlock &MBB) {
@@ -44,6 +40,7 @@ void RegScavenger::init(MachineBasicBlock &MBB) {
   TII = MF.getSubtarget().getInstrInfo();
   TRI = MF.getSubtarget().getRegisterInfo();
   MRI = &MF.getRegInfo();
+  LiveUnits.init(*TRI);
 
   assert((NumRegUnits == 0 || NumRegUnits == TRI->getNumRegUnits()) &&
          "Target changed?");
@@ -51,7 +48,6 @@ void RegScavenger::init(MachineBasicBlock &MBB) {
   // Self-initialize.
   if (!this->MBB) {
     NumRegUnits = TRI->getNumRegUnits();
-    RegUnitsAvailable.resize(NumRegUnits);
     KillRegUnits.resize(NumRegUnits);
     DefRegUnits.resize(NumRegUnits);
     TmpRegUnits.resize(NumRegUnits);
@@ -64,32 +60,17 @@ void RegScavenger::init(MachineBasicBlock &MBB) {
     I->Restore = nullptr;
   }
 
-  // All register units start out unused.
-  RegUnitsAvailable.set();
-
-  // Pristine CSRs are not available.
-  BitVector PR = MF.getFrameInfo().getPristineRegs(MF);
-  for (int I = PR.find_first(); I>0; I = PR.find_next(I))
-    setRegUsed(I);
-
   Tracking = false;
 }
 
-void RegScavenger::setLiveInsUsed(const MachineBasicBlock &MBB) {
-  for (const auto &LI : MBB.liveins())
-    setRegUsed(LI.PhysReg, LI.LaneMask);
-}
-
 void RegScavenger::enterBasicBlock(MachineBasicBlock &MBB) {
   init(MBB);
-  setLiveInsUsed(MBB);
+  LiveUnits.addLiveIns(MBB);
 }
 
 void RegScavenger::enterBasicBlockEnd(MachineBasicBlock &MBB) {
   init(MBB);
-  // Merge live-ins of successors to get live-outs.
-  for (const MachineBasicBlock *Succ : MBB.successors())
-    setLiveInsUsed(*Succ);
+  LiveUnits.addLiveOuts(MBB);
 
   // Move internal iterator at the last instruction of the block.
   if (MBB.begin() != MBB.end()) {
@@ -263,36 +244,7 @@ void RegScavenger::backward() {
   assert(Tracking && "Must be tracking to determine kills and defs");
 
   const MachineInstr &MI = *MBBI;
-  // Defined or clobbered registers are available now.
-  for (const MachineOperand &MO : MI.operands()) {
-    if (MO.isRegMask()) {
-      for (unsigned RU = 0, RUEnd = TRI->getNumRegUnits(); RU != RUEnd;
-           ++RU) {
-        for (MCRegUnitRootIterator RURI(RU, TRI); RURI.isValid(); ++RURI) {
-          if (MO.clobbersPhysReg(*RURI)) {
-            RegUnitsAvailable.set(RU);
-            break;
-          }
-        }
-      }
-    } else if (MO.isReg() && MO.isDef()) {
-      unsigned Reg = MO.getReg();
-      if (!Reg || TargetRegisterInfo::isVirtualRegister(Reg) ||
-          isReserved(Reg))
-        continue;
-      addRegUnits(RegUnitsAvailable, Reg);
-    }
-  }
-  // Mark read registers as unavailable.
-  for (const MachineOperand &MO : MI.uses()) {
-    if (MO.isReg() && MO.readsReg()) {
-      unsigned Reg = MO.getReg();
-      if (!Reg || TargetRegisterInfo::isVirtualRegister(Reg) ||
-          isReserved(Reg))
-        continue;
-      removeRegUnits(RegUnitsAvailable, Reg);
-    }
-  }
+  LiveUnits.stepBackward(MI);
 
   if (MBBI == MBB->begin()) {
     MBBI = MachineBasicBlock::iterator(nullptr);
@@ -302,12 +254,9 @@ void RegScavenger::backward() {
 }
 
 bool RegScavenger::isRegUsed(unsigned Reg, bool includeReserved) const {
-  if (includeReserved && isReserved(Reg))
-    return true;
-  for (MCRegUnitIterator RUI(Reg, TRI); RUI.isValid(); ++RUI)
-    if (!RegUnitsAvailable.test(*RUI))
-      return true;
-  return false;
+  if (isReserved(Reg))
+    return includeReserved;
+  return !LiveUnits.available(Reg);
 }
 
 unsigned RegScavenger::FindUnusedReg(const TargetRegisterClass *RC) const {