[AArch64] Extend AArch64RedundantCopyElimination to do simple copy propagation.

[gberry]

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).

[qcolombet]

test/CodeGen/AArch64/machine-copy-remove.mir

274

Could you add a positive and negative test case when the copy we want to eliminate is at the beginning of a block with two or more predecessors.

[MatzeB]

LGTM with nitpicks. Please wait for qcolombet though.

lib/Target/AArch64/AArch64RedundantCopyElimination.cpp

80

Don't repeat the function name in the doxygen comments.

84–85

could add a continue; here

121

Use MCPhysReg instead of unsigned, possibly at a few more places below as well.

Why is this called TargetRegs? How about calling this ZeroRegs or NullRegs if you don't want it confused with xzr/wzr?

127–128

could this be

if (!guaranteesZeroRegInBlock(...))
  continue;

to reduce indentation?

130–157

Should this have an additional check for situations in which all target regs are clobbered (similar to the one a few lines earlier?)

195–196

Drive-by-comment (even though the patch doesn't touch it), this can be done simpler as:
dbgs() << "Remove redundant Copy : " << *MI;

test/CodeGen/AArch64/machine-copy-remove.mir

4–13

I recently found out that you can even leave out the whole IR string literal and MIRParser will create dummy functions for you, that should work here.

(Unfortunately as soon as you need to define other IR elements like globals or declaring external functions you need the IR block and MIRParser won't create dummy function for you right now)

[gberry]

lib/Target/AArch64/AArch64RedundantCopyElimination.cpp

121

I changed all of these except for the SmallSetVector element type for MCPhysReg. That one can't be changed without also adding a DenseMapInfo<uint16_t> specialization, which doesn't seem worth doing for this change to me.

As for 'TargetRegs': Yeah, I considered renaming this variable as well. I went with 'KnownZeroRegs' in the latest version.

130–157

I don't think so. The clobbered reg tracking is for regs that are clobbered as we move backwards from the e.g. CBZ instruction in PredBB. At this point in the code we start looking forward at the start of MBB and process clobbers by removing the known zero regs from the list instead.

test/CodeGen/AArch64/machine-copy-remove.mir

4–13

Good to know.

274

I added a test (test9) to make sure we aren't removing a COPY from a block with multiple predecessors. I'm not sure what you mean by a positive test in this case though.

[gberry]

Update to address comments from Matthias and Quentin.

lib/Target/AArch64/AArch64RedundantCopyElimination.cpp

 #define DEBUG_TYPE "aarch64-copyelim"
 
 STATISTIC(NumCopiesRemoved, "Number of copies removed.");
 
 namespace {
 class AArch64RedundantCopyElimination : public MachineFunctionPass {
   const MachineRegisterInfo *MRI;
   const TargetRegisterInfo *TRI;
+  BitVector ClobberedRegs;
 
 public:
   static char ID;
   AArch64RedundantCopyElimination() : MachineFunctionPass(ID) {
     initializeAArch64RedundantCopyEliminationPass(
         *PassRegistry::getPassRegistry());
   }
   bool optimizeCopy(MachineBasicBlock *MBB);
static bool guaranteesZeroRegInBlock(MachineInstr &MI, MachineBasicBlock *MBB) {
   // Check if the current basic block is the target block to which the
   // CBZ/CBNZ instruction jumps when its Wt/Xt is zero.
   return ((Opc == AArch64::CBZW || Opc == AArch64::CBZX) &&
           MBB == MI.getOperand(1).getMBB()) ||
          ((Opc == AArch64::CBNZW || Opc == AArch64::CBNZX) &&
           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)
     return false;
 
   // Check if the predecessor has two successors, implying the block ends in a
   // conditional branch.
   MachineBasicBlock *PredMBB = *MBB->pred_begin();
   if (PredMBB->succ_size() != 2)
     return false;
 
   MachineBasicBlock::iterator CompBr = PredMBB->getLastNonDebugInstr();
   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))
+    if (!guaranteesZeroRegInBlock(*CompBr, MBB))
+      continue;
+
+    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;
-  } while (CompBr != PredMBB->begin() && CompBr->isTerminator());
+          }
+          // 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;
+          }
+        }
+      }
 
-  // We've not found a CBZ/CBNZ, time to bail out.
-  if (!guaranteesZeroRegInBlock(*CompBr, MBB))
-    return false;
+      // Stop if we get to the beginning of PredMBB.
+      if (PredI == PredMBB->begin())
+        break;
 
-  unsigned TargetReg = CompBr->getOperand(0).getReg();
-  if (!TargetReg)
-    return false;
-  assert(TargetRegisterInfo::isPhysicalRegister(TargetReg) &&
-         "Expect physical register");
+      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;
 
-  // Remember all registers aliasing with TargetReg.
-  SmallSetVector<unsigned, 8> TargetRegs;
-  for (MCRegAliasIterator AI(TargetReg, TRI, true); AI.isValid(); ++AI)
-    TargetRegs.insert(*AI);
+  } 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()));
+          !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++;
-        SmallestDef =
-            TRI->isSubRegister(SmallestDef, DefReg) ? DefReg : SmallestDef;
-        continue;
-      }
+            UsedKnownZeroRegs.insert(KnownZeroReg);
+            RemovedCopy = true;
+            break;
+          }
+        }
+      }
+
+    // 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;
       }
 
-    if (MI->modifiesRegister(TargetReg, TRI))
+    // 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);
-
-  // Clear any kills of TargetReg between CompBr and the last removed COPY.
+  // 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 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;
 }
 
 bool AArch64RedundantCopyElimination::runOnMachineFunction(
     MachineFunction &MF) {
   if (skipFunction(*MF.getFunction()))
     return false;
   TRI = MF.getSubtarget().getRegisterInfo();
   MRI = &MF.getRegInfo();
+  ClobberedRegs.resize(TRI->getNumRegs());
   bool Changed = false;
   for (MachineBasicBlock &MBB : MF)
     Changed |= optimizeCopy(&MBB);
   return Changed;
 }
 
 FunctionPass *llvm::createAArch64RedundantCopyEliminationPass() {
   return new AArch64RedundantCopyElimination();
 }

test/CodeGen/AArch64/machine-copy-remove.mir

+# RUN: llc -mtriple=aarch64--linux-gnu -run-pass=aarch64-copyelim %s -verify-machineinstrs -o - 2>&1 | FileCheck %s
+---
+# Check that bb.0 COPY is seen through to allow the bb.1 COPY of XZR to be removed.
+# CHECK-LABEL: name: test1
+# CHECK-NOT: COPY %xzr
+name:            test1
+tracksRegLiveness: true
+body:             |
+  bb.0:
+    successors: %bb.1, %bb.2
+    liveins: %x0, %x1
+
+    %x0 = COPY %x1
+    CBNZX %x1, %bb.2
+
+  bb.1:
+    successors: %bb.3
+
+    %x0 = COPY %xzr
+    B %bb.3
+
+  bb.2:
+    successors: %bb.3
+    liveins: %x1
+
+    %x0 = LDRXui %x1, 0
+
+  bb.3:
+    liveins: %x0
+
+    RET_ReallyLR implicit %x0
+
+...
+# Similar to test1, but with reversed COPY.
+# CHECK-LABEL: name: test2
+# CHECK-NOT: COPY %xzr
+name:            test2
+tracksRegLiveness: true
+body:             |
+  bb.0:
+    successors: %bb.1, %bb.2
+    liveins: %x0, %x1
+
+    %x1 = COPY %x0
+    CBNZX %x1, %bb.2
+
+  bb.1:
+    successors: %bb.3
+
+    %x0 = COPY %xzr
+    B %bb.3
+
+  bb.2:
+    successors: %bb.3
+    liveins: %x1
+
+    %x0 = LDRXui %x1, 0
+
+  bb.3:
+    liveins: %x0
+
+    RET_ReallyLR implicit %x0
+
+...
+# Similar to test1, but with a clobber that prevents removal of the XZR COPY.
+# CHECK-LABEL: name: test3
+# CHECK: COPY %xzr
+name:            test3
+tracksRegLiveness: true
+body:             |
+  bb.0:
+    successors: %bb.1, %bb.2
+    liveins: %x0, %x1, %x2
+
+    %x0 = COPY %x1
+    %x1 = LDRXui %x1, 0
+    CBNZX %x1, %bb.2
+
+  bb.1:
+    successors: %bb.3
+
+    %x0 = COPY %xzr
+    B %bb.3
+
+  bb.2:
+    successors: %bb.3
+    liveins: %x1
+
+    %x0 = LDRXui %x1, 0
+
+  bb.3:
+    liveins: %x0
+
+    RET_ReallyLR implicit %x0
+
+...
+# Similar to test2, but with a clobber that prevents removal of the XZR COPY.
+# CHECK-LABEL: name: test4
+# CHECK: COPY %xzr
+name:            test4
+tracksRegLiveness: true
+body:             |
+  bb.0:
+    successors: %bb.1, %bb.2
+    liveins: %x0, %x1, %x2
+
+    %x1 = COPY %x0
+    %x1 = LDRXui %x1, 0
+    CBNZX %x1, %bb.2
+
+  bb.1:
+    successors: %bb.3
+
+    %x0 = COPY %xzr
+    B %bb.3
+
+  bb.2:
+    successors: %bb.3
+    liveins: %x1
+
+    %x0 = LDRXui %x1, 0
+
+  bb.3:
+    liveins: %x0
+
+    RET_ReallyLR implicit %x0
+
+...
+# Similar to test2, but with a clobber that prevents removal of the XZR COPY.
+# CHECK-LABEL: name: test5
+# CHECK: COPY %xzr
+name:            test5
+tracksRegLiveness: true
+body:             |
+  bb.0:
+    successors: %bb.1, %bb.2
+    liveins: %x0, %x1, %x2
+
+    %x1 = COPY %x0
+    %x0 = LDRXui %x1, 0
+    CBNZX %x1, %bb.2
+
+  bb.1:
+    successors: %bb.3
+
+    %x0 = COPY %xzr
+    B %bb.3
+
+  bb.2:
+    successors: %bb.3
+    liveins: %x1
+
+    %x0 = LDRXui %x1, 0
+
+  bb.3:
+    liveins: %x0
+
+    RET_ReallyLR implicit %x0
+
+...
+# Similar to test1, but with two levels of COPYs.
+# CHECK-LABEL: name: test6
+# CHECK-NOT: COPY %xzr
+name:            test6
+tracksRegLiveness: true
+body:             |
+  bb.0:
+    successors: %bb.1, %bb.2
+    liveins: %x0, %x1, %x2
+
+    %x2 = COPY %x0
+    %x1 = COPY %x2
+    CBNZX %x1, %bb.2
+
+  bb.1:
+    successors: %bb.3
+
+    %x0 = COPY %xzr
+    B %bb.3
+
+  bb.2:
+    successors: %bb.3
+    liveins: %x1
+
+    %x0 = LDRXui %x1, 0
+
+  bb.3:
+    liveins: %x0
+
+    RET_ReallyLR implicit %x0
+
+...
+# Similar to test1, but with two levels of COPYs and a clobber preventing COPY of XZR removal.
+# CHECK-LABEL: name: test7
+# CHECK: COPY %xzr
+name:            test7
+tracksRegLiveness: true
+body:             |
+  bb.0:
+    successors: %bb.1, %bb.2
+    liveins: %x0, %x1, %x2
+
+    %x2 = COPY %x0
+    %x0 = LDRXui %x1, 0
+    %x1 = COPY %x2
+    CBNZX %x1, %bb.2
+
+  bb.1:
+    successors: %bb.3
+
+    %x0 = COPY %xzr
+    B %bb.3
+
+  bb.2:
+    successors: %bb.3
+    liveins: %x1
+
+    %x0 = LDRXui %x1, 0
+
+  bb.3:
+    liveins: %x0
+
+    RET_ReallyLR implicit %x0
+
+...
+# Check that the TargetRegs vector clobber update loop in
+#  AArch64RedundantCopyElimination::optimizeCopy works correctly.
+# CHECK-LABEL: name: test8
+# CHECK: x0 = COPY %xzr
+# CHECK: x1 = COPY %xzr
+name:            test8
+tracksRegLiveness: true
+body:             |
+  bb.0:
+    successors: %bb.1, %bb.2
+    liveins: %x0, %x1
+
+    %x1 = COPY %x0
+    CBNZX %x1, %bb.2
+
+  bb.1:
+    successors: %bb.3
+    liveins: %x0, %x2
+
+    %x0, %x1 = LDPXi %x2, 0
+    %x0 = COPY %xzr
+    %x1 = COPY %xzr
+    B %bb.3
+
+  bb.2:
+    successors: %bb.3
+    liveins: %x1
+
+    %x0 = LDRXui %x1, 0
+
+  bb.3:
+    liveins: %x0
+
+    RET_ReallyLR implicit %x0
+
+...
+# Check that copy isn't removed from a block with multiple predecessors.
+# CHECK-LABEL: name: test9
+# CHECK: x0 = COPY %xzr
+# CHECK-NEXT: B %bb.3
+name:            test9
+tracksRegLiveness: true
+body:             |
+  bb.0:
+    successors: %bb.1, %bb.2
+    liveins: %x0, %x1
+
+    CBNZX %x0, %bb.2
+
+  bb.1:
+    successors: %bb.3
+    liveins: %x0, %x2
+
+    %x0 = COPY %xzr
+    B %bb.3
+
+  bb.2:
+    successors: %bb.1, %bb.3
+    liveins: %x1
+
+    %x0 = LDRXui %x1, 0
+
+    CBNZX %x1, %bb.1
+
+  bb.3:
+    liveins: %x0
+
+    RET_ReallyLR implicit %x0
+
+...