[AArch64] Extend AArch64RedundantCopyElimination to do simple copy propagation.

Summary:
Extend AArch64RedundantCopyElimination to catch cases where the register
that is known to be zero is COPY'd in the predecessor block.  Before
this change, this pass would catch cases like:

      CBZW %W0, <BB#1>
  BB#1:
      %W0 = COPY %WZR // removed

After this change, cases like the one below are also caught:

      %W0 = COPY %W1
      CBZW %W1, <BB#1>
  BB#1:
      %W0 = COPY %WZR // removed

This change results in a 4% increase in static copies removed by this
pass when compiling the llvm test-suite.  It also fixes regressions
caused by doing post-RA copy propagation (a separate change to be put up
for review shortly).

Reviewers: junbuml, mcrosier, t.p.northover, qcolombet, MatzeB

Subscribers: aemerson, rengolin, llvm-commits

Differential Revision: https://reviews.llvm.org/D30113

llvm-svn: 295863

1 files changed, 127 insertions, 43 deletions

diff --git a/llvm/lib/Target/AArch64/AArch64RedundantCopyElimination.cpp b/llvm/lib/Target/AArch64/AArch64RedundantCopyElimination.cpp
index 45fd358a99b..7092061cbdf 100644
--- a/llvm/lib/Target/AArch64/AArch64RedundantCopyElimination.cpp
+++ b/llvm/lib/Target/AArch64/AArch64RedundantCopyElimination.cpp
@@ -43,6 +43,7 @@ namespace {
 class AArch64RedundantCopyElimination : public MachineFunctionPass {
   const MachineRegisterInfo *MRI;
   const TargetRegisterInfo *TRI;
+  BitVector ClobberedRegs;
 
 public:
   static char ID;
@@ -76,6 +77,28 @@ static bool guaranteesZeroRegInBlock(MachineInstr &MI, MachineBasicBlock *MBB) {
           MBB != MI.getOperand(1).getMBB());
 }
 
+/// Remember what registers the specified instruction modifies.
+static void trackRegDefs(const MachineInstr &MI, BitVector &ClobberedRegs,
+                         const TargetRegisterInfo *TRI) {
+  for (const MachineOperand &MO : MI.operands()) {
+    if (MO.isRegMask()) {
+      ClobberedRegs.setBitsNotInMask(MO.getRegMask());
+      continue;
+    }
+
+    if (!MO.isReg())
+      continue;
+    unsigned Reg = MO.getReg();
+    if (!Reg)
+      continue;
+    if (!MO.isDef())
+      continue;
+
+    for (MCRegAliasIterator AI(Reg, TRI, true); AI.isValid(); ++AI)
+      ClobberedRegs.set(*AI);
+  }
+}
+
 bool AArch64RedundantCopyElimination::optimizeCopy(MachineBasicBlock *MBB) {
   // Check if the current basic block has a single predecessor.
   if (MBB->pred_size() != 1)
@@ -91,74 +114,134 @@ bool AArch64RedundantCopyElimination::optimizeCopy(MachineBasicBlock *MBB) {
   if (CompBr == PredMBB->end())
     return false;
 
+  // Keep track of the earliest point in the PredMBB block where kill markers
+  // need to be removed if a COPY is removed.
+  MachineBasicBlock::iterator FirstUse;
+  // Registers that are known to contain zeros at the start of MBB.
+  SmallVector<MCPhysReg, 4> KnownZeroRegs;
+  // Registers clobbered in PredMBB between CompBr instruction and current
+  // instruction being checked in loop.
+  ClobberedRegs.reset();
   ++CompBr;
   do {
     --CompBr;
-    if (guaranteesZeroRegInBlock(*CompBr, MBB))
-      break;
-  } while (CompBr != PredMBB->begin() && CompBr->isTerminator());
+    if (!guaranteesZeroRegInBlock(*CompBr, MBB))
+      continue;
 
-  // We've not found a CBZ/CBNZ, time to bail out.
-  if (!guaranteesZeroRegInBlock(*CompBr, MBB))
-    return false;
+    KnownZeroRegs.push_back(CompBr->getOperand(0).getReg());
+    FirstUse = CompBr;
+    // Look backward in PredMBB for COPYs from the known zero reg to
+    // find other registers that are known to be zero.
+    for (auto PredI = CompBr;; --PredI) {
+      if (PredI->isCopy()) {
+        MCPhysReg CopyDstReg = PredI->getOperand(0).getReg();
+        MCPhysReg CopySrcReg = PredI->getOperand(1).getReg();
+        for (MCPhysReg KnownZeroReg : KnownZeroRegs) {
+          if (ClobberedRegs[KnownZeroReg])
+            continue;
+          // If we have X = COPY Y, and Y is known to be zero, then now X is
+          // known to be zero.
+          if (CopySrcReg == KnownZeroReg && !ClobberedRegs[CopyDstReg]) {
+            KnownZeroRegs.push_back(CopyDstReg);
+            FirstUse = PredI;
+            break;
+          }
+          // If we have X = COPY Y, and X is known to be zero, then now Y is
+          // known to be zero.
+          if (CopyDstReg == KnownZeroReg && !ClobberedRegs[CopySrcReg]) {
+            KnownZeroRegs.push_back(CopySrcReg);
+            FirstUse = PredI;
+            break;
+          }
+        }
+      }
 
-  unsigned TargetReg = CompBr->getOperand(0).getReg();
-  if (!TargetReg)
-    return false;
-  assert(TargetRegisterInfo::isPhysicalRegister(TargetReg) &&
-         "Expect physical register");
+      // Stop if we get to the beginning of PredMBB.
+      if (PredI == PredMBB->begin())
+        break;
 
-  // Remember all registers aliasing with TargetReg.
-  SmallSetVector<unsigned, 8> TargetRegs;
-  for (MCRegAliasIterator AI(TargetReg, TRI, true); AI.isValid(); ++AI)
-    TargetRegs.insert(*AI);
+      trackRegDefs(*PredI, ClobberedRegs, TRI);
+      // Stop if all of the known-zero regs have been clobbered.
+      if (all_of(KnownZeroRegs, [&](MCPhysReg KnownZeroReg) {
+            return ClobberedRegs[KnownZeroReg];
+          }))
+        break;
+    }
+    break;
+
+  } while (CompBr != PredMBB->begin() && CompBr->isTerminator());
+
+  // We've not found a known zero register, time to bail out.
+  if (KnownZeroRegs.empty())
+    return false;
 
   bool Changed = false;
+  // UsedKnownZeroRegs is the set of KnownZeroRegs that have had uses added to MBB.
+  SmallSetVector<unsigned, 4> UsedKnownZeroRegs;
   MachineBasicBlock::iterator LastChange = MBB->begin();
-  unsigned SmallestDef = TargetReg;
-  // Remove redundant Copy instructions unless TargetReg is modified.
+  // Remove redundant Copy instructions unless KnownZeroReg is modified.
   for (MachineBasicBlock::iterator I = MBB->begin(), E = MBB->end(); I != E;) {
     MachineInstr *MI = &*I;
     ++I;
-    if (MI->isCopy() && MI->getOperand(0).isReg() &&
-        MI->getOperand(1).isReg()) {
-
-      unsigned DefReg = MI->getOperand(0).getReg();
-      unsigned SrcReg = MI->getOperand(1).getReg();
+    bool RemovedCopy = false;
+    if (MI->isCopy()) {
+      MCPhysReg DefReg = MI->getOperand(0).getReg();
+      MCPhysReg SrcReg = MI->getOperand(1).getReg();
 
       if ((SrcReg == AArch64::XZR || SrcReg == AArch64::WZR) &&
-          !MRI->isReserved(DefReg) &&
-          (TargetReg == DefReg || TRI->isSuperRegister(DefReg, TargetReg))) {
-        DEBUG(dbgs() << "Remove redundant Copy : ");
-        DEBUG((MI)->print(dbgs()));
-
-        MI->eraseFromParent();
-        Changed = true;
-        LastChange = I;
-        NumCopiesRemoved++;
-        SmallestDef =
-            TRI->isSubRegister(SmallestDef, DefReg) ? DefReg : SmallestDef;
-        continue;
+          !MRI->isReserved(DefReg)) {
+        for (MCPhysReg KnownZeroReg : KnownZeroRegs) {
+          if (KnownZeroReg == DefReg ||
+              TRI->isSuperRegister(DefReg, KnownZeroReg)) {
+            DEBUG(dbgs() << "Remove redundant Copy : " << *MI);
+
+            MI->eraseFromParent();
+            Changed = true;
+            LastChange = I;
+            NumCopiesRemoved++;
+            UsedKnownZeroRegs.insert(KnownZeroReg);
+            RemovedCopy = true;
+            break;
+          }
+        }
       }
     }
 
-    if (MI->modifiesRegister(TargetReg, TRI))
+    // Skip to the next instruction if we removed the COPY from WZR/XZR.
+    if (RemovedCopy)
+      continue;
+
+    // Remove any regs the MI clobbers from the KnownZeroRegs set.
+    for (unsigned RI = 0; RI < KnownZeroRegs.size();)
+      if (MI->modifiesRegister(KnownZeroRegs[RI], TRI)) {
+        std::swap(KnownZeroRegs[RI], KnownZeroRegs[KnownZeroRegs.size() - 1]);
+        KnownZeroRegs.pop_back();
+        // Don't increment RI since we need to now check the swapped-in
+        // KnownZeroRegs[RI].
+      } else {
+        ++RI;
+      }
+
+    // Continue until the KnownZeroRegs set is empty.
+    if (KnownZeroRegs.empty())
       break;
   }
 
   if (!Changed)
     return false;
 
-  // Otherwise, we have to fixup the use-def chain, starting with the
-  // CBZ/CBNZ. Conservatively mark as much as we can live.
-  CompBr->clearRegisterKills(SmallestDef, TRI);
-
-  if (none_of(TargetRegs, [&](unsigned Reg) { return MBB->isLiveIn(Reg); }))
-    MBB->addLiveIn(TargetReg);
+  // Add newly used regs to the block's live-in list if they aren't there
+  // already.
+  for (MCPhysReg KnownZeroReg : UsedKnownZeroRegs)
+    if (!MBB->isLiveIn(KnownZeroReg))
+      MBB->addLiveIn(KnownZeroReg);
 
-  // Clear any kills of TargetReg between CompBr and the last removed COPY.
+  // Clear kills in the range where changes were made.  This is conservative,
+  // but should be okay since kill markers are being phased out.
+  for (MachineInstr &MMI : make_range(FirstUse, PredMBB->end()))
+    MMI.clearKillInfo();
   for (MachineInstr &MMI : make_range(MBB->begin(), LastChange))
-    MMI.clearRegisterKills(SmallestDef, TRI);
+    MMI.clearKillInfo();
 
   return true;
 }
@@ -169,6 +252,7 @@ bool AArch64RedundantCopyElimination::runOnMachineFunction(
     return false;
   TRI = MF.getSubtarget().getRegisterInfo();
   MRI = &MF.getRegInfo();
+  ClobberedRegs.resize(TRI->getNumRegs());
   bool Changed = false;
   for (MachineBasicBlock &MBB : MF)
     Changed |= optimizeCopy(&MBB);