[llvm] [SCEV] Speed up forgetLoop by avoiding def-use walk for loop-header PHIs (PR #201572)

[Mingjie Xu via llvm-commits]

https://github.com/Enna1 created https://github.com/llvm/llvm-project/pull/201572

None

>From 35f4b10f84f0c2a9f487599269b810801af44e0a Mon Sep 17 00:00:00 2001
From: Enna1 <xumingjie.enna1 at bytedance.com>
Date: Thu, 4 Jun 2026 17:37:00 +0800
Subject: [PATCH] [SCEV] Speed up forgetLoop by avoiding def-use walk for
 loop-header PHIs

---
 llvm/lib/Analysis/ScalarEvolution.cpp         | 22 +++--------
 .../Analysis/ScalarEvolutionTest.cpp          | 39 +++++++++++++++++++
 2 files changed, 45 insertions(+), 16 deletions(-)

diff --git a/llvm/lib/Analysis/ScalarEvolution.cpp b/llvm/lib/Analysis/ScalarEvolution.cpp
index ddf583082b2de..8bc009f389b9e 100644
--- a/llvm/lib/Analysis/ScalarEvolution.cpp
+++ b/llvm/lib/Analysis/ScalarEvolution.cpp
@@ -8677,18 +8677,6 @@ bool ScalarEvolution::isBackedgeTakenCountMaxOrZero(const Loop *L) {
   return getBackedgeTakenInfo(L).isConstantMaxOrZero(this);
 }
 
-/// Push PHI nodes in the header of the given loop onto the given Worklist.
-static void PushLoopPHIs(const Loop *L,
-                         SmallVectorImpl<Instruction *> &Worklist,
-                         SmallPtrSetImpl<Instruction *> &Visited) {
-  BasicBlock *Header = L->getHeader();
-
-  // Push all Loop-header PHIs onto the Worklist stack.
-  for (PHINode &PN : Header->phis())
-    if (Visited.insert(&PN).second)
-      Worklist.push_back(&PN);
-}
-
 ScalarEvolution::BackedgeTakenInfo &
 ScalarEvolution::getPredicatedBackedgeTakenInfo(const Loop *L) {
   auto &BTI = getBackedgeTakenInfo(L);
@@ -8798,8 +8786,6 @@ void ScalarEvolution::visitAndClearUsers(
 
 void ScalarEvolution::forgetLoop(const Loop *L) {
   SmallVector<const Loop *, 16> LoopWorklist(1, L);
-  SmallVector<Instruction *, 32> Worklist;
-  SmallPtrSet<Instruction *, 16> Visited;
   SmallVector<SCEVUse, 16> ToForget;
 
   // Iterate over all the loops and sub-loops to drop SCEV information.
@@ -8819,8 +8805,12 @@ void ScalarEvolution::forgetLoop(const Loop *L) {
       llvm::append_range(ToForget, LoopUsersItr->second);
 
     // Drop information about expressions based on loop-header PHIs.
-    PushLoopPHIs(CurrL, Worklist, Visited);
-    visitAndClearUsers(Worklist, Visited, ToForget);
+    for (PHINode &PN : CurrL->getHeader()->phis()) {
+      ConstantEvolutionLoopExitValue.erase(&PN);
+      auto VIt = ValueExprMap.find_as(static_cast<Value *>(&PN));
+      if (VIt != ValueExprMap.end())
+        ToForget.push_back(VIt->second);
+    }
 
     LoopPropertiesCache.erase(CurrL);
     // Forget all contained loops too, to avoid dangling entries in the
diff --git a/llvm/unittests/Analysis/ScalarEvolutionTest.cpp b/llvm/unittests/Analysis/ScalarEvolutionTest.cpp
index 8504534693281..dcffed39d9381 100644
--- a/llvm/unittests/Analysis/ScalarEvolutionTest.cpp
+++ b/llvm/unittests/Analysis/ScalarEvolutionTest.cpp
@@ -1669,6 +1669,45 @@ TEST_F(ScalarEvolutionsTest, ForgetValueWithOverflowInst) {
   });
 }
 
+TEST_F(ScalarEvolutionsTest, ForgetLoopPreservesUnrelatedCachesInLoopBody) {
+  LLVMContext C;
+  SMDiagnostic Err;
+  std::unique_ptr<Module> M =
+      parseAssemblyString("define void @foo(i32 %n) { "
+                          "entry: "
+                          "  br label %loop "
+                          "loop: "
+                          "  %iv = phi i32 [ 0, %entry ], [ %iv.next, %loop ] "
+                          "  %iv.next = add nsw i32 %iv, 1 "
+                          "  %cmp = icmp slt i32 %iv, %n "
+                          "  br i1 %cmp, label %loop, label %exit "
+                          "exit: "
+                          "  ret void "
+                          "} ",
+                          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 *IV = getInstructionByName(F, "iv");
+    auto *Cmp = getInstructionByName(F, "cmp");
+
+    const SCEV *IVScev = SE.getSCEV(IV);
+    EXPECT_NE(IVScev, nullptr);
+    EXPECT_TRUE(isa<SCEVAddRecExpr>(IVScev));
+
+    const SCEV *CmpScev = SE.getSCEV(Cmp);
+    EXPECT_NE(CmpScev, nullptr);
+    EXPECT_TRUE(isa<SCEVUnknown>(CmpScev));
+
+    Loop *L = *LI.begin();
+    SE.forgetLoop(L);
+    EXPECT_EQ(SE.getExistingSCEV(IV), nullptr);
+    EXPECT_EQ(SE.getExistingSCEV(Cmp), CmpScev);
+  });
+}
+
 TEST_F(ScalarEvolutionsTest, ComplexityComparatorIsStrictWeakOrdering) {
   // Regression test for a case where caching of equivalent values caused the
   // comparator to get inconsistent.