[CodeGen] Use RegUnits to track register aliases (NFC)

[junbuml]

Summary

Use RegUnits to track register aliases in PostRASink and AArch64LoadStoreOptimizer.

[thegameg]

Thanks for working on this! This looks good to me. Adding Tim for the AArch64LDST part.

[MatzeB]

Have you seen LiveRegUnits::accumulate() glancing over the code here (and earlier) it seems to me that TII::tracksRegDefsUses() was just invented to do the same thing: figuring out which registers are free/usable over a range of instructions...

[junbuml]

Have you seen LiveRegUnits::accumulate() glancing over the code here (and earlier) it seems to me that TII::tracksRegDefsUses() was just invented to do the same thing: figuring out which registers are free/usable over a range of instructions...

Yes, actually my current change in TII::tracksRegDefsUses() was inspired by LiveRegUnits::accumulate(). The only difference is that we need to track UsedRegs and DefReg separately in TII::tracksRegDefsUses(), while LiveRegUnits::accumulate() track both used/defed registers together. Based on your comment in https://reviews.llvm.org/D41463#inline-400109, would it make sense to introduce new static member function in LiveRegUnits like :

static void LiveRegUnits::accumulateUseDef(MachineInstr &MI, LiveRegUnits &ModifiedRegUnits, LiveRegUnits &UsedRegUnits)  ?

[MatzeB]

In D45695#1069868, @junbuml wrote:

Have you seen LiveRegUnits::accumulate() glancing over the code here (and earlier) it seems to me that TII::tracksRegDefsUses() was just invented to do the same thing: figuring out which registers are free/usable over a range of instructions...

Yes, actually my current change in TII::tracksRegDefsUses() was inspired by LiveRegUnits::accumulate(). The only difference is that we need to track UsedRegs and DefReg separately in TII::tracksRegDefsUses(), while LiveRegUnits::accumulate() track both used/defed registers together. Based on your comment in https://reviews.llvm.org/D41463#inline-400109, would it make sense to introduce new static member function in LiveRegUnits like :

static void LiveRegUnits::accumulateUseDef(MachineInstr &MI, LiveRegUnits &ModifiedRegUnits, LiveRegUnits &UsedRegUnits)  ?

Yes that sounds better!

[junbuml]

Addressed Matthias' comment by moving TargetInstrInfo::trackRegDefsUses into LiveRegUnits::accumulateUsedDefed.

[junbuml]

Hi Matthias,
Please let me know if my change in LiveRegUnits is okay with you.
Thanks,
Jun

[junbuml]

Just a minor change in accumulateUsedDefed(). Please let me know any comment.

[junbuml]

Kindly ping.
Matthias, please let me know if this change is still okay with you.

[sebpop]

looks good, thanks!

[junbuml]

Thanks Sebastian for your review!

include/llvm/CodeGen/LiveRegUnits.h

 /// A set of register units. It is intended for register liveness tracking.
 //
 //===----------------------------------------------------------------------===//
 
 #ifndef LLVM_CODEGEN_LIVEREGUNITS_H
 #define LLVM_CODEGEN_LIVEREGUNITS_H
 
 #include "llvm/ADT/BitVector.h"
+#include "llvm/CodeGen/MachineRegisterInfo.h"
 #include "llvm/CodeGen/TargetRegisterInfo.h"
 #include "llvm/MC/LaneBitmask.h"
 #include "llvm/MC/MCRegisterInfo.h"
 #include <cstdint>
 
 namespace llvm {
 
 class MachineInstr;
 class MachineBasicBlock;
 
 /// A set of register units used to track register liveness.
 class LiveRegUnits {
   const TargetRegisterInfo *TRI = nullptr;
   BitVector Units;
 
 public:
   /// Constructs a new empty LiveRegUnits set.
   LiveRegUnits() = default;
 
   /// Constructs and initialize an empty LiveRegUnits set.
   LiveRegUnits(const TargetRegisterInfo &TRI) {
     init(TRI);
   }
 
+  /// For a machine instruction \p MI, adds all register units used in
+  /// \p UsedRegUnits and defined or clobbered in \p ModifiedRegUnits. This is
+  /// useful when walking over a range of instructions to track registers
+  /// used or defined seperately.
+  static void accumulateUsedDefed(const MachineInstr &MI,
+                                  LiveRegUnits &ModifiedRegUnits,
+                                  LiveRegUnits &UsedRegUnits,
+                                  const TargetRegisterInfo *TRI) {
+    for (ConstMIBundleOperands O(MI); O.isValid(); ++O) {
+      if (O->isRegMask())
+        ModifiedRegUnits.addRegsInMask(O->getRegMask());
+      if (!O->isReg())
+        continue;
+      unsigned Reg = O->getReg();
+      if (!Reg || !TargetRegisterInfo::isPhysicalRegister(Reg))
+        continue;
+      if (O->isDef()) {
+        // Some architectures (e.g. AArch64 XZR/WZR) have registers that are
+        // constant and may be used as destinations to indicate the generated
+        // value is discarded. No need to track such case as a def.
+        if (!TRI->isConstantPhysReg(Reg))
+          ModifiedRegUnits.addReg(Reg);
+      } else {
+        assert(O->isUse() && "Reg operand not a def and not a use");
+        UsedRegUnits.addReg(Reg);
+      }
+    }
+    return;
+  }
+
   /// Initialize and clear the set.
   void init(const TargetRegisterInfo &TRI) {
     this->TRI = &TRI;
     Units.reset();
     Units.resize(TRI.getNumRegUnits());
   }
 
   /// Clears the set.

include/llvm/CodeGen/TargetInstrInfo.h

 
 #ifndef LLVM_TARGET_TARGETINSTRINFO_H
 #define LLVM_TARGET_TARGETINSTRINFO_H
 
 #include "llvm/ADT/ArrayRef.h"
 #include "llvm/ADT/DenseMap.h"
 #include "llvm/ADT/DenseMapInfo.h"
 #include "llvm/ADT/None.h"
+#include "llvm/CodeGen/LiveRegUnits.h"
 #include "llvm/CodeGen/MachineBasicBlock.h"
 #include "llvm/CodeGen/MachineCombinerPattern.h"
 #include "llvm/CodeGen/MachineFunction.h"
 #include "llvm/CodeGen/MachineInstr.h"
 #include "llvm/CodeGen/MachineLoopInfo.h"
 #include "llvm/CodeGen/MachineOperand.h"
 #include "llvm/CodeGen/PseudoSourceValue.h"
 #include "llvm/MC/MCInstrInfo.h"
public:
 
   /// Return true when a target supports MachineCombiner.
   virtual bool useMachineCombiner() const { return false; }
 
   /// Return true if the given SDNode can be copied during scheduling
   /// even if it has glue.
   virtual bool canCopyGluedNodeDuringSchedule(SDNode *N) const { return false; }
 
-  /// Remember what registers the specified instruction uses and modifies.
-  virtual void trackRegDefsUses(const MachineInstr &MI, BitVector &ModifiedRegs,
-                                BitVector &UsedRegs,
-                                const TargetRegisterInfo *TRI) const;
-
 protected:
   /// Target-dependent implementation for foldMemoryOperand.
   /// Target-independent code in foldMemoryOperand will
   /// take care of adding a MachineMemOperand to the newly created instruction.
   /// The instruction and any auxiliary instructions necessary will be inserted
   /// at InsertPt.
   virtual MachineInstr *
   foldMemoryOperandImpl(MachineFunction &MF, MachineInstr &MI,

lib/CodeGen/MachineSink.cpp

public:
   }
 
   MachineFunctionProperties getRequiredProperties() const override {
     return MachineFunctionProperties().set(
         MachineFunctionProperties::Property::NoVRegs);
   }
 
 private:
-  /// Track which registers have been modified and used.
-  BitVector ModifiedRegs, UsedRegs;
+  /// Track which register units have been modified and used.
+  LiveRegUnits ModifiedRegUnits, UsedRegUnits;
 
   /// Sink Copy instructions unused in the same block close to their uses in
   /// successors.
   bool tryToSinkCopy(MachineBasicBlock &BB, MachineFunction &MF,
                      const TargetRegisterInfo *TRI, const TargetInstrInfo *TII);
 };
 } // namespace
 
 char PostRAMachineSinking::ID = 0;
 char &llvm::PostRAMachineSinkingID = PostRAMachineSinking::ID;
 
 INITIALIZE_PASS(PostRAMachineSinking, "postra-machine-sink",
                 "PostRA Machine Sink", false, false)
 
-static bool aliasWithRegsInLiveIn(MachineBasicBlock *SI,
-                                  SmallSet<unsigned, 8> &AliasedRegs) {
-  for (const auto LI : SI->liveins())
-    if (AliasedRegs.count(LI.PhysReg))
-      return true;
-  return false;
+static bool aliasWithRegsInLiveIn(MachineBasicBlock &MBB, unsigned Reg,
+                                  const TargetRegisterInfo *TRI) {
+  LiveRegUnits LiveInRegUnits(*TRI);
+  LiveInRegUnits.addLiveIns(MBB);
+  return !LiveInRegUnits.available(Reg);
 }
 
 static MachineBasicBlock *
 getSingleLiveInSuccBB(MachineBasicBlock &CurBB,
                       ArrayRef<MachineBasicBlock *> SinkableBBs, unsigned Reg,
                       const TargetRegisterInfo *TRI) {
-  SmallSet<unsigned, 8> AliasedRegs;
-  for (MCRegAliasIterator AI(Reg, TRI, true); AI.isValid(); ++AI)
-    AliasedRegs.insert(*AI);
-
   // Try to find a single sinkable successor in which Reg is live-in.
   MachineBasicBlock *BB = nullptr;
   for (auto *SI : SinkableBBs) {
     if (SI->isLiveIn(Reg)) {
       // If BB is set here, Reg is live-in to at least two sinkable successors,
       // so quit.
       if (BB)
         return nullptr;
       BB = SI;
     }
   }
   // Reg is not live-in to any sinkable successors.
   if (!BB)
     return nullptr;
 
   // Check if any register aliased with Reg is live-in in other successors.
   for (auto *SI : CurBB.successors()) {
     if (SI == BB)
       continue;
-    if (aliasWithRegsInLiveIn(SI, AliasedRegs))
+    if (aliasWithRegsInLiveIn(*SI, Reg, TRI))
       return nullptr;
   }
   return BB;
 }
 
 static MachineBasicBlock *getSingleLiveInSuccBB(
     MachineBasicBlock &CurBB, ArrayRef<MachineBasicBlock *> SinkableBBs,
     ArrayRef<unsigned> DefedRegsInCopy, const TargetRegisterInfo *TRI) {
   MachineBasicBlock *SingleBB = nullptr;
   for (auto DefReg : DefedRegsInCopy) {
     MachineBasicBlock *BB =
         getSingleLiveInSuccBB(CurBB, SinkableBBs, DefReg, TRI);
     if (!BB || (SingleBB && SingleBB != BB))
       return nullptr;
     SingleBB = BB;
   }
   return SingleBB;
 }
 
 static void clearKillFlags(MachineInstr *MI, MachineBasicBlock &CurBB,
                            SmallVectorImpl<unsigned> &UsedOpsInCopy,
-                           BitVector &UsedRegs, const TargetRegisterInfo *TRI) {
+                           LiveRegUnits &UsedRegUnits,
+                           const TargetRegisterInfo *TRI) {
   for (auto U : UsedOpsInCopy) {
     MachineOperand &MO = MI->getOperand(U);
     unsigned SrcReg = MO.getReg();
-    if (UsedRegs[SrcReg]) {
+    if (!UsedRegUnits.available(SrcReg)) {
       MachineBasicBlock::iterator NI = std::next(MI->getIterator());
       for (MachineInstr &UI : make_range(NI, CurBB.end())) {
         if (UI.killsRegister(SrcReg, TRI)) {
           UI.clearRegisterKills(SrcReg, TRI);
           MO.setIsKill(true);
           break;
         }
       }
for (auto U : UsedOpsInCopy) {
     if (!SuccBB->isLiveIn(Reg))
       SuccBB->addLiveIn(Reg);
   }
 }
 
 static bool hasRegisterDependency(MachineInstr *MI,
                                   SmallVectorImpl<unsigned> &UsedOpsInCopy,
                                   SmallVectorImpl<unsigned> &DefedRegsInCopy,
-                                  BitVector &ModifiedRegs,
-                                  BitVector &UsedRegs) {
+                                  LiveRegUnits &ModifiedRegUnits,
+                                  LiveRegUnits &UsedRegUnits) {
   bool HasRegDependency = false;
   for (unsigned i = 0, e = MI->getNumOperands(); i != e; ++i) {
     MachineOperand &MO = MI->getOperand(i);
     if (!MO.isReg())
       continue;
     unsigned Reg = MO.getReg();
     if (!Reg)
       continue;
     if (MO.isDef()) {
-      if (ModifiedRegs[Reg] || UsedRegs[Reg]) {
+      if (!ModifiedRegUnits.available(Reg) || !UsedRegUnits.available(Reg)) {
         HasRegDependency = true;
         break;
       }
       DefedRegsInCopy.push_back(Reg);
 
       // FIXME: instead of isUse(), readsReg() would be a better fix here,
       // For example, we can ignore modifications in reg with undef. However,
       // it's not perfectly clear if skipping the internal read is safe in all
       // other targets.
     } else if (MO.isUse()) {
-      if (ModifiedRegs[Reg]) {
+      if (!ModifiedRegUnits.available(Reg)) {
         HasRegDependency = true;
         break;
       }
       UsedOpsInCopy.push_back(i);
     }
   }
   return HasRegDependency;
 }
bool PostRAMachineSinking::tryToSinkCopy(MachineBasicBlock &CurBB,
 
   if (SinkableBBs.empty())
     return false;
 
   bool Changed = false;
 
   // Track which registers have been modified and used between the end of the
   // block and the current instruction.
-  ModifiedRegs.reset();
-  UsedRegs.reset();
+  ModifiedRegUnits.clear();
+  UsedRegUnits.clear();
 
   for (auto I = CurBB.rbegin(), E = CurBB.rend(); I != E;) {
     MachineInstr *MI = &*I;
     ++I;
 
     // Do not move any instruction across function call.
     if (MI->isCall())
       return false;
 
     if (!MI->isCopy() || !MI->getOperand(0).isRenamable()) {
-      TII->trackRegDefsUses(*MI, ModifiedRegs, UsedRegs, TRI);
+      LiveRegUnits::accumulateUsedDefed(*MI, ModifiedRegUnits, UsedRegUnits,
+                                        TRI);
       continue;
     }
 
     // Track the operand index for use in Copy.
     SmallVector<unsigned, 2> UsedOpsInCopy;
     // Track the register number defed in Copy.
     SmallVector<unsigned, 2> DefedRegsInCopy;
 
     // Don't sink the COPY if it would violate a register dependency.
-    if (hasRegisterDependency(MI, UsedOpsInCopy, DefedRegsInCopy, ModifiedRegs,
-                              UsedRegs)) {
-      TII->trackRegDefsUses(*MI, ModifiedRegs, UsedRegs, TRI);
+    if (hasRegisterDependency(MI, UsedOpsInCopy, DefedRegsInCopy,
+                              ModifiedRegUnits, UsedRegUnits)) {
+      LiveRegUnits::accumulateUsedDefed(*MI, ModifiedRegUnits, UsedRegUnits,
+                                        TRI);
       continue;
     }
     assert((!UsedOpsInCopy.empty() && !DefedRegsInCopy.empty()) &&
            "Unexpect SrcReg or DefReg");
     MachineBasicBlock *SuccBB =
         getSingleLiveInSuccBB(CurBB, SinkableBBs, DefedRegsInCopy, TRI);
     // Don't sink if we cannot find a single sinkable successor in which Reg
     // is live-in.
     if (!SuccBB) {
-      TII->trackRegDefsUses(*MI, ModifiedRegs, UsedRegs, TRI);
+      LiveRegUnits::accumulateUsedDefed(*MI, ModifiedRegUnits, UsedRegUnits,
+                                        TRI);
       continue;
     }
     assert((SuccBB->pred_size() == 1 && *SuccBB->pred_begin() == &CurBB) &&
            "Unexpected predecessor");
 
     // Clear the kill flag if SrcReg is killed between MI and the end of the
     // block.
-    clearKillFlags(MI, CurBB, UsedOpsInCopy, UsedRegs, TRI);
+    clearKillFlags(MI, CurBB, UsedOpsInCopy, UsedRegUnits, TRI);
     MachineBasicBlock::iterator InsertPos = SuccBB->getFirstNonPHI();
     SuccBB->splice(InsertPos, &CurBB, MI);
     updateLiveIn(MI, SuccBB, UsedOpsInCopy, DefedRegsInCopy);
 
     Changed = true;
     ++NumPostRACopySink;
   }
   return Changed;
 }
 
 bool PostRAMachineSinking::runOnMachineFunction(MachineFunction &MF) {
   bool Changed = false;
   const TargetRegisterInfo *TRI = MF.getSubtarget().getRegisterInfo();
   const TargetInstrInfo *TII = MF.getSubtarget().getInstrInfo();
-  ModifiedRegs.resize(TRI->getNumRegs());
-  UsedRegs.resize(TRI->getNumRegs());
 
+  ModifiedRegUnits.init(*TRI);
+  UsedRegUnits.init(*TRI);
   for (auto &BB : MF)
     Changed |= tryToSinkCopy(BB, MF, TRI, TII);
 
   return Changed;
 }

lib/CodeGen/TargetInstrInfo.cpp

default:
     break;
   }
 
   assert(Prev && "Unknown pattern for machine combiner");
 
   reassociateOps(Root, *Prev, Pattern, InsInstrs, DelInstrs, InstIdxForVirtReg);
 }
 
-void TargetInstrInfo::trackRegDefsUses(const MachineInstr &MI,
-                                       BitVector &ModifiedRegs,
-                                       BitVector &UsedRegs,
-                                       const TargetRegisterInfo *TRI) const {
-  for (const MachineOperand &MO : MI.operands()) {
-    if (MO.isRegMask())
-      ModifiedRegs.setBitsNotInMask(MO.getRegMask());
-    if (!MO.isReg())
-      continue;
-    unsigned Reg = MO.getReg();
-    if (!Reg)
-      continue;
-    if (MO.isDef()) {
-      // Some architectures (e.g. AArch64 XZR/WZR) have registers that are
-      // constant and may be used as destinations to indicate the generated
-      // value is discarded. No need to track such case as a def.
-      if (!TRI->isConstantPhysReg(Reg))
-        for (MCRegAliasIterator AI(Reg, TRI, true); AI.isValid(); ++AI)
-          ModifiedRegs.set(*AI);
-    } else {
-      assert(MO.isUse() && "Reg operand not a def and not a use");
-      for (MCRegAliasIterator AI(Reg, TRI, true); AI.isValid(); ++AI)
-        UsedRegs.set(*AI);
-    }
-  }
-}
-
 bool TargetInstrInfo::isReallyTriviallyReMaterializableGeneric(
     const MachineInstr &MI, AliasAnalysis *AA) const {
   const MachineFunction &MF = *MI.getMF();
   const MachineRegisterInfo &MRI = MF.getRegInfo();
 
   // Remat clients assume operand 0 is the defined register.
   if (!MI.getNumOperands() || !MI.getOperand(0).isReg())
     return false;

lib/Target/AArch64/AArch64LoadStoreOptimizer.cpp

AArch64LoadStoreOpt() : MachineFunctionPass(ID) {
     initializeAArch64LoadStoreOptPass(*PassRegistry::getPassRegistry());
   }
 
   AliasAnalysis *AA;
   const AArch64InstrInfo *TII;
   const TargetRegisterInfo *TRI;
   const AArch64Subtarget *Subtarget;
 
-  // Track which registers have been modified and used.
-  BitVector ModifiedRegs, UsedRegs;
+  // Track which register units have been modified and used.
+  LiveRegUnits ModifiedRegUnits, UsedRegUnits;
 
   void getAnalysisUsage(AnalysisUsage &AU) const override {
     AU.addRequired<AAResultsWrapperPass>();
     MachineFunctionPass::getAnalysisUsage(AU);
   }
 
   // Scan the instructions looking for a load/store that can be combined
   // with the current instruction into a load/store pair.
bool AArch64LoadStoreOpt::findMatchingStore(
   MachineInstr &LoadMI = *I;
   unsigned BaseReg = getLdStBaseOp(LoadMI).getReg();
 
   // If the load is the first instruction in the block, there's obviously
   // not any matching store.
   if (MBBI == B)
     return false;
 
-  // Track which registers have been modified and used between the first insn
-  // and the second insn.
-  ModifiedRegs.reset();
-  UsedRegs.reset();
+  // Track which register units have been modified and used between the first
+  // insn and the second insn.
+  ModifiedRegUnits.clear();
+  UsedRegUnits.clear();
 
   unsigned Count = 0;
   do {
     --MBBI;
     MachineInstr &MI = *MBBI;
 
     // Don't count transient instructions towards the search limit since there
     // may be different numbers of them if e.g. debug information is present.
     if (!MI.isTransient())
       ++Count;
 
     // If the load instruction reads directly from the address to which the
     // store instruction writes and the stored value is not modified, we can
     // promote the load. Since we do not handle stores with pre-/post-index,
     // it's unnecessary to check if BaseReg is modified by the store itself.
     if (MI.mayStore() && isMatchingStore(LoadMI, MI) &&
         BaseReg == getLdStBaseOp(MI).getReg() &&
         isLdOffsetInRangeOfSt(LoadMI, MI, TII) &&
-        !ModifiedRegs[getLdStRegOp(MI).getReg()]) {
+        ModifiedRegUnits.available(getLdStRegOp(MI).getReg())) {
       StoreI = MBBI;
       return true;
     }
 
     if (MI.isCall())
       return false;
 
-    // Update modified / uses register lists.
-    TII->trackRegDefsUses(MI, ModifiedRegs, UsedRegs, TRI);
+    // Update modified / uses register units.
+    LiveRegUnits::accumulateUsedDefed(MI, ModifiedRegUnits, UsedRegUnits, TRI);
 
     // Otherwise, if the base register is modified, we have no match, so
     // return early.
-    if (ModifiedRegs[BaseReg])
+    if (!ModifiedRegUnits.available(BaseReg))
       return false;
 
     // If we encounter a store aliased with the load, return early.
     if (MI.mayStore() && mayAlias(LoadMI, MI, AA))
       return false;
   } while (MBBI != B && Count < Limit);
   return false;
 }
AArch64LoadStoreOpt::findMatchingInsn(MachineBasicBlock::iterator I,
   bool MayLoad = FirstMI.mayLoad();
   bool IsUnscaled = TII->isUnscaledLdSt(FirstMI);
   unsigned Reg = getLdStRegOp(FirstMI).getReg();
   unsigned BaseReg = getLdStBaseOp(FirstMI).getReg();
   int Offset = getLdStOffsetOp(FirstMI).getImm();
   int OffsetStride = IsUnscaled ? getMemScale(FirstMI) : 1;
   bool IsPromotableZeroStore = isPromotableZeroStoreInst(FirstMI);
 
-  // Track which registers have been modified and used between the first insn
-  // (inclusive) and the second insn.
-  ModifiedRegs.reset();
-  UsedRegs.reset();
+  // Track which register units have been modified and used between the first
+  // insn (inclusive) and the second insn.
+  ModifiedRegUnits.clear();
+  UsedRegUnits.clear();
 
   // Remember any instructions that read/write memory between FirstMI and MI.
   SmallVector<MachineInstr *, 4> MemInsns;
 
   for (unsigned Count = 0; MBBI != E && Count < Limit; ++MBBI) {
     MachineInstr &MI = *MBBI;
 
     // Don't count transient instructions towards the search limit since there
if (areCandidatesToMergeOrPair(FirstMI, MI, Flags, TII) &&
         // We're trying to pair instructions that differ in how they are scaled.
         // If FirstMI is scaled then scale the offset of MI accordingly.
         // Otherwise, do the opposite (i.e., make MI's offset unscaled).
         int MemSize = getMemScale(MI);
         if (MIIsUnscaled) {
           // If the unscaled offset isn't a multiple of the MemSize, we can't
           // pair the operations together: bail and keep looking.
           if (MIOffset % MemSize) {
-            TII->trackRegDefsUses(MI, ModifiedRegs, UsedRegs, TRI);
+            LiveRegUnits::accumulateUsedDefed(MI, ModifiedRegUnits,
+                                              UsedRegUnits, TRI);
             MemInsns.push_back(&MI);
             continue;
           }
           MIOffset /= MemSize;
         } else {
           MIOffset *= MemSize;
         }
       }
 
       if (BaseReg == MIBaseReg && ((Offset == MIOffset + OffsetStride) ||
                                    (Offset + OffsetStride == MIOffset))) {
         int MinOffset = Offset < MIOffset ? Offset : MIOffset;
         if (FindNarrowMerge) {
           // If the alignment requirements of the scaled wide load/store
           // instruction can't express the offset of the scaled narrow input,
           // bail and keep looking. For promotable zero stores, allow only when
           // the stored value is the same (i.e., WZR).
           if ((!IsUnscaled && alignTo(MinOffset, 2) != MinOffset) ||
               (IsPromotableZeroStore && Reg != getLdStRegOp(MI).getReg())) {
-            TII->trackRegDefsUses(MI, ModifiedRegs, UsedRegs, TRI);
+            LiveRegUnits::accumulateUsedDefed(MI, ModifiedRegUnits,
+                                              UsedRegUnits, TRI);
             MemInsns.push_back(&MI);
             continue;
           }
         } else {
           // Pairwise instructions have a 7-bit signed offset field. Single
           // insns have a 12-bit unsigned offset field.  If the resultant
           // immediate offset of merging these instructions is out of range for
           // a pairwise instruction, bail and keep looking.
           if (!inBoundsForPair(IsUnscaled, MinOffset, OffsetStride)) {
-            TII->trackRegDefsUses(MI, ModifiedRegs, UsedRegs, TRI);
+            LiveRegUnits::accumulateUsedDefed(MI, ModifiedRegUnits,
+                                              UsedRegUnits, TRI);
             MemInsns.push_back(&MI);
             continue;
           }
           // If the alignment requirements of the paired (scaled) instruction
           // can't express the offset of the unscaled input, bail and keep
           // looking.
           if (IsUnscaled && (alignTo(MinOffset, OffsetStride) != MinOffset)) {
-            TII->trackRegDefsUses(MI, ModifiedRegs, UsedRegs, TRI);
+            LiveRegUnits::accumulateUsedDefed(MI, ModifiedRegUnits,
+                                              UsedRegUnits, TRI);
             MemInsns.push_back(&MI);
             continue;
           }
         }
         // If the destination register of the loads is the same register, bail
         // and keep looking. A load-pair instruction with both destination
         // registers the same is UNPREDICTABLE and will result in an exception.
         if (MayLoad && Reg == getLdStRegOp(MI).getReg()) {
-          TII->trackRegDefsUses(MI, ModifiedRegs, UsedRegs, TRI);
+          LiveRegUnits::accumulateUsedDefed(MI, ModifiedRegUnits, UsedRegUnits,
+                                            TRI);
           MemInsns.push_back(&MI);
           continue;
         }
 
         // If the Rt of the second instruction was not modified or used between
         // the two instructions and none of the instructions between the second
         // and first alias with the second, we can combine the second into the
         // first.
-        if (!ModifiedRegs[getLdStRegOp(MI).getReg()] &&
-            !(MI.mayLoad() && UsedRegs[getLdStRegOp(MI).getReg()]) &&
+        if (ModifiedRegUnits.available(getLdStRegOp(MI).getReg()) &&
+            !(MI.mayLoad() &&
+              !UsedRegUnits.available(getLdStRegOp(MI).getReg())) &&
             !mayAlias(MI, MemInsns, AA)) {
           Flags.setMergeForward(false);
           return MBBI;
         }
 
         // Likewise, if the Rt of the first instruction is not modified or used
         // between the two instructions and none of the instructions between the
         // first and the second alias with the first, we can combine the first
         // into the second.
-        if (!ModifiedRegs[getLdStRegOp(FirstMI).getReg()] &&
-            !(MayLoad && UsedRegs[getLdStRegOp(FirstMI).getReg()]) &&
+        if (ModifiedRegUnits.available(getLdStRegOp(FirstMI).getReg()) &&
+            !(MayLoad &&
+              !UsedRegUnits.available(getLdStRegOp(FirstMI).getReg())) &&
             !mayAlias(FirstMI, MemInsns, AA)) {
           Flags.setMergeForward(true);
           return MBBI;
         }
         // Unable to combine these instructions due to interference in between.
         // Keep looking.
       }
     }
 
     // If the instruction wasn't a matching load or store.  Stop searching if we
     // encounter a call instruction that might modify memory.
     if (MI.isCall())
       return E;
 
-    // Update modified / uses register lists.
-    TII->trackRegDefsUses(MI, ModifiedRegs, UsedRegs, TRI);
+    // Update modified / uses register units.
+    LiveRegUnits::accumulateUsedDefed(MI, ModifiedRegUnits, UsedRegUnits, TRI);
 
     // Otherwise, if the base register is modified, we have no match, so
     // return early.
-    if (ModifiedRegs[BaseReg])
+    if (!ModifiedRegUnits.available(BaseReg))
       return E;
 
     // Update list of instructions that read/write memory.
     if (MI.mayLoadOrStore())
       MemInsns.push_back(&MI);
   }
   return E;
 }
MachineBasicBlock::iterator AArch64LoadStoreOpt::findMatchingUpdateInsnForward(
   // merge the update.
   bool IsPairedInsn = isPairedLdSt(MemMI);
   for (unsigned i = 0, e = IsPairedInsn ? 2 : 1; i != e; ++i) {
     unsigned DestReg = getLdStRegOp(MemMI, i).getReg();
     if (DestReg == BaseReg || TRI->isSubRegister(BaseReg, DestReg))
       return E;
   }
 
-  // Track which registers have been modified and used between the first insn
-  // (inclusive) and the second insn.
-  ModifiedRegs.reset();
-  UsedRegs.reset();
+  // Track which register units have been modified and used between the first
+  // insn (inclusive) and the second insn.
+  ModifiedRegUnits.clear();
+  UsedRegUnits.clear();
   ++MBBI;
   for (unsigned Count = 0; MBBI != E && Count < Limit; ++MBBI) {
     MachineInstr &MI = *MBBI;
 
     // Don't count transient instructions towards the search limit since there
     // may be different numbers of them if e.g. debug information is present.
     if (!MI.isTransient())
       ++Count;
 
     // If we found a match, return it.
     if (isMatchingUpdateInsn(*I, MI, BaseReg, UnscaledOffset))
       return MBBI;
 
     // Update the status of what the instruction clobbered and used.
-    TII->trackRegDefsUses(MI, ModifiedRegs, UsedRegs, TRI);
+    LiveRegUnits::accumulateUsedDefed(MI, ModifiedRegUnits, UsedRegUnits, TRI);
 
     // Otherwise, if the base register is used or modified, we have no match, so
     // return early.
-    if (ModifiedRegs[BaseReg] || UsedRegs[BaseReg])
+    if (!ModifiedRegUnits.available(BaseReg) ||
+        !UsedRegUnits.available(BaseReg))
       return E;
   }
   return E;
 }
 
 MachineBasicBlock::iterator AArch64LoadStoreOpt::findMatchingUpdateInsnBackward(
     MachineBasicBlock::iterator I, unsigned Limit) {
   MachineBasicBlock::iterator B = I->getParent()->begin();
MachineBasicBlock::iterator AArch64LoadStoreOpt::findMatchingUpdateInsnBackward(
   // merge the update.
   bool IsPairedInsn = isPairedLdSt(MemMI);
   for (unsigned i = 0, e = IsPairedInsn ? 2 : 1; i != e; ++i) {
     unsigned DestReg = getLdStRegOp(MemMI, i).getReg();
     if (DestReg == BaseReg || TRI->isSubRegister(BaseReg, DestReg))
       return E;
   }
 
-  // Track which registers have been modified and used between the first insn
-  // (inclusive) and the second insn.
-  ModifiedRegs.reset();
-  UsedRegs.reset();
+  // Track which register units have been modified and used between the first
+  // insn (inclusive) and the second insn.
+  ModifiedRegUnits.clear();
+  UsedRegUnits.clear();
   unsigned Count = 0;
   do {
     --MBBI;
     MachineInstr &MI = *MBBI;
 
     // Don't count transient instructions towards the search limit since there
     // may be different numbers of them if e.g. debug information is present.
     if (!MI.isTransient())
       ++Count;
 
     // If we found a match, return it.
     if (isMatchingUpdateInsn(*I, MI, BaseReg, Offset))
       return MBBI;
 
     // Update the status of what the instruction clobbered and used.
-    TII->trackRegDefsUses(MI, ModifiedRegs, UsedRegs, TRI);
+    LiveRegUnits::accumulateUsedDefed(MI, ModifiedRegUnits, UsedRegUnits, TRI);
 
     // Otherwise, if the base register is used or modified, we have no match, so
     // return early.
-    if (ModifiedRegs[BaseReg] || UsedRegs[BaseReg])
+    if (!ModifiedRegUnits.available(BaseReg) ||
+        !UsedRegUnits.available(BaseReg))
       return E;
   } while (MBBI != B && Count < Limit);
   return E;
 }
 
 bool AArch64LoadStoreOpt::tryToPromoteLoadFromStore(
     MachineBasicBlock::iterator &MBBI) {
   MachineInstr &MI = *MBBI;
bool AArch64LoadStoreOpt::runOnMachineFunction(MachineFunction &Fn) {
   if (skipFunction(Fn.getFunction()))
     return false;
 
   Subtarget = &static_cast<const AArch64Subtarget &>(Fn.getSubtarget());
   TII = static_cast<const AArch64InstrInfo *>(Subtarget->getInstrInfo());
   TRI = Subtarget->getRegisterInfo();
   AA = &getAnalysis<AAResultsWrapperPass>().getAAResults();
 
-  // Resize the modified and used register bitfield trackers.  We do this once
-  // per function and then clear the bitfield each time we optimize a load or
-  // store.
-  ModifiedRegs.resize(TRI->getNumRegs());
-  UsedRegs.resize(TRI->getNumRegs());
+  // Resize the modified and used register unit trackers.  We do this once
+  // per function and then clear the register units each time we optimize a load
+  // or store.
+  ModifiedRegUnits.init(*TRI);
+  UsedRegUnits.init(*TRI);
 
   bool Modified = false;
   bool enableNarrowZeroStOpt = !Subtarget->requiresStrictAlign();
   for (auto &MBB : Fn)
     Modified |= optimizeBlock(MBB, enableNarrowZeroStOpt);
 
   return Modified;
 }