[SCEV] Track and invalidate ValuesAtScopes users

ValuesAtScopes maps a SCEV and a Loop to another SCEV. While we
invalidate entries if the left-hand SCEV is invalidated, we
currently don't do this for the right-hand SCEV. Fix this by
tracking users in a reverse map and using it for invalidation.

This is conceptually the same change as D114738, but using the
reverse map to avoid performance issues.

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

llvm/include/llvm/Analysis/ScalarEvolution.h

@@ -1492,6 +1492,11 @@ private:
   DenseMap<const SCEV *, SmallVector<std::pair<const Loop *, const SCEV *>, 2>>
       ValuesAtScopes;
 
+  /// Reverse map for invalidation purposes: Stores of which SCEV and which
+  /// loop this is the value-at-scope of.
+  DenseMap<const SCEV *, SmallVector<std::pair<const Loop *, const SCEV *>, 2>>
+      ValuesAtScopesUsers;
+
   /// Memoized computeLoopDisposition results.
   DenseMap<const SCEV *,
            SmallVector<PointerIntPair<const Loop *, 2, LoopDisposition>, 2>>

llvm/lib/Analysis/ScalarEvolution.cpp

@@ -7607,6 +7607,7 @@ void ScalarEvolution::forgetAllLoops() {
   ConstantEvolutionLoopExitValue.clear();
   ValueExprMap.clear();
   ValuesAtScopes.clear();
+  ValuesAtScopesUsers.clear();
   LoopDispositions.clear();
   BlockDispositions.clear();
   UnsignedRanges.clear();
@@ -8847,6 +8848,9 @@ const SCEV *ScalarEvolution::getSCEVAtScope(const SCEV *V, const Loop *L) {
       LS.second = C;
       break;
     }
+
+  if (!isa<SCEVConstant>(C))
+    ValuesAtScopesUsers[C].push_back({L, V});
   return C;
 }
 
@@ -12465,6 +12469,7 @@ ScalarEvolution::ScalarEvolution(ScalarEvolution &&Arg)
       ConstantEvolutionLoopExitValue(
           std::move(Arg.ConstantEvolutionLoopExitValue)),
       ValuesAtScopes(std::move(Arg.ValuesAtScopes)),
+      ValuesAtScopesUsers(std::move(Arg.ValuesAtScopesUsers)),
       LoopDispositions(std::move(Arg.LoopDispositions)),
       LoopPropertiesCache(std::move(Arg.LoopPropertiesCache)),
       BlockDispositions(std::move(Arg.BlockDispositions)),
@@ -12919,7 +12924,6 @@ void ScalarEvolution::forgetMemoizedResults(ArrayRef<const SCEV *> SCEVs) {
 }
 
 void ScalarEvolution::forgetMemoizedResultsImpl(const SCEV *S) {
-  ValuesAtScopes.erase(S);
   LoopDispositions.erase(S);
   BlockDispositions.erase(S);
   UnsignedRanges.erase(S);
@@ -12938,6 +12942,22 @@ void ScalarEvolution::forgetMemoizedResultsImpl(const SCEV *S) {
     }
     ExprValueMap.erase(ExprIt);
   }
+
+  auto ScopeIt = ValuesAtScopes.find(S);
+  if (ScopeIt != ValuesAtScopes.end()) {
+    for (const auto &Pair : ScopeIt->second)
+      if (!isa_and_nonnull<SCEVConstant>(Pair.second))
+        erase_value(ValuesAtScopesUsers[Pair.second],
+                    std::make_pair(Pair.first, S));
+    ValuesAtScopes.erase(ScopeIt);
+  }
+
+  auto ScopeUserIt = ValuesAtScopesUsers.find(S);
+  if (ScopeUserIt != ValuesAtScopesUsers.end()) {
+    for (const auto &Pair : ScopeUserIt->second)
+      erase_value(ValuesAtScopes[Pair.second], std::make_pair(Pair.first, S));
+    ValuesAtScopesUsers.erase(ScopeUserIt);
+  }
 }
 
 void

llvm/test/Transforms/IndVarSimplify/bbi-63564.ll

@@ -0,0 +1,58 @@
+; NOTE: Assertions have been autogenerated by utils/update_test_checks.py
+; RUN: opt -passes='loop-mssa(indvars,indvars)' -S < %s | FileCheck %s
+
+target triple = "x86_64-unknown-linux-gnu"
+
+@c = external global i16, align 1
+@a = external global i16, align 1
+
+; When we delete %c.promoted, we need to ensure we clear the getSCEVAtScope
+; cache or we'll crash trying to access loop disposition of the exit value
+; for the remaining IV.
+define void @f() {
+; CHECK-LABEL: @f(
+; CHECK-NEXT:  entry:
+; CHECK-NEXT:    br label [[FOR_COND:%.*]]
+; CHECK:       for.cond:
+; CHECK-NEXT:    br i1 false, label [[FOR_BODY:%.*]], label [[FOR_END4:%.*]]
+; CHECK:       for.body:
+; CHECK-NEXT:    br label [[FOR_BODY2:%.*]]
+; CHECK:       for.body2:
+; CHECK-NEXT:    [[INC2:%.*]] = phi i16 [ undef, [[FOR_BODY]] ], [ [[INC:%.*]], [[FOR_BODY2]] ]
+; CHECK-NEXT:    [[INC]] = add nsw i16 [[INC2]], 1
+; CHECK-NEXT:    store i16 [[INC]], i16* undef, align 1
+; CHECK-NEXT:    br i1 true, label [[FOR_BODY2]], label [[CRIT_EDGE:%.*]]
+; CHECK:       crit_edge:
+; CHECK-NEXT:    [[INC_LCSSA:%.*]] = phi i16 [ [[INC]], [[FOR_BODY2]] ]
+; CHECK-NEXT:    store i16 [[INC_LCSSA]], i16* @a, align 1
+; CHECK-NEXT:    unreachable
+; CHECK:       for.end4:
+; CHECK-NEXT:    ret void
+;
+entry:
+  br label %for.cond
+
+for.cond:                                         ; preds = %entry
+  br i1 false, label %for.body, label %for.end4
+
+for.body:                                         ; preds = %for.cond
+  %c.promoted = load i16, i16* @c, align 1
+  br label %for.body2
+
+for.body2:                                        ; preds = %for.body2, %for.body
+  %inc33 = phi i16 [ %c.promoted, %for.body ], [ %inc3, %for.body2 ]
+  %inc2 = phi i16 [ undef, %for.body ], [ %inc, %for.body2 ]
+  %inc = add nsw i16 %inc2, 1
+  store i16 %inc, i16* undef, align 1
+  %inc3 = add nsw i16 %inc33, 1
+  %tobool = icmp ne i16 %inc3, 0
+  br i1 %tobool, label %for.body2, label %crit_edge
+
+crit_edge:                                        ; preds = %for.body2
+  %inc.lcssa = phi i16 [ %inc, %for.body2 ]
+  store i16 %inc.lcssa, i16* @a, align 1
+  unreachable
+
+for.end4:                                         ; preds = %for.cond
+  ret void
+}