[IPT] Restructure cache to allow lazy update following invalidation [NFC]

This change restructures the cache used in IPT to point not to the first special instruction, but to the first instruction which *could* be special. That is, the cached reference is always equal to the first special, or comes before it in the block.

This avoids expensive block scans when we are removing special instructions from the beginning of the block. At the moment, this case is not heavily used, though it does trigger in GVN when doing CSE of calls. The main motivation was a change I'm no longer planning to move forward with, but the cache optimization seemed worthwhile as a minor perf win at low cost.

Differential Revision: https://reviews.llvm.org/D111768

Author: preames

llvm/include/llvm/Analysis/InstructionPrecedenceTracking.h

@@ -28,14 +28,13 @@ class BasicBlock;
 class Instruction;
 
 class InstructionPrecedenceTracking {
-  // Maps a block to the topmost special instruction in it. If the value is
-  // nullptr, it means that it is known that this block does not contain any
-  // special instructions.
+  // Maps a block to the topmost instruction which *might* be a special
+  // instruction in it. This value is lazily updated on query to point to the
+  // topmost special instruction, but we allow it to point before that for
+  // efficient invalidation.  If the value is nullptr, it means that it is
+  // known that this block does not contain any special instructions.
   DenseMap<const BasicBlock *, const Instruction *> FirstSpecialInsts;
 
-  // Fills information about the given block's special instructions.
-  void fill(const BasicBlock *BB);
-
 #ifndef NDEBUG
   /// Asserts that the cached info for \p BB is up-to-date. This helps to catch
   /// the usage error of accessing a block without properly invalidating after a

llvm/lib/Analysis/InstructionPrecedenceTracking.cpp

@@ -47,11 +47,29 @@ const Instruction *InstructionPrecedenceTracking::getFirstSpecialInstruction(
     validate(BB);
 #endif
 
-  if (FirstSpecialInsts.find(BB) == FirstSpecialInsts.end()) {
-    fill(BB);
-    assert(FirstSpecialInsts.find(BB) != FirstSpecialInsts.end() && "Must be!");
-  }
-  return FirstSpecialInsts[BB];
+  if (!FirstSpecialInsts.count(BB))
+    // Seed the lazy scan
+    FirstSpecialInsts[BB] = &*BB->begin();
+
+  auto *CurI = FirstSpecialInsts[BB];
+  if (!CurI || isSpecialInstruction(CurI))
+    // We found a cached definite result
+    return CurI;
+
+  // Otherwise, scan forward until we find a definite result, then cache that.
+  auto *Res = [&]() -> const Instruction * {
+    for (auto &I : make_range(CurI->getIterator(), BB->end())) {
+      NumInstScanned++;
+      if (isSpecialInstruction(&I))
+        // Found next special instruction
+        return &I;
+    }
+    // Mark this block as having no special instructions.
+    return nullptr;
+  }();
+
+  FirstSpecialInsts[BB] = Res;
+  return Res;
 }
 
 bool InstructionPrecedenceTracking::hasSpecialInstructions(
@@ -66,33 +84,20 @@ bool InstructionPrecedenceTracking::isPreceededBySpecialInstruction(
   return MaybeFirstSpecial && MaybeFirstSpecial->comesBefore(Insn);
 }
 
-void InstructionPrecedenceTracking::fill(const BasicBlock *BB) {
-  FirstSpecialInsts.erase(BB);
-  for (auto &I : *BB) {
-    NumInstScanned++;
-    if (isSpecialInstruction(&I)) {
-      FirstSpecialInsts[BB] = &I;
-      return;
-    }
-  }
-
-  // Mark this block as having no special instructions.
-  FirstSpecialInsts[BB] = nullptr;
-}
-
 #ifndef NDEBUG
 void InstructionPrecedenceTracking::validate(const BasicBlock *BB) const {
   auto It = FirstSpecialInsts.find(BB);
   // Bail if we don't have anything cached for this block.
   if (It == FirstSpecialInsts.end())
     return;
 
-  for (const Instruction &Insn : *BB)
-    if (isSpecialInstruction(&Insn)) {
-      assert(It->second == &Insn &&
-             "Cached first special instruction is wrong!");
+  for (const Instruction &I : *BB) {
+    if (&I == It->second)
+      // No special instruction before cached result
       return;
-    }
+    assert(!isSpecialInstruction(&I) &&
+           "Cached first special instruction is wrong!");
+  }
 
   assert(It->second == nullptr &&
          "Block is marked as having special instructions but in fact it  has "
@@ -115,8 +120,12 @@ void InstructionPrecedenceTracking::insertInstructionTo(const Instruction *Inst,
 void InstructionPrecedenceTracking::removeInstruction(const Instruction *Inst) {
   auto *BB = Inst->getParent();
   assert(BB && "must be called before instruction is actually removed");
-  if (FirstSpecialInsts.count(BB) && FirstSpecialInsts[BB] == Inst)
-    FirstSpecialInsts.erase(BB);
+  if (FirstSpecialInsts.count(BB) && FirstSpecialInsts[BB] == Inst) {
+    if (Inst->isTerminator())
+      FirstSpecialInsts[BB] = nullptr;
+    else
+      FirstSpecialInsts[BB] = &*std::next(Inst->getIterator());
+  }
 }
 
 void InstructionPrecedenceTracking::removeUsersOf(const Instruction *Inst) {