[GlobalISel] Add MachineInstNumbering to CSEInfo and propagate CSE throughout AArch64 pipeline.

[aemerson]

Summary

MachineInstNumbering is a wrapper around SlotIndexes, which exposes an API to allow fast dominance checks in the CSEMIRBuilder, and to speed up intra-block instruction sinking in the Localizer.

The numbering is stored within the CSEInfo analysis. This change also sets the AArch64 and generic GlobalISel passes to preserve the CSEInfo, and therefore use the CSEMIRBuilder within the RegBankSelect and other combiner passes.

Overall this seems to impact compile time negatively by about ~0.5% on -O0 and -Os but prevents us from having unreasonably long compile times in the worst cases.

[aemerson]

Improve the performance of the insertion queue flushing.

For huge basic blocks, this was having hanging in the renumbering. This new version now tries to not flush on removal of an instruction.

The second issue was that in those huge blocks, any time there was multiple insertions, we'd end up renumbering the entire block. For such blocks, this was way too slow. I've added a more intelligent flushing algorithm to instead search up from each instruction in the block's queue, so that we find an instruction that allows us to use the SlotIndexes to do a fast comparison of ordering. This lets us find the begin iterator for a numbering repair, and likewise walk down each queued instruction until we find the post dominator instruction that gives us an "end" iterator.

These two improvements don't really help on CTMark, they only help on very large pathological blocks.

I think it may be worth exploring using some heuristics early in the pipeline like IRTranslation to detect huge blocks, and switch to using this numbering analysis. Otherwise, we can continue to use the simple linked list block walk, which is slightly faster in the common cases.

[aemerson]

Fix clang-format warnings.

[paquette]

How big is a huge block?

[aemerson]

In D94264#2549325, @paquette wrote:

How big is a huge block?

On the order of millions of instructions.

[paquette]

llvm/include/llvm/CodeGen/GlobalISel/CSEInfo.h

109

Maybe good to note that this also removes MI from the queue if it is present?

llvm/include/llvm/CodeGen/GlobalISel/MachineIRBuilder.h

372

Why is this necessary?

llvm/lib/CodeGen/GlobalISel/CSEInfo.cpp

215

Does this bug have a FIXME or something somewhere?

356

I think you can remove the check for CurrInst->isDebugInstr() if you use skipDebugInstructionsBackward.

380

This function and findEarlierNonQueuedInst are the same aside from the direction they iterate in.

Is it possible to refactor them somehow so they can share more code?

findEarlierNonQueuedInst looks like it's the same loop if you use a reverse_iterator.

380

I think you can remove the check for CurrInst->isDebugInstr() if you use skipDebugInstructionsForward.

[tschuett]

Could you split large basic blocks into smaller ones?

[arsenm]

I think legalizing/combining will do way, way more insertions than would happen in regalloc. Is this losing time to renumberings? Would it make sense to increase the increment size?

[aemerson]

In D94264#2579954, @arsenm wrote:

I think legalizing/combining will do way, way more insertions than would happen in regalloc. Is this losing time to renumberings?

On average, this does cost more in renumbering/flushing than it saves. It only provides a net benefit on the pathological cases in legalization & localization.

In D94264#2579954, @arsenm wrote:

Would it make sense to increase the increment size?

What do you mean by increment size?

In D94264#2579294, @tschuett wrote:

Could you split large basic blocks into smaller ones?

The issue does get resolved if we keep the constants entry block separate after IR translation, instead of merging them. However @qcolombet believes that's not the right approach.

[arsenm]

In D94264#2579990, @aemerson wrote:

In D94264#2579954, @arsenm wrote:

I think legalizing/combining will do way, way more insertions than would happen in regalloc. Is this losing time to renumberings?

On average, this does cost more in renumbering/flushing than it saves. It only provides a net benefit on the pathological cases in legalization & localization.

In D94264#2579954, @arsenm wrote:

Would it make sense to increase the increment size?

What do you mean by increment size?

The instruction index values initially increment by 16 so you only renumber for every 15 or so insertions between instructions:

enum {
  /// The default distance between instructions as returned by distance().
  /// This may vary as instructions are inserted and removed.
  InstrDist = 4 * Slot_Count
};

[aemerson]

In D94264#2580058, @arsenm wrote:

In D94264#2579990, @aemerson wrote:

In D94264#2579954, @arsenm wrote:

I think legalizing/combining will do way, way more insertions than would happen in regalloc. Is this losing time to renumberings?

On average, this does cost more in renumbering/flushing than it saves. It only provides a net benefit on the pathological cases in legalization & localization.

In D94264#2579954, @arsenm wrote:

Would it make sense to increase the increment size?

What do you mean by increment size?

The instruction index values initially increment by 16 so you only renumber for every 15 or so insertions between instructions:

enum {
  /// The default distance between instructions as returned by distance().
  /// This may vary as instructions are inserted and removed.
  InstrDist = 4 * Slot_Count
};

Ok, I'll try that and see what the effect is.

llvm/include/llvm/CodeGen/GlobalISel/MachineIRBuilder.h

372

Must be vestigial.

llvm/lib/CodeGen/GlobalISel/CSEInfo.cpp

215

I think Matt fixed this by removing the observer from the MIRBuilder a while ago, which landed after I wrote this.

380

I'll see what I can do here. I did try to factor things out originally but ran into a a few issues.

[jroelofs]

llvm/include/llvm/CodeGen/GlobalISel/CSEInfo.h

70

It would eliminate the explicit dtor and simplify the logic in init() if this were std::unique_ptr<SlotIndexes> SI;

[arsenm]

Is this still relevant? I thought this was done already

[arsenm]

Thought there was a previously committed solution for this

llvm/include/llvm/CodeGen/GlobalISel/CSEInfo.h

 //===- llvm/CodeGen/GlobalISel/CSEInfo.h ------------------*- C++ -*-===//
 //
 // Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
 // See https://llvm.org/LICENSE.txt for license information.
 // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
 //
 //===----------------------------------------------------------------------===//
 //
 /// Provides analysis for continuously CSEing during GISel passes.
 //
 //===----------------------------------------------------------------------===//
 #ifndef LLVM_CODEGEN_GLOBALISEL_CSEINFO_H
 #define LLVM_CODEGEN_GLOBALISEL_CSEINFO_H
 
 #include "llvm/ADT/FoldingSet.h"
+#include "llvm/ADT/SmallPtrSet.h"
+#include "llvm/ADT/SmallVector.h"
+#include "llvm/ADT/iterator_range.h"
 #include "llvm/CodeGen/CSEConfigBase.h"
 #include "llvm/CodeGen/GlobalISel/GISelChangeObserver.h"
 #include "llvm/CodeGen/GlobalISel/GISelWorkList.h"
 #include "llvm/CodeGen/MachineFunctionPass.h"
+#include "llvm/CodeGen/MachineMemOperand.h"
+#include "llvm/CodeGen/SlotIndexes.h"
 #include "llvm/Support/Allocator.h"
 #include "llvm/Support/CodeGen.h"
 
 namespace llvm {
 class MachineBasicBlock;
 
 /// A class that wraps MachineInstrs and derives from FoldingSetNode in order to
 /// be uniqued in a CSEMap. The tradeoff here is extra memory allocations for

 
 // Returns the standard expected CSEConfig for the given optimization level.
 // We have this logic here so targets can make use of it from their derived
 // TargetPassConfig, but can't put this logic into TargetPassConfig directly
 // because the CodeGen library can't depend on GlobalISel.
 std::unique_ptr<CSEConfigBase>
 getStandardCSEConfigForOpt(CodeGenOpt::Level Level);
 
+/// A wrapper for SlotIndexes that holds yet-to-be inserted instructions.
+/// This is used so we can lazily renumber the indices when the next
+/// query is made. The created() observer callback actually executes
+/// *before* the MachineInstr is inserted, but SlotIndexes needs to have
+/// the instruction fully inserted before it starts renumbering.
+class MachineInstNumbering {
+  SlotIndexes *SI = nullptr;
+  MachineFunction *MF = nullptr;
+  bool IsEnabled = true;
+  /// Contains all instructions that need renumbering.
+  SmallVector<MachineInstr *> InsertionQueue;
+
+  /// Extract instructions from the queue that have the same parent block,
+  /// and null them out in the queue.
+  void
+  extractInstsFromSameBlock(SmallPtrSetImpl<MachineInstr *> &ExtractedInsts);
+  /// Find the nearest predecessor instruction that's not queued.
+  MachineBasicBlock::iterator
+  findEarlierNonQueuedInst(MachineBasicBlock::iterator Start,
+                           SmallPtrSetImpl<MachineInstr *> &InstQueue);
+  /// Find the nearest successor instruction that's not queued.
+  MachineBasicBlock::iterator
+  findLaterNonQueuedInst(MachineBasicBlock::iterator Start,
+                         SmallPtrSetImpl<MachineInstr *> &InstQueue);
+
+  /// Returns true if the instruction \p A is earlier than the SlotIndex \p BIdx
+  /// in the block using the instruction numbering. If \p NoFlush is true, then
+  /// do not flush any queued insertions before performing the query. This can
+  /// be used for efficiency reasons if both \A and \B are known to be already
+  /// in the SlotIndexes.
+  bool isEarlier(const MachineInstr &A, const SlotIndex &BIdx,
+                 bool NoFlush = false);
+
+public:
+  MachineInstNumbering() = default;
+  void init(MachineFunction &MF);
+
+  void setMF(MachineFunction &MF) { this->MF = &MF; }
+  bool isEnabled() const { return SI != nullptr; }
+  /// Enable numbering of the given MF, unless the MF has already been analyzed.
+  void enable(MachineFunction *MF = nullptr);
+  void disable();
+
+  /// Flush the queue and start a renumbering for the regions that need it.
+  void flushAndRenumber();
+  /// Remove an instruction from SlotIndexes.
+  void removeInstruction(MachineInstr &MI);
+  /// Queue an instruction to be inserted and trigger renumbering
+  /// on the next query event.
+  void queueInsertion(MachineInstr &MI);
+  /// Repair numbering in the given range.
+  void repairNumberingInRange(MachineBasicBlock::iterator Begin,
+                              MachineBasicBlock::iterator End);
+  /// Returns true if \p A is earlier than \p B in the block using the
+  /// instruction numbering.
+  bool isEarlier(const MachineInstr &A, const MachineInstr &B);
+
+  /// Returns the earliest MachineInstr in the given set \p Instrs.
+  /// If \p NoFlush is true, then do not flush any queued insertions before
+  /// performing the query. This can be used for efficiency reasons if both
+  /// \A and \B are known to be already in the SlotIndexes.
+  MachineInstr *getEarliestInstr(const SmallPtrSetImpl<MachineInstr *> &Instrs,
+                                 bool NoFlush = false);
+
+  /// Deleted copy ctor and assignment operator to prevent copying these.
+  MachineInstNumbering(const MachineInstNumbering &_) = delete;
+  MachineInstNumbering &operator=(const MachineInstNumbering &_) = delete;
+
+  ~MachineInstNumbering();
+};
+
 /// The CSE Analysis object.
 /// This installs itself as a delegate to the MachineFunction to track
 /// new instructions as well as deletions. It however will not be able to
 /// track instruction mutations. In such cases, recordNewInstruction should be
 /// called (for eg inside MachineIRBuilder::recordInsertion).
 /// Also because of how just the instruction can be inserted without adding any
 /// operands to the instruction, instructions are uniqued and inserted lazily.
 /// CSEInfo should assert when trying to enter an incomplete instruction into
 /// the CSEMap. There is Opcode level granularity on which instructions can be
 /// CSE'd and for now, only Generic instructions are CSEable.
 class GISelCSEInfo : public GISelChangeObserver {
   // Make it accessible only to CSEMIRBuilder.
   friend class CSEMIRBuilder;
 
   BumpPtrAllocator UniqueInstrAllocator;
   FoldingSet<UniqueMachineInstr> CSEMap;
   MachineRegisterInfo *MRI = nullptr;
   MachineFunction *MF = nullptr;
+  MachineInstNumbering MINumbering;
   std::unique_ptr<CSEConfigBase> CSEOpt;
   /// Keep a cache of UniqueInstrs for each MachineInstr. In GISel,
   /// often instructions are mutated (while their ID has completely changed).
   /// Whenever mutation happens, invalidate the UniqueMachineInstr for the
   /// MachineInstr
   DenseMap<const MachineInstr *, UniqueMachineInstr *> InstrMapping;
 
   /// Store instructions that are not fully formed in TemporaryInsts.
public:
   GISelCSEInfo() = default;
 
   virtual ~GISelCSEInfo();
 
   void setMF(MachineFunction &MF);
 
   Error verify();
 
+  MachineInstNumbering &getInstNumbering() { return MINumbering; }
+  const MachineInstNumbering &getInstNumbering() const { return MINumbering; }
+
   /// Records a newly created inst in a list and lazily insert it to the CSEMap.
   /// Sometimes, this method might be called with a partially constructed
   /// MachineInstr,
   //  (right after BuildMI without adding any operands) - and in such cases,
   //  defer the hashing of the instruction to a later stage.
   void recordNewInstruction(MachineInstr *MI);
 
   /// Use this callback to inform CSE about a newly fully created instruction.

llvm/include/llvm/CodeGen/GlobalISel/CSEMIRBuilder.h

class CSEMIRBuilder : public MachineIRBuilder {
   //
   // TODO: Another approach for checking dominance is having two iterators and
   // making them go towards each other until they meet or reach begin/end. Which
   // approach is better? Should this even change dynamically? For G_CONSTANTS
   // most of which will be at the top of the BB, the top down approach would be
   // a better choice. Does IRTranslator placing constants at the beginning still
   // make sense? Should this change based on Opcode?
   bool dominates(MachineBasicBlock::const_iterator A,
-                 MachineBasicBlock::const_iterator B) const;
+                 MachineBasicBlock::const_iterator B);
 
   /// For given ID, find a machineinstr in the CSE Map. If found, check if it
   /// dominates the current insertion point and if not, move it just before the
   /// current insertion point and return it. If not found, return Null
   /// MachineInstrBuilder.
   MachineInstrBuilder getDominatingInstrForID(FoldingSetNodeID &ID,
                                               void *&NodeInsertPos);
   /// Simple check if we can CSE (we have the CSEInfo) or if this Opcode is

llvm/include/llvm/CodeGen/GlobalISel/Legalizer.h

 /// instruction.
 //
 //===----------------------------------------------------------------------===//
 
 #ifndef LLVM_CODEGEN_GLOBALISEL_LEGALIZEMACHINEIRPASS_H
 #define LLVM_CODEGEN_GLOBALISEL_LEGALIZEMACHINEIRPASS_H
 
 #include "llvm/CodeGen/GlobalISel/MachineIRBuilder.h"
+#include "llvm/CodeGen/MachineBasicBlock.h"
 #include "llvm/CodeGen/MachineFunctionPass.h"
 
 namespace llvm {
 
 class MachineRegisterInfo;
 class LostDebugLocObserver;
 
 class Legalizer : public MachineFunctionPass {

llvm/include/llvm/CodeGen/GlobalISel/Localizer.h

 #include "llvm/ADT/SetVector.h"
 #include "llvm/CodeGen/GlobalISel/MachineIRBuilder.h"
 #include "llvm/CodeGen/MachineFunctionPass.h"
 
 namespace llvm {
 // Forward declarations.
 class MachineRegisterInfo;
 class TargetTransformInfo;
+class CSEInfo;
 
 /// This pass implements the localization mechanism described at the
 /// top of this file. One specificity of the implementation is that
 /// it will materialize one and only one instance of a constant per
 /// basic block, thus enabling reuse of that constant within that block.
 /// Moreover, it only materializes constants in blocks where they
 /// are used. PHI uses are considered happening at the end of the
 /// related predecessor.
 class Localizer : public MachineFunctionPass {
 public:
   static char ID;
 
 private:
   /// An input function to decide if the pass should run or not
   /// on the given MachineFunction.
   std::function<bool(const MachineFunction &)> DoNotRunPass;
 
   /// MRI contains all the register class/bank information that this
   /// pass uses and updates.
   MachineRegisterInfo *MRI;
   /// TTI used for getting remat costs for instructions.
   TargetTransformInfo *TTI;
+  /// Used for accessing the instruction numbering to speed up localizing.
+  GISelCSEInfo *CSEInfo = nullptr;
 
   /// Check if \p MOUse is used in the same basic block as \p Def.
   /// If the use is in the same block, we say it is local.
   /// When the use is not local, \p InsertMBB will contain the basic
   /// block when to insert \p Def to have a local use.
   static bool isLocalUse(MachineOperand &MOUse, const MachineInstr &Def,
                          MachineBasicBlock *&InsertMBB);
 

llvm/include/llvm/CodeGen/GlobalISel/MachineIRBuilder.h

 #ifndef LLVM_CODEGEN_GLOBALISEL_MACHINEIRBUILDER_H
 #define LLVM_CODEGEN_GLOBALISEL_MACHINEIRBUILDER_H
 
 #include "llvm/CodeGen/GlobalISel/CSEInfo.h"
 #include "llvm/CodeGen/LowLevelType.h"
 #include "llvm/CodeGen/MachineBasicBlock.h"
 #include "llvm/CodeGen/MachineInstrBuilder.h"
 #include "llvm/CodeGen/MachineRegisterInfo.h"
+#include "llvm/CodeGen/SlotIndexes.h"
 #include "llvm/CodeGen/TargetOpcodes.h"
 #include "llvm/IR/Constants.h"
 #include "llvm/IR/DebugLoc.h"
 #include "llvm/IR/Module.h"
 
 namespace llvm {
 
 // Forward declarations.
public:
   /// Build and insert <empty> = \p Opcode <empty>.
   /// The insertion point is the one set by the last call of either
   /// setBasicBlock or setMI.
   ///
   /// \pre setBasicBlock or setMI must have been called.
   ///
   /// \return a MachineInstrBuilder for the newly created instruction.
   MachineInstrBuilder buildInstr(unsigned Opcode) {
-    return insertInstr(buildInstrNoInsert(Opcode));
+    auto MIB = insertInstr(buildInstrNoInsert(Opcode));
+    return MIB;
   }
 
   /// Build but don't insert <empty> = \p Opcode <empty>.
   ///
   /// \pre setMF, setBasicBlock or setMI  must have been called.
   ///
   /// \return a MachineInstrBuilder for the newly created instruction.
   MachineInstrBuilder buildInstrNoInsert(unsigned Opcode);

llvm/include/llvm/CodeGen/GlobalISel/RegBankSelect.h

private:
   /// Get the frequency of the edges.
   /// This is required for non-fast mode.
   MachineBranchProbabilityInfo *MBPI = nullptr;
 
   /// Current optimization remark emitter. Used to report failures.
   std::unique_ptr<MachineOptimizationRemarkEmitter> MORE;
 
   /// Helper class used for every code morphing.
-  MachineIRBuilder MIRBuilder;
+  std::unique_ptr<MachineIRBuilder> MIRBuilder;
 
   /// Optimization mode of the pass.
   Mode OptMode;
 
   /// Current target configuration. Controls how the pass handles errors.
   const TargetPassConfig *TPC;
 
   /// Assign the register bank of each operand of \p MI.

llvm/lib/CodeGen/GlobalISel/CSEInfo.cpp

 //===- CSEInfo.cpp ------------------------------===//
 //
 // Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
 // See https://llvm.org/LICENSE.txt for license information.
 // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
 //
 //===----------------------------------------------------------------------===//
 //
 //
 //===----------------------------------------------------------------------===//
 #include "llvm/CodeGen/GlobalISel/CSEInfo.h"
+#include "llvm/ADT/STLExtras.h"
+#include "llvm/ADT/SmallVector.h"
+#include "llvm/CodeGen/MachineBasicBlock.h"
+#include "llvm/CodeGen/MachineInstr.h"
 #include "llvm/CodeGen/MachineRegisterInfo.h"
+#include "llvm/CodeGen/SlotIndexes.h"
 #include "llvm/InitializePasses.h"
 
 #define DEBUG_TYPE "cseinfo"
 
 using namespace llvm;
 char llvm::GISelCSEAnalysisWrapperPass::ID = 0;
 GISelCSEAnalysisWrapperPass::GISelCSEAnalysisWrapperPass()
     : MachineFunctionPass(ID) {
#ifndef NDEBUG
   if (OpcodeHitTable.count(Opc))
     OpcodeHitTable[Opc] += 1;
   else
     OpcodeHitTable[Opc] = 1;
 #endif
   // Else do nothing.
 }
 
+void MachineInstNumbering::init(MachineFunction &MF) {
+  if (!IsEnabled)
+    return;
+  this->MF = &MF;
+  if (SI) {
+    SI->releaseMemory();
+    SI->runOnMachineFunction(MF);
+    return;
+  }
+  SI = new class SlotIndexes();
+  SI->runOnMachineFunction(MF);
+}
+
+void MachineInstNumbering::enable(MachineFunction *MF) {
+  if (IsEnabled && this->MF == MF)
+    return;
+  IsEnabled = true;
+  if (MF)
+    init(*MF);
+}
+
+void MachineInstNumbering::disable() {
+  if (SI) {
+    delete SI;
+    SI = nullptr;
+  }
+  InsertionQueue.clear();
+  IsEnabled = false;
+  return;
+}
+
+void MachineInstNumbering::removeInstruction(MachineInstr &MI) {
+  assert(IsEnabled && "Expected Inst Numbering to be enabled");
+  // Remove the instruction from the queue if it exists. Due to a bug that
+  // inserts instructions twice due to duplicate observers, we check the entire
+  // queue instead of quitting on the first hit.
+  bool Found = false;
+  for (unsigned I = 0, E = InsertionQueue.size(); I < E; ++I) {
+    if (InsertionQueue[I] == &MI) {
+      InsertionQueue[I] = nullptr;
+      Found = true;
+    }
+  }
+  if (!Found)
+    SI->removeSingleMachineInstrFromMaps(MI);
+}
+
+void MachineInstNumbering::queueInsertion(MachineInstr &MI) {
+  assert(IsEnabled && "Expected Inst Numbering to be enabled");
+  InsertionQueue.emplace_back(&MI);
+}
+
+void MachineInstNumbering::repairNumberingInRange(
+    MachineBasicBlock::iterator Begin, MachineBasicBlock::iterator End) {
+  assert(IsEnabled && "Expected Inst Numbering to be enabled");
+  flushAndRenumber();
+  SI->repairIndexesInRange(Begin->getParent(), Begin, End);
+}
+
+void MachineInstNumbering::flushAndRenumber() {
+  assert(IsEnabled && "Expected Inst Numbering to be enabled");
+
+  if (InsertionQueue.empty())
+    return;
+
+  // There may be multiple queued instructions for a given block. Because of
+  // this we can't trivially find the iterator ranges to do a local index
+  // repair since the neighbouring instructions may not have valid indexes yet.
+  // So, we have to find the nearest non-queued predecessor instruction in the
+  // block that dominates all of the queued instructions.
+  //
+  // If the instruction is not queued, then we should have a valid index, and
+  // therefore can quickly check if any other non-queued instruction is earlier
+  // or later.
+
+  SmallPtrSet<MachineInstr *, 8> SameBlockSubset;
+  while (true) {
+    // First we have to extract all the instructions that share a common basic
+    // block.
+    extractInstsFromSameBlock(SameBlockSubset);
+    if (SameBlockSubset.empty()) {
+      // We have no more instructions to process.
+      InsertionQueue.clear();
+      return;
+    }
+
+    // Now walk up from each instruction position until we hit an unqueued
+    // instruction, and compare the SlotIndex for that to the current "earliest"
+    // index.
+    SlotIndex DominatingInstSlotIdx;
+    MachineBasicBlock::iterator StartIt;
+    for (auto *MI : SameBlockSubset) {
+      auto It = MI->getIterator();
+      auto TerminalIt = MI->getParent()->begin();
+      if (It == TerminalIt) {
+        // This is the first instruction, end search.
+        StartIt = It;
+        break;
+      }
+
+      auto NonQueuedIt = findEarlierNonQueuedInst(It, SameBlockSubset);
+
+      // There's an edge case where the iterator returned by
+      // findEarlierNonQueuedInst() is begin(), but is also in the queue.
+      if (NonQueuedIt == TerminalIt) {
+        StartIt = NonQueuedIt;
+        break;
+      }
+
+      SlotIndex NonQueuedIndex = SI->getInstructionIndex(*NonQueuedIt);
+      if (!DominatingInstSlotIdx.isValid()) {
+        // This is the first instruction checked.
+        DominatingInstSlotIdx = NonQueuedIndex;
+        StartIt = NonQueuedIt;
+      } else if (SlotIndex::isEarlierInstr(NonQueuedIndex,
+                                           DominatingInstSlotIdx)) {
+        DominatingInstSlotIdx = NonQueuedIndex;
+        StartIt = NonQueuedIt;
+      }
+
+      // If the current iterator is the first instruction, then end the search
+      // early.
+      if (StartIt == TerminalIt)
+        break;
+    }
+
+    // Now we have the nearest predecessor which we can use as a repair range
+    // start point. Now do the same thing for the end iterator.
+
+    SlotIndex PostDomSlotIdx;
+    MachineBasicBlock::iterator EndIt;
+    for (auto *MI : SameBlockSubset) {
+      auto It = MI->getIterator();
+      auto TerminalIt = MI->getParent()->end();
+      if (std::next(It) == TerminalIt) {
+        EndIt = TerminalIt;
+        break;
+      }
+
+      auto NonQueuedIt = findLaterNonQueuedInst(It, SameBlockSubset);
+
+      // There's an edge case where the iterator returned by
+      // findLaterNonQueuedInst() is end(), but is also in the queue.
+      if (NonQueuedIt == TerminalIt) {
+        EndIt = NonQueuedIt;
+        break;
+      }
+
+      SlotIndex NonQueuedIndex = SI->getInstructionIndex(*NonQueuedIt);
+      if (!PostDomSlotIdx.isValid()) {
+        // This is the first instruction checked.
+        PostDomSlotIdx = NonQueuedIndex;
+        EndIt = NonQueuedIt;
+      } else if (SlotIndex::isEarlierInstr(PostDomSlotIdx, NonQueuedIndex)) {
+        PostDomSlotIdx = NonQueuedIndex;
+        EndIt = NonQueuedIt;
+      }
+
+      // If the current iterator is the first instruction, then end the search
+      // early.
+      if (EndIt == TerminalIt)
+        break;
+    }
+    SI->repairIndexesInRange(StartIt->getParent(), StartIt, EndIt);
+    SameBlockSubset.clear();
+  }
+}
+
+MachineBasicBlock::iterator MachineInstNumbering::findEarlierNonQueuedInst(
+    MachineBasicBlock::iterator Start,
+    SmallPtrSetImpl<MachineInstr *> &InstQueue) {
+  MachineBasicBlock *ParentBB = Start->getParent();
+  assert(Start != ParentBB->begin() && "Iterator should not be the start");
+  MachineBasicBlock::iterator II = Start;
+  do {
+    --II;
+    MachineInstr *CurrInst = &*II;
+
+    // DBG_VALUE instructions aren't inserted into SlotIndexes, so skip them.
+    if (CurrInst->isDebugInstr()) {
+      if (II == ParentBB->begin())
+        return II;
+      continue;
+    }
+    if (!InstQueue.contains(CurrInst)) {
+      // We have a non-queued instruction. Return the iterator.
+      return CurrInst->getIterator();
+    }
+  } while (II != ParentBB->begin());
+  return II;
+}
+
+MachineBasicBlock::iterator MachineInstNumbering::findLaterNonQueuedInst(
+    MachineBasicBlock::iterator Start,
+    SmallPtrSetImpl<MachineInstr *> &InstQueue) {
+  MachineBasicBlock *ParentBB = Start->getParent();
+  assert(Start != ParentBB->end() && "Iterator should not be the start");
+  MachineBasicBlock::iterator II = Start;
+  do {
+    ++II;
+    if (II == ParentBB->end())
+      return II;
+
+    MachineInstr *CurrInst = &*II;
+
+    // DBG_VALUE instructions aren't inserted into SlotIndexes, so skip them.
+    if (CurrInst->isDebugInstr()) {
+      if (II == ParentBB->end())
+        return II;
+      continue;
+    }
+    if (!InstQueue.contains(CurrInst)) {
+      // We have a non-queued instruction. Return the iterator.
+      return CurrInst->getIterator();
+    }
+  } while (true);
+  return II;
+}
+
+void MachineInstNumbering::extractInstsFromSameBlock(
+    SmallPtrSetImpl<MachineInstr *> &ExtractedInsts) {
+  MachineBasicBlock *ParentBB = nullptr;
+  for (unsigned I = 0, E = InsertionQueue.size(); I < E; ++I) {
+    MachineInstr *MI = InsertionQueue[I];
+    if (!MI)
+      continue;
+    if (!ParentBB) {
+      ParentBB = MI->getParent();
+      ExtractedInsts.insert(MI);
+      InsertionQueue[I] = nullptr;
+      continue;
+    }
+    if (MI->getParent() == ParentBB) {
+      ExtractedInsts.insert(MI);
+      InsertionQueue[I] = nullptr;
+    }
+  }
+}
+
+bool MachineInstNumbering::isEarlier(const MachineInstr &A,
+                                     const MachineInstr &B) {
+  assert(IsEnabled && "Expected Inst Numbering to be enabled");
+  // If we have outstanding queued instructions, process them first.
+  flushAndRenumber();
+
+  SlotIndex SlotA = SI->getInstructionIndex(A);
+  SlotIndex SlotB = SI->getInstructionIndex(B);
+  return SlotIndex::isEarlierInstr(SlotA, SlotB);
+}
+
+bool MachineInstNumbering::isEarlier(const MachineInstr &A,
+                                     const SlotIndex &BIdx, bool NoFlush) {
+  assert(BIdx.isValid() && "Expected valid index");
+  // If we have outstanding queued instructions, process them first.
+  if (!NoFlush)
+    flushAndRenumber();
+
+  SlotIndex SlotA = SI->getInstructionIndex(A);
+  return SlotIndex::isEarlierInstr(SlotA, BIdx);
+}
+
+MachineInstr *MachineInstNumbering::getEarliestInstr(
+    const SmallPtrSetImpl<MachineInstr *> &Instrs, bool NoFlush) {
+  MachineInstr *Earliest = nullptr;
+  SlotIndex EarliestIdx;
+  for (auto *Other : Instrs) {
+    if (!Earliest) {
+      Earliest = Other;
+      EarliestIdx = SI->getInstructionIndex(*Earliest);
+      continue;
+    }
+    assert(Earliest != Other);
+    if (isEarlier(*Other, EarliestIdx, NoFlush)) {
+      Earliest = Other;
+      EarliestIdx = SI->getInstructionIndex(*Other);
+    }
+  }
+  return Earliest;
+}
+
+MachineInstNumbering::~MachineInstNumbering() {
+  if (SI)
+    delete SI;
+}
+
 void GISelCSEInfo::recordNewInstruction(MachineInstr *MI) {
   if (shouldCSE(MI->getOpcode())) {
     TemporaryInsts.insert(MI);
     LLVM_DEBUG(dbgs() << "CSEInfo::Recording new MI " << *MI);
   }
 }
 
 void GISelCSEInfo::handleRecordedInst(MachineInstr *MI) {
void GISelCSEInfo::handleRecordedInsts() {
   }
 }
 
 bool GISelCSEInfo::shouldCSE(unsigned Opc) const {
   assert(CSEOpt.get() && "CSEConfig not set");
   return CSEOpt->shouldCSEOpc(Opc);
 }
 
-void GISelCSEInfo::erasingInstr(MachineInstr &MI) { handleRemoveInst(&MI); }
-void GISelCSEInfo::createdInstr(MachineInstr &MI) { recordNewInstruction(&MI); }
+void GISelCSEInfo::erasingInstr(MachineInstr &MI) {
+  handleRemoveInst(&MI);
+  if (MINumbering.isEnabled())
+    MINumbering.removeInstruction(MI);
+}
+
+void GISelCSEInfo::createdInstr(MachineInstr &MI) {
+  recordNewInstruction(&MI);
+  if (MINumbering.isEnabled())
+    MINumbering.queueInsertion(MI);
+}
 void GISelCSEInfo::changingInstr(MachineInstr &MI) {
-  // For now, perform erase, followed by insert.
-  erasingInstr(MI);
-  createdInstr(MI);
+  // We effectively do an erase + insert, but we don't call erasingInstr() or
+  // createdInstr() because this we don't need to do anything for the
+  // instruction numbering for instruction mutations.
+  handleRemoveInst(&MI);
+  recordNewInstruction(&MI);
 }
+
 void GISelCSEInfo::changedInstr(MachineInstr &MI) { changingInstr(MI); }
 
 void GISelCSEInfo::analyze(MachineFunction &MF) {
   setMF(MF);
   for (auto &MBB : MF) {
     if (MBB.empty())
       continue;
     for (MachineInstr &MI : MBB) {
void GISelCSEInfo::releaseMemory() {
   print();
   CSEMap.clear();
   InstrMapping.clear();
   UniqueInstrAllocator.Reset();
   TemporaryInsts.clear();
   CSEOpt.reset();
   MRI = nullptr;
   MF = nullptr;
+  MINumbering.disable();
 #ifndef NDEBUG
   OpcodeHitTable.clear();
 #endif
 }
 
 Error GISelCSEInfo::verify() {
 #ifndef NDEBUG
   handleRecordedInsts();

 }
 
 GISelCSEInfo &
 GISelCSEAnalysisWrapper::get(std::unique_ptr<CSEConfigBase> CSEOpt,
                              bool Recompute) {
   if (!AlreadyComputed || Recompute) {
     Info.releaseMemory();
     Info.setCSEConfig(std::move(CSEOpt));
+    Info.getInstNumbering().init(*MF);
     Info.analyze(*MF);
     AlreadyComputed = true;
   }
   return Info;
 }
 void GISelCSEAnalysisWrapperPass::getAnalysisUsage(AnalysisUsage &AU) const {
   AU.setPreservesAll();
   MachineFunctionPass::getAnalysisUsage(AU);
 }
 
 bool GISelCSEAnalysisWrapperPass::runOnMachineFunction(MachineFunction &MF) {
   releaseMemory();
   Wrapper.setMF(MF);
   return false;
 }

llvm/lib/CodeGen/GlobalISel/CSEMIRBuilder.cpp

 //===-- llvm/CodeGen/GlobalISel/CSEMIRBuilder.cpp - MIBuilder--*- C++ -*-==//
 //
 // Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
 // See https://llvm.org/LICENSE.txt for license information.
 // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
 //
 //===----------------------------------------------------------------------===//
 /// \file
 /// This file implements the CSEMIRBuilder class which CSEs as it builds
 /// instructions.
 //===----------------------------------------------------------------------===//
 //
 
 #include "llvm/CodeGen/GlobalISel/CSEMIRBuilder.h"
+#include "llvm/CodeGen/GlobalISel/CSEInfo.h"
 #include "llvm/CodeGen/GlobalISel/GISelChangeObserver.h"
+#include "llvm/CodeGen/MachineBasicBlock.h"
 #include "llvm/IR/DebugInfoMetadata.h"
 
 using namespace llvm;
 
 bool CSEMIRBuilder::dominates(MachineBasicBlock::const_iterator A,
-                              MachineBasicBlock::const_iterator B) const {
+                              MachineBasicBlock::const_iterator B) {
   auto MBBEnd = getMBB().end();
   if (B == MBBEnd)
     return true;
   assert(A->getParent() == B->getParent() &&
          "Iterators should be in same block");
+
+  auto &InstNumbering = getCSEInfo()->getInstNumbering();
+  if (InstNumbering.isEnabled())
+    return InstNumbering.isEarlier(*A, *B);
+
   const MachineBasicBlock *BBA = A->getParent();
   MachineBasicBlock::const_iterator I = BBA->begin();
   for (; &*I != A && &*I != B; ++I)
     ;
   return &*I == A;
 }
 
 MachineInstrBuilder
 CSEMIRBuilder::getDominatingInstrForID(FoldingSetNodeID &ID,
                                        void *&NodeInsertPos) {
   GISelCSEInfo *CSEInfo = getCSEInfo();
   assert(CSEInfo && "Can't get here without setting CSEInfo");
   MachineBasicBlock *CurMBB = &getMBB();
   MachineInstr *MI =
       CSEInfo->getMachineInstrIfExists(ID, CurMBB, NodeInsertPos);
   if (MI) {
     CSEInfo->countOpcodeHit(MI->getOpcode());
     auto CurrPos = getInsertPt();
     auto MII = MachineBasicBlock::iterator(MI);
     if (MII == CurrPos) {
       // Move the insert point ahead of the instruction so any future uses of
       // this builder will have the def ready.
       setInsertPt(*CurMBB, std::next(MII));
     } else if (!dominates(MI, CurrPos)) {
+      ++MII;
       CurMBB->splice(CurrPos, CurMBB, MI);
+      auto &InstNumbering = CSEInfo->getInstNumbering();
+      if (InstNumbering.isEnabled()) {
+        InstNumbering.removeInstruction(*MI);
+        // Update numbering for the modified range.
+        InstNumbering.repairNumberingInRange(MachineBasicBlock::iterator(MI),
+                                             MII);
+      }
     }
     return MachineInstrBuilder(getMF(), MI);
   }
   return MachineInstrBuilder();
 }
 
 bool CSEMIRBuilder::canPerformCSEForOpc(unsigned Opc) const {
   const GISelCSEInfo *CSEInfo = getCSEInfo();

llvm/lib/CodeGen/GlobalISel/IRTranslator.cpp

bool EnableCSE = EnableCSEInIRTranslator.getNumOccurrences()
                        ? EnableCSEInIRTranslator
                        : TPC->isGISelCSEEnabled();
 
   if (EnableCSE) {
     EntryBuilder = std::make_unique<CSEMIRBuilder>(CurMF);
     CSEInfo = &Wrapper.get(TPC->getCSEConfig());
     EntryBuilder->setCSEInfo(CSEInfo);
     CurBuilder = std::make_unique<CSEMIRBuilder>(CurMF);
+    // Disable instruction numbering while building the MF.
+    CSEInfo->getInstNumbering().disable();
     CurBuilder->setCSEInfo(CSEInfo);
   } else {
     EntryBuilder = std::make_unique<MachineIRBuilder>();
     CurBuilder = std::make_unique<MachineIRBuilder>();
   }
   CLI = MF->getSubtarget().getCallLowering();
   CurBuilder->setMF(*MF);
   EntryBuilder->setMF(*MF);
#endif
   MF->DeleteMachineBasicBlock(EntryBB);
 
   assert(&MF->front() == &NewEntryBB &&
          "New entry wasn't next in the list of basic block!");
 
   // Initialize stack protector information.
   StackProtector &SP = getAnalysis<StackProtector>();
   SP.copyToMachineFrameInfo(MF->getFrameInfo());
-
   return false;
 }

llvm/lib/CodeGen/GlobalISel/Legalizer.cpp

 #include "llvm/CodeGen/GlobalISel/CSEInfo.h"
 #include "llvm/CodeGen/GlobalISel/CSEMIRBuilder.h"
 #include "llvm/CodeGen/GlobalISel/GISelChangeObserver.h"
 #include "llvm/CodeGen/GlobalISel/GISelWorkList.h"
 #include "llvm/CodeGen/GlobalISel/LegalizationArtifactCombiner.h"
 #include "llvm/CodeGen/GlobalISel/LegalizerHelper.h"
 #include "llvm/CodeGen/GlobalISel/LostDebugLocObserver.h"
 #include "llvm/CodeGen/GlobalISel/Utils.h"
+#include "llvm/CodeGen/MachineBasicBlock.h"
 #include "llvm/CodeGen/MachineOptimizationRemarkEmitter.h"
 #include "llvm/CodeGen/MachineRegisterInfo.h"
 #include "llvm/CodeGen/TargetPassConfig.h"
 #include "llvm/CodeGen/TargetSubtargetInfo.h"
 #include "llvm/InitializePasses.h"
 #include "llvm/Support/Debug.h"
 #include "llvm/Support/Error.h"
 #include "llvm/Target/TargetMachine.h"
bool Legalizer::runOnMachineFunction(MachineFunction &MF) {
   GISelCSEInfo *CSEInfo = nullptr;
   bool EnableCSE = EnableCSEInLegalizer.getNumOccurrences()
                        ? EnableCSEInLegalizer
                        : TPC.isGISelCSEEnabled();
   if (EnableCSE) {
     MIRBuilder = std::make_unique<CSEMIRBuilder>();
     CSEInfo = &Wrapper.get(TPC.getCSEConfig());
     MIRBuilder->setCSEInfo(CSEInfo);
+    CSEInfo->getInstNumbering().enable(&MF);
   } else
     MIRBuilder = std::make_unique<MachineIRBuilder>();
 
   SmallVector<GISelChangeObserver *, 1> AuxObservers;
   if (EnableCSE && CSEInfo) {
     // We want CSEInfo in addition to WorkListObserver to observe all changes.
     AuxObservers.push_back(CSEInfo);
   }

llvm/lib/CodeGen/GlobalISel/Localizer.cpp

 //===- Localizer.cpp ---------------------- Localize some instrs -*- C++ -*-==//
 //
 // Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
 // See https://llvm.org/LICENSE.txt for license information.
 // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
 //
 //===----------------------------------------------------------------------===//
 /// \file
 /// This file implements the Localizer class.
 //===----------------------------------------------------------------------===//
 
 #include "llvm/CodeGen/GlobalISel/Localizer.h"
 #include "llvm/ADT/DenseMap.h"
 #include "llvm/Analysis/TargetTransformInfo.h"
+#include "llvm/CodeGen/MachineBasicBlock.h"
 #include "llvm/CodeGen/MachineRegisterInfo.h"
 #include "llvm/CodeGen/TargetLowering.h"
+#include "llvm/CodeGen/TargetPassConfig.h"
 #include "llvm/InitializePasses.h"
 #include "llvm/Support/Debug.h"
 
 #define DEBUG_TYPE "localizer"
 
 using namespace llvm;
 
+static cl::opt<bool> UseInstNumbering(
+    "localizer-use-numbering", cl::Hidden,
+    cl::desc("Use the instruction numbering in the localizer."),
+    cl::init(true));
+
 char Localizer::ID = 0;
 INITIALIZE_PASS_BEGIN(Localizer, DEBUG_TYPE,
                       "Move/duplicate certain instructions close to their use",
                       false, false)
 INITIALIZE_PASS_DEPENDENCY(TargetTransformInfoWrapperPass)
 INITIALIZE_PASS_END(Localizer, DEBUG_TYPE,
                     "Move/duplicate certain instructions close to their use",
                     false, false)
 
 Localizer::Localizer(std::function<bool(const MachineFunction &)> F)
     : MachineFunctionPass(ID), DoNotRunPass(F) {}
 
 Localizer::Localizer()
     : Localizer([](const MachineFunction &) { return false; }) {}
 
 void Localizer::init(MachineFunction &MF) {
   MRI = &MF.getRegInfo();
   TTI = &getAnalysis<TargetTransformInfoWrapperPass>().getTTI(MF.getFunction());
+
+  if (!UseInstNumbering) {
+    CSEInfo = nullptr;
+    return;
+  }
+  // Enable inst numbering, we need to require the CSEInfo analysis since that
+  // holds the numbering.
+  auto &TPC = getAnalysis<TargetPassConfig>();
+  GISelCSEAnalysisWrapper &Wrapper =
+      getAnalysis<GISelCSEAnalysisWrapperPass>().getCSEWrapper();
+  auto *FnCSEInfo = &Wrapper.get(TPC.getCSEConfig());
+  FnCSEInfo->getInstNumbering().enable(&MF);
+  CSEInfo = FnCSEInfo;
 }
 
 void Localizer::getAnalysisUsage(AnalysisUsage &AU) const {
   AU.addRequired<TargetTransformInfoWrapperPass>();
+  AU.addRequired<TargetPassConfig>();
   getSelectionDAGFallbackAnalysisUsage(AU);
+  AU.addRequired<GISelCSEAnalysisWrapperPass>();
+  AU.addPreserved<GISelCSEAnalysisWrapperPass>();
   MachineFunctionPass::getAnalysisUsage(AU);
 }
 
 bool Localizer::isLocalUse(MachineOperand &MOUse, const MachineInstr &Def,
                            MachineBasicBlock *&InsertMBB) {
   MachineInstr &MIUse = *MOUse.getParent();
   InsertMBB = MIUse.getParent();
   if (MIUse.isPHI())
for (auto MOIt = MRI->use_begin(Reg), MOItEnd = MRI->use_end();
         LocalizedMI->getOperand(0).setReg(NewReg);
         NewVRegIt =
             MBBWithLocalDef.insert(std::make_pair(MBBAndReg, NewReg)).first;
         LLVM_DEBUG(dbgs() << "Inserted: " << *LocalizedMI);
       }
       LLVM_DEBUG(dbgs() << "Update use with: " << printReg(NewVRegIt->second)
                         << '\n');
       // Update the user reg.
+      if (CSEInfo)
+        CSEInfo->changedInstr(*MOUse.getParent());
       MOUse.setReg(NewVRegIt->second);
+      if (CSEInfo)
+        CSEInfo->changedInstr(*MOUse.getParent());
     }
   }
   return Changed;
 }
 
 bool Localizer::localizeIntraBlock(LocalizedSetVecT &LocalizedInstrs) {
   bool Changed = false;
 
for (MachineInstr &UseMI : MRI->use_nodbg_instructions(Reg)) {
         Users.insert(&UseMI);
     }
     // If all the users were PHIs then they're not going to be in our block,
     // don't try to move this instruction.
     if (Users.empty())
       continue;
 
     MachineBasicBlock::iterator II(MI);
+    if (UseInstNumbering) {
+      auto &InstNumbering = CSEInfo->getInstNumbering();
+      assert(InstNumbering.isEnabled() && "Inst number is not enabled!");
+      // We have a set of instructions which we know are users within the same
+      // block. We want to find the first user of MI.
+      MachineInstr *Nearest =
+          InstNumbering.getEarliestInstr(Users, /* NoFlush */ true);
+      II = Nearest->getIterator();
+    } else {
       ++II;
       while (II != MBB.end() && !Users.count(&*II))
         ++II;
-
+    }
     LLVM_DEBUG(dbgs() << "Intra-block: moving " << *MI << " before " << *&*II
                       << "\n");
     assert(II != MBB.end() && "Didn't find the user in the MBB");
     MI->removeFromParent();
     MBB.insert(II, MI);
+
+    if (UseInstNumbering)
+      CSEInfo->getInstNumbering().flushAndRenumber();
+
     Changed = true;
   }
   return Changed;
 }
 
 bool Localizer::runOnMachineFunction(MachineFunction &MF) {
   // If the ISel pipeline failed, do not bother running that pass.
   if (MF.getProperties().hasProperty(
           MachineFunctionProperties::Property::FailedISel))
     return false;
 
   // Don't run the pass if the target asked so.
   if (DoNotRunPass(MF))
     return false;
 
   LLVM_DEBUG(dbgs() << "Localize instructions for: " << MF.getName() << '\n');
 
   init(MF);
 
+  GISelObserverWrapper WrapperObserver(CSEInfo);
+  RAIIMFObsDelInstaller Installer(MF, WrapperObserver);
+
   // Keep track of the instructions we localized. We'll do a second pass of
   // intra-block localization to further reduce live ranges.
   LocalizedSetVecT LocalizedInstrs;
 
+  CSEInfo->getInstNumbering().flushAndRenumber();
+
   bool Changed = localizeInterBlock(MF, LocalizedInstrs);
   Changed |= localizeIntraBlock(LocalizedInstrs);
   return Changed;
 }

llvm/lib/CodeGen/GlobalISel/RegBankSelect.cpp

 //==- llvm/CodeGen/GlobalISel/RegBankSelect.cpp - RegBankSelect --*- C++ -*-==//
 //
 // Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
 // See https://llvm.org/LICENSE.txt for license information.
 // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
 //
 //===----------------------------------------------------------------------===//
 /// \file
 /// This file implements the RegBankSelect class.
 //===----------------------------------------------------------------------===//
 
 #include "llvm/CodeGen/GlobalISel/RegBankSelect.h"
 #include "llvm/ADT/PostOrderIterator.h"
 #include "llvm/ADT/STLExtras.h"
 #include "llvm/ADT/SmallVector.h"
+#include "llvm/CodeGen/GlobalISel/CSEMIRBuilder.h"
 #include "llvm/CodeGen/GlobalISel/LegalizerInfo.h"
 #include "llvm/CodeGen/GlobalISel/RegisterBank.h"
 #include "llvm/CodeGen/GlobalISel/RegisterBankInfo.h"
 #include "llvm/CodeGen/GlobalISel/Utils.h"
 #include "llvm/CodeGen/MachineBasicBlock.h"
 #include "llvm/CodeGen/MachineBlockFrequencyInfo.h"
 #include "llvm/CodeGen/MachineBranchProbabilityInfo.h"
 #include "llvm/CodeGen/MachineFunction.h"
void RegBankSelect::init(MachineFunction &MF) {
   TPC = &getAnalysis<TargetPassConfig>();
   if (OptMode != Mode::Fast) {
     MBFI = &getAnalysis<MachineBlockFrequencyInfo>();
     MBPI = &getAnalysis<MachineBranchProbabilityInfo>();
   } else {
     MBFI = nullptr;
     MBPI = nullptr;
   }
-  MIRBuilder.setMF(MF);
+  // Enable CSE and inst numbering.
+  GISelCSEAnalysisWrapper &Wrapper =
+      getAnalysis<GISelCSEAnalysisWrapperPass>().getCSEWrapper();
+  auto *CSEInfo = &Wrapper.get(TPC->getCSEConfig());
+  CSEInfo->getInstNumbering().enable(&MF);
+  MIRBuilder = std::make_unique<CSEMIRBuilder>();
+  MIRBuilder->setMF(MF);
+  MIRBuilder->setCSEInfo(CSEInfo);
+
   MORE = std::make_unique<MachineOptimizationRemarkEmitter>(MF, MBFI);
 }
 
 void RegBankSelect::getAnalysisUsage(AnalysisUsage &AU) const {
   if (OptMode != Mode::Fast) {
     // We could preserve the information from these two analysis but
     // the APIs do not allow to do so yet.
     AU.addRequired<MachineBlockFrequencyInfo>();
     AU.addRequired<MachineBranchProbabilityInfo>();
   }
   AU.addRequired<TargetPassConfig>();
   getSelectionDAGFallbackAnalysisUsage(AU);
+  AU.addRequired<GISelCSEAnalysisWrapperPass>();
+  AU.addPreserved<GISelCSEAnalysisWrapperPass>();
   MachineFunctionPass::getAnalysisUsage(AU);
 }
 
 bool RegBankSelect::assignmentMatch(
     Register Reg, const RegisterBankInfo::ValueMapping &ValMapping,
     bool &OnlyAssign) const {
   // By default we assume we will have to repair something.
   OnlyAssign = false;
if (MO.isDef())
       std::swap(Src, Dst);
 
     assert((RepairPt.getNumInsertPoints() == 1 ||
             Register::isPhysicalRegister(Dst)) &&
            "We are about to create several defs for Dst");
 
     // Build the instruction used to repair, then clone it at the right
     // places. Avoiding buildCopy bypasses the check that Src and Dst have the
-    // same types because the type is a placeholder when this function is called.
-    MI = MIRBuilder.buildInstrNoInsert(TargetOpcode::COPY)
+    // same types because the type is a placeholder when this function is
+    // called.
+    MI = MIRBuilder->buildInstrNoInsert(TargetOpcode::COPY)
              .addDef(Dst)
              .addUse(Src);
     LLVM_DEBUG(dbgs() << "Copy: " << printReg(Src) << " to: " << printReg(Dst)
                << '\n');
   } else {
     // TODO: Support with G_IMPLICIT_DEF + G_INSERT sequence or G_EXTRACT
     // sequence.
     assert(ValMapping.partsAllUniform() && "irregular breakdowns not supported");
 
     LLT RegTy = MRI->getType(MO.getReg());
if (MO.isDef()) {
               "don't understand this value breakdown");
 
           MergeOp = TargetOpcode::G_CONCAT_VECTORS;
         }
       } else
         MergeOp = TargetOpcode::G_MERGE_VALUES;
 
       auto MergeBuilder =
-        MIRBuilder.buildInstrNoInsert(MergeOp)
-        .addDef(MO.getReg());
+          MIRBuilder->buildInstrNoInsert(MergeOp).addDef(MO.getReg());
 
       for (Register SrcReg : NewVRegs)
         MergeBuilder.addUse(SrcReg);
 
       MI = MergeBuilder;
     } else {
       MachineInstrBuilder UnMergeBuilder =
-        MIRBuilder.buildInstrNoInsert(TargetOpcode::G_UNMERGE_VALUES);
+          MIRBuilder->buildInstrNoInsert(TargetOpcode::G_UNMERGE_VALUES);
       for (Register DefReg : NewVRegs)
         UnMergeBuilder.addDef(DefReg);
 
       UnMergeBuilder.addUse(MO.getReg());
       MI = UnMergeBuilder;
     }
   }
 
   if (RepairPt.getNumInsertPoints() != 1)
     report_fatal_error("need testcase to support multiple insertion points");
 
   // TODO:
   // Check if MI is legal. if not, we need to legalize all the
   // instructions we are going to insert.
   std::unique_ptr<MachineInstr *[]> NewInstrs(
       new MachineInstr *[RepairPt.getNumInsertPoints()]);
   bool IsFirst = true;
   unsigned Idx = 0;
   for (const std::unique_ptr<InsertPoint> &InsertPt : RepairPt) {
     MachineInstr *CurMI;
     if (IsFirst)
       CurMI = MI;
     else
-      CurMI = MIRBuilder.getMF().CloneMachineInstr(MI);
+      CurMI = MIRBuilder->getMF().CloneMachineInstr(MI);
     InsertPt->insert(*CurMI);
     NewInstrs[Idx++] = CurMI;
     IsFirst = false;
   }
   // TODO:
   // Legalize NewInstrs if need be.
   return true;
 }
bool RegBankSelect::runOnMachineFunction(MachineFunction &MF) {
 
   LLVM_DEBUG(dbgs() << "Assign register banks for: " << MF.getName() << '\n');
   const Function &F = MF.getFunction();
   Mode SaveOptMode = OptMode;
   if (F.hasOptNone())
     OptMode = Mode::Fast;
   init(MF);
 
+  GISelObserverWrapper WrapperObserver(MIRBuilder->getCSEInfo());
+  RAIIMFObsDelInstaller Installer(MF, WrapperObserver);
+
 #ifndef NDEBUG
   // Check that our input is fully legal: we require the function to have the
   // Legalized property, so it should be.
   // FIXME: This should be in the MachineVerifier.
   if (!DisableGISelLegalityCheck)
     if (const MachineInstr *MI = machineFunctionIsIllegal(MF)) {
       reportGISelFailure(MF, *TPC, *MORE, "gisel-regbankselect",
                          "instruction is not legal", *MI);
       return false;
     }
 #endif
 
   // Walk the function and assign register banks to all operands.
   // Use a RPOT to make sure all registers are assigned before we choose
   // the best mapping of the current instruction.
   ReversePostOrderTraversal<MachineFunction*> RPOT(&MF);
   for (MachineBasicBlock *MBB : RPOT) {
     // Set a sensible insertion point so that subsequent calls to
     // MIRBuilder.
-    MIRBuilder.setMBB(*MBB);
+    MIRBuilder->setMBB(*MBB);
     for (MachineBasicBlock::iterator MII = MBB->begin(), End = MBB->end();
          MII != End;) {
       // MI might be invalidated by the assignment, so move the
       // iterator before hand.
       MachineInstr &MI = *MII++;
 
       // Ignore target-specific post-isel instructions: they should use proper
       // regclasses.
for (MachineBasicBlock::iterator MII = MBB->begin(), End = MBB->end();
 
       // It's possible the mapping changed control flow, and moved the following
       // instruction to a new block, so figure out the new parent.
       if (MII != End) {
         MachineBasicBlock *NextInstBB = MII->getParent();
         if (NextInstBB != MBB) {
           LLVM_DEBUG(dbgs() << "Instruction mapping changed control flow\n");
           MBB = NextInstBB;
-          MIRBuilder.setMBB(*MBB);
+          MIRBuilder->setMBB(*MBB);
           End = MBB->end();
         }
       }
     }
   }
 
   OptMode = SaveOptMode;
   return false;

llvm/lib/Target/AArch64/GISel/AArch64PostLegalizerCombiner.cpp

void AArch64PostLegalizerCombiner::getAnalysisUsage(AnalysisUsage &AU) const {
   AU.setPreservesCFG();
   getSelectionDAGFallbackAnalysisUsage(AU);
   AU.addRequired<GISelKnownBitsAnalysis>();
   AU.addPreserved<GISelKnownBitsAnalysis>();
   if (!IsOptNone) {
     AU.addRequired<MachineDominatorTree>();
     AU.addPreserved<MachineDominatorTree>();
   }
+  AU.addRequired<GISelCSEAnalysisWrapperPass>();
+  AU.addPreserved<GISelCSEAnalysisWrapperPass>();
   MachineFunctionPass::getAnalysisUsage(AU);
 }
 
 AArch64PostLegalizerCombiner::AArch64PostLegalizerCombiner(bool IsOptNone)
     : MachineFunctionPass(ID), IsOptNone(IsOptNone) {
   initializeAArch64PostLegalizerCombinerPass(*PassRegistry::getPassRegistry());
 }
 
 bool AArch64PostLegalizerCombiner::runOnMachineFunction(MachineFunction &MF) {
   if (MF.getProperties().hasProperty(
           MachineFunctionProperties::Property::FailedISel))
     return false;
   assert(MF.getProperties().hasProperty(
              MachineFunctionProperties::Property::Legalized) &&
          "Expected a legalized function?");
-  auto *TPC = &getAnalysis<TargetPassConfig>();
+  auto &TPC = getAnalysis<TargetPassConfig>();
+
+  // Enable CSE and inst numbering.
+  GISelCSEAnalysisWrapper &Wrapper =
+      getAnalysis<GISelCSEAnalysisWrapperPass>().getCSEWrapper();
+  auto *CSEInfo = &Wrapper.get(TPC.getCSEConfig());
+  CSEInfo->getInstNumbering().enable(&MF);
+
   const Function &F = MF.getFunction();
   bool EnableOpt =
       MF.getTarget().getOptLevel() != CodeGenOpt::None && !skipFunction(F);
   GISelKnownBits *KB = &getAnalysis<GISelKnownBitsAnalysis>().get(MF);
   MachineDominatorTree *MDT =
       IsOptNone ? nullptr : &getAnalysis<MachineDominatorTree>();
   AArch64PostLegalizerCombinerInfo PCInfo(EnableOpt, F.hasOptSize(),
                                           F.hasMinSize(), KB, MDT);
-  Combiner C(PCInfo, TPC);
-  return C.combineMachineInstrs(MF, /*CSEInfo*/ nullptr);
+  Combiner C(PCInfo, &TPC);
+  return C.combineMachineInstrs(MF, CSEInfo);
 }
 
 char AArch64PostLegalizerCombiner::ID = 0;
 INITIALIZE_PASS_BEGIN(AArch64PostLegalizerCombiner, DEBUG_TYPE,
                       "Combine AArch64 MachineInstrs after legalization", false,
                       false)
 INITIALIZE_PASS_DEPENDENCY(TargetPassConfig)
 INITIALIZE_PASS_DEPENDENCY(GISelKnownBitsAnalysis)

llvm/lib/Target/AArch64/GISel/AArch64PostLegalizerLowering.cpp

public:
   void getAnalysisUsage(AnalysisUsage &AU) const override;
 };
 } // end anonymous namespace
 
 void AArch64PostLegalizerLowering::getAnalysisUsage(AnalysisUsage &AU) const {
   AU.addRequired<TargetPassConfig>();
   AU.setPreservesCFG();
   getSelectionDAGFallbackAnalysisUsage(AU);
+  AU.addRequired<GISelCSEAnalysisWrapperPass>();
+  AU.addPreserved<GISelCSEAnalysisWrapperPass>();
   MachineFunctionPass::getAnalysisUsage(AU);
 }
 
 AArch64PostLegalizerLowering::AArch64PostLegalizerLowering()
     : MachineFunctionPass(ID) {
   initializeAArch64PostLegalizerLoweringPass(*PassRegistry::getPassRegistry());
 }
 
 bool AArch64PostLegalizerLowering::runOnMachineFunction(MachineFunction &MF) {
   if (MF.getProperties().hasProperty(
           MachineFunctionProperties::Property::FailedISel))
     return false;
   assert(MF.getProperties().hasProperty(
              MachineFunctionProperties::Property::Legalized) &&
          "Expected a legalized function?");
-  auto *TPC = &getAnalysis<TargetPassConfig>();
+  auto &TPC = getAnalysis<TargetPassConfig>();
+
+  // Enable CSE and inst numbering.
+  GISelCSEAnalysisWrapper &Wrapper =
+      getAnalysis<GISelCSEAnalysisWrapperPass>().getCSEWrapper();
+  auto *CSEInfo = &Wrapper.get(TPC.getCSEConfig());
+  CSEInfo->getInstNumbering().enable(&MF);
+
   const Function &F = MF.getFunction();
   AArch64PostLegalizerLoweringInfo PCInfo(F.hasOptSize(), F.hasMinSize());
-  Combiner C(PCInfo, TPC);
-  return C.combineMachineInstrs(MF, /*CSEInfo*/ nullptr);
+  Combiner C(PCInfo, &TPC);
+  return C.combineMachineInstrs(MF, CSEInfo);
 }
 
 char AArch64PostLegalizerLowering::ID = 0;
 INITIALIZE_PASS_BEGIN(AArch64PostLegalizerLowering, DEBUG_TYPE,
                       "Lower AArch64 MachineInstrs after legalization", false,
                       false)
 INITIALIZE_PASS_DEPENDENCY(TargetPassConfig)
 INITIALIZE_PASS_END(AArch64PostLegalizerLowering, DEBUG_TYPE,
                     "Lower AArch64 MachineInstrs after legalization", false,
                     false)
 
 namespace llvm {
 FunctionPass *createAArch64PostLegalizerLowering() {
   return new AArch64PostLegalizerLowering();
 }
 } // end namespace llvm

llvm/lib/Target/AArch64/GISel/AArch64PreLegalizerCombiner.cpp

void AArch64PreLegalizerCombiner::getAnalysisUsage(AnalysisUsage &AU) const {
   AU.setPreservesCFG();
   getSelectionDAGFallbackAnalysisUsage(AU);
   AU.addRequired<GISelKnownBitsAnalysis>();
   AU.addPreserved<GISelKnownBitsAnalysis>();
   if (!IsOptNone) {
     AU.addRequired<MachineDominatorTree>();
     AU.addPreserved<MachineDominatorTree>();
   }
+  AU.addRequired<GISelCSEAnalysisWrapperPass>();
+  AU.addPreserved<GISelCSEAnalysisWrapperPass>();
   MachineFunctionPass::getAnalysisUsage(AU);
 }
 
 AArch64PreLegalizerCombiner::AArch64PreLegalizerCombiner(bool IsOptNone)
     : MachineFunctionPass(ID), IsOptNone(IsOptNone) {
   initializeAArch64PreLegalizerCombinerPass(*PassRegistry::getPassRegistry());
 }
 
 bool AArch64PreLegalizerCombiner::runOnMachineFunction(MachineFunction &MF) {
   if (MF.getProperties().hasProperty(
           MachineFunctionProperties::Property::FailedISel))
     return false;
-  auto *TPC = &getAnalysis<TargetPassConfig>();
+  auto &TPC = getAnalysis<TargetPassConfig>();
+
+  // Enable CSE and inst numbering.
+  GISelCSEAnalysisWrapper &Wrapper =
+      getAnalysis<GISelCSEAnalysisWrapperPass>().getCSEWrapper();
+  auto *CSEInfo = &Wrapper.get(TPC.getCSEConfig());
+  CSEInfo->getInstNumbering().enable(&MF);
+
   const Function &F = MF.getFunction();
   bool EnableOpt =
       MF.getTarget().getOptLevel() != CodeGenOpt::None && !skipFunction(F);
   GISelKnownBits *KB = &getAnalysis<GISelKnownBitsAnalysis>().get(MF);
   MachineDominatorTree *MDT =
       IsOptNone ? nullptr : &getAnalysis<MachineDominatorTree>();
   AArch64PreLegalizerCombinerInfo PCInfo(EnableOpt, F.hasOptSize(),
                                          F.hasMinSize(), KB, MDT);
-  Combiner C(PCInfo, TPC);
-  return C.combineMachineInstrs(MF, /*CSEInfo*/ nullptr);
+  Combiner C(PCInfo, &TPC);
+  return C.combineMachineInstrs(MF, CSEInfo);
 }
 
 char AArch64PreLegalizerCombiner::ID = 0;
 INITIALIZE_PASS_BEGIN(AArch64PreLegalizerCombiner, DEBUG_TYPE,
                       "Combine AArch64 machine instrs before legalization",
                       false, false)
 INITIALIZE_PASS_DEPENDENCY(TargetPassConfig)
 INITIALIZE_PASS_DEPENDENCY(GISelKnownBitsAnalysis)
+INITIALIZE_PASS_DEPENDENCY(GISelCSEAnalysisWrapperPass)
 INITIALIZE_PASS_END(AArch64PreLegalizerCombiner, DEBUG_TYPE,
                     "Combine AArch64 machine instrs before legalization", false,
                     false)
 
 
 namespace llvm {
 FunctionPass *createAArch64PreLegalizerCombiner(bool IsOptNone) {
   return new AArch64PreLegalizerCombiner(IsOptNone);
 }
 } // end namespace llvm

llvm/test/CodeGen/AArch64/GlobalISel/combine-shift-immed-mismatch-crash.mir

 # NOTE: Assertions have been autogenerated by utils/update_mir_test_checks.py
 # RUN: llc -run-pass=aarch64-prelegalizer-combiner -verify-machineinstrs -mtriple aarch64-unknown-unknown %s -o - | FileCheck %s
 ---
 name:            shift_immed_chain_mismatch_size_crash
 alignment:       4
 tracksRegLiveness: true
 liveins:
   - { reg: '$x0' }
 body:             |
   ; CHECK-LABEL: name: shift_immed_chain_mismatch_size_crash
   ; CHECK: bb.0:
   ; CHECK:   successors: %bb.1(0x40000000), %bb.2(0x40000000)
   ; CHECK:   liveins: $x0
-  ; CHECK:   [[DEF:%[0-9]+]]:_(p0) = G_IMPLICIT_DEF
+  ; CHECK:   [[DEF:%[0-9]+]]:_(s1) = G_IMPLICIT_DEF
+  ; CHECK:   [[DEF1:%[0-9]+]]:_(p0) = G_IMPLICIT_DEF
   ; CHECK:   [[C:%[0-9]+]]:_(s32) = G_CONSTANT i32 16
   ; CHECK:   [[C1:%[0-9]+]]:_(s32) = G_CONSTANT i32 9
-  ; CHECK:   [[DEF1:%[0-9]+]]:_(s1) = G_IMPLICIT_DEF
-  ; CHECK:   G_BRCOND [[DEF1]](s1), %bb.2
+  ; CHECK:   G_BRCOND [[DEF]](s1), %bb.2
   ; CHECK:   G_BR %bb.1
   ; CHECK: bb.1:
   ; CHECK:   successors:
   ; CHECK: bb.2:
-  ; CHECK:   [[LOAD:%[0-9]+]]:_(s32) = G_LOAD [[DEF]](p0) :: (load 4 from `i32* undef`, align 8)
+  ; CHECK:   [[LOAD:%[0-9]+]]:_(s32) = G_LOAD [[DEF1]](p0) :: (load 4 from `i32* undef`, align 8)
   ; CHECK:   [[MUL:%[0-9]+]]:_(s32) = nsw G_MUL [[C]], [[LOAD]]
   ; CHECK:   [[MUL1:%[0-9]+]]:_(s32) = nsw G_MUL [[MUL]], [[C1]]
   ; CHECK:   [[C2:%[0-9]+]]:_(s64) = G_CONSTANT i64 2
   ; CHECK:   [[SHL:%[0-9]+]]:_(s32) = G_SHL [[MUL1]], [[C2]](s64)
   ; CHECK:   $w0 = COPY [[SHL]](s32)
   ; CHECK:   RET_ReallyLR implicit $w0
   bb.1:
     liveins: $x0

llvm/test/CodeGen/AArch64/GlobalISel/gisel-commandline-option.ll

 ; RUN:   --debugify-and-strip-all-safe=0 \
 ; RUN:   -debug-pass=Structure %s -o /dev/null 2>&1 -verify-machineinstrs=0 \
 ; RUN:   | FileCheck %s --check-prefix DISABLED
 
 ; ENABLED:       IRTranslator
 ; VERIFY-NEXT:   Verify generated machine code
 ; ENABLED-NEXT:  Analysis for ComputingKnownBits
 ; ENABLED-O1-NEXT:  MachineDominator Tree Construction
+; ENABLED-NEXT:  Analysis containing CSE Info
 ; ENABLED-NEXT:  PreLegalizerCombiner
 ; VERIFY-NEXT:   Verify generated machine code
-; ENABLED-NEXT:  Analysis containing CSE Info
 ; ENABLED-NEXT:  Legalizer
 ; VERIFY-NEXT:   Verify generated machine code
 ; ENABLED:  RegBankSelect
 ; VERIFY-NEXT:   Verify generated machine code
 ; ENABLED-NEXT:  Localizer
 ; VERIFY-O0-NEXT:   Verify generated machine code
 ; ENABLED-NEXT: Analysis for ComputingKnownBits
 ; ENABLED-NEXT:  InstructionSelect

llvm/test/CodeGen/AArch64/O0-pipeline.ll

 ; CHECK-NEXT:     FunctionPass Manager
 ; CHECK-NEXT:       Exception handling preparation
 ; CHECK-NEXT:       Safe Stack instrumentation pass
 ; CHECK-NEXT:       Insert stack protectors
 ; CHECK-NEXT:       Module Verifier
 ; CHECK-NEXT:       Analysis containing CSE Info
 ; CHECK-NEXT:       IRTranslator
 ; CHECK-NEXT:       Analysis for ComputingKnownBits
-; CHECK-NEXT:       AArch64PreLegalizerCombiner
 ; CHECK-NEXT:       Analysis containing CSE Info
+; CHECK-NEXT:       AArch64PreLegalizerCombiner
 ; CHECK-NEXT:       Legalizer
 ; CHECK-NEXT:       AArch64PostLegalizerLowering
 ; CHECK-NEXT:       RegBankSelect
 ; CHECK-NEXT:       Localizer
 ; CHECK-NEXT:       Analysis for ComputingKnownBits
 ; CHECK-NEXT:       InstructionSelect
 ; CHECK-NEXT:       ResetMachineFunction
 ; CHECK-NEXT:       AArch64 Instruction Selection