[SCEV] Track and invalidate ValuesAtScopes users

[nikic]

Summary

This is the same change as D114738, but tracks ValuesAtScopes users with a reverse map, so we can invalidate values without scanning the whole map.

[reames]

LGTM

llvm/include/llvm/Analysis/ScalarEvolution.h

private:
   DenseMap<PHINode *, Constant *> ConstantEvolutionLoopExitValue;
 
   /// This map contains entries for all the expressions that we attempt to
   /// compute getSCEVAtScope information for, which can be expensive in
   /// extreme cases.
   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>>
       LoopDispositions;
 
   struct LoopProperties {
     /// Set to true if the loop contains no instruction that can abnormally exit
     /// the loop (i.e. via throwing an exception, by terminating the thread

llvm/lib/Analysis/ScalarEvolution.cpp

void ScalarEvolution::forgetAllLoops() {
   // - Every llvm::Value has been updated in place to produce a different
   // result.
   BackedgeTakenCounts.clear();
   PredicatedBackedgeTakenCounts.clear();
   LoopPropertiesCache.clear();
   ConstantEvolutionLoopExitValue.clear();
   ValueExprMap.clear();
   ValuesAtScopes.clear();
+  ValuesAtScopesUsers.clear();
   LoopDispositions.clear();
   BlockDispositions.clear();
   UnsignedRanges.clear();
   SignedRanges.clear();
   ExprValueMap.clear();
   HasRecMap.clear();
   MinTrailingZerosCache.clear();
   PredicatedSCEVRewrites.clear();
const SCEV *ScalarEvolution::getSCEVAtScope(const SCEV *V, const Loop *L) {
 
   // Otherwise compute it.
   const SCEV *C = computeSCEVAtScope(V, L);
   for (auto &LS : reverse(ValuesAtScopes[V]))
     if (LS.first == L) {
       LS.second = C;
       break;
     }
+
+  if (!isa<SCEVConstant>(C))
+    ValuesAtScopesUsers[C].push_back({L, V});
   return C;
 }
 
 /// This builds up a Constant using the ConstantExpr interface.  That way, we
 /// will return Constants for objects which aren't represented by a
 /// SCEVConstant, because SCEVConstant is restricted to ConstantInt.
 /// Returns NULL if the SCEV isn't representable as a Constant.
 static Constant *BuildConstantFromSCEV(const SCEV *V) {
: F(Arg.F), HasGuards(Arg.HasGuards), TLI(Arg.TLI), AC(Arg.AC), DT(Arg.DT),
       PendingMerges(std::move(Arg.PendingMerges)),
       MinTrailingZerosCache(std::move(Arg.MinTrailingZerosCache)),
       BackedgeTakenCounts(std::move(Arg.BackedgeTakenCounts)),
       PredicatedBackedgeTakenCounts(
           std::move(Arg.PredicatedBackedgeTakenCounts)),
       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)),
       SCEVUsers(std::move(Arg.SCEVUsers)),
       UnsignedRanges(std::move(Arg.UnsignedRanges)),
       SignedRanges(std::move(Arg.SignedRanges)),
       UniqueSCEVs(std::move(Arg.UniqueSCEVs)),
       UniquePreds(std::move(Arg.UniquePreds)),
auto RemoveSCEVFromBackedgeMap = [&ToForget](
         }
   };
 
   RemoveSCEVFromBackedgeMap(BackedgeTakenCounts);
   RemoveSCEVFromBackedgeMap(PredicatedBackedgeTakenCounts);
 }
 
 void ScalarEvolution::forgetMemoizedResultsImpl(const SCEV *S) {
-  ValuesAtScopes.erase(S);
   LoopDispositions.erase(S);
   BlockDispositions.erase(S);
   UnsignedRanges.erase(S);
   SignedRanges.erase(S);
   HasRecMap.erase(S);
   MinTrailingZerosCache.erase(S);
 
   auto ExprIt = ExprValueMap.find(S);
   if (ExprIt != ExprValueMap.end()) {
     for (auto &ValueAndOffset : ExprIt->second) {
       if (ValueAndOffset.second == nullptr) {
         auto ValueIt = ValueExprMap.find_as(ValueAndOffset.first);
         if (ValueIt != ValueExprMap.end())
           ValueExprMap.erase(ValueIt);
       }
     }
     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
 ScalarEvolution::getUsedLoops(const SCEV *S,
                               SmallPtrSetImpl<const Loop *> &LoopsUsed) {
   struct FindUsedLoops {
     FindUsedLoops(SmallPtrSetImpl<const Loop *> &LoopsUsed)
         : LoopsUsed(LoopsUsed) {}

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

+; 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
+}