Add order bit to instructions to lazily track dominance queries. This improves the performance of dominance queries, which are used quite often within the compiler(especially within the verifier).

This reduced the execution time of a few internal tests from ~2 minutes to ~4 seconds.

PiperOrigin-RevId: 230819723
diff --git a/include/mlir/IR/Block.h b/include/mlir/IR/Block.h
index 2d22efe..1b14d92 100644
--- a/include/mlir/IR/Block.h
+++ b/include/mlir/IR/Block.h
@@ -51,7 +51,7 @@
   }
 
   /// Blocks are maintained in a list of BlockList type.
-  BlockList *getParent() const { return parent; }
+  BlockList *getParent() const { return parentValidInstOrderPair.getPointer(); }
 
   /// Returns the closest surrounding instruction that contains this block or
   /// nullptr if this is a top-level block.
@@ -171,6 +171,24 @@
   /// they are to be deleted.
   void dropAllReferences();
 
+  /// Returns true if the ordering of the child instructions is valid, false
+  /// otherwise.
+  bool isInstOrderValid() const { return parentValidInstOrderPair.getInt(); }
+
+  /// Invalidates the current ordering of instructions.
+  void invalidateInstOrder() {
+    // Validate the current ordering.
+    assert(!verifyInstOrder());
+    parentValidInstOrderPair.setInt(false);
+  }
+
+  /// Verifies the current ordering of child instructions matches the
+  /// validInstOrder flag. Returns false if the order is valid, true otherwise.
+  bool verifyInstOrder() const;
+
+  /// Recomputes the ordering of child instructions within the block.
+  void recomputeInstOrder();
+
   //===--------------------------------------------------------------------===//
   // Terminator management
   //===--------------------------------------------------------------------===//
@@ -264,8 +282,10 @@
   void printAsOperand(raw_ostream &os, bool printType = true);
 
 private:
-  /// This is the parent object that owns this block.
-  BlockList *parent = nullptr;
+  /// Pair of the parent object that owns this block and a bit that signifies if
+  /// the instructions within this block have a valid ordering.
+  llvm::PointerIntPair<BlockList *, /*IntBits=*/1, bool>
+      parentValidInstOrderPair;
 
   /// This is the list of instructions in the block.
   InstListType instructions;
diff --git a/include/mlir/IR/Instruction.h b/include/mlir/IR/Instruction.h
index 9f4f592..6a296b7 100644
--- a/include/mlir/IR/Instruction.h
+++ b/include/mlir/IR/Instruction.h
@@ -121,7 +121,14 @@
   /// right before `iterator` in the specified block.
   void moveBefore(Block *block, llvm::iplist<Instruction>::iterator iterator);
 
-  // Returns whether the Instruction is a terminator.
+  /// Given an instruction 'other' that is within the same parent block, return
+  /// whether the current instruction is before 'other' in the instruction list
+  /// of the parent block.
+  /// Note: This function has an average complexity of O(1), but worst case may
+  /// take O(N) where N is the number of instructions within the parent block.
+  bool isBeforeInBlock(const Instruction *other) const;
+
+  /// Returns whether the instruction is a terminator.
   bool isTerminator() const;
 
   void print(raw_ostream &os) const;
@@ -199,8 +206,21 @@
   /// The instruction block that containts this instruction.
   Block *block = nullptr;
 
+  /// Relative order of this instruction in its parent block. Used for
+  /// O(1) local dominance checks between instructions.
+  mutable unsigned orderIndex = 0;
+
+  // Provide a 'getParent' method for ilist_node_with_parent methods.
+  const Block *getParent() const { return getBlock(); }
+
   // allow ilist_traits access to 'block' field.
   friend struct llvm::ilist_traits<Instruction>;
+
+  // allow block to access the 'orderIndex' field.
+  friend class Block;
+
+  // allow ilist_node_with_parent to access the 'getParent' method.
+  friend class llvm::ilist_node_with_parent<Instruction, Block>;
 };
 
 inline raw_ostream &operator<<(raw_ostream &os, const Instruction &inst) {
diff --git a/lib/Analysis/Dominance.cpp b/lib/Analysis/Dominance.cpp
index a5f3be4..b98efa7 100644
--- a/lib/Analysis/Dominance.cpp
+++ b/lib/Analysis/Dominance.cpp
@@ -79,31 +79,9 @@
   auto *aBlock = a->getBlock();
   auto *bBlock = b->getBlock();
 
-  // If the blocks are the same, then we do a linear scan.
-  if (aBlock == bBlock) {
-    // If a/b are the same, then they don't properly dominate each other.
-    if (a == b)
-      return false;
-
-    // If one is a terminator, then the other dominates it.
-    if (a->isTerminator())
-      return false;
-
-    if (b->isTerminator())
-      return true;
-
-    // Otherwise, do a linear scan to determine whether B comes after A.
-    // TODO: This is an O(n) scan that can be bad for very large blocks.
-    auto aIter = Block::const_iterator(a);
-    auto bIter = Block::const_iterator(b);
-    auto fIter = aBlock->begin();
-    while (bIter != fIter) {
-      --bIter;
-      if (aIter == bIter)
-        return true;
-    }
-    return false;
-  }
+  // If the blocks are the same, then check if b is before a in the block.
+  if (aBlock == bBlock)
+    return a->isBeforeInBlock(b);
 
   // If the blocks are different, but in the same function-level block list,
   // then a standard block dominance query is sufficient.
@@ -141,34 +119,12 @@
 /// Returns true if statement 'a' properly postdominates statement b.
 bool PostDominanceInfo::properlyPostDominates(const Instruction *a,
                                               const Instruction *b) {
-  // If a/b are the same, then they don't properly dominate each other.
-  if (a == b)
-    return false;
-
   auto *aBlock = a->getBlock();
   auto *bBlock = b->getBlock();
 
-  // If the blocks are the same, then we do a linear scan.
-  if (aBlock == bBlock) {
-    // If one is a terminator, it postdominates the other.
-    if (a->isTerminator())
-      return true;
-
-    if (b->isTerminator())
-      return false;
-
-    // Otherwise, do a linear scan to determine whether A comes after B.
-    // TODO: This is an O(n) scan that can be bad for very large blocks.
-    auto aIter = Block::const_iterator(a);
-    auto bIter = Block::const_iterator(b);
-    auto fIter = bBlock->begin();
-    while (aIter != fIter) {
-      --aIter;
-      if (aIter == bIter)
-        return true;
-    }
-    return false;
-  }
+  // If the blocks are the same, check if b is before a in the block.
+  if (aBlock == bBlock)
+    return b->isBeforeInBlock(a);
 
   // If the blocks are different, but in the same function-level block list,
   // then a standard block dominance query is sufficient.
diff --git a/lib/IR/Block.cpp b/lib/IR/Block.cpp
index 3d3150a..b6095bf 100644
--- a/lib/IR/Block.cpp
+++ b/lib/IR/Block.cpp
@@ -19,6 +19,7 @@
 #include "mlir/IR/BlockAndValueMapping.h"
 #include "mlir/IR/Builders.h"
 #include "mlir/IR/InstVisitor.h"
+#include "mlir/IR/Instruction.h"
 using namespace mlir;
 
 //===----------------------------------------------------------------------===//
@@ -37,6 +38,7 @@
 //===----------------------------------------------------------------------===//
 
 Block::~Block() {
+  assert(!verifyInstOrder() && "Expected valid instruction ordering.");
   clear();
 
   llvm::DeleteContainerPointers(arguments);
@@ -45,7 +47,7 @@
 /// Returns the closest surrounding instruction that contains this block or
 /// nullptr if this is a top-level instruction block.
 Instruction *Block::getContainingInst() {
-  return parent ? parent->getContainingInst() : nullptr;
+  return getParent() ? getParent()->getContainingInst() : nullptr;
 }
 
 Function *Block::getFunction() {
@@ -97,6 +99,39 @@
     i.dropAllReferences();
 }
 
+/// Verifies the current ordering of child instructions. Returns false if the
+/// order is valid, true otherwise.
+bool Block::verifyInstOrder() const {
+  // The order is already known to be invalid.
+  if (!isInstOrderValid())
+    return false;
+  // The order is valid if there are less than 2 instructions.
+  if (instructions.empty() ||
+      std::next(instructions.begin()) == instructions.end())
+    return false;
+
+  const Instruction *prev = nullptr;
+  for (auto &i : *this) {
+    // The previous instruction must have a smaller order index than the next as
+    // it appears earlier in the list.
+    if (prev && prev->orderIndex >= i.orderIndex)
+      return true;
+    prev = &i;
+  }
+  return false;
+}
+
+/// Recomputes the ordering of child instructions within the block.
+void Block::recomputeInstOrder() {
+  parentValidInstOrderPair.setInt(true);
+
+  // TODO(riverriddle) Have non-congruent indices to reduce the number of times
+  // an insert invalidates the list.
+  unsigned orderIndex = 0;
+  for (auto &inst : *this)
+    inst.orderIndex = orderIndex++;
+}
+
 //===----------------------------------------------------------------------===//
 // Argument list management.
 //===----------------------------------------------------------------------===//
@@ -287,15 +322,15 @@
 /// This is a trait method invoked when a basic block is added to a function.
 /// We keep the function pointer up to date.
 void llvm::ilist_traits<::mlir::Block>::addNodeToList(Block *block) {
-  assert(!block->parent && "already in a function!");
-  block->parent = getContainingBlockList();
+  assert(!block->getParent() && "already in a function!");
+  block->parentValidInstOrderPair.setPointer(getContainingBlockList());
 }
 
 /// This is a trait method invoked when an instruction is removed from a
 /// function.  We keep the function pointer up to date.
 void llvm::ilist_traits<::mlir::Block>::removeNodeFromList(Block *block) {
-  assert(block->parent && "not already in a function!");
-  block->parent = nullptr;
+  assert(block->getParent() && "not already in a function!");
+  block->parentValidInstOrderPair.setPointer(nullptr);
 }
 
 /// This is a trait method invoked when an instruction is moved from one block
@@ -310,5 +345,5 @@
 
   // Update the 'parent' member of each Block.
   for (; first != last; ++first)
-    first->parent = curParent;
+    first->parentValidInstOrderPair.setPointer(curParent);
 }
diff --git a/lib/IR/Instruction.cpp b/lib/IR/Instruction.cpp
index c75e141..c491d10 100644
--- a/lib/IR/Instruction.cpp
+++ b/lib/IR/Instruction.cpp
@@ -179,7 +179,22 @@
   return getContext()->emitError(getLoc(), message);
 }
 
-// Returns whether the Instruction is a terminator.
+/// Given an instruction 'other' that is within the same parent block, return
+/// whether the current instruction is before 'other' in the instruction list
+/// of the parent block.
+/// Note: This function has an average complexity of O(1), but worst case may
+/// take O(N) where N is the number of instructions within the parent block.
+bool Instruction::isBeforeInBlock(const Instruction *other) const {
+  assert(block && "Instructions without parent blocks have no order.");
+  assert(other && other->block == block &&
+         "Expected other instruction to have the same parent block.");
+  // Recompute the parent ordering if necessary.
+  if (!block->isInstOrderValid())
+    block->recomputeInstOrder();
+  return orderIndex < other->orderIndex;
+}
+
+/// Returns whether the Instruction is a terminator.
 bool Instruction::isTerminator() const {
   if (auto *op = dyn_cast<OperationInst>(this))
     return op->isTerminator();
@@ -205,6 +220,9 @@
 void llvm::ilist_traits<::mlir::Instruction>::addNodeToList(Instruction *inst) {
   assert(!inst->getBlock() && "already in a instruction block!");
   inst->block = getContainingBlock();
+
+  // Invalidate the block ordering.
+  inst->block->invalidateInstOrder();
 }
 
 /// This is a trait method invoked when a instruction is removed from a block.
@@ -220,9 +238,13 @@
 void llvm::ilist_traits<::mlir::Instruction>::transferNodesFromList(
     ilist_traits<Instruction> &otherList, inst_iterator first,
     inst_iterator last) {
+  Block *curParent = getContainingBlock();
+
+  // Invalidate the ordering of the parent block.
+  curParent->invalidateInstOrder();
+
   // If we are transferring instructions within the same block, the block
   // pointer doesn't need to be updated.
-  Block *curParent = getContainingBlock();
   if (curParent == otherList.getContainingBlock())
     return;