AMDGPU/SILoadStoreOptimizer: Optimize scanning for mergeable instructions

[tstellar]

Summary

This adds a pre-pass to this optimization that scans through the basic
block and generates lists of mergeable instructions with one list per unique
address.

In the optimization phase instead of scanning through the basic block for mergeable
instructions, we now iterate over the lists generated by the pre-pass.

The decision to re-optimize a block is now made per list, so if we fail to merge any
instructions with the same address, then we do not attempt to optimize them in
future passes over the block. This will help to reduce the time this pass
spends re-optimizing instructions.

In one pathological test case, this change reduces the time spent in the
SILoadStoreOptimizer from 0.2s to 0.03s.

This restructuring will also make it possible to implement further solutions in
this pass, because we can now add less expensive checks to the pre-pass and
filter instructions out early which will avoid the need to do the expensive
scanning during the optimization pass. For example, checking for adjacent
offsets is an inexpensive test we can move to the pre-pass.

[arsenm]

llvm/lib/Target/AMDGPU/SILoadStoreOptimizer.cpp

181

Typo s/on/no

1589

Why std::list, and a std::list of lists?

[tstellar]

llvm/lib/Target/AMDGPU/SILoadStoreOptimizer.cpp

1589

The main reason to use lists is so I can remove items without invalidating iterators.

[arsenm]

LGTM but I would rather avoid the list usage

[tstellar]

In D65961#1675758, @arsenm wrote:

LGTM but I would rather avoid the list usage

What do you think would be a better alternative? The operations used iterate, emplace_back, size, and erase.

[arsenm]

In D65961#1692092, @tstellar wrote:

In D65961#1675758, @arsenm wrote:

LGTM but I would rather avoid the list usage

What do you think would be a better alternative? The operations used iterate, emplace_back, size, and erase.

ok

[nhaehnle]

Thank you for doing this, it seems quite useful. As a follow-up to this change, do you think it makes sense to refactor CombineInfo a bit? We have a list of mergeable instructions, but the CombineInfo structure also has fields for a second instruction, which are only for temporary use, which is a bit odd.

[foad]

Hi @tstellar, I'm looking into a case where this patch slowed down a shader by 10%. Before I go too far, was this patch supposed to change the behaviour at all, or was it supposed to be purely a compile time improvement?

In the case I'm looking at it seems to do the same amount of load merging as before, but the merged loads are inserted at different places in the basic block.

[tstellar]

In D65961#1763143, @foad wrote:

Hi @tstellar, I'm looking into a case where this patch slowed down a shader by 10%. Before I go too far, was this patch supposed to change the behaviour at all, or was it supposed to be purely a compile time improvement?

The intention was to not change the behavior at all.

In the case I'm looking at it seems to do the same amount of load merging as before, but the merged loads are inserted at different places in the basic block.

Do you have a MIR or .ll dump of the shader I could look at ? Also, does https://reviews.llvm.org/D65966 help?

[foad]

In D65961#1765026, @tstellar wrote:

In D65961#1763143, @foad wrote:

Hi @tstellar, I'm looking into a case where this patch slowed down a shader by 10%. Before I go too far, was this patch supposed to change the behaviour at all, or was it supposed to be purely a compile time improvement?

The intention was to not change the behavior at all.

In the case I'm looking at it seems to do the same amount of load merging as before, but the merged loads are inserted at different places in the basic block.

Do you have a MIR or .ll dump of the shader I could look at ? Also, does https://reviews.llvm.org/D65966 help?

P8173 is a MIR test case. See the RUN line for how to run it. I see significant differences in the placing of the merged BUFFER_LOAD instructions with/without D65961 (or before/after it was committed).

I tried applying D65966 on top of rGe6f51713054f but it made no difference to the output.

[tstellar]

In D65961#1766760, @foad wrote:

In D65961#1765026, @tstellar wrote:

In D65961#1763143, @foad wrote:

Hi @tstellar, I'm looking into a case where this patch slowed down a shader by 10%. Before I go too far, was this patch supposed to change the behaviour at all, or was it supposed to be purely a compile time improvement?

The intention was to not change the behavior at all.

In the case I'm looking at it seems to do the same amount of load merging as before, but the merged loads are inserted at different places in the basic block.

Do you have a MIR or .ll dump of the shader I could look at ? Also, does https://reviews.llvm.org/D65966 help?

P8173 is a MIR test case. See the RUN line for how to run it. I see significant differences in the placing of the merged BUFFER_LOAD instructions with/without D65961 (or before/after it was committed).

I tried applying D65966 on top of rGe6f51713054f but it made no difference to the output.

You need to apply it on top of 3a8d80944b7766449e2c8784a8fb30d19a2ba16c or newer for it to have an impact. When I do that, D65966 does help improve the merging.

llvm/lib/Target/AMDGPU/SILoadStoreOptimizer.cpp

bool hasSameBaseAddress(const MachineInstr &MI) {
         if (AddrReg[i]->getReg() != AddrRegNext.getReg() ||
             AddrReg[i]->getSubReg() != AddrRegNext.getSubReg()) {
          return false;
         }
       }
       return true;
     }
 
+    bool hasMergeableAddress(const MachineRegisterInfo &MRI) {
+      for (unsigned i = 0; i < NumAddresses; ++i) {
+        const MachineOperand *AddrOp = AddrReg[i];
+        // Immediates are always OK.
+        if (AddrOp->isImm())
+          continue;
+
+        // Don't try to merge addresses that aren't either immediates or registers.
+        // TODO: Should be possible to merge FrameIndexes and maybe some other
+        // non-register
+        if (!AddrOp->isReg())
+          return false;
+
+        // TODO: We should be able to merge physical reg addreses.
+        if (Register::isPhysicalRegister(AddrOp->getReg()))
+          return false;
+
+        // If an address has only one use then there will be on other
+        // instructions with the same address, so we can't merge this one.
+        if (MRI.hasOneNonDBGUse(AddrOp->getReg()))
+          return false;
+      }
+      return true;
+    }
+
     void setMI(MachineBasicBlock::iterator MI, const SIInstrInfo &TII,
                const GCNSubtarget &STM);
     void setPaired(MachineBasicBlock::iterator MI, const SIInstrInfo &TII);
   };
 
   struct BaseRegisters {
     unsigned LoReg = 0;
     unsigned HiReg = 0;
private:
   Optional<int32_t> extractConstOffset(const MachineOperand &Op) const;
   void processBaseWithConstOffset(const MachineOperand &Base, MemAddress &Addr) const;
   /// Promotes constant offset to the immediate by adjusting the base. It
   /// tries to use a base from the nearby instructions that allows it to have
   /// a 13bit constant offset which gets promoted to the immediate.
   bool promoteConstantOffsetToImm(MachineInstr &CI,
                                   MemInfoMap &Visited,
                                   SmallPtrSet<MachineInstr *, 4> &Promoted) const;
+  void addInstToMergeableList(const CombineInfo &CI,
+                  std::list<std::list<CombineInfo> > &MergeableInsts) const;
+  bool collectMergeableInsts(MachineBasicBlock &MBB,
+                  std::list<std::list<CombineInfo> > &MergeableInsts) const;
 
 public:
   static char ID;
 
   SILoadStoreOptimizer() : MachineFunctionPass(ID) {
     initializeSILoadStoreOptimizerPass(*PassRegistry::getPassRegistry());
   }
 
-  bool optimizeBlock(MachineBasicBlock &MBB);
+  void removeCombinedInst(std::list<CombineInfo> &MergeList,
+                                         const MachineInstr &MI);
+  bool optimizeInstsWithSameBaseAddr(std::list<CombineInfo> &MergeList,
+                                     bool &OptimizeListAgain);
+  bool optimizeBlock(std::list<std::list<CombineInfo> > &MergeableInsts);
 
   bool runOnMachineFunction(MachineFunction &MF) override;
 
   StringRef getPassName() const override { return "SI Load Store Optimizer"; }
 
   void getAnalysisUsage(AnalysisUsage &AU) const override {
     AU.setPreservesCFG();
     AU.addRequired<AAResultsWrapperPass>();
void SILoadStoreOptimizer::CombineInfo::setMI(MachineBasicBlock::iterator MI,
   if (Regs & VADDR) {
     AddrOpName[NumAddresses++] = AMDGPU::OpName::vaddr;
   }
 
   for (unsigned i = 0; i < NumAddresses; i++) {
     AddrIdx[i] = AMDGPU::getNamedOperandIdx(I->getOpcode(), AddrOpName[i]);
     AddrReg[i] = &I->getOperand(AddrIdx[i]);
   }
+
+  InstsToMove.clear();
 }
 
 void SILoadStoreOptimizer::CombineInfo::setPaired(MachineBasicBlock::iterator MI,
                                                   const SIInstrInfo &TII) {
   Paired = MI;
   assert(InstClass == getInstClass(Paired->getOpcode(), TII));
   int OffsetIdx =
       AMDGPU::getNamedOperandIdx(I->getOpcode(), AMDGPU::OpName::offset);
bool SILoadStoreOptimizer::findMatchingInst(CombineInfo &CI) {
 
   const unsigned Opc = CI.I->getOpcode();
   const InstClassEnum InstClass = getInstClass(Opc, *TII);
 
   if (InstClass == UNKNOWN) {
     return false;
   }
 
-  for (unsigned i = 0; i < CI.NumAddresses; i++) {
-    // We only ever merge operations with the same base address register, so
-    // don't bother scanning forward if there are no other uses.
-    if (CI.AddrReg[i]->isReg() &&
-        (Register::isPhysicalRegister(CI.AddrReg[i]->getReg()) ||
-         MRI->hasOneNonDBGUse(CI.AddrReg[i]->getReg())))
-      return false;
-  }
-
   ++MBBI;
 
   DenseSet<unsigned> RegDefsToMove;
   DenseSet<unsigned> PhysRegUsesToMove;
   addDefsUsesToList(*CI.I, RegDefsToMove, PhysRegUsesToMove);
 
   for (; MBBI != E; ++MBBI) {
     const bool IsDS = (InstClass == DS_READ) || (InstClass == DS_WRITE);
BuildMI(*MBB, CI.Paired, DL, CopyDesc)
       .add(*Dest0) // Copy to same destination including flags and sub reg.
       .addReg(DestReg, 0, SubRegIdx0);
   MachineInstr *Copy1 = BuildMI(*MBB, CI.Paired, DL, CopyDesc)
                             .add(*Dest1)
                             .addReg(DestReg, RegState::Kill, SubRegIdx1);
 
   moveInstsAfter(Copy1, CI.InstsToMove);
 
-  MachineBasicBlock::iterator Next = std::next(CI.I);
   CI.I->eraseFromParent();
   CI.Paired->eraseFromParent();
 
   LLVM_DEBUG(dbgs() << "Inserted read2: " << *Read2 << '\n');
-  return Next;
+  return Read2;
 }
 
 unsigned SILoadStoreOptimizer::write2Opcode(unsigned EltSize) const {
   if (STM->ldsRequiresM0Init())
     return (EltSize == 4) ? AMDGPU::DS_WRITE2_B32 : AMDGPU::DS_WRITE2_B64;
   return (EltSize == 4) ? AMDGPU::DS_WRITE2_B32_gfx9
                         : AMDGPU::DS_WRITE2_B64_gfx9;
 }
MachineInstrBuilder Write2 =
           .add(*Data1)                               // data1
           .addImm(NewOffset0)                        // offset0
           .addImm(NewOffset1)                        // offset1
           .addImm(0)                                 // gds
           .cloneMergedMemRefs({&*CI.I, &*CI.Paired});
 
   moveInstsAfter(Write2, CI.InstsToMove);
 
-  MachineBasicBlock::iterator Next = std::next(CI.I);
   CI.I->eraseFromParent();
   CI.Paired->eraseFromParent();
 
   LLVM_DEBUG(dbgs() << "Inserted write2 inst: " << *Write2 << '\n');
-  return Next;
+  return Write2;
 }
 
 MachineBasicBlock::iterator
 SILoadStoreOptimizer::mergeSBufferLoadImmPair(CombineInfo &CI) {
   MachineBasicBlock *MBB = CI.I->getParent();
   DebugLoc DL = CI.I->getDebugLoc();
   const unsigned Opcode = getNewOpcode(CI);
 
   const TargetRegisterClass *SuperRC = getTargetRegisterClass(CI);
 
   unsigned DestReg = MRI->createVirtualRegister(SuperRC);
   unsigned MergedOffset = std::min(CI.Offset0, CI.Offset1);
 
   // It shouldn't be possible to get this far if the two instructions
   // don't have a single memoperand, because MachineInstr::mayAlias()
   // will return true if this is the case.
   assert(CI.I->hasOneMemOperand() && CI.Paired->hasOneMemOperand());
 
   const MachineMemOperand *MMOa = *CI.I->memoperands_begin();
   const MachineMemOperand *MMOb = *CI.Paired->memoperands_begin();
 
+  MachineInstr *New =
     BuildMI(*MBB, CI.Paired, DL, TII->get(Opcode), DestReg)
         .add(*TII->getNamedOperand(*CI.I, AMDGPU::OpName::sbase))
         .addImm(MergedOffset) // offset
         .addImm(CI.GLC0)      // glc
         .addImm(CI.DLC0)      // dlc
         .addMemOperand(combineKnownAdjacentMMOs(*MBB->getParent(), MMOa, MMOb));
 
   std::pair<unsigned, unsigned> SubRegIdx = getSubRegIdxs(CI);
   const unsigned SubRegIdx0 = std::get<0>(SubRegIdx);
   const unsigned SubRegIdx1 = std::get<1>(SubRegIdx);
 
   // Copy to the old destination registers.
   const MCInstrDesc &CopyDesc = TII->get(TargetOpcode::COPY);
   const auto *Dest0 = TII->getNamedOperand(*CI.I, AMDGPU::OpName::sdst);
   const auto *Dest1 = TII->getNamedOperand(*CI.Paired, AMDGPU::OpName::sdst);
 
   BuildMI(*MBB, CI.Paired, DL, CopyDesc)
       .add(*Dest0) // Copy to same destination including flags and sub reg.
       .addReg(DestReg, 0, SubRegIdx0);
   MachineInstr *Copy1 = BuildMI(*MBB, CI.Paired, DL, CopyDesc)
                             .add(*Dest1)
                             .addReg(DestReg, RegState::Kill, SubRegIdx1);
 
   moveInstsAfter(Copy1, CI.InstsToMove);
 
-  MachineBasicBlock::iterator Next = std::next(CI.I);
   CI.I->eraseFromParent();
   CI.Paired->eraseFromParent();
-  return Next;
+  return New;
 }
 
 MachineBasicBlock::iterator
 SILoadStoreOptimizer::mergeBufferLoadPair(CombineInfo &CI) {
   MachineBasicBlock *MBB = CI.I->getParent();
   DebugLoc DL = CI.I->getDebugLoc();
 
   const unsigned Opcode = getNewOpcode(CI);
SILoadStoreOptimizer::mergeBufferLoadPair(CombineInfo &CI) {
   // It shouldn't be possible to get this far if the two instructions
   // don't have a single memoperand, because MachineInstr::mayAlias()
   // will return true if this is the case.
   assert(CI.I->hasOneMemOperand() && CI.Paired->hasOneMemOperand());
 
   const MachineMemOperand *MMOa = *CI.I->memoperands_begin();
   const MachineMemOperand *MMOb = *CI.Paired->memoperands_begin();
 
+  MachineInstr *New =
     MIB.add(*TII->getNamedOperand(*CI.I, AMDGPU::OpName::srsrc))
         .add(*TII->getNamedOperand(*CI.I, AMDGPU::OpName::soffset))
         .addImm(MergedOffset) // offset
         .addImm(CI.GLC0)      // glc
         .addImm(CI.SLC0)      // slc
         .addImm(0)            // tfe
         .addImm(CI.DLC0)      // dlc
         .addMemOperand(combineKnownAdjacentMMOs(*MBB->getParent(), MMOa, MMOb));
 
   std::pair<unsigned, unsigned> SubRegIdx = getSubRegIdxs(CI);
   const unsigned SubRegIdx0 = std::get<0>(SubRegIdx);
   const unsigned SubRegIdx1 = std::get<1>(SubRegIdx);
 
   // Copy to the old destination registers.
   const MCInstrDesc &CopyDesc = TII->get(TargetOpcode::COPY);
   const auto *Dest0 = TII->getNamedOperand(*CI.I, AMDGPU::OpName::vdata);
   const auto *Dest1 = TII->getNamedOperand(*CI.Paired, AMDGPU::OpName::vdata);
 
   BuildMI(*MBB, CI.Paired, DL, CopyDesc)
       .add(*Dest0) // Copy to same destination including flags and sub reg.
       .addReg(DestReg, 0, SubRegIdx0);
   MachineInstr *Copy1 = BuildMI(*MBB, CI.Paired, DL, CopyDesc)
                             .add(*Dest1)
                             .addReg(DestReg, RegState::Kill, SubRegIdx1);
 
   moveInstsAfter(Copy1, CI.InstsToMove);
 
-  MachineBasicBlock::iterator Next = std::next(CI.I);
   CI.I->eraseFromParent();
   CI.Paired->eraseFromParent();
-  return Next;
+  return New;
 }
 
 unsigned SILoadStoreOptimizer::getNewOpcode(const CombineInfo &CI) {
   const unsigned Width = CI.Width0 + CI.Width1;
 
   switch (CI.InstClass) {
   default:
     return AMDGPU::getMUBUFOpcode(CI.InstClass, Width);
SILoadStoreOptimizer::mergeBufferStorePair(CombineInfo &CI) {
   // It shouldn't be possible to get this far if the two instructions
   // don't have a single memoperand, because MachineInstr::mayAlias()
   // will return true if this is the case.
   assert(CI.I->hasOneMemOperand() && CI.Paired->hasOneMemOperand());
 
   const MachineMemOperand *MMOa = *CI.I->memoperands_begin();
   const MachineMemOperand *MMOb = *CI.Paired->memoperands_begin();
 
+  MachineInstr *New =
     MIB.add(*TII->getNamedOperand(*CI.I, AMDGPU::OpName::srsrc))
         .add(*TII->getNamedOperand(*CI.I, AMDGPU::OpName::soffset))
         .addImm(std::min(CI.Offset0, CI.Offset1)) // offset
         .addImm(CI.GLC0)      // glc
         .addImm(CI.SLC0)      // slc
         .addImm(0)            // tfe
         .addImm(CI.DLC0)      // dlc
         .addMemOperand(combineKnownAdjacentMMOs(*MBB->getParent(), MMOa, MMOb));
 
   moveInstsAfter(MIB, CI.InstsToMove);
 
-  MachineBasicBlock::iterator Next = std::next(CI.I);
   CI.I->eraseFromParent();
   CI.Paired->eraseFromParent();
-  return Next;
+  return New;
 }
 
 MachineOperand
 SILoadStoreOptimizer::createRegOrImm(int32_t Val, MachineInstr &MI) const {
   APInt V(32, Val, true);
   if (TII->isInlineConstant(V))
     return MachineOperand::CreateImm(Val);
 
if (AnchorInst) {
     }
     AnchorList.insert(AnchorInst);
     return true;
   }
 
   return false;
 }
 
-// Scan through looking for adjacent LDS operations with constant offsets from
-// the same base register. We rely on the scheduler to do the hard work of
-// clustering nearby loads, and assume these are all adjacent.
-bool SILoadStoreOptimizer::optimizeBlock(MachineBasicBlock &MBB) {
-  bool Modified = false;
+void SILoadStoreOptimizer::addInstToMergeableList(const CombineInfo &CI,
+                 std::list<std::list<CombineInfo> > &MergeableInsts) const {
+  for (std::list<CombineInfo> &AddrList : MergeableInsts) {
+    if (AddrList.front().hasSameBaseAddress(*CI.I) &&
+        AddrList.front().InstClass == CI.InstClass) {
+      AddrList.emplace_back(CI);
+      return;
+    }
+  }
+
+  // Base address not found, so add a new list.
+  MergeableInsts.emplace_back(1, CI);
+}
 
+bool SILoadStoreOptimizer::collectMergeableInsts(MachineBasicBlock &MBB,
+                 std::list<std::list<CombineInfo> > &MergeableInsts) const {
+  bool Modified = false;
   // Contain the list
   MemInfoMap Visited;
   // Contains the list of instructions for which constant offsets are being
   // promoted to the IMM.
   SmallPtrSet<MachineInstr *, 4> AnchorList;
 
-  for (MachineBasicBlock::iterator I = MBB.begin(), E = MBB.end(); I != E;) {
-    MachineInstr &MI = *I;
-
+  // Sort potential mergeable instructions into lists.  One list per base address.
+  for (MachineInstr &MI : MBB.instrs()) {
+    // We run this before checking if an address is mergeable, because it can produce
+    // better code even if the instructions aren't mergeable.
     if (promoteConstantOffsetToImm(MI, Visited, AnchorList))
       Modified = true;
 
+    const InstClassEnum InstClass = getInstClass(MI.getOpcode(), *TII);
+    if (InstClass == UNKNOWN)
+      continue;
+
     // Don't combine if volatile.
-    if (MI.hasOrderedMemoryRef()) {
-      ++I;
+    if (MI.hasOrderedMemoryRef())
       continue;
-    }
 
     CombineInfo CI;
-    CI.setMI(I, *TII, *STM);
+    CI.setMI(MI, *TII, *STM);
+
+    if (!CI.hasMergeableAddress(*MRI))
+      continue;
+
+    addInstToMergeableList(CI, MergeableInsts);
+  }
+  return Modified;
+}
+
+// Scan through looking for adjacent LDS operations with constant offsets from
+// the same base register. We rely on the scheduler to do the hard work of
+// clustering nearby loads, and assume these are all adjacent.
+bool SILoadStoreOptimizer::optimizeBlock(
+                       std::list<std::list<CombineInfo> > &MergeableInsts) {
+  bool Modified = false;
+
+  for (std::list<CombineInfo> &MergeList : MergeableInsts) {
+    if (MergeList.size() < 2)
+      continue;
+
+    bool OptimizeListAgain = false;
+    if (!optimizeInstsWithSameBaseAddr(MergeList, OptimizeListAgain)) {
+      // We weren't able to make any changes, so clear the list so we don't
+      // process the same instructions the next time we try to optimize this
+      // block.
+      MergeList.clear();
+      continue;
+    }
+
+    // We made changes, but also determined that there were no more optimization
+    // opportunities, so we don't need to reprocess the list
+    if (!OptimizeListAgain)
+      MergeList.clear();
+
+    OptimizeAgain |= OptimizeListAgain;
+    Modified = true;
+  }
+  return Modified;
+}
+
+void
+SILoadStoreOptimizer::removeCombinedInst(std::list<CombineInfo> &MergeList,
+                                         const MachineInstr &MI) {
+
+  for (auto CI = MergeList.begin(), E = MergeList.end(); CI != E; ++CI) {
+    if (&*CI->I == &MI) {
+      MergeList.erase(CI);
+      return;
+    }
+  }
+}
+
+bool
+SILoadStoreOptimizer::optimizeInstsWithSameBaseAddr(
+                                          std::list<CombineInfo> &MergeList,
+                                          bool &OptimizeListAgain) {
+  bool Modified = false;
+  for (auto I = MergeList.begin(); I != MergeList.end(); ++I) {
+    CombineInfo &CI = *I;
 
     switch (CI.InstClass) {
     default:
       break;
     case DS_READ:
       if (findMatchingInst(CI)) {
         Modified = true;
-        I = mergeRead2Pair(CI);
-      } else {
-        ++I;
+        removeCombinedInst(MergeList, *CI.Paired);
+        MachineBasicBlock::iterator NewMI = mergeRead2Pair(CI);
+        CI.setMI(NewMI, *TII, *STM);
       }
-      continue;
+      break;
     case DS_WRITE:
       if (findMatchingInst(CI)) {
         Modified = true;
-        I = mergeWrite2Pair(CI);
-      } else {
-        ++I;
+        removeCombinedInst(MergeList, *CI.Paired);
+        MachineBasicBlock::iterator NewMI = mergeWrite2Pair(CI);
+        CI.setMI(NewMI, *TII, *STM);
       }
-      continue;
+      break;
     case S_BUFFER_LOAD_IMM:
       if (findMatchingInst(CI)) {
         Modified = true;
-        I = mergeSBufferLoadImmPair(CI);
-        OptimizeAgain |= (CI.Width0 + CI.Width1) < 16;
-      } else {
-        ++I;
+        removeCombinedInst(MergeList, *CI.Paired);
+        MachineBasicBlock::iterator NewMI = mergeSBufferLoadImmPair(CI);
+        CI.setMI(NewMI, *TII, *STM);
+        OptimizeListAgain |= (CI.Width0 + CI.Width1) < 16;
       }
-      continue;
+      break;
     case BUFFER_LOAD_OFFEN:
     case BUFFER_LOAD_OFFSET:
     case BUFFER_LOAD_OFFEN_exact:
     case BUFFER_LOAD_OFFSET_exact:
       if (findMatchingInst(CI)) {
         Modified = true;
-        I = mergeBufferLoadPair(CI);
-        OptimizeAgain |= (CI.Width0 + CI.Width1) < 4;
-      } else {
-        ++I;
+        removeCombinedInst(MergeList, *CI.Paired);
+        MachineBasicBlock::iterator NewMI = mergeBufferLoadPair(CI);
+        CI.setMI(NewMI, *TII, *STM);
+        OptimizeListAgain |= (CI.Width0 + CI.Width1) < 4;
       }
-      continue;
+      break;
     case BUFFER_STORE_OFFEN:
     case BUFFER_STORE_OFFSET:
     case BUFFER_STORE_OFFEN_exact:
     case BUFFER_STORE_OFFSET_exact:
       if (findMatchingInst(CI)) {
         Modified = true;
-        I = mergeBufferStorePair(CI);
-        OptimizeAgain |= (CI.Width0 + CI.Width1) < 4;
-      } else {
-        ++I;
+        removeCombinedInst(MergeList, *CI.Paired);
+        MachineBasicBlock::iterator NewMI = mergeBufferStorePair(CI);
+        CI.setMI(NewMI, *TII, *STM);
+        OptimizeListAgain |= (CI.Width0 + CI.Width1) < 4;
       }
-      continue;
+      break;
     }
-
-    ++I;
+    // Clear the InstsToMove after we have finished searching so we don't have
+    // stale values left over if we search for this CI again in another pass
+    // over the block.
+    CI.InstsToMove.clear();
   }
 
   return Modified;
 }
 
 bool SILoadStoreOptimizer::runOnMachineFunction(MachineFunction &MF) {
   if (skipFunction(MF.getFunction()))
     return false;
bool SILoadStoreOptimizer::runOnMachineFunction(MachineFunction &MF) {
   AA = &getAnalysis<AAResultsWrapperPass>().getAAResults();
 
   assert(MRI->isSSA() && "Must be run on SSA");
 
   LLVM_DEBUG(dbgs() << "Running SILoadStoreOptimizer\n");
 
   bool Modified = false;
 
+
   for (MachineBasicBlock &MBB : MF) {
+    std::list<std::list<CombineInfo> > MergeableInsts;
+    // First pass: Collect list of all instructions we know how to merge.
+    Modified |= collectMergeableInsts(MBB, MergeableInsts);
     do {
       OptimizeAgain = false;
-      Modified |= optimizeBlock(MBB);
+      Modified |= optimizeBlock(MergeableInsts);
     } while (OptimizeAgain);
   }
 
   return Modified;
 }