[AArch64] Teach Load/Store optimizier to rename store operands for pairing.

In some cases, we can rename a store operand, in order to enable pairing
of stores.  For store pairs, that cannot be merged because the first
tored register is defined in between the second store, we try to find
suitable rename register.

First, we check if we can rename the given register:

1. The first store register must be killed at the store, which means we
   do not have to rename instructions after the first store.
2. We scan backwards from the first store, to find the definition of the
   stored register and check all uses in between are renamable. Along
   they way, we collect the minimal register classes of the uses for
   overlapping (sub/super)registers.

Second, we try to find an available register from the minimal physical
register class of the original register. A suitable register must not be

1. defined before FirstMI
2. between the previous definition of the register to rename
3. a callee saved register.

We use KILL flags to clear defined registers while scanning from the
beginning to the end of the block.

This triggers quite often, here are the top changes for MultiSource,
SPEC2000, SPEC2006 compiled with -O3 for iOS:

Metric: aarch64-ldst-opt.NumPairCreated

Program                                        base     patch    diff
 test-suite...nch/fourinarow/fourinarow.test     2.00    39.00   1850.0%
 test-suite...s/ASC_Sequoia/IRSmk/IRSmk.test    46.00    80.00   73.9%
 test-suite...chmarks/Olden/power/power.test    70.00    96.00   37.1%
 test-suite...cations/hexxagon/hexxagon.test    29.00    39.00   34.5%
 test-suite...nchmarks/McCat/05-eks/eks.test   100.00   132.00   32.0%
 test-suite.../Trimaran/enc-rc4/enc-rc4.test    46.00    59.00   28.3%
 test-suite...T2006/473.astar/473.astar.test   160.00   200.00   25.0%
 test-suite.../Trimaran/enc-md5/enc-md5.test     8.00    10.00   25.0%
 test-suite...telecomm-gsm/telecomm-gsm.test   113.00   139.00   23.0%
 test-suite...ediabench/gsm/toast/toast.test   113.00   139.00   23.0%
 test-suite...Source/Benchmarks/sim/sim.test    91.00   111.00   22.0%
 test-suite...C/CFP2000/179.art/179.art.test    41.00    49.00   19.5%
 test-suite...peg2/mpeg2dec/mpeg2decode.test   245.00   279.00   13.9%
 test-suite...marks/Olden/health/health.test    16.00    18.00   12.5%
 test-suite...ks/Prolangs-C/cdecl/cdecl.test    90.00   101.00   12.2%
 test-suite...fice-ispell/office-ispell.test    91.00   100.00    9.9%
 test-suite...oxyApps-C/miniGMG/miniGMG.test   430.00   465.00    8.1%
 test-suite...lowfish/security-blowfish.test    39.00    42.00    7.7%
 test-suite.../Applications/spiff/spiff.test    42.00    45.00    7.1%
 test-suite...arks/mafft/pairlocalalign.test   2473.00  2646.00   7.0%
 test-suite.../VersaBench/ecbdes/ecbdes.test    29.00    31.00    6.9%
 test-suite...nch/beamformer/beamformer.test   220.00   235.00    6.8%
 test-suite...CFP2000/177.mesa/177.mesa.test   2110.00  2252.00   6.7%
 test-suite...ve-susan/automotive-susan.test   109.00   116.00    6.4%
 test-suite...s-C/unix-smail/unix-smail.test    65.00    69.00    6.2%
 test-suite...CI_Purple/SMG2000/smg2000.test   1194.00  1265.00   5.9%
 test-suite.../Benchmarks/nbench/nbench.test   472.00   500.00    5.9%
 test-suite...oxyApps-C/miniAMR/miniAMR.test   248.00   262.00    5.6%
 test-suite...quoia/CrystalMk/CrystalMk.test    18.00    19.00    5.6%
 test-suite...rks/tramp3d-v4/tramp3d-v4.test   7331.00  7710.00   5.2%
 test-suite.../Benchmarks/Bullet/bullet.test   5651.00  5938.00   5.1%
 test-suite...ternal/HMMER/hmmcalibrate.test   750.00   788.00    5.1%
 test-suite...T2006/456.hmmer/456.hmmer.test   764.00   802.00    5.0%
 test-suite...ications/JM/ldecod/ldecod.test   1028.00  1079.00   5.0%
 test-suite...CFP2006/444.namd/444.namd.test   1368.00  1434.00   4.8%
 test-suite...marks/7zip/7zip-benchmark.test   4471.00  4685.00   4.8%
 test-suite...6/464.h264ref/464.h264ref.test   3122.00  3271.00   4.8%
 test-suite...pplications/oggenc/oggenc.test   1497.00  1565.00   4.5%
 test-suite...T2000/300.twolf/300.twolf.test   742.00   774.00    4.3%
 test-suite.../Prolangs-C/loader/loader.test    24.00    25.00    4.2%
 test-suite...0.perlbench/400.perlbench.test   1983.00  2058.00   3.8%
 test-suite...ications/JM/lencod/lencod.test   4612.00  4785.00   3.8%
 test-suite...yApps-C++/PENNANT/PENNANT.test   995.00   1032.00   3.7%
 test-suite...arks/VersaBench/dbms/dbms.test    54.00    56.00    3.7%

Reviewers: efriedma, thegameg, samparker, dmgreen, paquette, evandro

Reviewed By: paquette

Differential Revision: https://reviews.llvm.org/D70450
diff --git a/llvm/lib/Target/AArch64/AArch64LoadStoreOptimizer.cpp b/llvm/lib/Target/AArch64/AArch64LoadStoreOptimizer.cpp
index a0c4a25..d60dc43 100644
--- a/llvm/lib/Target/AArch64/AArch64LoadStoreOptimizer.cpp
+++ b/llvm/lib/Target/AArch64/AArch64LoadStoreOptimizer.cpp
@@ -32,10 +32,12 @@
 #include "llvm/Pass.h"
 #include "llvm/Support/CommandLine.h"
 #include "llvm/Support/Debug.h"
+#include "llvm/Support/DebugCounter.h"
 #include "llvm/Support/ErrorHandling.h"
 #include "llvm/Support/raw_ostream.h"
 #include <cassert>
 #include <cstdint>
+#include <functional>
 #include <iterator>
 #include <limits>
 
@@ -51,6 +53,9 @@
 STATISTIC(NumZeroStoresPromoted, "Number of narrow zero stores promoted");
 STATISTIC(NumLoadsFromStoresPromoted, "Number of loads from stores promoted");
 
+DEBUG_COUNTER(RegRenamingCounter, DEBUG_TYPE "-reg-renaming",
+              "Controls which pairs are considered for renaming");
+
 // The LdStLimit limits how far we search for load/store pairs.
 static cl::opt<unsigned> LdStLimit("aarch64-load-store-scan-limit",
                                    cl::init(20), cl::Hidden);
@@ -76,6 +81,11 @@
   // to be extended, 0 means I, and 1 means the returned iterator.
   int SExtIdx = -1;
 
+  // If not none, RenameReg can be used to rename the result register of the
+  // first store in a pair. Currently this only works when merging stores
+  // forward.
+  Optional<MCPhysReg> RenameReg = None;
+
   LdStPairFlags() = default;
 
   void setMergeForward(bool V = true) { MergeForward = V; }
@@ -83,6 +93,10 @@
 
   void setSExtIdx(int V) { SExtIdx = V; }
   int getSExtIdx() const { return SExtIdx; }
+
+  void setRenameReg(MCPhysReg R) { RenameReg = R; }
+  void clearRenameReg() { RenameReg = None; }
+  Optional<MCPhysReg> getRenameReg() const { return RenameReg; }
 };
 
 struct AArch64LoadStoreOpt : public MachineFunctionPass {
@@ -99,6 +113,7 @@
 
   // Track which register units have been modified and used.
   LiveRegUnits ModifiedRegUnits, UsedRegUnits;
+  LiveRegUnits DefinedInBB;
 
   void getAnalysisUsage(AnalysisUsage &AU) const override {
     AU.addRequired<AAResultsWrapperPass>();
@@ -599,8 +614,8 @@
   }
 }
 
-static const MachineOperand &getLdStRegOp(const MachineInstr &MI,
-                                          unsigned PairedRegOp = 0) {
+static MachineOperand &getLdStRegOp(MachineInstr &MI,
+                                    unsigned PairedRegOp = 0) {
   assert(PairedRegOp < 2 && "Unexpected register operand idx.");
   unsigned Idx = isPairedLdSt(MI) ? PairedRegOp : 0;
   return MI.getOperand(Idx);
@@ -783,6 +798,44 @@
   return NextI;
 }
 
+// Apply Fn to all instructions between MI and the beginning of the block, until
+// a def for DefReg is reached. Returns true, iff Fn returns true for all
+// visited instructions. Stop after visiting Limit iterations.
+static bool forAllMIsUntilDef(MachineInstr &MI, MCPhysReg DefReg,
+                              const TargetRegisterInfo *TRI, unsigned Limit,
+                              std::function<bool(MachineInstr &, bool)> &Fn) {
+  auto MBB = MI.getParent();
+  for (MachineBasicBlock::reverse_iterator I = MI.getReverseIterator(),
+                                           E = MBB->rend();
+       I != E; I++) {
+    if (!Limit)
+      return false;
+    --Limit;
+
+    bool isDef = any_of(I->operands(), [DefReg, TRI](MachineOperand &MOP) {
+      return MOP.isReg() && MOP.isDef() &&
+             TRI->regsOverlap(MOP.getReg(), DefReg);
+    });
+    if (!Fn(*I, isDef))
+      return false;
+    if (isDef)
+      break;
+  }
+  return true;
+}
+
+static void updateDefinedRegisters(MachineInstr &MI, LiveRegUnits &Units,
+                                   const TargetRegisterInfo *TRI) {
+
+  for (const MachineOperand &MOP : phys_regs_and_masks(MI))
+    if (MOP.isReg() && MOP.isKill())
+      Units.removeReg(MOP.getReg());
+
+  for (const MachineOperand &MOP : phys_regs_and_masks(MI))
+    if (MOP.isReg() && !MOP.isKill())
+      Units.addReg(MOP.getReg());
+}
+
 MachineBasicBlock::iterator
 AArch64LoadStoreOpt::mergePairedInsns(MachineBasicBlock::iterator I,
                                       MachineBasicBlock::iterator Paired,
@@ -803,6 +856,70 @@
   int OffsetStride = IsUnscaled ? getMemScale(*I) : 1;
 
   bool MergeForward = Flags.getMergeForward();
+
+  Optional<MCPhysReg> RenameReg = Flags.getRenameReg();
+  if (MergeForward && RenameReg) {
+    MCRegister RegToRename = getLdStRegOp(*I).getReg();
+    DefinedInBB.addReg(*RenameReg);
+
+    // Return the sub/super register for RenameReg, matching the size of
+    // OriginalReg.
+    auto GetMatchingSubReg = [this,
+                              RenameReg](MCPhysReg OriginalReg) -> MCPhysReg {
+      for (MCPhysReg SubOrSuper : TRI->sub_and_superregs_inclusive(*RenameReg))
+        if (TRI->getMinimalPhysRegClass(OriginalReg) ==
+            TRI->getMinimalPhysRegClass(SubOrSuper))
+          return SubOrSuper;
+      llvm_unreachable("Should have found matching sub or super register!");
+    };
+
+    std::function<bool(MachineInstr &, bool)> UpdateMIs =
+        [this, RegToRename, GetMatchingSubReg](MachineInstr &MI, bool IsDef) {
+          if (IsDef) {
+            bool SeenDef = false;
+            for (auto &MOP : MI.operands()) {
+              // Rename the first explicit definition and all implicit
+              // definitions matching RegToRename.
+              if (MOP.isReg() &&
+                  (!SeenDef || (MOP.isDef() && MOP.isImplicit())) &&
+                  TRI->regsOverlap(MOP.getReg(), RegToRename)) {
+                assert((MOP.isImplicit() ||
+                        (MOP.isRenamable() && !MOP.isEarlyClobber())) &&
+                       "Need renamable operands");
+                MOP.setReg(GetMatchingSubReg(MOP.getReg()));
+                SeenDef = true;
+              }
+            }
+          } else {
+            for (auto &MOP : MI.operands()) {
+              if (MOP.isReg() && TRI->regsOverlap(MOP.getReg(), RegToRename)) {
+                assert(MOP.isImplicit() ||
+                       (MOP.isRenamable() && !MOP.isEarlyClobber()) &&
+                           "Need renamable operands");
+                MOP.setReg(GetMatchingSubReg(MOP.getReg()));
+              }
+            }
+          }
+          LLVM_DEBUG(dbgs() << "Renamed " << MI << "\n");
+          return true;
+        };
+    forAllMIsUntilDef(*I, RegToRename, TRI, LdStLimit, UpdateMIs);
+
+    // Make sure the register used for renaming is not used between the paired
+    // instructions. That would trash the content before the new paired
+    // instruction.
+    for (auto &MI :
+         iterator_range<MachineInstrBundleIterator<llvm::MachineInstr>>(
+             std::next(I), std::next(Paired)))
+      assert(all_of(MI.operands(),
+                    [this, &RenameReg](const MachineOperand &MOP) {
+                      return !MOP.isReg() ||
+                             !TRI->regsOverlap(MOP.getReg(), *RenameReg);
+                    }) &&
+             "Rename register used between paired instruction, trashing the "
+             "content");
+  }
+
   // Insert our new paired instruction after whichever of the paired
   // instructions MergeForward indicates.
   MachineBasicBlock::iterator InsertionPoint = MergeForward ? Paired : I;
@@ -931,6 +1048,11 @@
   }
   LLVM_DEBUG(dbgs() << "\n");
 
+  if (MergeForward)
+    for (const MachineOperand &MOP : phys_regs_and_masks(*I))
+      if (MOP.isReg() && MOP.isKill())
+        DefinedInBB.addReg(MOP.getReg());
+
   // Erase the old instructions.
   I->eraseFromParent();
   Paired->eraseFromParent();
@@ -1207,6 +1329,144 @@
   // FIXME: Can we also match a mixed sext/zext unscaled/scaled pair?
 }
 
+static bool
+canRenameUpToDef(MachineInstr &FirstMI, LiveRegUnits &UsedInBetween,
+                 SmallPtrSetImpl<const TargetRegisterClass *> &RequiredClasses,
+                 const TargetRegisterInfo *TRI) {
+  if (!FirstMI.mayStore())
+    return false;
+
+  // Check if we can find an unused register which we can use to rename
+  // the register used by the first load/store.
+  auto *RegClass = TRI->getMinimalPhysRegClass(getLdStRegOp(FirstMI).getReg());
+  MachineFunction &MF = *FirstMI.getParent()->getParent();
+  if (!RegClass || !MF.getRegInfo().tracksLiveness())
+    return false;
+
+  auto RegToRename = getLdStRegOp(FirstMI).getReg();
+  // For now, we only rename if the store operand gets killed at the store.
+  if (!getLdStRegOp(FirstMI).isKill() &&
+      !any_of(FirstMI.operands(),
+              [TRI, RegToRename](const MachineOperand &MOP) {
+                return MOP.isReg() && MOP.isImplicit() && MOP.isKill() &&
+                       TRI->regsOverlap(RegToRename, MOP.getReg());
+              })) {
+    LLVM_DEBUG(dbgs() << "  Operand not killed at " << FirstMI << "\n");
+    return false;
+  }
+  auto canRenameMOP = [](const MachineOperand &MOP) {
+    return MOP.isImplicit() ||
+           (MOP.isRenamable() && !MOP.isEarlyClobber() && !MOP.isTied());
+  };
+
+  bool FoundDef = false;
+
+  // For each instruction between FirstMI and the previous def for RegToRename,
+  // we
+  // * check if we can rename RegToRename in this instruction
+  // * collect the registers used and required register classes for RegToRename.
+  std::function<bool(MachineInstr &, bool)> CheckMIs = [&](MachineInstr &MI,
+                                                           bool IsDef) {
+    LLVM_DEBUG(dbgs() << "Checking " << MI << "\n");
+    // Currently we do not try to rename across frame-setup instructions.
+    if (MI.getFlag(MachineInstr::FrameSetup)) {
+      LLVM_DEBUG(dbgs() << "  Cannot rename framesetup instructions currently ("
+                        << MI << ")\n");
+      return false;
+    }
+
+    UsedInBetween.accumulate(MI);
+
+    // For a definition, check that we can rename the definition and exit the
+    // loop.
+    FoundDef = IsDef;
+
+    // For defs, check if we can rename the first def of RegToRename.
+    if (FoundDef) {
+      for (auto &MOP : MI.operands()) {
+        if (!MOP.isReg() || !MOP.isDef() ||
+            !TRI->regsOverlap(MOP.getReg(), RegToRename))
+          continue;
+        if (!canRenameMOP(MOP)) {
+          LLVM_DEBUG(dbgs()
+                     << "  Cannot rename " << MOP << " in " << MI << "\n");
+          return false;
+        }
+        RequiredClasses.insert(TRI->getMinimalPhysRegClass(MOP.getReg()));
+      }
+      return true;
+    } else {
+      for (auto &MOP : MI.operands()) {
+        if (!MOP.isReg() || !TRI->regsOverlap(MOP.getReg(), RegToRename))
+          continue;
+
+        if (!canRenameMOP(MOP)) {
+          LLVM_DEBUG(dbgs()
+                     << "  Cannot rename " << MOP << " in " << MI << "\n");
+          return false;
+        }
+        RequiredClasses.insert(TRI->getMinimalPhysRegClass(MOP.getReg()));
+      }
+    }
+    return true;
+  };
+
+  if (!forAllMIsUntilDef(FirstMI, RegToRename, TRI, LdStLimit, CheckMIs))
+    return false;
+
+  if (!FoundDef) {
+    LLVM_DEBUG(dbgs() << "  Did not find definition for register in BB\n");
+    return false;
+  }
+  return true;
+}
+
+// Check if we can find a physical register for renaming. This register must:
+// * not be defined up to FirstMI (checking DefinedInBB)
+// * not used between the MI and the defining instruction of the register to
+//   rename (checked using UsedInBetween).
+// * is available in all used register classes (checked using RequiredClasses).
+static Optional<MCPhysReg> tryToFindRegisterToRename(
+    MachineInstr &FirstMI, MachineInstr &MI, LiveRegUnits &DefinedInBB,
+    LiveRegUnits &UsedInBetween,
+    SmallPtrSetImpl<const TargetRegisterClass *> &RequiredClasses,
+    const TargetRegisterInfo *TRI) {
+  auto &MF = *FirstMI.getParent()->getParent();
+
+  // Checks if any sub- or super-register of PR is callee saved.
+  auto AnySubOrSuperRegCalleePreserved = [&MF, TRI](MCPhysReg PR) {
+    return any_of(TRI->sub_and_superregs_inclusive(PR),
+                  [&MF, TRI](MCPhysReg SubOrSuper) {
+                    return TRI->isCalleeSavedPhysReg(SubOrSuper, MF);
+                  });
+  };
+
+  // Check if PR or one of its sub- or super-registers can be used for all
+  // required register classes.
+  auto CanBeUsedForAllClasses = [&RequiredClasses, TRI](MCPhysReg PR) {
+    return all_of(RequiredClasses, [PR, TRI](const TargetRegisterClass *C) {
+      return any_of(TRI->sub_and_superregs_inclusive(PR),
+                    [C, TRI](MCPhysReg SubOrSuper) {
+                      return C == TRI->getMinimalPhysRegClass(SubOrSuper);
+                    });
+    });
+  };
+
+  auto *RegClass = TRI->getMinimalPhysRegClass(getLdStRegOp(FirstMI).getReg());
+  for (const MCPhysReg &PR : *RegClass) {
+    if (DefinedInBB.available(PR) && UsedInBetween.available(PR) &&
+        !AnySubOrSuperRegCalleePreserved(PR) && CanBeUsedForAllClasses(PR)) {
+      DefinedInBB.addReg(PR);
+      LLVM_DEBUG(dbgs() << "Found rename register " << printReg(PR, TRI)
+                        << "\n");
+      return {PR};
+    }
+  }
+  LLVM_DEBUG(dbgs() << "No rename register found from "
+                    << TRI->getRegClassName(RegClass) << "\n");
+  return None;
+}
+
 /// Scan the instructions looking for a load/store that can be combined with the
 /// current instruction into a wider equivalent or a load/store pair.
 MachineBasicBlock::iterator
@@ -1215,6 +1475,7 @@
                                       bool FindNarrowMerge) {
   MachineBasicBlock::iterator E = I->getParent()->end();
   MachineBasicBlock::iterator MBBI = I;
+  MachineBasicBlock::iterator MBBIWithRenameReg;
   MachineInstr &FirstMI = *I;
   ++MBBI;
 
@@ -1226,6 +1487,13 @@
   int OffsetStride = IsUnscaled ? getMemScale(FirstMI) : 1;
   bool IsPromotableZeroStore = isPromotableZeroStoreInst(FirstMI);
 
+  Optional<bool> MaybeCanRename = None;
+  SmallPtrSet<const TargetRegisterClass *, 5> RequiredClasses;
+  LiveRegUnits UsedInBetween;
+  UsedInBetween.init(*TRI);
+
+  Flags.clearRenameReg();
+
   // Track which register units have been modified and used between the first
   // insn (inclusive) and the second insn.
   ModifiedRegUnits.clear();
@@ -1237,6 +1505,8 @@
   for (unsigned Count = 0; MBBI != E && Count < Limit; ++MBBI) {
     MachineInstr &MI = *MBBI;
 
+    UsedInBetween.accumulate(MI);
+
     // 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())
@@ -1329,7 +1599,9 @@
             !(MI.mayLoad() &&
               !UsedRegUnits.available(getLdStRegOp(MI).getReg())) &&
             !mayAlias(MI, MemInsns, AA)) {
+
           Flags.setMergeForward(false);
+          Flags.clearRenameReg();
           return MBBI;
         }
 
@@ -1337,18 +1609,41 @@
         // 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 (ModifiedRegUnits.available(getLdStRegOp(FirstMI).getReg()) &&
-            !(MayLoad &&
+        if (!(MayLoad &&
               !UsedRegUnits.available(getLdStRegOp(FirstMI).getReg())) &&
             !mayAlias(FirstMI, MemInsns, AA)) {
-          Flags.setMergeForward(true);
-          return MBBI;
+
+          if (ModifiedRegUnits.available(getLdStRegOp(FirstMI).getReg())) {
+            Flags.setMergeForward(true);
+            Flags.clearRenameReg();
+            return MBBI;
+          }
+
+          if (DebugCounter::shouldExecute(RegRenamingCounter)) {
+            if (!MaybeCanRename)
+              MaybeCanRename = {canRenameUpToDef(FirstMI, UsedInBetween,
+                                                 RequiredClasses, TRI)};
+
+            if (*MaybeCanRename) {
+              Optional<MCPhysReg> MaybeRenameReg = tryToFindRegisterToRename(
+                  FirstMI, MI, DefinedInBB, UsedInBetween, RequiredClasses,
+                  TRI);
+              if (MaybeRenameReg) {
+                Flags.setRenameReg(*MaybeRenameReg);
+                Flags.setMergeForward(true);
+                MBBIWithRenameReg = MBBI;
+              }
+            }
+          }
         }
         // Unable to combine these instructions due to interference in between.
         // Keep looking.
       }
     }
 
+    if (Flags.getRenameReg())
+      return MBBIWithRenameReg;
+
     // 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())
@@ -1680,7 +1975,13 @@
       ++NumUnscaledPairCreated;
     // Keeping the iterator straight is a pain, so we let the merge routine tell
     // us what the next instruction is after it's done mucking about.
+    auto Prev = std::prev(MBBI);
     MBBI = mergePairedInsns(MBBI, Paired, Flags);
+    // Collect liveness info for instructions between Prev and the new position
+    // MBBI.
+    for (auto I = std::next(Prev); I != MBBI; I++)
+      updateDefinedRegisters(*I, DefinedInBB, TRI);
+
     return true;
   }
   return false;
@@ -1742,6 +2043,7 @@
 
 bool AArch64LoadStoreOpt::optimizeBlock(MachineBasicBlock &MBB,
                                         bool EnableNarrowZeroStOpt) {
+
   bool Modified = false;
   // Four tranformations to do here:
   // 1) Find loads that directly read from stores and promote them by
@@ -1786,8 +2088,17 @@
   //        ldr x1, [x2, #8]
   //        ; becomes
   //        ldp x0, x1, [x2]
+
+  if (MBB.getParent()->getRegInfo().tracksLiveness()) {
+    DefinedInBB.clear();
+    DefinedInBB.addLiveIns(MBB);
+  }
+
   for (MachineBasicBlock::iterator MBBI = MBB.begin(), E = MBB.end();
        MBBI != E;) {
+    // Track currently live registers up to this point, to help with
+    // searching for a rename register on demand.
+    updateDefinedRegisters(*MBBI, DefinedInBB, TRI);
     if (TII->isPairableLdStInst(*MBBI) && tryToPairLdStInst(MBBI))
       Modified = true;
     else
@@ -1825,11 +2136,14 @@
   // or store.
   ModifiedRegUnits.init(*TRI);
   UsedRegUnits.init(*TRI);
+  DefinedInBB.init(*TRI);
 
   bool Modified = false;
   bool enableNarrowZeroStOpt = !Subtarget->requiresStrictAlign();
-  for (auto &MBB : Fn)
-    Modified |= optimizeBlock(MBB, enableNarrowZeroStOpt);
+  for (auto &MBB : Fn) {
+    auto M = optimizeBlock(MBB, enableNarrowZeroStOpt);
+    Modified |= M;
+  }
 
   return Modified;
 }