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

This reverts commit 92f698f01f.

The updated version of the patch includes handling for non-SCEVable
types. A test case has been added in ec86e9a99b.

llvm/include/llvm/Analysis/ScalarEvolution.h

@@ -944,7 +944,7 @@ public:
   ///
   /// 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,39 @@ 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;
+  }
+
+  if (!isSCEVable(V->getType()))
+    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 @@ ReprocessLoop:
         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