[llvm] [SCEV] Avoid redundant ValueExprMap lookups in forgetValue (PR #201573)

[Mingjie Xu via llvm-commits]

https://github.com/Enna1 updated https://github.com/llvm/llvm-project/pull/201573

>From 21a5664f792522ad9944f76a75cc977ecb8eafb8 Mon Sep 17 00:00:00 2001
From: Enna1 <xumingjie.enna1 at bytedance.com>
Date: Thu, 4 Jun 2026 18:31:54 +0800
Subject: [PATCH 1/3] [SCEV] Avoid redundant ValueExprMap lookups in
 forgetValue

---
 llvm/lib/Analysis/ScalarEvolution.cpp | 28 ++++++++++++++++++++++-----
 1 file changed, 23 insertions(+), 5 deletions(-)

diff --git a/llvm/lib/Analysis/ScalarEvolution.cpp b/llvm/lib/Analysis/ScalarEvolution.cpp
index ddf583082b2de..af093d6c8cbb4 100644
--- a/llvm/lib/Analysis/ScalarEvolution.cpp
+++ b/llvm/lib/Analysis/ScalarEvolution.cpp
@@ -8838,13 +8838,31 @@ void ScalarEvolution::forgetValue(Value *V) {
   Instruction *I = dyn_cast<Instruction>(V);
   if (!I) return;
 
-  // Drop information about expressions based on loop-header PHIs.
-  SmallVector<Instruction *, 16> Worklist;
+  SmallVector<std::pair<Instruction *, const SCEV *>, 16> Worklist;
   SmallPtrSet<Instruction *, 8> Visited;
   SmallVector<SCEVUse, 8> ToForget;
-  Worklist.push_back(I);
-  Visited.insert(I);
-  visitAndClearUsers(Worklist, Visited, ToForget);
+
+  auto PushIntoWorklist = [&](Instruction *UI) {
+    bool IsSCEVable = isSCEVable(UI->getType());
+    bool IsWO = isa<WithOverflowInst>(UI);
+    if ((!IsSCEVable && !IsWO) || !Visited.insert(UI).second)
+      return;
+    const SCEV *S = IsSCEVable ? getExistingSCEV(UI) : nullptr;
+    Worklist.emplace_back(UI, S);
+  };
+
+  PushIntoWorklist(I);
+
+  while (!Worklist.empty()) {
+    auto [Cur, S] = Worklist.pop_back_val();
+    if (S) {
+      ToForget.push_back(S);
+      if (auto *PN = dyn_cast<PHINode>(Cur))
+        ConstantEvolutionLoopExitValue.erase(PN);
+    }
+    for (User *U : Cur->users())
+      PushIntoWorklist(cast<Instruction>(U));
+  }
 
   forgetMemoizedResults(ToForget);
 }

>From 2652c5c44ba5ae2217243c90c9ec5bdf1b2e0f48 Mon Sep 17 00:00:00 2001
From: Enna1 <xumingjie.enna1 at bytedance.com>
Date: Thu, 4 Jun 2026 20:41:09 +0800
Subject: [PATCH 2/3] add test

---
 .../Analysis/ScalarEvolutionTest.cpp          | 30 +++++++++++++++++++
 1 file changed, 30 insertions(+)

diff --git a/llvm/unittests/Analysis/ScalarEvolutionTest.cpp b/llvm/unittests/Analysis/ScalarEvolutionTest.cpp
index 8504534693281..443a510e88223 100644
--- a/llvm/unittests/Analysis/ScalarEvolutionTest.cpp
+++ b/llvm/unittests/Analysis/ScalarEvolutionTest.cpp
@@ -1669,6 +1669,36 @@ TEST_F(ScalarEvolutionsTest, ForgetValueWithOverflowInst) {
   });
 }
 
+TEST_F(ScalarEvolutionsTest, ForgetValueAddChainFolding) {
+  LLVMContext C;
+  SMDiagnostic Err;
+  std::unique_ptr<Module> M = parseAssemblyString(
+      "define i32 @foo(i32 %x) { "
+      "  %add1 = add i32 %x, 1 "
+      "  %add2 = add i32 %add1, 1 "
+      "  %add3 = add i32 %add2, 1 "
+      "  ret i32 %add3 "
+      "} ",
+      Err, C);
+
+  ASSERT_TRUE(M && "Could not parse module?");
+  ASSERT_TRUE(!verifyModule(*M) && "Must have been well formed!");
+
+  runWithSE(*M, "foo", [](Function &F, LoopInfo &LI, ScalarEvolution &SE) {
+    auto *Add1 = getInstructionByName(F, "add1");
+    auto *Add2 = getInstructionByName(F, "add2");
+    auto *Add3 = getInstructionByName(F, "add3");
+
+    const SCEV *Add3Scev = SE.getSCEV(Add3);
+    EXPECT_NE(Add3Scev, nullptr);
+    EXPECT_EQ(SE.getExistingSCEV(Add1), nullptr);
+    EXPECT_EQ(SE.getExistingSCEV(Add2), nullptr);
+
+    SE.forgetValue(Add1);
+    EXPECT_EQ(SE.getExistingSCEV(Add3), nullptr);
+  });
+}
+
 TEST_F(ScalarEvolutionsTest, ComplexityComparatorIsStrictWeakOrdering) {
   // Regression test for a case where caching of equivalent values caused the
   // comparator to get inconsistent.

>From ac367618b6209eb40a725b518f38d0ee477d4cc8 Mon Sep 17 00:00:00 2001
From: Enna1 <xumingjie.enna1 at bytedance.com>
Date: Thu, 4 Jun 2026 20:47:15 +0800
Subject: [PATCH 3/3] format

---
 llvm/unittests/Analysis/ScalarEvolutionTest.cpp | 15 +++++++--------
 1 file changed, 7 insertions(+), 8 deletions(-)

diff --git a/llvm/unittests/Analysis/ScalarEvolutionTest.cpp b/llvm/unittests/Analysis/ScalarEvolutionTest.cpp
index 443a510e88223..165845042c3bd 100644
--- a/llvm/unittests/Analysis/ScalarEvolutionTest.cpp
+++ b/llvm/unittests/Analysis/ScalarEvolutionTest.cpp
@@ -1672,14 +1672,13 @@ TEST_F(ScalarEvolutionsTest, ForgetValueWithOverflowInst) {
 TEST_F(ScalarEvolutionsTest, ForgetValueAddChainFolding) {
   LLVMContext C;
   SMDiagnostic Err;
-  std::unique_ptr<Module> M = parseAssemblyString(
-      "define i32 @foo(i32 %x) { "
-      "  %add1 = add i32 %x, 1 "
-      "  %add2 = add i32 %add1, 1 "
-      "  %add3 = add i32 %add2, 1 "
-      "  ret i32 %add3 "
-      "} ",
-      Err, C);
+  std::unique_ptr<Module> M = parseAssemblyString("define i32 @foo(i32 %x) { "
+                                                  "  %add1 = add i32 %x, 1 "
+                                                  "  %add2 = add i32 %add1, 1 "
+                                                  "  %add3 = add i32 %add2, 1 "
+                                                  "  ret i32 %add3 "
+                                                  "} ",
+                                                  Err, C);
 
   ASSERT_TRUE(M && "Could not parse module?");
   ASSERT_TRUE(!verifyModule(*M) && "Must have been well formed!");