[CodeGen] MachinePostRAUpdater

[samparker]

Summary

Introduce a new class that acts as a wrapper around ReachingDefAnalysis that can be used to track changes and still query use-def chains. It's primary purpose is to enable the insertion, removal and movement of machine instructions. This replaces how we were using RDA in ARMLowOverheadLoops.
The initialisation code for RDA has been split up a bit so that we can re-run the traversal after making modifications to the function.

[samparker]

• isSafeToRemove now looks at the use-def chain of the block and returns yes/no along with all the instructions that should be removed. For low-overhead loops, this is means that we don't have to explicitly check for LR users within the loop and the updater will look at copies so we can remove them too! The new test case is an example where LR copies previously prevented the transform.

[SjoerdMeijer]

Good stuff!

I haven't looked at details yet, but wanted to ask a high-level dumb question first. Essentially, this is DCE on MachineFunctions/MachineLoops. Just floating an idea, and you can correct me where I go wrong. Do you perhaps see any benefits in making this is a separate pass? I was thinking that by separating the concerns (i.e. LowOverhead loops, and cleaning-up and DCE), we can have a clean-up pass running after low-overhead loops, with possibly the benefit being that low-overhead loops does not have to deal with any of the cleanup and bookkeeping.

[samparker]

No, another pass won't help here. The point of this class is to allow transforms to use ReachingDefAnalysis while changing the code at the same time, not something a pass is useful for. But it should indeed be useful for DCE. ReachingDefAnalysis could(?) be modified to re-evaluate new instructions, but I don't think general queries about being able to move instructions fits into the analysis. In the context of low-overhead loops, we still need specific information from that analysis, such as instructions to ignore, to be able to do anything useful. For instance, we have cases where we want to remove 'redundant' mov instructions but, without knowing that the loop control instructions get removed, then those movs wouldn't be dead according to this class.

[samparker]

Rebased.

[samparker]

• Rebased.
• Broke up some of the initialisation code in RDA so we can re-run after finalization.

[samparker]

I've extracted the immediately useful changes into D73460, but it doesn't include the RDA initialisation changes. With those changes, I think being able to re-run the analysis is also likely to make this wrapper redundant.

llvm/include/llvm/CodeGen/MachinePostRAUpdater.h

+#include "llvm/ADT/SmallPtrSet.h"
+#include "llvm/CodeGen/MachineInstr.h"
+#include "llvm/CodeGen/ReachingDefAnalysis.h"
+
+namespace llvm {
+
+class MachineBasicBlock;
+class MachinePostRAUpdater;
+class ReachingDefAnalysis;
+
+class ModifiedBlock {
+  friend class MachinePostRAUpdater;
+
+public:
+  ModifiedBlock() { }
+  ModifiedBlock(MachineInstr *MI) { add(MI); }
+
+private:
+  std::map<int, SmallPtrSet<MachineInstr*, 2>> NewDefs;
+  std::map<int, SmallPtrSet<MachineInstr*, 2>> NewUses;
+  SmallPtrSet<MachineInstr*, 2> DeadInsts;
+  SmallPtrSet<MachineInstr*, 2> NewInsts;
+
+  void remove(MachineInstr *MI);
+  void add(MachineInstr *MI);
+  void eraseDead();
+  bool hasNewDef(int PhysReg) const { return NewDefs.count(PhysReg); }
+  bool hasNewUse(int PhysReg) const { return NewUses.count(PhysReg); }
+  bool isDead(MachineInstr *MI) const { return DeadInsts.count(MI); }
+  const SmallPtrSetImpl<MachineInstr*> &getNewDefs(int PhysReg) const {
+    return NewDefs.at(PhysReg);
+  }
+  const SmallPtrSetImpl<MachineInstr*> &getNewUses(int PhysReg) const {
+    return NewUses.at(PhysReg);
+  }
+};
+
+class MachinePostRAUpdater {
+  ReachingDefAnalysis &RDA;
+  std::map<MachineBasicBlock*, ModifiedBlock> ModifiedBlocks;
+  SmallPtrSet<MachineInstr*, 4> AllNewInsts;
+
+public:
+  MachinePostRAUpdater() = delete;
+  MachinePostRAUpdater(ReachingDefAnalysis &RDA) : RDA(RDA) { }
+  /// Return whether From can be moved forwards to just before To.
+  bool IsSafeToMoveForwards(MachineInstr *From, MachineInstr *To);
+  /// Return whether From can be moved backwards to just after To.
+  bool IsSafeToMoveBackwards(MachineInstr *From, MachineInstr *To);
+  /// Return whether removing this instruction will have no effect on the
+  /// program, returning the redundant use-def chain.
+  bool IsSafeToRemove(MachineInstr *MI,
+                      SmallPtrSetImpl<MachineInstr*> &ToRemove) const;
+  /// Return whether removing this instruction will have no effect on the
+  /// program, ignoring the possible effects on some instructions, returning
+  /// the redundant use-def chain.
+  bool IsSafeToRemove(MachineInstr *MI,
+                      SmallPtrSetImpl<MachineInstr*> &ToRemove,
+                      SmallPtrSetImpl<MachineInstr*> &Ignore) const;
+  /// Return whether defining the given register at MI would have no effect
+  /// on the program.
+  bool IsSafeToDefRegAt(MachineInstr *MI, int PhysReg) const;
+  /// Return whether defining the given register at MI would have no effect
+  /// on the program, ignoring the possible effects on instructions in Ignore.
+  bool IsSafeToDefRegAt(MachineInstr *MI, int PhysReg,
+                        SmallPtrSetImpl<MachineInstr*> &Ignore) const;
+  /// Return whether the given register is defined within the block.
+  bool isDefinedInBlock(MachineBasicBlock *MBB, int PhysReg) const;
+  /// Return whether the given register is defined within MIs parent block
+  /// but before MI.
+  bool hasLocalDefBefore(MachineInstr *MI, int PhysReg) const;
+  /// Return whether we can prove that A and B use the same value in the
+  /// given register.
+  bool hasSameReachingDef(MachineInstr *A, MachineInstr *B,
+                          int PhysReg) const;
+  /// Return the instruction which defines the register and provides the value
+  /// that lives out of the block, or nullptr for a non-local or newly added
+  /// definition.
+  MachineInstr *getLocalLiveOutDef(MachineBasicBlock *MBB, int PhysReg) const;
+  /// Return whether the value in PhysReg, used by MI, also lives out of the
+  /// basic block.
+  bool isReachingDefLiveOut(MachineInstr *MI, int PhysReg) const;
+  /// Return the local MachineInstr that is the reaching def for MI, or nullptr
+  /// for a non-local definition.
+  MachineInstr *getReachingDef(MachineInstr *MI, int PhysReg);
+
+  void MoveBefore(MachineInstr *From, MachineInstr *To);
+  void MoveAfter(MachineInstr *From, MachineInstr *To);
+  void Replace(MachineInstr *From, MachineInstr *To);
+  void Insert(MachineInstr *MI);
+  void Remove(MachineInstr *MI);
+  void Finalize();
+
+private:
+  template<typename Iterator>
+  bool IsSafeToMove(MachineInstr *From, MachineInstr *To);
+  bool IsSafeToRemove(MachineInstr *MI,
+                      SmallPtrSetImpl<MachineInstr*> &Visited,
+                      SmallPtrSetImpl<MachineInstr*> &ToRemove,
+                      SmallPtrSetImpl<MachineInstr*> &Ignore) const;
+  /// Return whether an instruction has been created and inserted before
+  /// MI which defines the given register.
+  bool hasNewDefBefore(MachineInstr *MI, int PhysReg) const;
+
+};
+
+} // end namespace llvm
+

llvm/include/llvm/CodeGen/ReachingDefAnalysis.h

 
 class MachineBasicBlock;
 class MachineInstr;
 
 /// This class provides the reaching def analysis.
 class ReachingDefAnalysis : public MachineFunctionPass {
 private:
   MachineFunction *MF;
+  LoopTraversal::TraversalOrder TraversedMBBOrder;
   const TargetRegisterInfo *TRI;
   unsigned NumRegUnits;
   /// Instruction that defined each register, relative to the beginning of the
   /// current basic block.  When a LiveRegsDefInfo is used to represent a
   /// live-out register, this value is relative to the end of the basic block,
   /// so it will be a negative number.
   using LiveRegsDefInfo = std::vector<int>;
   LiveRegsDefInfo LiveRegs;
public:
   bool runOnMachineFunction(MachineFunction &MF) override;
 
   MachineFunctionProperties getRequiredProperties() const override {
     return MachineFunctionProperties().set(
         MachineFunctionProperties::Property::NoVRegs).set(
           MachineFunctionProperties::Property::TracksLiveness);
   }
 
+  /// Re-run the analysis.
+  void reset();
+
+  /// Initialize data structures.
+  void init();
+
+  /// Traverse the machine function, mapping definitions.
+  void traverse();
+
   /// Provides the instruction id of the closest reaching def instruction of
   /// PhysReg that reaches MI, relative to the begining of MI's basic block.
   int getReachingDef(MachineInstr *MI, int PhysReg) const;
 
   /// Provides the instruction of the closest reaching def instruction of
   /// PhysReg that reaches MI, relative to the begining of MI's basic block.
   MachineInstr *getReachingMIDef(MachineInstr *MI, int PhysReg) const;
 
public:
   /// nullptr for a non-live out or non-local def.
   MachineInstr *getLocalLiveOutMIDef(MachineBasicBlock *MBB,
                                      int PhysReg) const;
 
   /// Return whether the given register is used after MI, whether it's a local
   /// use or a live out.
   bool isRegUsedAfter(MachineInstr *MI, int PhysReg) const;
 
+  /// Return whether the given register is defined after MI.
+  bool isRegDefinedAfter(MachineInstr *MI, int PhysReg) const;
+
   /// Provides the clearance - the number of instructions since the closest
   /// reaching def instuction of PhysReg that reaches MI.
   int getClearance(MachineInstr *MI, MCPhysReg PhysReg) const;
 
   /// Provides the uses, in the same block as MI, of register that MI defines.
   /// This does not consider live-outs.
   void getReachingLocalUses(MachineInstr *MI, int PhysReg,
                             SmallPtrSetImpl<MachineInstr*> &Uses) const;

llvm/lib/CodeGen/CMakeLists.txt

add_llvm_component_library(LLVMCodeGen
   MachineLoopUtils.cpp
   MachineModuleInfo.cpp
   MachineModuleInfoImpls.cpp
   MachineOperand.cpp
   MachineOptimizationRemarkEmitter.cpp
   MachineOutliner.cpp
   MachinePipeliner.cpp
   MachinePostDominators.cpp
+  MachinePostRAUpdater.cpp
   MachineRegionInfo.cpp
   MachineRegisterInfo.cpp
   MachineScheduler.cpp
   MachineSink.cpp
   MachineSizeOpts.cpp
   MachineSSAUpdater.cpp
   MachineTraceMetrics.cpp
   MachineVerifier.cpp

llvm/lib/CodeGen/MachinePostRAUpdater.cpp

+#include "llvm/ADT/SmallSet.h"
+#include "llvm/CodeGen/MachineBasicBlock.h"
+#include "llvm/CodeGen/MachinePostRAUpdater.h"
+#include "llvm/Support/Debug.h"
+
+#define DEBUG_TYPE "post-ra-updater"
+
+using namespace llvm;
+
+void ModifiedBlock::eraseDead() {
+  for (auto *MI : DeadInsts) {
+    LLVM_DEBUG(dbgs() << "PostRA Updater: Erasing " << *MI);
+    MI->eraseFromParent();
+  }
+}
+
+void ModifiedBlock::add(MachineInstr *MI) {
+  LLVM_DEBUG(dbgs() << "PostRA Updater: Adding " << *MI);
+  NewInsts.insert(MI);
+  for (auto &MO : MI->operands()) {
+    if (!MO.isReg())
+      continue;
+    if (MO.isDef())
+      NewDefs[MO.getReg()].insert(MI);
+    if (MO.isUse())
+      NewUses[MO.getReg()].insert(MI);
+  }
+}
+
+void ModifiedBlock::remove(MachineInstr *MI) {
+  LLVM_DEBUG(dbgs() << "PostRA Updater: Removing " << *MI);
+  DeadInsts.insert(MI);
+  if (NewInsts.count(MI))
+    NewInsts.erase(MI);
+}
+
+// Can we safely move 'From' to just before 'To'? To satisfy this, 'From' must
+// not define a register that is used by any instructions, after and including,
+// 'To'. These instructions also must not redefine any of Froms operands.
+template<typename Iterator>
+bool MachinePostRAUpdater::IsSafeToMove(MachineInstr *From, MachineInstr *To) {
+  if (From->getParent() != To->getParent())
+    return false;
+
+  // TODO: Handle modified blocks.
+  if (ModifiedBlocks.count(From->getParent()))
+    return false;
+
+  SmallSet<int, 2> Defs;
+  // First check that From would compute the same value if moved.
+  for (auto &MO : From->operands()) {
+    if (!MO.isReg() || MO.isUndef() || !MO.getReg())
+      continue;
+    if (MO.isDef())
+      Defs.insert(MO.getReg());
+    else if (!RDA.hasSameReachingDef(From, To, MO.getReg()))
+      return false;
+  }
+
+  // Now walk checking that the rest of the instructions will compute the same
+  // value.
+  for (auto I = ++Iterator(From), E = Iterator(To); I != E; ++I) {
+    for (auto &MO : I->operands())
+      if (MO.isReg() && MO.getReg() && MO.isUse() && Defs.count(MO.getReg()))
+        return false;
+  }
+  return true;
+}
+
+bool MachinePostRAUpdater::IsSafeToMoveForwards(MachineInstr *From,
+                                                MachineInstr *To) {
+  return IsSafeToMove<MachineBasicBlock::reverse_iterator>(From, To);
+}
+
+bool MachinePostRAUpdater::IsSafeToMoveBackwards(MachineInstr *From,
+                                                 MachineInstr *To) {
+  return IsSafeToMove<MachineBasicBlock::iterator>(From, To);
+}
+
+bool
+MachinePostRAUpdater::IsSafeToRemove(MachineInstr *MI,
+                            SmallPtrSetImpl<MachineInstr*> &ToRemove) const {
+  SmallPtrSet<MachineInstr*, 1> Ignore;
+  SmallPtrSet<MachineInstr*, 2> Visited;
+  return IsSafeToRemove(MI, Visited, ToRemove, Ignore);
+}
+
+bool
+MachinePostRAUpdater::IsSafeToRemove(MachineInstr *MI,
+                                     SmallPtrSetImpl<MachineInstr*> &ToRemove,
+                                SmallPtrSetImpl<MachineInstr*> &Ignore) const {
+  SmallPtrSet<MachineInstr*, 2> Visited;
+  return IsSafeToRemove(MI, Visited, ToRemove, Ignore);
+}
+
+bool
+MachinePostRAUpdater::IsSafeToRemove(MachineInstr *MI,
+                                     SmallPtrSetImpl<MachineInstr*> &Visited,
+                                     SmallPtrSetImpl<MachineInstr*> &ToRemove,
+                                SmallPtrSetImpl<MachineInstr*> &Ignore) const {
+  LLVM_DEBUG(dbgs() << "PostRA Updater: Is safe to remove " << *MI);
+  if (Visited.count(MI) || Ignore.count(MI))
+    return true;
+  else if (MI->mayLoadOrStore() || MI->hasUnmodeledSideEffects() ||
+           MI->isBranch() || MI->isTerminator() || MI->isReturn()) {
+    // Unless told to ignore the instruction, don't remove anything which has
+    // side effects.
+    LLVM_DEBUG(dbgs() << "PostRA Updater: Has side effects.\n");
+    return false;
+  } else if (AllNewInsts.count(MI)) {
+    LLVM_DEBUG(dbgs() << "PostRA Updater: Instruction not known to RDA.\n");
+    return false;
+  }
+
+  Visited.insert(MI);
+  for (auto &MO : MI->operands()) {
+    if (!MO.isReg() || MO.isUse() || MO.getReg() == 0)
+      continue;
+
+    SmallPtrSet<MachineInstr*, 4> Uses;
+    RDA.getGlobalUses(MI, MO.getReg(), Uses);
+
+    for (auto I : Uses) {
+      if (Ignore.count(I) || ToRemove.count(I))
+        continue;
+      if (!IsSafeToRemove(I, Visited, ToRemove, Ignore)) {
+        LLVM_DEBUG(dbgs() << "PostRA Updater: Unable to remove " << *I);
+        return false;
+      }
+    }
+  }
+  ToRemove.insert(MI);
+  LLVM_DEBUG(dbgs() << "PostRA Updater: Can remove: " << *MI);
+  return true;
+}
+
+void MachinePostRAUpdater::MoveAfter(MachineInstr *From, MachineInstr *To) {
+  assert(From->getParent() == To->getParent() &&
+         "Can't move instructions between blocks");
+  MachineBasicBlock *MBB = From->getParent();
+  From->removeFromParent();
+  MBB->insertAfter(MachineBasicBlock::iterator(To), From);
+  LLVM_DEBUG(dbgs() << "PostRA Updater: Moved " << *From);
+}
+
+void MachinePostRAUpdater::MoveBefore(MachineInstr *From, MachineInstr *To) {
+  assert(From->getParent() == To->getParent() &&
+         "Can't move instructions between blocks");
+  MachineBasicBlock *MBB = From->getParent();
+  From->removeFromParent();
+  MBB->insert(MachineBasicBlock::iterator(To), From);
+  LLVM_DEBUG(dbgs() << "PostRA Updater: Moved " << *From);
+}
+
+bool
+MachinePostRAUpdater::IsSafeToDefRegAt(MachineInstr *MI, int PhysReg) const {
+  SmallPtrSet<MachineInstr*, 1> Ignore;
+  return IsSafeToDefRegAt(MI, PhysReg, Ignore);
+}
+
+bool
+MachinePostRAUpdater::IsSafeToDefRegAt(MachineInstr *MI, int PhysReg,
+                                SmallPtrSetImpl<MachineInstr*> &Ignore) const {
+  MachineBasicBlock *MBB = MI->getParent();
+
+  // TODO: Figure out if the new definition really prevents another def.
+  if (ModifiedBlocks.count(MBB) && ModifiedBlocks.at(MBB).hasNewDef(PhysReg))
+    return false;
+
+  // Check for any uses of the register after MI.
+  if (RDA.isRegUsedAfter(MI, PhysReg)) {
+    if (auto *Def = RDA.getReachingMIDef(MI, PhysReg)) {
+      SmallPtrSet<MachineInstr*, 2> Uses;
+      RDA.getReachingLocalUses(Def, PhysReg, Uses);
+      for (auto *Use : Uses)
+        if (!Ignore.count(Use))
+          return false;
+    } else
+      return false;
+  }
+
+  // Check for any defs after MI.
+  if (RDA.isRegDefinedAfter(MI, PhysReg)) {
+    auto I = MachineBasicBlock::iterator(MI);
+    for (auto E = MBB->end(); I != E; ++I) {
+      if (Ignore.count(&*I))
+        continue;
+      for (auto &MO : I->operands())
+        if (MO.isReg() && MO.isDef() && MO.getReg() == PhysReg)
+          return false;
+    }
+  }
+  LLVM_DEBUG(dbgs() << "PostRA Updater: Safe to define register at: " << *MI);
+  return true;
+}
+
+bool
+MachinePostRAUpdater::hasNewDefBefore(MachineInstr *MI, int PhysReg) const {
+  auto isBefore = [](MachineInstr *A, MachineInstr *B) {
+    return std::distance(MachineBasicBlock::iterator(A),
+                         MachineBasicBlock::iterator(B)) > 0;
+  };
+
+  MachineBasicBlock *MBB = MI->getParent();
+  if (!ModifiedBlocks.count(MBB))
+    return false;
+  if (!ModifiedBlocks.at(MBB).hasNewDef(PhysReg))
+    return false;
+  for (auto *NewDef : ModifiedBlocks.at(MBB).getNewDefs(PhysReg))
+    if (isBefore(NewDef, MI))
+      return true;
+  return false;
+}
+
+bool
+MachinePostRAUpdater::hasLocalDefBefore(MachineInstr *MI, int PhysReg) const {
+  if (hasNewDefBefore(MI, PhysReg))
+    return true;
+  return RDA.getReachingDef(MI, PhysReg) >= 0;
+}
+
+bool MachinePostRAUpdater::isDefinedInBlock(MachineBasicBlock *MBB,
+                                            int PhysReg) const {
+  return hasLocalDefBefore(&MBB->back(), PhysReg) ||
+         !isReachingDefLiveOut(&MBB->back(), PhysReg);
+}
+
+void MachinePostRAUpdater::Finalize() {
+  LLVM_DEBUG(dbgs() << "PostRA Updater: Finalizing changes...\n");
+  for (auto &Pair : ModifiedBlocks)
+    Pair.second.eraseDead();
+  ModifiedBlocks.clear();
+  AllNewInsts.clear();
+  RDA.reset();
+}
+
+void MachinePostRAUpdater::Replace(MachineInstr *Old, MachineInstr *New) {
+  assert(Old->getParent() == New->getParent() &&
+         "Can't replace instructions in different blocks");
+  LLVM_DEBUG(dbgs() << "PostRA Updater:\n- Replacing " << *Old
+             << "- With " << *New);
+  Insert(New);
+  Remove(Old);
+}
+
+void MachinePostRAUpdater::Insert(MachineInstr *MI) {
+  AllNewInsts.insert(MI);
+  MachineBasicBlock *MBB = MI->getParent();
+  if (ModifiedBlocks.count(MBB))
+    ModifiedBlocks.at(MBB).add(MI);
+  else
+    ModifiedBlocks.emplace(MBB, ModifiedBlock(MI));
+}
+
+void MachinePostRAUpdater::Remove(MachineInstr *MI) {
+  MachineBasicBlock *MBB = MI->getParent();
+  if (ModifiedBlocks.count(MBB))
+    ModifiedBlocks.at(MBB).remove(MI);
+  else {
+    ModifiedBlocks.emplace(MBB, ModifiedBlock());
+    ModifiedBlocks.at(MBB).remove(MI);
+  }
+}
+
+bool
+MachinePostRAUpdater::hasSameReachingDef(MachineInstr *A, MachineInstr *B,
+                                         int PhysReg) const {
+  // For now, just bail if we find a new def.
+  // TODO: Check the new definition.
+  if ((ModifiedBlocks.count(A->getParent()) &&
+       ModifiedBlocks.at(A->getParent()).hasNewDef(PhysReg)) ||
+      (ModifiedBlocks.count(B->getParent()) &&
+       ModifiedBlocks.at(B->getParent()).hasNewDef(PhysReg)))
+    return false;
+  return RDA.hasSameReachingDef(A, B, PhysReg);
+}
+
+MachineInstr* MachinePostRAUpdater::getReachingDef(MachineInstr *MI,
+                                                   int PhysReg) {
+  MachineBasicBlock *MBB = MI->getParent();
+  // TODO: Find the most recent new definition and return it.
+  if (hasNewDefBefore(MI, PhysReg))
+    return nullptr;
+
+  MachineInstr *Def = RDA.getReachingMIDef(MI, PhysReg);
+  if (ModifiedBlocks.count(MBB) && ModifiedBlocks.at(MBB).isDead(Def))
+    return nullptr;
+  return Def;
+}
+
+MachineInstr* MachinePostRAUpdater::getLocalLiveOutDef(MachineBasicBlock *MBB,
+                                                       int PhysReg) const {
+  // TODO: Allow a new definition to be returned.
+  if (ModifiedBlocks.count(MBB) && ModifiedBlocks.at(MBB).hasNewDef(PhysReg))
+    return nullptr;
+  return RDA.getLocalLiveOutMIDef(MBB, PhysReg);
+}
+
+bool MachinePostRAUpdater::isReachingDefLiveOut(MachineInstr *MI,
+                                                int PhysReg) const {
+  MachineBasicBlock *MBB = MI->getParent();
+  // TODO: Allow a new definition to be returned.
+  if (ModifiedBlocks.count(MBB) &&
+      ModifiedBlocks.at(MBB).hasNewDef(PhysReg))
+    return false;
+  return RDA.isReachingDefLiveOut(MI, PhysReg);
+}

llvm/lib/CodeGen/ReachingDefAnalysis.cpp

void ReachingDefAnalysis::processBasicBlock(
   }
   leaveBasicBlock(TraversedMBB);
 }
 
 bool ReachingDefAnalysis::runOnMachineFunction(MachineFunction &mf) {
   MF = &mf;
   TRI = MF->getSubtarget().getRegisterInfo();
 
-  LiveRegs.clear();
-  NumRegUnits = TRI->getNumRegUnits();
+  LLVM_DEBUG(dbgs() << "********** REACHING DEFINITION ANALYSIS **********\n");
 
-  MBBReachingDefs.resize(mf.getNumBlockIDs());
+  init();
+  traverse();
+  return false;
+}
 
-  LLVM_DEBUG(dbgs() << "********** REACHING DEFINITION ANALYSIS **********\n");
+void ReachingDefAnalysis::releaseMemory() {
+  // Clear the internal vectors.
+  MBBOutRegsInfos.clear();
+  MBBReachingDefs.clear();
+  InstIds.clear();
+}
+
+void ReachingDefAnalysis::reset() {
+  InstIds.clear();
+  init();
+  traverse();
+}
 
+void ReachingDefAnalysis::init() {
+  LiveRegs.clear();
+  NumRegUnits = TRI->getNumRegUnits();
+  MBBReachingDefs.resize(MF->getNumBlockIDs());
   // Initialize the MBBOutRegsInfos
-  MBBOutRegsInfos.resize(mf.getNumBlockIDs());
+  MBBOutRegsInfos.resize(MF->getNumBlockIDs());
 
-  // Traverse the basic blocks.
+  // Get the traversal order.
   LoopTraversal Traversal;
-  LoopTraversal::TraversalOrder TraversedMBBOrder = Traversal.traverse(mf);
-  for (LoopTraversal::TraversedMBBInfo TraversedMBB : TraversedMBBOrder) {
-    processBasicBlock(TraversedMBB);
-  }
+  TraversedMBBOrder = Traversal.traverse(*MF);
 
+void ReachingDefAnalysis::traverse() {
+  // Traverse the basic blocks.
+  for (LoopTraversal::TraversedMBBInfo TraversedMBB : TraversedMBBOrder)
+    processBasicBlock(TraversedMBB);
+
   // Sorting all reaching defs found for a ceartin reg unit in a given BB.
   for (MBBDefsInfo &MBBDefs : MBBReachingDefs) {
     for (MBBRegUnitDefs &RegUnitDefs : MBBDefs)
       llvm::sort(RegUnitDefs);
   }
-
-  return false;
-}
-
-void ReachingDefAnalysis::releaseMemory() {
-  // Clear the internal vectors.
-  MBBOutRegsInfos.clear();
-  MBBReachingDefs.clear();
-  InstIds.clear();
 }
 
 int ReachingDefAnalysis::getReachingDef(MachineInstr *MI, int PhysReg) const {
   assert(InstIds.count(MI) && "Unexpected machine instuction.");
   int InstId = InstIds.lookup(MI);
   int DefRes = ReachingDefDefaultVal;
   unsigned MBBNumber = MI->getParent()->getNumber();
   assert(MBBNumber < MBBReachingDefs.size() &&
bool ReachingDefAnalysis::isRegUsedAfter(MachineInstr *MI, int PhysReg) const {
   for (auto Last = MBB->rbegin(), End = MBB->rend(); Last != End; ++Last) {
     LiveRegs.stepBackward(*Last);
     if (LiveRegs.contains(PhysReg))
       return InstIds.lookup(&*Last) > InstIds.lookup(MI);
   }
   return false;
 }
 
+bool ReachingDefAnalysis::isRegDefinedAfter(MachineInstr *MI,
+                                            int PhysReg) const {
+  MachineBasicBlock *MBB = MI->getParent();
+  if (getReachingDef(MI, PhysReg) != getReachingDef(&MBB->back(), PhysReg))
+    return true;
+
+  if (auto *Def = getLocalLiveOutMIDef(MBB, PhysReg))
+    return Def == getReachingMIDef(MI, PhysReg);
+
+  return false;
+}
+
 bool
 ReachingDefAnalysis::isReachingDefLiveOut(MachineInstr *MI, int PhysReg) const {
+
   MachineBasicBlock *MBB = MI->getParent();
   LivePhysRegs LiveRegs(*TRI);
   LiveRegs.addLiveOuts(*MBB);
   if (!LiveRegs.contains(PhysReg))
     return false;
 
   MachineInstr *Last = &MBB->back();
   int Def = getReachingDef(MI, PhysReg);

llvm/lib/Target/ARM/ARMLowOverheadLoops.cpp

 #include "ARMSubtarget.h"
 #include "Thumb2InstrInfo.h"
 #include "llvm/ADT/SetOperations.h"
 #include "llvm/ADT/SmallSet.h"
 #include "llvm/CodeGen/LivePhysRegs.h"
 #include "llvm/CodeGen/MachineFunctionPass.h"
 #include "llvm/CodeGen/MachineLoopInfo.h"
 #include "llvm/CodeGen/MachineLoopUtils.h"
+#include "llvm/CodeGen/MachinePostRAUpdater.h"
 #include "llvm/CodeGen/MachineRegisterInfo.h"
 #include "llvm/CodeGen/Passes.h"
-#include "llvm/CodeGen/ReachingDefAnalysis.h"
 #include "llvm/MC/MCInstrDesc.h"
 
 using namespace llvm;
 
 #define DEBUG_TYPE "arm-low-overhead-loops"
 #define ARM_LOW_OVERHEAD_LOOPS_NAME "ARM Low Overhead Loops pass"
 
 namespace {
public:
     MachineInstr *getVPST() const { return VPST->MI; }
     PredicatedMI *getDivergent() const { return Divergent; }
   };
 
   struct LowOverheadLoop {
 
     MachineLoop *ML = nullptr;
     MachineLoopInfo *MLI = nullptr;
-    ReachingDefAnalysis *RDA = nullptr;
+    MachinePostRAUpdater *Updater = nullptr;
     MachineFunction *MF = nullptr;
     MachineInstr *InsertPt = nullptr;
     MachineInstr *Start = nullptr;
     MachineInstr *Dec = nullptr;
     MachineInstr *End = nullptr;
     MachineInstr *VCTP = nullptr;
     VPTBlock *CurrentBlock = nullptr;
     SetVector<MachineInstr*> CurrentPredicate;
     SmallVector<VPTBlock, 4> VPTBlocks;
-    SmallPtrSet<MachineInstr*, 4> ToRemove;
     bool Revert = false;
     bool CannotTailPredicate = false;
 
     LowOverheadLoop(MachineLoop *ML, MachineLoopInfo *MLI,
-                    ReachingDefAnalysis *RDA) : ML(ML), MLI(MLI), RDA(RDA) {
+                    MachinePostRAUpdater *Updater) :
+      ML(ML), MLI(MLI), Updater(Updater) {
       MF = ML->getHeader()->getParent();
     }
 
     // If this is an MVE instruction, check that we know how to use tail
     // predication with it. Record VPT blocks and return whether the
     // instruction is valid for tail predication.
     bool ValidateMVEInst(MachineInstr *MI);
 
void dump() const {
       else if (!(Start && Dec && End))
         dbgs() << "ARM Loops: Failed to find all loop components.\n";
     }
   };
 
   class ARMLowOverheadLoops : public MachineFunctionPass {
     MachineFunction           *MF = nullptr;
     MachineLoopInfo           *MLI = nullptr;
-    ReachingDefAnalysis       *RDA = nullptr;
     const ARMBaseInstrInfo    *TII = nullptr;
     MachineRegisterInfo       *MRI = nullptr;
     const TargetRegisterInfo  *TRI = nullptr;
     std::unique_ptr<ARMBasicBlockUtils> BBUtils = nullptr;
+    std::unique_ptr<MachinePostRAUpdater> Updater;
 
   public:
     static char ID;
 
     ARMLowOverheadLoops() : MachineFunctionPass(ID) { }
 
     void getAnalysisUsage(AnalysisUsage &AU) const override {
       AU.setPreservesCFG();
namespace {
 
   private:
     bool ProcessLoop(MachineLoop *ML);
 
     bool RevertNonLoops();
 
     void RevertWhile(MachineInstr *MI) const;
 
-    bool RevertLoopDec(MachineInstr *MI, bool AllowFlags = false) const;
+    bool RevertLoopDec(MachineInstr *MI) const;
 
     void RevertLoopEnd(MachineInstr *MI, bool SkipCmp = false) const;
 
     void ConvertVPTBlocks(LowOverheadLoop &LoLoop);
 
     MachineInstr *ExpandLoopStart(LowOverheadLoop &LoLoop);
 
     void Expand(LowOverheadLoop &LoLoop);
return MI->getOpcode() == ARM::tMOVr &&
            MI->getOperand(2).getImm() == ARMCC::AL;
    };
 
   MachineBasicBlock *MBB = Start->getParent();
 
   // Find an insertion point:
   // - Is there a (mov lr, Count) before Start? If so, and nothing else writes
   //   to Count before Start, we can insert at that mov.
-  if (auto *LRDef = RDA->getReachingMIDef(Start, ARM::LR))
-    if (IsMoveLR(LRDef) && RDA->hasSameReachingDef(Start, LRDef, CountReg))
+  if (auto *LRDef = Updater->getReachingDef(Start, ARM::LR))
+    if (IsMoveLR(LRDef) && Updater->hasSameReachingDef(Start, LRDef, CountReg))
       return LRDef;
 
   // - Is there a (mov lr, Count) after Start? If so, and nothing else writes
   //   to Count after Start, we can insert at that mov.
-  if (auto *LRDef = RDA->getLocalLiveOutMIDef(MBB, ARM::LR))
-    if (IsMoveLR(LRDef) && RDA->hasSameReachingDef(Start, LRDef, CountReg))
+  if (auto *LRDef = Updater->getLocalLiveOutDef(MBB, ARM::LR))
+    if (IsMoveLR(LRDef) && Updater->hasSameReachingDef(Start, LRDef, CountReg))
       return LRDef;
 
   // We've found no suitable LR def and Start doesn't use LR directly. Can we
   // just define LR anyway?
-  if (!RDA->isRegUsedAfter(Start, ARM::LR))
-    return Start;
-
-  return nullptr;
-}
-
-// Can we safely move 'From' to just before 'To'? To satisfy this, 'From' must
-// not define a register that is used by any instructions, after and including,
-// 'To'. These instructions also must not redefine any of Froms operands.
-template<typename Iterator>
-static bool IsSafeToMove(MachineInstr *From, MachineInstr *To, ReachingDefAnalysis *RDA) {
-  SmallSet<int, 2> Defs;
-  // First check that From would compute the same value if moved.
-  for (auto &MO : From->operands()) {
-    if (!MO.isReg() || MO.isUndef() || !MO.getReg())
-      continue;
-    if (MO.isDef())
-      Defs.insert(MO.getReg());
-    else if (!RDA->hasSameReachingDef(From, To, MO.getReg()))
-      return false;
-  }
-
-  // Now walk checking that the rest of the instructions will compute the same
-  // value.
-  for (auto I = ++Iterator(From), E = Iterator(To); I != E; ++I) {
-    for (auto &MO : I->operands())
-      if (MO.isReg() && MO.getReg() && MO.isUse() && Defs.count(MO.getReg()))
-        return false;
-  }
-  return true;
-}
-
-static bool IsSafeToRemove(MachineInstr *MI, ReachingDefAnalysis *RDA,
-                           SmallPtrSetImpl<MachineInstr*> &Visited,
-                           SmallPtrSetImpl<MachineInstr*> &ToRemove,
-                           SmallPtrSetImpl<MachineInstr*> &Ignore) {
-  if (Visited.count(MI) || Ignore.count(MI))
-    return true;
-  else if (MI->mayLoadOrStore() || MI->hasUnmodeledSideEffects() ||
-           MI->isBranch() || MI->isTerminator() || MI->isReturn()) {
-    // Unless told to ignore the instruction, don't remove anything which has
-    // side effects.
-    LLVM_DEBUG(dbgs() << "ARM Loops: Has side effects: " << *MI);
-    return false;
-  }
-
-  Visited.insert(MI);
-  for (auto &MO : MI->operands()) {
-    if (!MO.isReg() || MO.isUse() || MO.getReg() == 0)
-      continue;
-
-    SmallPtrSet<MachineInstr*, 4> Uses;
-    RDA->getGlobalUses(MI, MO.getReg(), Uses);
-
-    for (auto I : Uses) {
-      if (Ignore.count(I) || ToRemove.count(I))
-        continue;
-      if (!IsSafeToRemove(I, RDA, Visited, ToRemove, Ignore)) {
-        LLVM_DEBUG(dbgs() << "ARM Loops: Unable to remove " << *I);
-        return false;
-      }
-    }
-  }
-  ToRemove.insert(MI);
-  LLVM_DEBUG(dbgs() << "ARM Loops: Can remove: " << *MI);
-  return true;
+  return Updater->IsSafeToDefRegAt(Start, ARM::LR) ? Start : nullptr;
 }
 
 bool LowOverheadLoop::ValidateTailPredicate(MachineInstr *StartInsertPt) {
   assert(VCTP && "VCTP instruction expected but is not set");
   // All predication within the loop should be based on vctp. If the block
   // isn't predicated on entry, check whether the vctp is within the block
   // and that all other instructions are then predicated on it.
   for (auto &Block : VPTBlocks) {
bool LowOverheadLoop::ValidateTailPredicate(MachineInstr *StartInsertPt) {
   // of the iteration count, to the loop start instruction. The number of
   // elements is provided to the vctp instruction, so we need to check that
   // we can use this register at InsertPt.
   Register NumElements = VCTP->getOperand(1).getReg();
 
   // If the register is defined within loop, then we can't perform TP.
   // TODO: Check whether this is just a mov of a register that would be
   // available.
-  if (RDA->getReachingDef(VCTP, NumElements) >= 0) {
+  if (Updater->hasLocalDefBefore(VCTP, NumElements)) {
     LLVM_DEBUG(dbgs() << "ARM Loops: VCTP operand is defined in the loop.\n");
     return false;
   }
 
   // The element count register maybe defined after InsertPt, in which case we
   // need to try to move either InsertPt or the def so that the [w|d]lstp can
   // use the value.
   MachineBasicBlock *InsertBB = StartInsertPt->getParent();
-  if (!RDA->isReachingDefLiveOut(StartInsertPt, NumElements)) {
-    if (auto *ElemDef = RDA->getLocalLiveOutMIDef(InsertBB, NumElements)) {
-      if (IsSafeToMove<MachineBasicBlock::reverse_iterator>(
-          ElemDef, StartInsertPt, RDA)) {
-        ElemDef->removeFromParent();
-        InsertBB->insert(MachineBasicBlock::iterator(StartInsertPt), ElemDef);
-        LLVM_DEBUG(dbgs() << "ARM Loops: Moved element count def: "
-                   << *ElemDef);
-      } else if (IsSafeToMove<MachineBasicBlock::iterator>(
-          StartInsertPt, ElemDef, RDA)) {
-        StartInsertPt->removeFromParent();
-        InsertBB->insertAfter(MachineBasicBlock::iterator(ElemDef),
-                              StartInsertPt);
-        LLVM_DEBUG(dbgs() << "ARM Loops: Moved start past: " << *ElemDef);
-      } else {
+  if (!Updater->isReachingDefLiveOut(StartInsertPt, NumElements)) {
+    if (auto *ElemDef = Updater->getLocalLiveOutDef(InsertBB, NumElements)) {
+      if (Updater->IsSafeToMoveForwards(ElemDef, StartInsertPt))
+        Updater->MoveBefore(ElemDef, StartInsertPt);
+      else if (Updater->IsSafeToMoveBackwards(StartInsertPt, ElemDef))
+        Updater->MoveAfter(StartInsertPt, ElemDef);
+      else {
         LLVM_DEBUG(dbgs() << "ARM Loops: Unable to move element count to loop "
                    << "start instruction.\n");
         return false;
       }
     }
   }
 
   // Especially in the case of while loops, InsertBB may not be the
   // preheader, so we need to check that the register isn't redefined
   // before entering the loop.
   auto CannotProvideElements = [this](MachineBasicBlock *MBB,
                                       Register NumElements) {
     // NumElements is redefined in this block.
-    if (RDA->getReachingDef(&MBB->back(), NumElements) >= 0)
+    if (Updater->hasLocalDefBefore(&MBB->back(), NumElements))
       return true;
 
     // Don't continue searching up through multiple predecessors.
     if (MBB->pred_size() > 1)
       return true;
 
     return false;
   };
auto IsValidSub = [](MachineInstr *MI, unsigned ExpectedVecWidth) {
     case ARM::t2SUBri12:
       ImmOpIdx = 2;
       break;
     }
     return MI->getOperand(ImmOpIdx).getImm() == ExpectedVecWidth;
   };
 
   MBB = VCTP->getParent();
-  if (MachineInstr *Def = RDA->getReachingMIDef(&MBB->back(), NumElements)) {
-    SmallPtrSet<MachineInstr*, 2> Visited;
+  if (MachineInstr *Def = Updater->getReachingDef(&MBB->back(), NumElements)) {
     SmallPtrSet<MachineInstr*, 2> ElementChain;
     SmallPtrSet<MachineInstr*, 2> Ignore = { VCTP };
     unsigned ExpectedVectorWidth = getTailPredVectorWidth(VCTP->getOpcode());
 
-    if (IsSafeToRemove(Def, RDA, Visited, ElementChain, Ignore)) {
+    if (Updater->IsSafeToRemove(Def, ElementChain, Ignore)) {
       bool FoundSub = false;
 
       for (auto *MI : ElementChain) {
         if (isMovRegOpcode(MI->getOpcode()))
           continue;
 
         if (isSubImmOpcode(MI->getOpcode())) {
           if (FoundSub || !IsValidSub(MI, ExpectedVectorWidth))
             return false;
           FoundSub = true;
         } else
           return false;
       }
 
       LLVM_DEBUG(dbgs() << "ARM Loops: Will remove element count chain:\n";
                  for (auto *MI : ElementChain)
                    dbgs() << " - " << *MI);
-      ToRemove.insert(ElementChain.begin(), ElementChain.end());
+      for (auto *MI : ElementChain)
+        Updater->Remove(MI);
     }
   }
   return true;
 }
 
 void LowOverheadLoop::CheckLegality(ARMBasicBlockUtils *BBUtils) {
   if (Revert)
     return;
bool ARMLowOverheadLoops::runOnMachineFunction(MachineFunction &mf) {
   const ARMSubtarget &ST = static_cast<const ARMSubtarget&>(mf.getSubtarget());
   if (!ST.hasLOB())
     return false;
 
   MF = &mf;
   LLVM_DEBUG(dbgs() << "ARM Loops on " << MF->getName() << " ------------- \n");
 
   MLI = &getAnalysis<MachineLoopInfo>();
-  RDA = &getAnalysis<ReachingDefAnalysis>();
+  Updater = std::make_unique<MachinePostRAUpdater>(
+        getAnalysis<ReachingDefAnalysis>());
   MF->getProperties().set(MachineFunctionProperties::Property::TracksLiveness);
   MRI = &MF->getRegInfo();
   TII = static_cast<const ARMBaseInstrInfo*>(ST.getInstrInfo());
   TRI = ST.getRegisterInfo();
   BBUtils = std::unique_ptr<ARMBasicBlockUtils>(new ARMBasicBlockUtils(*MF));
   BBUtils->computeAllBlockSizes();
   BBUtils->adjustBBOffsetsAfter(&MF->front());
 
for (auto &MI : *MBB) {
       if (isLoopStart(MI))
         return &MI;
     }
     if (MBB->pred_size() == 1)
       return SearchForStart(*MBB->pred_begin());
     return nullptr;
   };
 
-  LowOverheadLoop LoLoop(ML, MLI, RDA);
+  LowOverheadLoop LoLoop(ML, MLI, Updater.get());
   // Search the preheader for the start intrinsic.
   // FIXME: I don't see why we shouldn't be supporting multiple predecessors
   // with potentially multiple set.loop.iterations, so we need to enable this.
   if (auto *Preheader = ML->getLoopPreheader())
     LoLoop.Start = SearchForStart(Preheader);
   else if (auto *Preheader = MLI->findLoopPreheader(ML, true))
     LoLoop.Start = SearchForStart(Preheader);
   else
bool ARMLowOverheadLoops::ProcessLoop(MachineLoop *ML) {
   }
 
   LLVM_DEBUG(LoLoop.dump());
   if (!LoLoop.FoundAllComponents()) {
     LLVM_DEBUG(dbgs() << "ARM Loops: Didn't find loop start, update, end\n");
     return false;
   }
 
-  SmallPtrSet<MachineInstr*, 2> Visited;
   SmallPtrSet<MachineInstr*, 2> Ignore = { LoLoop.End };
   SmallPtrSet<MachineInstr*, 4> Remove;
-  if (!IsSafeToRemove(LoLoop.Dec, RDA, Visited, Remove, Ignore)) {
+  if (!Updater->IsSafeToRemove(LoLoop.Dec, Remove, Ignore)) {
     LLVM_DEBUG(dbgs() << "ARM Loops: Unable to remove loop count chain.\n");
     LoLoop.Revert = true;
   } else {
     LLVM_DEBUG(dbgs() << "ARM Loops: Will need to remove:\n";
                for (auto *I : Remove)
                  dbgs() << " - " << *I);
-    LoLoop.ToRemove.insert(Remove.begin(), Remove.end());
+    for (auto *MI : Remove)
+      Updater->Remove(MI);
   }
 
   LoLoop.CheckLegality(BBUtils.get());
   Expand(LoLoop);
   return true;
 }
 
 // WhileLoopStart holds the exit block, so produce a cmp lr, 0 and then a
void ARMLowOverheadLoops::RevertWhile(MachineInstr *MI) const {
 
   MIB = BuildMI(*MBB, MI, MI->getDebugLoc(), TII->get(BrOpc));
   MIB.add(MI->getOperand(1));   // branch target
   MIB.addImm(ARMCC::EQ);        // condition code
   MIB.addReg(ARM::CPSR);
   MI->eraseFromParent();
 }
 
-bool ARMLowOverheadLoops::RevertLoopDec(MachineInstr *MI,
-                                        bool SetFlags) const {
+bool ARMLowOverheadLoops::RevertLoopDec(MachineInstr *MI) const {
   LLVM_DEBUG(dbgs() << "ARM Loops: Reverting to sub: " << *MI);
   MachineBasicBlock *MBB = MI->getParent();
+  MachineInstr *Last = &MBB->back();
+
+  SmallPtrSet<MachineInstr*, 1> Ignore;
+  if (Last->getOpcode() == ARM::t2LoopEnd)
+    Ignore.insert(Last);
 
   // If nothing defines CPSR between LoopDec and LoopEnd, use a t2SUBS.
-  if (SetFlags &&
-      (RDA->isRegUsedAfter(MI, ARM::CPSR) ||
-       !RDA->hasSameReachingDef(MI, &MBB->back(), ARM::CPSR)))
-      SetFlags = false;
+  bool SetFlags = Updater->IsSafeToDefRegAt(MI, ARM::CPSR, Ignore);
 
   MachineInstrBuilder MIB = BuildMI(*MBB, MI, MI->getDebugLoc(),
                                     TII->get(ARM::t2SUBri));
   MIB.addDef(ARM::LR);
   MIB.add(MI->getOperand(1));
   MIB.add(MI->getOperand(2));
   MIB.addImm(ARMCC::AL);
   MIB.addReg(0);

 
 MachineInstr* ARMLowOverheadLoops::ExpandLoopStart(LowOverheadLoop &LoLoop) {
   LLVM_DEBUG(dbgs() << "ARM Loops: Expanding LoopStart.\n");
   // When using tail-predication, try to delete the dead code that was used to
   // calculate the number of loop iterations.
   if (LoLoop.IsTailPredicationLegal()) {
     SmallVector<MachineInstr*, 4> Killed;
     SmallVector<MachineInstr*, 4> Dead;
-    if (auto *Def = RDA->getReachingMIDef(LoLoop.Start,
+    if (auto *Def = Updater->getReachingDef(LoLoop.Start,
                                       LoLoop.Start->getOperand(0).getReg())) {
-      SmallPtrSet<MachineInstr*, 4> Visited;
       SmallPtrSet<MachineInstr*, 4> Remove;
       SmallPtrSet<MachineInstr*, 4> Ignore = { LoLoop.Start, LoLoop.Dec,
                                                LoLoop.End, LoLoop.VCTP,
                                                LoLoop.InsertPt };
       SmallVector<MachineInstr*, 4> Chain = { Def };
       while (!Chain.empty()) {
         MachineInstr *MI = Chain.back();
         Chain.pop_back();
-        if (IsSafeToRemove(MI, RDA, Visited, Remove, Ignore)) {
+        if (Updater->IsSafeToRemove(MI, Remove, Ignore)) {
           for (auto &MO : MI->operands()) {
             if (!MO.isReg() || !MO.isUse() || MO.getReg() == 0)
               continue;
-            if (auto *Op = RDA->getReachingMIDef(MI, MO.getReg()))
+            if (auto *Op = Updater->getReachingDef(MI, MO.getReg()))
               Chain.push_back(Op);
           }
           Ignore.insert(MI);
         }
       }
-      LoLoop.ToRemove.insert(Remove.begin(), Remove.end());
+      for (auto *MI : Remove)
+        Updater->Remove(MI);
     }
   }
 
   MachineInstr *InsertPt = LoLoop.InsertPt;
   MachineInstr *Start = LoLoop.Start;
   MachineBasicBlock *MBB = InsertPt->getParent();
   bool IsDo = Start->getOpcode() == ARM::t2DoLoopStart;
   unsigned Opc = LoLoop.getStartOpcode();
   MachineOperand &Count = LoLoop.getCount();
 
   MachineInstrBuilder MIB =
     BuildMI(*MBB, InsertPt, InsertPt->getDebugLoc(), TII->get(Opc));
 
   MIB.addDef(ARM::LR);
   MIB.add(Count);
   if (!IsDo)
     MIB.add(Start->getOperand(1));
 
   // If we're inserting at a mov lr, then remove it as it's redundant.
   if (InsertPt != Start)
-    LoLoop.ToRemove.insert(InsertPt);
-  LoLoop.ToRemove.insert(Start);
+    Updater->Remove(InsertPt);
+  Updater->Remove(Start);
   LLVM_DEBUG(dbgs() << "ARM Loops: Inserted start: " << *MIB);
   return &*MIB;
 }
 
 void ARMLowOverheadLoops::ConvertVPTBlocks(LowOverheadLoop &LoLoop) {
   auto RemovePredicate = [](MachineInstr *MI) {
     LLVM_DEBUG(dbgs() << "ARM Loops: Removing predicate from: " << *MI);
     if (int PIdx = llvm::findFirstVPTPredOperandIdx(*MI)) {
if (Block.HasNonUniformPredicate()) {
           ++I;
         }
         MachineInstrBuilder MIB = BuildMI(*InsertAt->getParent(), InsertAt,
                                           InsertAt->getDebugLoc(),
                                           TII->get(ARM::MVE_VPST));
         MIB.addImm(getARMVPTBlockMask(Size));
         LLVM_DEBUG(dbgs() << "ARM Loops: Removing VPST: " << *Block.getVPST());
         LLVM_DEBUG(dbgs() << "ARM Loops: Created VPST: " << *MIB);
-        LoLoop.ToRemove.insert(Block.getVPST());
+        Updater->Remove(Block.getVPST());
       }
     } else if (Block.IsOnlyPredicatedOn(LoLoop.VCTP)) {
       // A vpt block which is only predicated upon vctp and has no internal vpr
       // defs:
       // - Remove vpst.
       // - Unpredicate the remaining instructions.
       LLVM_DEBUG(dbgs() << "ARM Loops: Removing VPST: " << *Block.getVPST());
-      LoLoop.ToRemove.insert(Block.getVPST());
+      Updater->Remove(Block.getVPST());
       for (auto &PredMI : Insts)
         RemovePredicate(PredMI.MI);
     }
   }
   LLVM_DEBUG(dbgs() << "ARM Loops: Removing VCTP: " << *LoLoop.VCTP);
-  LoLoop.ToRemove.insert(LoLoop.VCTP);
+  Updater->Remove(LoLoop.VCTP);
 }
 
 void ARMLowOverheadLoops::Expand(LowOverheadLoop &LoLoop) {
 
   // Combine the LoopDec and LoopEnd instructions into LE(TP).
   auto ExpandLoopEnd = [this](LowOverheadLoop &LoLoop) {
     MachineInstr *End = LoLoop.End;
     MachineBasicBlock *MBB = End->getParent();
     unsigned Opc = LoLoop.IsTailPredicationLegal() ?
       ARM::MVE_LETP : ARM::t2LEUpdate;
     MachineInstrBuilder MIB = BuildMI(*MBB, End, End->getDebugLoc(),
                                       TII->get(Opc));
     MIB.addDef(ARM::LR);
     MIB.add(End->getOperand(0));
     MIB.add(End->getOperand(1));
     LLVM_DEBUG(dbgs() << "ARM Loops: Inserted LE: " << *MIB);
-    End->eraseFromParent();
+    Updater->Remove(End);
     return &*MIB;
   };
 
   // TODO: We should be able to automatically remove these branches before we
   // get here - probably by teaching analyzeBranch about the pseudo
   // instructions.
   // If there is an unconditional branch, after I, that just branches to the
   // next block, remove it.
-  auto RemoveDeadBranch = [](MachineInstr *I) {
+  auto RemoveDeadBranch = [this](MachineInstr *I) {
     MachineBasicBlock *BB = I->getParent();
     MachineInstr *Terminator = &BB->instr_back();
     if (Terminator->isUnconditionalBranch() && I != Terminator) {
       MachineBasicBlock *Succ = Terminator->getOperand(0).getMBB();
       if (BB->isLayoutSuccessor(Succ)) {
         LLVM_DEBUG(dbgs() << "ARM Loops: Removing branch: " << *Terminator);
-        Terminator->eraseFromParent();
+        Updater->Remove(Terminator);
       }
     }
   };
 
   if (LoLoop.Revert) {
     if (LoLoop.Start->getOpcode() == ARM::t2WhileLoopStart)
       RevertWhile(LoLoop.Start);
     else
-      LoLoop.Start->eraseFromParent();
-    bool FlagsAlreadySet = RevertLoopDec(LoLoop.Dec, true);
+      Updater->Remove(LoLoop.Start);
+    bool FlagsAlreadySet = RevertLoopDec(LoLoop.Dec);
     RevertLoopEnd(LoLoop.End, FlagsAlreadySet);
   } else {
     LoLoop.Start = ExpandLoopStart(LoLoop);
     RemoveDeadBranch(LoLoop.Start);
     LoLoop.End = ExpandLoopEnd(LoLoop);
     RemoveDeadBranch(LoLoop.End);
     if (LoLoop.IsTailPredicationLegal())
       ConvertVPTBlocks(LoLoop);
-    for (auto *I : LoLoop.ToRemove) {
-      LLVM_DEBUG(dbgs() << "ARM Loops: Erasing " << *I);
-      I->eraseFromParent();
-    }
+    Updater->Finalize();
   }
 
   PostOrderLoopTraversal DFS(*LoLoop.ML, *MLI);
   DFS.ProcessLoop();
   const SmallVectorImpl<MachineBasicBlock*> &PostOrder = DFS.getOrder();
   for (auto *MBB : PostOrder) {
     recomputeLiveIns(*MBB);
     // FIXME: For some reason, the live-in print order is non-deterministic for

llvm/test/CodeGen/Thumb2/LowOverheadLoops/unsafe-use-after.mir

- { id: 1, name: '', type: spill-slot, offset: -8, size: 4, alignment: 4,
       debug-info-variable: '', debug-info-expression: '', debug-info-location: '' }
 callSites:       []
 constants:       []
 machineFunctionInfo: {}
 body:             |
   ; CHECK-LABEL: name: do_copy
   ; CHECK: bb.0.entry:
   ; CHECK:   successors: %bb.1(0x80000000)
-  ; CHECK:   liveins: $lr, $r2, $r7
+  ; CHECK:   liveins: $lr, $r0, $r2, $r7
   ; CHECK:   frame-setup tPUSH 14, $noreg, killed $r7, implicit-def $sp, implicit $sp
   ; CHECK:   frame-setup CFI_INSTRUCTION def_cfa_offset 8
   ; CHECK:   frame-setup CFI_INSTRUCTION offset $lr, -4
   ; CHECK:   frame-setup CFI_INSTRUCTION offset $r7, -8
   ; CHECK:   renamable $r0 = t2SUBri killed renamable $lr, 4, 14, $noreg, def dead $cpsr
   ; CHECK:   renamable $r1, dead $cpsr = tSUBi3 killed renamable $r2, 4, 14, $noreg
   ; CHECK: bb.1.preheader:
   ; CHECK:   successors: %bb.2(0x80000000)