[SCEV] Support clearing Block/LoopDispositions for a single value.

Extend forgetBlockAndLoopDisposition to allow clearing information for a
single value. This can be useful when only a single value is changed,
e.g. because the instruction is moved.

We also need to clear the cached values for all SCEV users, because they
may depend on the starting value's disposition.

Reviewed By: nikic

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

llvm/include/llvm/Analysis/ScalarEvolution.h

@@ -944,7 +944,7 @@ class ScalarEvolution {
   ///
   /// We don't have a way to invalidate per-loop/per-block dispositions. Clear
   /// and recompute is simpler.
-  void forgetBlockAndLoopDispositions();
+  void forgetBlockAndLoopDispositions(Value *V = nullptr);
 
   /// Determine the minimum number of zero bits that S is guaranteed to end in
   /// (at every loop iteration).  It is, at the same time, the minimum number

llvm/lib/Analysis/ScalarEvolution.cpp

@@ -8384,9 +8384,36 @@ void ScalarEvolution::forgetValue(Value *V) {
 
 void ScalarEvolution::forgetLoopDispositions() { LoopDispositions.clear(); }
 
-void ScalarEvolution::forgetBlockAndLoopDispositions() {
-  BlockDispositions.clear();
-  LoopDispositions.clear();
+void ScalarEvolution::forgetBlockAndLoopDispositions(Value *V) {
+  // Unless a specific value is passed to invalidation, completely clear both
+  // caches.
+  if (!V) {
+    BlockDispositions.clear();
+    LoopDispositions.clear();
+    return;
+  }
+
+  const SCEV *S = getExistingSCEV(V);
+  if (!S)
+    return;
+
+  // Invalidate the block and loop dispositions cached for S. Dispositions of
+  // S's users may change if S's disposition changes (i.e. a user may change to
+  // loop-invariant, if S changes to loop invariant), so also invalidate
+  // dispositions of S's users recursively.
+  SmallVector<const SCEV *, 8> Worklist = {S};
+  SmallPtrSet<const SCEV *, 8> Seen = {S};
+  while (!Worklist.empty()) {
+    const SCEV *Curr = Worklist.pop_back_val();
+    if (!LoopDispositions.erase(Curr) && !BlockDispositions.erase(S))
+      continue;
+
+    auto Users = SCEVUsers.find(Curr);
+    if (Users != SCEVUsers.end())
+      for (const auto *User : Users->second)
+        if (Seen.insert(User).second)
+          Worklist.push_back(User);
+  }
 }
 
 /// Get the exact loop backedge taken count considering all loop exits. A

llvm/lib/Transforms/Scalar/LoopDeletion.cpp

@@ -90,22 +90,21 @@ static bool isLoopDead(Loop *L, ScalarEvolution &SE,
         break;
 
       if (Instruction *I = dyn_cast<Instruction>(incoming)) {
-        if (!L->makeLoopInvariant(I, Changed, Preheader->getTerminator())) {
+        bool InstrMoved = false;
+        if (!L->makeLoopInvariant(I, InstrMoved, Preheader->getTerminator())) {
           AllEntriesInvariant = false;
           break;
         }
-        if (Changed) {
+        Changed |= InstrMoved;
+        if (InstrMoved) {
           // Moving I to a different location may change its block disposition,
           // so invalidate its SCEV.
-          SE.forgetValue(I);
+          SE.forgetBlockAndLoopDispositions(I);
         }
       }
     }
   }
 
-  if (Changed)
-    SE.forgetLoopDispositions();
-
   if (!AllEntriesInvariant || !AllOutgoingValuesSame)
     return false;
 

llvm/lib/Transforms/Scalar/LoopSink.cpp

@@ -312,12 +312,13 @@ static bool sinkLoopInvariantInstructions(Loop &L, AAResults &AA, LoopInfo &LI,
     if (!canSinkOrHoistInst(I, &AA, &DT, &L, MSSAU, false, LICMFlags))
       continue;
     if (sinkInstruction(L, I, ColdLoopBBs, LoopBlockNumber, LI, DT, BFI,
-                        &MSSAU))
+                        &MSSAU)) {
       Changed = true;
+      if (SE)
+        SE->forgetBlockAndLoopDispositions(&I);
+    }
   }
 
-  if (Changed && SE)
-    SE->forgetLoopDispositions();
   return Changed;
 }
 

llvm/lib/Transforms/Utils/LoopSimplify.cpp

@@ -647,20 +647,22 @@ static bool simplifyOneLoop(Loop *L, SmallVectorImpl<Loop *> &Worklist,
         Instruction *Inst = &*I++;
         if (Inst == CI)
           continue;
+        bool InstInvariant = false;
         if (!L->makeLoopInvariant(
-                Inst, AnyInvariant,
+                Inst, InstInvariant,
                 Preheader ? Preheader->getTerminator() : nullptr, MSSAU)) {
           AllInvariant = false;
           break;
         }
+        if (InstInvariant && SE) {
+          // The loop disposition of all SCEV expressions that depend on any
+          // hoisted values have also changed.
+          SE->forgetBlockAndLoopDispositions(Inst);
+        }
+        AnyInvariant |= InstInvariant;
       }
-      if (AnyInvariant) {
+      if (AnyInvariant)
         Changed = true;
-        // The loop disposition of all SCEV expressions that depend on any
-        // hoisted values have also changed.
-        if (SE)
-          SE->forgetLoopDispositions();
-      }
       if (!AllInvariant) continue;
 
       // The block has now been cleared of all instructions except for