[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

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

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

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

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

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 hasAvailableCopies() { return !AvailCopyMap.empty(); }
+  bool hasAvailableCopies() {
+    // TODO: This is conservatively correct, but is it useful?
+    return !Copies.empty();
+  }
 
-  MachineInstr *findAvailCopy(unsigned Reg) {
-    auto CI = AvailCopyMap.find(Reg);
-    if (CI != AvailCopyMap.end())
-      return CI->second;
+  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 *findCopy(unsigned Reg) {
-    auto CI = CopyMap.find(Reg);
-    if (CI != CopyMap.end())
-      return CI->second;
+  MachineInstr *findCopy(unsigned Reg, const TargetRegisterInfo &TRI,
+                         bool MustBeAvailable = false) {
+    MCRegUnitIterator RUI(Reg, &TRI);
+    MachineInstr *MI = findCopyForUnit(*RUI, TRI, MustBeAvailable);
+    if (!MI || !TRI.isSubRegisterEq(MI->getOperand(0).getReg(), Reg))
       return nullptr;
+    return MI;
+  }
+
+  MachineInstr *findAvailCopy(unsigned Reg, const TargetRegisterInfo &TRI) {
+    return findCopy(Reg, TRI, /*MustBeAvailable=*/true);
   }
 
   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(Def);
+  MachineInstr *PrevCopy = Tracker.findAvailCopy(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 (!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(MOUse.getReg());
+    MachineInstr *Copy = Tracker.findAvailCopy(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.