[SCEV] Properly clean up duplicated FoldCacheUser ID entries.

The current code did not properly handled duplicated FoldCacheUser ID entries when overwriting an existing entry in the FoldCache. This triggered verification failures reported by @uabelho and #59721. The patch fixes that by removing stale IDs when overwriting an existing entry in the cache. Fixes #59721.
llvm · Dec 28, 2022 · a564048 · a564048
1 parent 58906e4
commit a564048
Show file tree

Hide file tree

Showing 2 changed files with 31 additions and 13 deletions.
diff --git a/llvm/lib/Analysis/ScalarEvolution.cpp b/llvm/lib/Analysis/ScalarEvolution.cpp
@@ -1624,6 +1624,29 @@ static APInt extractConstantWithoutWrapping(ScalarEvolution &SE,
   return APInt(BitWidth, 0);
 }
 
+static void insertFoldCacheEntry(
+    const ScalarEvolution::FoldID &ID, const SCEV *S,
+    DenseMap<ScalarEvolution::FoldID, const SCEV *> &FoldCache,
+    DenseMap<const SCEV *, SmallVector<ScalarEvolution::FoldID, 2>>
+        &FoldCacheUser) {
+  auto I = FoldCache.insert({ID, S});
+  if (!I.second) {
+    // Remove FoldCacheUser entry for ID when replacing an existing FoldCache
+    // entry.
+    auto &UserIDs = FoldCacheUser[I.first->second];
+    assert(count(UserIDs, ID) == 1 && "unexpected duplicates in UserIDs");
+    for (unsigned I = 0; I != UserIDs.size(); ++I)
+      if (UserIDs[I] == ID) {
+        std::swap(UserIDs[I], UserIDs.back());
+        break;
+      }
+    UserIDs.pop_back();
+    I.first->second = S;
+  }
+  auto R = FoldCacheUser.insert({S, {}});
+  R.first->second.push_back(ID);
+}
+
 const SCEV *
 ScalarEvolution::getZeroExtendExpr(const SCEV *Op, Type *Ty, unsigned Depth) {
   assert(getTypeSizeInBits(Op->getType()) < getTypeSizeInBits(Ty) &&
@@ -1642,11 +1665,8 @@ ScalarEvolution::getZeroExtendExpr(const SCEV *Op, Type *Ty, unsigned Depth) {
     return Iter->second;
 
   const SCEV *S = getZeroExtendExprImpl(Op, Ty, Depth);
-  if (!isa<SCEVZeroExtendExpr>(S)) {
-    FoldCache.insert({ID, S});
-    auto R = FoldCacheUser.insert({S, {}});
-    R.first->second.push_back(ID);
-  }
+  if (!isa<SCEVZeroExtendExpr>(S))
+    insertFoldCacheEntry(ID, S, FoldCache, FoldCacheUser);
   return S;
 }
 
@@ -1968,11 +1988,8 @@ ScalarEvolution::getSignExtendExpr(const SCEV *Op, Type *Ty, unsigned Depth) {
     return Iter->second;
 
   const SCEV *S = getSignExtendExprImpl(Op, Ty, Depth);
-  if (!isa<SCEVSignExtendExpr>(S)) {
-    FoldCache.insert({ID, S});
-    auto R = FoldCacheUser.insert({S, {}});
-    R.first->second.push_back(ID);
-  }
+  if (!isa<SCEVSignExtendExpr>(S))
+    insertFoldCacheEntry(ID, S, FoldCache, FoldCacheUser);
   return S;
 }
 

diff --git a/llvm/test/Analysis/IVUsers/zext-fold-cache-invalidation.ll b/llvm/test/Analysis/IVUsers/zext-fold-cache-invalidation.ll
@@ -1,10 +1,11 @@
-; RUN: opt -verify-scev -passes='print<iv-users>' -disable-output %s | FileCheck %s
-
-; XFAIL: *
+; RUN: opt -verify-scev -passes='print<iv-users>' -disable-output %s 2>&1 | FileCheck %s
 
 target datalayout = "n16"
 
 define i16 @zext_cache_invalidation_1(i1 %c) {
+; CHECK:      IV Users for loop %loop with backedge-taken count 13:
+; CHECK-NEXT:   %iv = {-3,+,4}<nuw><nsw><%loop> in    %iv.ext = zext i16 %iv to i32
+;
 entry:
   br i1 false, label %loop, label %exit