[AArch64] Extend redundant copy elimination pass to handle non-zero stores.

[mcrosier]

Summary

This patch extends the current functionality of the AArch64 redundant copy elimination pass to handle non-zero cases such as:

BB#0:
  cmp x0, #1
  b.eq .LBB0_1
.LBB0_1:
  orr x0, xzr, #0x1  ; <-- redundant copy; x0 known to hold #1.

This removes redundant copies from most of the SPEC200X tests. When running on Kryo, I saw a 2.7% and 1% improvement in spec2000/gap and spec2006/povray, respectively. All other testing was within noise.

Chad

[junbuml]

lib/Target/AArch64/AArch64RedundantCopyElimination.cpp

121 ↗

(On Diff #86479)

Is it always safe to do this directly without any check for isConditionalBranch(), Opcode, or MI.getOperand(0).isImm()?

137 ↗

(On Diff #86479)

If we didn't see any instruction that sets NZCV between MI and PredMBB->begin() and NZCV is live-in from PredMBB's single predecessor, then we could continue scan the single predecessor of PredMBB. Maybe different patch if it make sense.

160 ↗

(On Diff #86479)

At this point we know that I->getOperand(0) is zero register, so this if statement will be taken only when I->getOperand(1) is zero register. I don't see any problem when both getOperand(0) and getOperand(1) are zero register.

222 ↗

(On Diff #86479)

Please update the comment as we continue checking non-zero case.

233 ↗

(On Diff #86479)

Is there any case where both ZeroRegVal and KnownRegVal is true? If not, adding assert() will be good.

264 ↗

(On Diff #86479)

We could move check for ZeroRegVal above if statement in line 259.

283 ↗

(On Diff #86479)

The comment in MachineInstr.h mention that isMoveImmediate() return true for a move immediate (including conditional moves). I think the conditional move is not the case for AArch64, but just want to double confirm.

283 ↗

(On Diff #86479)

We can move KnownRegVal in line 283. It seems that both ZeroRegVal and KnownRegVal cannot be true at the same time. Then we can use if-else, instead of if-if.

[mcrosier]

lib/Target/AArch64/AArch64RedundantCopyElimination.cpp

121 ↗

(On Diff #86479)

Yes, I believe so. We directly check that the opcode is Bcc above, so we know MI must be a conditional branch at this point. AFAIK, the 0 operand is always the CC for a Bcc.

137 ↗

(On Diff #86479)

True, but I'd prefer to defer this for another day as you suggest.

160 ↗

(On Diff #86479)

You're correct. I'll remove accordingly.

222 ↗

(On Diff #86479)

Will do.

233 ↗

(On Diff #86479)

No. I'll add an assertion.

283 ↗

(On Diff #86479)

For AArch64, only the MOVi32imm and MOVi64imm pseudos will return true for isMoveImmediate() and these instructions are unconditionally executed.

[mcrosier]

Address Jun's comments.

[mcrosier]

lib/Target/AArch64/AArch64RedundantCopyElimination.cpp

283 ↗

(On Diff #86479)

I hoisted the ZeroRegVal and KnownRegVal to the outer checks, but left the code as an if-if, rather thank if-else to reduce indention. Let me know if you strongly prefer the the if-else and I'll happily

[mcrosier]

Include test case that was accidentally dropped from previous diff.

[MatzeB]

This looks good to me when the comments below are addressed.

lib/Target/AArch64/AArch64RedundantCopyElimination.cpp

84–86 ↗

(On Diff #89374)

Looks like you synced that with https://reviews.llvm.org/D30113, good.

105–108 ↗

(On Diff #89374)

Use doxygen /// comment. It would help to talk about block \p MBB instead of current block, similarily to explain some of the other parameters.

120 ↗

(On Diff #89374)

Is that cast to (int) here necessary?

129–130 ↗

(On Diff #89374)

Add a comment that this is a shortcut for a block that is empty except for the branch.

136 ↗

(On Diff #89374)

This slightly confused me as most loops in LLVM first initialize/declare the iteration variable then the end variable.

This could also be written as a range based for: for (MachineInstr &PredMI : make_range(std::next(MI.getReverseIterator()), PredMBB->rend()).

146–148 ↗

(On Diff #89374)

does the DestReg matter here? I would assume this transformation to work as well when we write to an actual register...

155 ↗

(On Diff #89374)

Somewhat off-topic from this specific patch and I don't expect this to be fixed here: It feels bad to scatter all those magic values like the 3 here around the targets, this function is another example where we have a lot of magic input numbers. At some point we should add some getter/accessor functions somewhere to give them a name or extend tablegen to generate some enum constants for us...

170 ↗

(On Diff #89374)

typo but->that

182–183 ↗

(On Diff #89374)

You can use DEBUG(dbgs() << "Remove redundant copy: " << *MI);

218 ↗

(On Diff #89374)

Can't you just have a single variable KnownVal or similar here? You can then use KnownVal == 0 instead of ZeroRegVal.

271–280 ↗

(On Diff #89374)

Do we really need to split this into Case 1/Case 2? There seems to be more equal code than different one so I would expect this to be better merged (with a slightly more complicated condition to check).

328 ↗

(On Diff #89374)

You may even delay the resizing/allocation of the bitvector to the point where it is actually used (resizing it multiple times to the same size shouldn't hurt performance).

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

1 ↗

(On Diff #89374)

The 2>&1 is no longer necessary for .mir tests and is discouraged now as it hides abort()/assert() messages if they happen.

5–8 ↗

(On Diff #89374)

The IR block can be left out completely in this case the MIRParser will create the dummy functions.

14 ↗

(On Diff #89374)

This can be slightly simplified by leaving out the block frequencies (i.e. just successors: %bb.1, %bb.2 Similar for the other blocks.

38 ↗

(On Diff #89374)

Typo!
I also tend to use CHECK-LABEL: name: test1 to better match the mir output with the typo this probably matched the dummy IR function instead of the beginning of the actual test1 mir function.

[mcrosier]

lib/Target/AArch64/AArch64RedundantCopyElimination.cpp

120 ↗

(On Diff #89374)

It is not. Will remove.

146–148 ↗

(On Diff #89374)

No, this is a simplification not a requirement. I'll remove this check. I think the only necessary check is to make sure the compare doesn't clobber itself (e.g., SUBS x1, x1, 5).

155 ↗

(On Diff #89374)

I agree.

[mcrosier]

Address Matthias' comments as well as rebase/rework patch after r295863.

[junbuml]

LGTM

lib/Target/AArch64/AArch64RedundantCopyElimination.cpp

115 ↗

(On Diff #90051)

This return the compare instruction the branch is predicated upon.

214 ↗

(On Diff #90051)

Please remove it.

215–226 ↗

(On Diff #90051)

Isn't it possible to merge the "if" and "else if" block as these two block do almost the same thing ? Maybe move guaranteesZeroRegInBlock() in knownRegValInBlock() ?

[junbuml]

lib/Target/AArch64/AArch64RedundantCopyElimination.cpp

115 ↗

(On Diff #90051)

Please ignore this.

[mcrosier]

Address Jun's feedback. Thanks, Jun.

[gberry]

A few minor comments below. Two questions:

• can we avoid doing two backward scans of the pred block? perhaps just tracking the clobbers once
• can you add some negative tests that check that clobbers inhibit this optimization when expected?

lib/Target/AArch64/AArch64RedundantCopyElimination.cpp

57 ↗

(On Diff #90069)

This typedef seems unnecessary. Why not just struct RegImm?

59 ↗

(On Diff #90069)

This could be int32_t to save some space since SUBS can only reference 12-bit immediates (24 with lsl 12 which isn't currently handled by this change), but it probably isn't worth it since there aren't many RegImm objects.

171 ↗

(On Diff #90069)

Typo 'include' -> 'includes'

226 ↗

(On Diff #90069)

Perhaps invert the condition above and put the continue inside it to get rid of the 'else' here

114 ↗

(On Diff #90051)

Typo 'whos' -> 'whose'

[mcrosier]

-Address Geoff's inline comments.
-Add a few negative tests.
-Use a single ClobberReg bit vector to track clobbered registers.

While I think it's possible to only scan the predecessor instructions once, I believe it would add a fair amount of complexity. More importantly we would need to maintain a list of copies every time we scan the predecessors (even in those cases when we can't optimize the move immediate/copy, which is the common case). Regardless, I'm happy to discuss this offline and update everyone with the result of that discussion.

[gberry]

I'm convinced the two passes over PredBB makes sense.

Could you add a couple of test cases where the known register and the eliminated MOV are different register sizes? I feel like we may get into trouble in the case where we're depending on the MOVi32imm zeroing out the high 32-bits of a 64-bit value. For example:
b1:

CMP w0, 1
BCC EQ, b2

b2:

w0 = MOVi32imm 1
STORE x0 # depends on high 32-bits of w0 being zero'd by MOVi32imm

A few more questions/suggestions below.

lib/Target/AArch64/AArch64RedundantCopyElimination.cpp

62 ↗

(On Diff #90194)

Perhaps this could return an Optional<RegImm>, then you could get rid of the two out ref parameters?

173 ↗

(On Diff #90194)

I think you could get rid of lines 172 and 173 if you do the trackRegDefs before the ClobberedRegs check.

215 ↗

(On Diff #90194)

I believe this is no longer needed here.

355 ↗

(On Diff #90194)

I don't think FirstUse is being calculated correctly in the case of the non-zero immediate. E.g. if there are no COPYs to worry about, FirstUse should be the SUBS compare instruction since we want to remove the kill of the known value register in that case if it is present. Your test case test10 should show this problem, though I'm not sure why -verfiy-machineinstrs isn't catching it.

[mcrosier]

Address Geoff's latest comments.

[mcrosier]

I added the test cases you suggested (and cleaned up the others). In particular, I added the test (i.e., test 13) that has a 32 bit comparison and a move immediate that implicitly defines the upper bits. We can definitely get into trouble here. I added the necessary checks to make sure we don't remove the mov immediate in this case because that would leave the upper bits in some unknown state.

lib/Target/AArch64/AArch64RedundantCopyElimination.cpp

355 ↗

(On Diff #90194)

You are correct. I've updated the logic and test cases to make sure the kill flags are removed.

[junbuml]

lib/Target/AArch64/AArch64RedundantCopyElimination.cpp

314–315 ↗

(On Diff #90352)

If this is needed, it will be good to add comment why this is needed only for mov, not in Copy.

[gberry]

LGTM with a few nits

lib/Target/AArch64/AArch64RedundantCopyElimination.cpp

212 ↗

(On Diff #90352)

Typo 'inbetween' -> 'in between'

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

408 ↗

(On Diff #90352)

Typo 'eliminated' -> 'eliminate'. And another before test15.

llvm/trunk/lib/Target/AArch64/AArch64RedundantCopyElimination.cpp

 //=- AArch64RedundantCopyElimination.cpp - Remove useless copy for AArch64 -=//
 //
 //                     The LLVM Compiler Infrastructure
 //
 // This file is distributed under the University of Illinois Open Source
 // License. See LICENSE.TXT for details.
 //
 // This pass removes unnecessary zero copies in BBs that are targets of
 // cbz/cbnz instructions. For instance, the copy instruction in the code below
 // can be removed because the CBZW jumps to BB#2 when W0 is zero.
 //  BB#1:
 //    CBZW %W0, <BB#2>
 //  BB#2:
 //    %W0 = COPY %WZR
-// This pass should be run after register allocation.
+// Similarly, this pass also handles non-zero copies.
+//  BB#0:
+//    cmp x0, #1
+//    b.eq .LBB0_1
+//  .LBB0_1:
+//    orr x0, xzr, #0x1
 //
-// FIXME: This should be extended to handle any constant other than zero. E.g.,
-//   cmp w0, #1
-//     b.eq .BB1
-//   BB1:
-//     mov w0, #1
+// This pass should be run after register allocation.
 //
 // FIXME: This could also be extended to check the whole dominance subtree below
 // the comparison if the compile time regression is acceptable.
 //
 //===----------------------------------------------------------------------===//
 
 #include "AArch64.h"
+#include "llvm/ADT/Optional.h"
 #include "llvm/ADT/SetVector.h"
 #include "llvm/ADT/Statistic.h"
 #include "llvm/ADT/iterator_range.h"
 #include "llvm/CodeGen/MachineFunctionPass.h"
 #include "llvm/CodeGen/MachineRegisterInfo.h"
 #include "llvm/Support/Debug.h"
 
 using namespace llvm;
class AArch64RedundantCopyElimination : public MachineFunctionPass {
   BitVector ClobberedRegs;
 
 public:
   static char ID;
   AArch64RedundantCopyElimination() : MachineFunctionPass(ID) {
     initializeAArch64RedundantCopyEliminationPass(
         *PassRegistry::getPassRegistry());
   }
+
+  struct RegImm {
+    MCPhysReg Reg;
+    int32_t Imm;
+    RegImm(MCPhysReg Reg, int32_t Imm) : Reg(Reg), Imm(Imm) {}
+  };
+
+  Optional<RegImm> knownRegValInBlock(MachineInstr &CondBr,
+                                      MachineBasicBlock *MBB,
+                                      MachineBasicBlock::iterator &FirstUse);
   bool optimizeCopy(MachineBasicBlock *MBB);
   bool runOnMachineFunction(MachineFunction &MF) override;
   MachineFunctionProperties getRequiredProperties() const override {
     return MachineFunctionProperties().set(
         MachineFunctionProperties::Property::NoVRegs);
   }
   StringRef getPassName() const override {
     return "AArch64 Redundant Copy Elimination";
   }
 };
 char AArch64RedundantCopyElimination::ID = 0;
 }
 
 INITIALIZE_PASS(AArch64RedundantCopyElimination, "aarch64-copyelim",
                 "AArch64 redundant copy elimination pass", false, false)
 
-static bool guaranteesZeroRegInBlock(MachineInstr &MI, MachineBasicBlock *MBB) {
-  unsigned Opc = MI.getOpcode();
-  // 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);
   }
 }
 
+/// It's possible to determine the value of a register based on a dominating
+/// condition.  To do so, this function checks to see if the basic block \p MBB
+/// is the target to which a conditional branch \p CondBr jumps and whose
+/// equality comparison is against a constant.  If so, return a known physical
+/// register and constant value pair.  Otherwise, return None.
+Optional<AArch64RedundantCopyElimination::RegImm>
+AArch64RedundantCopyElimination::knownRegValInBlock(
+    MachineInstr &CondBr, MachineBasicBlock *MBB,
+    MachineBasicBlock::iterator &FirstUse) {
+  unsigned Opc = CondBr.getOpcode();
+
+  // Check if the current basic block is the target block to which the
+  // CBZ/CBNZ instruction jumps when its Wt/Xt is zero.
+  if (((Opc == AArch64::CBZW || Opc == AArch64::CBZX) &&
+       MBB == CondBr.getOperand(1).getMBB()) ||
+      ((Opc == AArch64::CBNZW || Opc == AArch64::CBNZX) &&
+       MBB != CondBr.getOperand(1).getMBB())) {
+    FirstUse = CondBr;
+    return RegImm(CondBr.getOperand(0).getReg(), 0);
+  }
+
+  // Otherwise, must be a conditional branch.
+  if (Opc != AArch64::Bcc)
+    return None;
+
+  // Must be an equality check (i.e., == or !=).
+  AArch64CC::CondCode CC = (AArch64CC::CondCode)CondBr.getOperand(0).getImm();
+  if (CC != AArch64CC::EQ && CC != AArch64CC::NE)
+    return None;
+
+  MachineBasicBlock *BrTarget = CondBr.getOperand(1).getMBB();
+  if ((CC == AArch64CC::EQ && BrTarget != MBB) ||
+      (CC == AArch64CC::NE && BrTarget == MBB))
+    return None;
+
+  // Stop if we get to the beginning of PredMBB.
+  MachineBasicBlock *PredMBB = *MBB->pred_begin();
+  assert(PredMBB == CondBr.getParent() &&
+         "Conditional branch not in predecessor block!");
+  if (CondBr == PredMBB->begin())
+    return None;
+
+  // Registers clobbered in PredMBB between CondBr instruction and current
+  // instruction being checked in loop.
+  ClobberedRegs.reset();
+
+  // Find compare instruction that sets NZCV used by CondBr.
+  MachineBasicBlock::reverse_iterator RIt = CondBr.getReverseIterator();
+  for (MachineInstr &PredI : make_range(std::next(RIt), PredMBB->rend())) {
+
+    // Track clobbered registers.
+    trackRegDefs(PredI, ClobberedRegs, TRI);
+
+    switch (PredI.getOpcode()) {
+    default:
+      break;
+
+    // CMP is an alias for SUBS with a dead destination register.
+    case AArch64::SUBSWri:
+    case AArch64::SUBSXri: {
+      unsigned SrcReg = PredI.getOperand(1).getReg();
+      // Must not be a symbolic immediate.
+      if (!PredI.getOperand(2).isImm())
+        return None;
+
+      // FIXME: For simplicity, give up on non-zero shifts.
+      if (PredI.getOperand(3).getImm())
+        return None;
+
+      // The src register must not be modified between the cmp and conditional
+      // branch.  This includes a self-clobbering compare.
+      if (ClobberedRegs[SrcReg])
+        return None;
+
+      // We've found the Cmp that sets NZCV.
+      FirstUse = PredI;
+      return RegImm(PredI.getOperand(1).getReg(), PredI.getOperand(2).getImm());
+    }
+    }
+
+    // Bail if we see an instruction that defines NZCV that we don't handle.
+    if (PredI.definesRegister(AArch64::NZCV))
+      return None;
+  }
+  return None;
+}
+
 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())
+  MachineBasicBlock::iterator CondBr = PredMBB->getLastNonDebugInstr();
+  if (CondBr == 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;
+  // After calling knownRegValInBlock, FirstUse will either point to a CBZ/CBNZ
+  // or a compare (i.e., SUBS).  In the latter case, we must take care when
+  // updating FirstUse when scanning for COPY instructions.  In particular, if
+  // there's a COPY in between the compare and branch the COPY should not
+  // update FirstUse.
+  bool SeenFirstUse = false;
+  // Registers that contain a known value at the start of MBB.
+  SmallVector<RegImm, 4> KnownRegs;
+
+  MachineBasicBlock::iterator Itr = std::next(CondBr);
   do {
-    --CompBr;
-    if (!guaranteesZeroRegInBlock(*CompBr, MBB))
+    --Itr;
+
+    Optional<RegImm> KnownRegImm = knownRegValInBlock(*Itr, MBB, FirstUse);
+    if (KnownRegImm == None)
       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) {
+    KnownRegs.push_back(*KnownRegImm);
+
+    // Reset the clobber list, which is used by knownRegValInBlock.
+    ClobberedRegs.reset();
+
+    // Look backward in PredMBB for COPYs from the known reg to find other
+    // registers that are known to be a constant value.
+    for (auto PredI = Itr;; --PredI) {
+      if (FirstUse == PredI)
+        SeenFirstUse = true;
+
       if (PredI->isCopy()) {
         MCPhysReg CopyDstReg = PredI->getOperand(0).getReg();
         MCPhysReg CopySrcReg = PredI->getOperand(1).getReg();
-        for (MCPhysReg KnownZeroReg : KnownZeroRegs) {
-          if (ClobberedRegs[KnownZeroReg])
+        for (auto &KnownReg : KnownRegs) {
+          if (ClobberedRegs[KnownReg.Reg])
             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);
+          if (CopySrcReg == KnownReg.Reg && !ClobberedRegs[CopyDstReg]) {
+            KnownRegs.push_back(RegImm(CopyDstReg, KnownReg.Imm));
+            if (SeenFirstUse)
               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);
+          if (CopyDstReg == KnownReg.Reg && !ClobberedRegs[CopySrcReg]) {
+            KnownRegs.push_back(RegImm(CopySrcReg, KnownReg.Imm));
+            if (SeenFirstUse)
               FirstUse = PredI;
             break;
           }
         }
       }
 
       // Stop if we get to the beginning of PredMBB.
       if (PredI == PredMBB->begin())
         break;
 
       trackRegDefs(*PredI, ClobberedRegs, TRI);
       // Stop if all of the known-zero regs have been clobbered.
-      if (all_of(KnownZeroRegs, [&](MCPhysReg KnownZeroReg) {
-            return ClobberedRegs[KnownZeroReg];
+      if (all_of(KnownRegs, [&](RegImm KnownReg) {
+            return ClobberedRegs[KnownReg.Reg];
           }))
         break;
     }
     break;
 
-  } while (CompBr != PredMBB->begin() && CompBr->isTerminator());
+  } while (Itr != PredMBB->begin() && Itr->isTerminator());
 
-  // We've not found a known zero register, time to bail out.
-  if (KnownZeroRegs.empty())
+  // We've not found a registers with a known value, time to bail out.
+  if (KnownRegs.empty())
     return false;
 
   bool Changed = false;
-  // UsedKnownZeroRegs is the set of KnownZeroRegs that have had uses added to MBB.
-  SmallSetVector<unsigned, 4> UsedKnownZeroRegs;
+  // UsedKnownRegs is the set of KnownRegs that have had uses added to MBB.
+  SmallSetVector<unsigned, 4> UsedKnownRegs;
   MachineBasicBlock::iterator LastChange = MBB->begin();
-  // Remove redundant Copy instructions unless KnownZeroReg is modified.
+  // Remove redundant Copy instructions unless KnownReg is modified.
   for (MachineBasicBlock::iterator I = MBB->begin(), E = MBB->end(); I != E;) {
     MachineInstr *MI = &*I;
     ++I;
-    bool RemovedCopy = false;
-    if (MI->isCopy()) {
+    bool RemovedMI = false;
+    bool IsCopy = MI->isCopy();
+    bool IsMoveImm = MI->isMoveImmediate();
+    if (IsCopy || IsMoveImm) {
       MCPhysReg DefReg = MI->getOperand(0).getReg();
-      MCPhysReg SrcReg = MI->getOperand(1).getReg();
+      MCPhysReg SrcReg = IsCopy ? MI->getOperand(1).getReg() : 0;
+      int64_t SrcImm = IsMoveImm ? MI->getOperand(1).getImm() : 0;
+      if (!MRI->isReserved(DefReg) &&
+          ((IsCopy && (SrcReg == AArch64::XZR || SrcReg == AArch64::WZR)) ||
+           IsMoveImm)) {
+        for (RegImm &KnownReg : KnownRegs) {
+          if (KnownReg.Reg != DefReg &&
+              !TRI->isSuperRegister(DefReg, KnownReg.Reg))
+            continue;
+
+          // For a copy, the known value must be a zero.
+          if (IsCopy && KnownReg.Imm != 0)
+            continue;
 
-      if ((SrcReg == AArch64::XZR || SrcReg == AArch64::WZR) &&
-          !MRI->isReserved(DefReg)) {
-        for (MCPhysReg KnownZeroReg : KnownZeroRegs) {
-          if (KnownZeroReg == DefReg ||
-              TRI->isSuperRegister(DefReg, KnownZeroReg)) {
+          if (IsMoveImm) {
+            // For a move immediate, the known immediate must match the source
+            // immediate.
+            if (KnownReg.Imm != SrcImm)
+              continue;
+
+            // Don't remove a move immediate that implicitly defines the upper
+            // bits when only the lower 32 bits are known.
+            MCPhysReg CmpReg = KnownReg.Reg;
+            if (any_of(MI->implicit_operands(), [CmpReg](MachineOperand &O) {
+                  return !O.isDead() && O.isReg() && O.isDef() &&
+                         O.getReg() != CmpReg;
+                }))
+              continue;
+          }
+
+          if (IsCopy)
             DEBUG(dbgs() << "Remove redundant Copy : " << *MI);
+          else
+            DEBUG(dbgs() << "Remove redundant Move : " << *MI);
 
-            MI->eraseFromParent();
-            Changed = true;
-            LastChange = I;
-            NumCopiesRemoved++;
-            UsedKnownZeroRegs.insert(KnownZeroReg);
-            RemovedCopy = true;
-            break;
-          }
-        }
-      }
-    }
+          UsedKnownRegs.insert(KnownReg.Reg);
+          RemovedMI = true;
 
-    // Skip to the next instruction if we removed the COPY from WZR/XZR.
-    if (RemovedCopy)
+    // Skip to the next instruction if we removed the COPY/MovImm.
+    if (RemovedMI)
       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();
+    // Remove any regs the MI clobbers from the KnownConstRegs set.
+    for (unsigned RI = 0; RI < KnownRegs.size();)
+      if (MI->modifiesRegister(KnownRegs[RI].Reg, TRI)) {
+        std::swap(KnownRegs[RI], KnownRegs[KnownRegs.size() - 1]);
+        KnownRegs.pop_back();
         // Don't increment RI since we need to now check the swapped-in
-        // KnownZeroRegs[RI].
+        // KnownRegs[RI].
       } else {
         ++RI;
       }
 
-    // Continue until the KnownZeroRegs set is empty.
-    if (KnownZeroRegs.empty())
+    // Continue until the KnownRegs set is empty.
+    if (KnownRegs.empty())
       break;
   }
 
   if (!Changed)
     return false;
 
   // 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);
+  for (MCPhysReg KnownReg : UsedKnownRegs)
+    if (!MBB->isLiveIn(KnownReg))
+      MBB->addLiveIn(KnownReg);
 
   // Clear kills in the range where changes were made.  This is conservative,
   // but should be okay since kill markers are being phased out.
+  DEBUG(dbgs() << "Clearing kill flags.\n\tFirstUse: " << *FirstUse
+               << "\tLastChange: " << *LastChange);
   for (MachineInstr &MMI : make_range(FirstUse, PredMBB->end()))
     MMI.clearKillInfo();
   for (MachineInstr &MMI : make_range(MBB->begin(), LastChange))
     MMI.clearKillInfo();
 
   return true;
 }
 
 bool AArch64RedundantCopyElimination::runOnMachineFunction(
     MachineFunction &MF) {
   if (skipFunction(*MF.getFunction()))
     return false;
   TRI = MF.getSubtarget().getRegisterInfo();
   MRI = &MF.getRegInfo();
+
+  // Resize the clobber register bitfield tracker.  We do this once per
+  // function and then clear the bitfield each time we optimize a copy.
   ClobberedRegs.resize(TRI->getNumRegs());
+
   bool Changed = false;
   for (MachineBasicBlock &MBB : MF)
     Changed |= optimizeCopy(&MBB);
   return Changed;
 }
 
 FunctionPass *llvm::createAArch64RedundantCopyEliminationPass() {
   return new AArch64RedundantCopyElimination();
 }

llvm/trunk/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
+# RUN: llc -mtriple=aarch64--linux-gnu -run-pass=aarch64-copyelim %s -verify-machineinstrs -o - | 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:
bb.2:
     CBNZX %x1, %bb.1
 
   bb.3:
     liveins: %x0
 
     RET_ReallyLR implicit %x0
 
 ...
+# Eliminate redundant MOVi32imm 7 in bb.1
+# Note: 32-bit compare/32-bit move imm
+# Kill marker should be removed from compare.
+# CHECK-LABEL: name: test10
+# CHECK: SUBSWri %w0, 7, 0, implicit-def %nzcv
+# CHECK: bb.1:
+# CHECK-NOT: MOVi32imm
+name:            test10
+tracksRegLiveness: true
+body:             |
+  bb.0.entry:
+    successors: %bb.1, %bb.2
+    liveins: %w0, %x1
+
+    dead %wzr = SUBSWri killed %w0, 7, 0, implicit-def %nzcv
+    Bcc 1, %bb.2, implicit killed %nzcv
+    B %bb.1
+
+  bb.1:
+    successors: %bb.2
+    liveins: %x1
+
+    %w0 = MOVi32imm 7
+    STRWui killed %w0, killed %x1, 0
+
+  bb.2:
+    RET_ReallyLR
+...
+# Eliminate redundant MOVi32imm 7 in bb.1
+# Note: 64-bit compare/32-bit move imm w/implicit def
+# Kill marker should be removed from compare.
+# CHECK-LABEL: name: test11
+# CHECK: SUBSXri %x0, 7, 0, implicit-def %nzcv
+# CHECK: bb.1:
+# CHECK-NOT: MOVi32imm
+name:            test11
+tracksRegLiveness: true
+body:             |
+  bb.0.entry:
+    successors: %bb.1, %bb.2
+    liveins: %x0, %x1
+
+    dead %xzr = SUBSXri killed %x0, 7, 0, implicit-def %nzcv
+    Bcc 1, %bb.2, implicit killed %nzcv
+    B %bb.1
+
+  bb.1:
+    successors: %bb.2
+    liveins: %x1
+
+    %w0 = MOVi32imm 7, implicit-def %x0
+    STRXui killed %x0, killed %x1, 0
+
+  bb.2:
+    RET_ReallyLR
+...
+# Eliminate redundant MOVi32imm 7 in bb.1
+# Note: 64-bit compare/32-bit move imm
+# Kill marker should be removed from compare.
+# CHECK-LABEL: name: test12
+# CHECK: SUBSXri %x0, 7, 0, implicit-def %nzcv
+# CHECK: bb.1:
+# CHECK-NOT: MOVi32imm
+name:            test12
+tracksRegLiveness: true
+body:             |
+  bb.0.entry:
+    successors: %bb.1, %bb.2
+    liveins: %x0, %x1
+
+    dead %xzr = SUBSXri killed %x0, 7, 0, implicit-def %nzcv
+    Bcc 1, %bb.2, implicit killed %nzcv
+    B %bb.1
+
+  bb.1:
+    successors: %bb.2
+    liveins: %x1
+
+    %w0 = MOVi32imm 7
+    STRWui killed %w0, killed %x1, 0
+
+  bb.2:
+    RET_ReallyLR
+...
+# Don't eliminate MOVi32imm 7 in bb.1 as we don't necessarily know the upper 32-bits.
+# Note: 32-bit compare/32-bit move imm w/implicit def
+# Kill marker should remain on compare.
+# CHECK-LABEL: name: test13
+# CHECK: SUBSWri killed %w0, 7, 0, implicit-def %nzcv
+# CHECK: bb.1:
+# CHECK: MOVi32imm
+name:            test13
+tracksRegLiveness: true
+body:             |
+  bb.0.entry:
+    successors: %bb.1, %bb.2
+    liveins: %w0, %x1
+
+    dead %wzr = SUBSWri killed %w0, 7, 0, implicit-def %nzcv
+    Bcc 1, %bb.2, implicit killed %nzcv
+    B %bb.1
+
+  bb.1:
+    successors: %bb.2
+    liveins: %x1
+
+    %w0 = MOVi32imm 7, implicit-def %x0
+    STRXui killed %x0, killed %x1, 0
+
+  bb.2:
+    RET_ReallyLR
+...
+# We can't eliminate the MOVi32imm because of the clobbering LDRWui.
+# CHECK-LABEL: name: test14
+# CHECK: bb.1:
+# CHECK: MOVi32imm
+name:            test14
+tracksRegLiveness: true
+body:             |
+  bb.0.entry:
+    successors: %bb.1, %bb.2
+    liveins: %w0, %x1, %x2
+
+    dead %wzr = SUBSWri killed %w0, 7, 0, implicit-def %nzcv
+    %w0 = LDRWui %x1, 0
+    STRWui killed %w0, killed %x2, 0
+    Bcc 1, %bb.2, implicit killed %nzcv
+    B %bb.1
+
+  bb.1:
+    successors: %bb.2
+    liveins: %x1
+
+    %w0 = MOVi32imm 7
+    STRWui killed %w0, killed %x1, 0
+
+  bb.2:
+    RET_ReallyLR
+...
+# We can't eliminate the MOVi32imm because of the clobbering LDRWui.
+# CHECK-LABEL: name: test15
+# CHECK: bb.1:
+# CHECK: MOVi32imm
+name:            test15
+tracksRegLiveness: true
+body:             |
+  bb.0.entry:
+    successors: %bb.1, %bb.2
+    liveins: %w0, %x1, %x2
+
+    dead %wzr = SUBSWri killed %w0, 7, 0, implicit-def %nzcv
+    Bcc 1, %bb.2, implicit killed %nzcv
+    B %bb.1
+
+  bb.1:
+    successors: %bb.2
+    liveins: %x1, %x2
+
+    %w0 = LDRWui %x1, 0
+    STRWui killed %w0, killed %x2, 0
+    %w0 = MOVi32imm 7
+    STRWui killed %w0, killed %x1, 0
+
+  bb.2:
+    RET_ReallyLR
+...
+# Check that bb.0 COPY is seen through to allow the bb.1 MOVi32imm to be removed.
+# CHECK-LABEL: name: test16
+# CHECK: bb.1:
+# CHECK-NOT: MOVi32imm
+name:            test16
+tracksRegLiveness: true
+body:             |
+  bb.0.entry:
+    successors: %bb.1, %bb.2
+    liveins: %w0, %x1
+
+    dead %wzr = SUBSWri %w0, 7, 0, implicit-def %nzcv
+    %w2 = COPY %w0
+    Bcc 1, %bb.2, implicit killed %nzcv
+    B %bb.1
+
+  bb.1:
+    successors: %bb.2
+    liveins: %x1
+
+    %w2 = MOVi32imm 7
+    STRWui killed %w2, killed %x1, 0
+
+  bb.2:
+    RET_ReallyLR
+...
+# Check that bb.1 MOVi32imm is not removed due to self clobbering compare.
+# CHECK-LABEL: name: test17
+# CHECK: bb.1:
+# CHECK: MOVi32imm
+name:            test17
+tracksRegLiveness: true
+body:             |
+  bb.0.entry:
+    successors: %bb.1, %bb.2
+    liveins: %w0, %x1
+
+    dead %w0 = SUBSWri killed %w0, 7, 0, implicit-def %nzcv
+    Bcc 1, %bb.2, implicit killed %nzcv
+    B %bb.1
+
+  bb.1:
+    successors: %bb.2
+    liveins: %x1
+
+    %w0 = MOVi32imm 7
+    STRWui killed %w0, killed %x1, 0
+
+  bb.2:
+    RET_ReallyLR
+...
+# Make sure the MOVi64imm is not removed.  In one version of this patch the
+# MOVi64imm immediate was truncated to 32 bits and incorrectly matched because
+# the low 32 bits of 4252017623040 are all zero.
+# CHECK-LABEL: name: test18
+# CHECK: bb.1:
+# CHECK: MOVi64imm
+name:            test18
+tracksRegLiveness: true
+body:             |
+  bb.0.entry:
+    successors: %bb.1, %bb.2
+    liveins: %x0, %x1
+
+    CBNZX killed %x0, %bb.2
+    B %bb.1
+
+  bb.1:
+    successors: %bb.2
+    liveins: %x1
+
+    %x0 = MOVi64imm 4252017623040
+    STRXui killed %x0, killed %x1, 0
+
+  bb.2:
+    RET_ReallyLR