[MachineCopyPropagation] Reimplement CopyTracker in terms of register units

[bogner]

Summary

Change the copy tracker to keep a single map of register units instead
of 3 maps of registers. This gives a very significant compile time
performance improvement to the pass. I measured a 30-40% decrease in
time spent in MCP on x86 and AArch64 and much more significant
improvements on an out of tree target with more registers.

[MatzeB]

lib/CodeGen/MachineCopyPropagation.cpp

133–136 ↗

(On Diff #166531)

This feels odd now, as in the new code you say Avail = false instead of removing elements from the Copies set... Would it make sense to maintain a counter of entries with Avail == true for this?

148–155 ↗

(On Diff #166531)

Is it correct to only check the first register unit here? Isn't it possible that only some of the units gots clobbered in the meantime, while the first one still makes it seem like the copy is available?

[bogner]

Clarify a bit based on Matthias's comments.

[bogner]

lib/CodeGen/MachineCopyPropagation.cpp

133–136 ↗

(On Diff #166531)

It's not really worth the bookkeeping to keep track of the number of available copies, but this is used to shortcut out of forwardUses and it's measurably slower if we don't do it at all.

I've renamed this to hasAnyCopies so it's more obvious what it actually checks.

148–155 ↗

(On Diff #166531)

Yes, as discussed offline this is correct. When we clobber units we end up marking registers unavailable based on the original instruction's defs (not just the unit itself), so we won't hit any problems here.

[MatzeB]

I'm convinced now the clobbering works correctly because as clobberRegister does not just clobber the register at hand but the whole destination register of the copy touched by the COPY at hand.
LGTM (this time with the right review :)

llvm/trunk/lib/CodeGen/MachineCopyPropagation.cpp

 STATISTIC(NumDeletes, "Number of dead copies deleted");
 STATISTIC(NumCopyForwards, "Number of copy uses forwarded");
 DEBUG_COUNTER(FwdCounter, "machine-cp-fwd",
               "Controls which register COPYs are forwarded");
 
 namespace {
 
 class CopyTracker {
-  using RegList = SmallVector<unsigned, 4>;
-  using SourceMap = DenseMap<unsigned, RegList>;
-  using Reg2MIMap = DenseMap<unsigned, MachineInstr *>;
-
-  /// Def -> available copies map.
-  Reg2MIMap AvailCopyMap;
-
-  /// Def -> copies map.
-  Reg2MIMap CopyMap;
+  struct CopyInfo {
+    MachineInstr *MI;
+    SmallVector<unsigned, 4> DefRegs;
+    bool Avail;
+  };
 
-  /// Src -> Def map
-  SourceMap SrcMap;
+  DenseMap<unsigned, CopyInfo> Copies;
 
 public:
   /// Mark all of the given registers and their subregisters as unavailable for
   /// copying.
-  void markRegsUnavailable(const RegList &Regs, const TargetRegisterInfo &TRI) {
+  void markRegsUnavailable(ArrayRef<unsigned> Regs,
+                           const TargetRegisterInfo &TRI) {
     for (unsigned Reg : Regs) {
       // Source of copy is no longer available for propagation.
-      for (MCSubRegIterator SR(Reg, &TRI, true); SR.isValid(); ++SR)
-        AvailCopyMap.erase(*SR);
+      for (MCRegUnitIterator RUI(Reg, &TRI); RUI.isValid(); ++RUI) {
+        auto CI = Copies.find(*RUI);
+        if (CI != Copies.end())
+          CI->second.Avail = false;
+      }
     }
   }
 
   /// Clobber a single register, removing it from the tracker's copy maps.
   void clobberRegister(unsigned Reg, const TargetRegisterInfo &TRI) {
-    for (MCRegAliasIterator AI(Reg, &TRI, true); AI.isValid(); ++AI) {
-      CopyMap.erase(*AI);
-      AvailCopyMap.erase(*AI);
-
-      SourceMap::iterator SI = SrcMap.find(*AI);
-      if (SI != SrcMap.end()) {
-        markRegsUnavailable(SI->second, TRI);
-        SrcMap.erase(SI);
+    for (MCRegUnitIterator RUI(Reg, &TRI); RUI.isValid(); ++RUI) {
+      auto I = Copies.find(*RUI);
+      if (I != Copies.end()) {
+        markRegsUnavailable(I->second.DefRegs, TRI);
+        Copies.erase(I);
       }
     }
   }
 
   /// Add this copy's registers into the tracker's copy maps.
   void trackCopy(MachineInstr *Copy, const TargetRegisterInfo &TRI) {
     assert(Copy->isCopy() && "Tracking non-copy?");
 
     unsigned Def = Copy->getOperand(0).getReg();
     unsigned Src = Copy->getOperand(1).getReg();
 
     // Remember Def is defined by the copy.
-    for (MCSubRegIterator SR(Def, &TRI, /*IncludeSelf=*/true); SR.isValid();
-         ++SR) {
-      CopyMap[*SR] = Copy;
-      AvailCopyMap[*SR] = Copy;
-    }
+    for (MCRegUnitIterator RUI(Def, &TRI); RUI.isValid(); ++RUI)
+      Copies[*RUI] = {Copy, {}, true};
 
     // Remember source that's copied to Def. Once it's clobbered, then
     // it's no longer available for copy propagation.
-    RegList &DestList = SrcMap[Src];
-    if (!is_contained(DestList, Def))
-      DestList.push_back(Def);
+    for (MCRegUnitIterator RUI(Src, &TRI); RUI.isValid(); ++RUI) {
+      auto I = Copies.insert({*RUI, {nullptr, {}, false}});
+      auto &Copy = I.first->second;
+      if (!is_contained(Copy.DefRegs, Def))
+        Copy.DefRegs.push_back(Def);
+    }
+  }
+
+  bool hasAnyCopies() {
+    return !Copies.empty();
   }
 
-  bool hasAvailableCopies() { return !AvailCopyMap.empty(); }
+  MachineInstr *findCopyForUnit(unsigned RegUnit, const TargetRegisterInfo &TRI,
+                         bool MustBeAvailable = false) {
+    auto CI = Copies.find(RegUnit);
+    if (CI == Copies.end())
+      return nullptr;
+    if (MustBeAvailable && !CI->second.Avail)
+      return nullptr;
+    return CI->second.MI;
+  }
 
-  MachineInstr *findAvailCopy(MachineInstr &DestCopy, unsigned Reg) {
-    auto CI = AvailCopyMap.find(Reg);
-    if (CI == AvailCopyMap.end())
+  MachineInstr *findAvailCopy(MachineInstr &DestCopy, unsigned Reg,
+                              const TargetRegisterInfo &TRI) {
+    // We check the first RegUnit here, since we'll only be interested in the
+    // copy if it copies the entire register anyway.
+    MCRegUnitIterator RUI(Reg, &TRI);
+    MachineInstr *AvailCopy =
+        findCopyForUnit(*RUI, TRI, /*MustBeAvailable=*/true);
+    if (!AvailCopy ||
+        !TRI.isSubRegisterEq(AvailCopy->getOperand(0).getReg(), Reg))
       return nullptr;
-    MachineInstr &AvailCopy = *CI->second;
 
     // Check that the available copy isn't clobbered by any regmasks between
     // itself and the destination.
-    unsigned AvailSrc = AvailCopy.getOperand(1).getReg();
-    unsigned AvailDef = AvailCopy.getOperand(0).getReg();
+    unsigned AvailSrc = AvailCopy->getOperand(1).getReg();
+    unsigned AvailDef = AvailCopy->getOperand(0).getReg();
     for (const MachineInstr &MI :
-         make_range(AvailCopy.getIterator(), DestCopy.getIterator()))
+         make_range(AvailCopy->getIterator(), DestCopy.getIterator()))
       for (const MachineOperand &MO : MI.operands())
         if (MO.isRegMask())
           if (MO.clobbersPhysReg(AvailSrc) || MO.clobbersPhysReg(AvailDef))
             return nullptr;
 
-    return &AvailCopy;
-  }
-
-  MachineInstr *findCopy(unsigned Reg) {
-    auto CI = CopyMap.find(Reg);
-    if (CI != CopyMap.end())
-      return CI->second;
-    return nullptr;
+    return AvailCopy;
   }
 
   void clear() {
-    AvailCopyMap.clear();
-    CopyMap.clear();
-    SrcMap.clear();
+    Copies.clear();
   }
 };
 
 class MachineCopyPropagation : public MachineFunctionPass {
   const TargetRegisterInfo *TRI;
   const TargetInstrInfo *TII;
   const MachineRegisterInfo *MRI;
 

 char &llvm::MachineCopyPropagationID = MachineCopyPropagation::ID;
 
 INITIALIZE_PASS(MachineCopyPropagation, DEBUG_TYPE,
                 "Machine Copy Propagation Pass", false, false)
 
 void MachineCopyPropagation::ReadRegister(unsigned Reg) {
   // If 'Reg' is defined by a copy, the copy is no longer a candidate
   // for elimination.
-  for (MCRegAliasIterator AI(Reg, TRI, true); AI.isValid(); ++AI) {
-    if (MachineInstr *Copy = Tracker.findCopy(*AI)) {
+  for (MCRegUnitIterator RUI(Reg, TRI); RUI.isValid(); ++RUI) {
+    if (MachineInstr *Copy = Tracker.findCopyForUnit(*RUI, *TRI)) {
       LLVM_DEBUG(dbgs() << "MCP: Copy is used - not dead: "; Copy->dump());
       MaybeDeadCopies.remove(Copy);
     }
   }
 }
 
 /// Return true if \p PreviousCopy did copy register \p Src to register \p Def.
 /// This fact may have been obscured by sub register usage or may not be true at

 bool MachineCopyPropagation::eraseIfRedundant(MachineInstr &Copy, unsigned Src,
                                               unsigned Def) {
   // Avoid eliminating a copy from/to a reserved registers as we cannot predict
   // the value (Example: The sparc zero register is writable but stays zero).
   if (MRI->isReserved(Src) || MRI->isReserved(Def))
     return false;
 
   // Search for an existing copy.
-  MachineInstr *PrevCopy = Tracker.findAvailCopy(Copy, Def);
+  MachineInstr *PrevCopy = Tracker.findAvailCopy(Copy, Def, *TRI);
   if (!PrevCopy)
     return false;
 
   // Check that the existing copy uses the correct sub registers.
   if (PrevCopy->getOperand(0).isDead())
     return false;
   if (!isNopCopy(*PrevCopy, Src, Def, TRI))
     return false;
if (&MIUse != &Use && MIUse.isReg() && MIUse.isImplicit() &&
       return true;
 
   return false;
 }
 
 /// Look for available copies whose destination register is used by \p MI and
 /// replace the use in \p MI with the copy's source register.
 void MachineCopyPropagation::forwardUses(MachineInstr &MI) {
-  if (!Tracker.hasAvailableCopies())
+  if (!Tracker.hasAnyCopies())
     return;
 
   // Look for non-tied explicit vreg uses that have an active COPY
   // instruction that defines the physical register allocated to them.
   // Replace the vreg with the source of the active COPY.
   for (unsigned OpIdx = 0, OpEnd = MI.getNumOperands(); OpIdx < OpEnd;
        ++OpIdx) {
     MachineOperand &MOUse = MI.getOperand(OpIdx);
if (!MOUse.getReg())
       continue;
 
     // Check that the register is marked 'renamable' so we know it is safe to
     // rename it without violating any constraints that aren't expressed in the
     // IR (e.g. ABI or opcode requirements).
     if (!MOUse.isRenamable())
       continue;
 
-    MachineInstr *Copy = Tracker.findAvailCopy(MI, MOUse.getReg());
+    MachineInstr *Copy = Tracker.findAvailCopy(MI, MOUse.getReg(), *TRI);
     if (!Copy)
       continue;
 
     unsigned CopyDstReg = Copy->getOperand(0).getReg();
     const MachineOperand &CopySrc = Copy->getOperand(1);
     unsigned CopySrcReg = CopySrc.getReg();
 
     // FIXME: Don't handle partial uses of wider COPYs yet.