[llvm] [SCEV] Find existing AR in cache in getAR (NFC) (PR #218927)

[Ramkumar Ramachandra via llvm-commits]

https://github.com/artagnon created https://github.com/llvm/llvm-project/pull/218927

Also use the extended findExistingSCEVInCache to clean up code in proveNoWrapByVaryingStart.

-- 8< --
Compile-time neutral: https://llvm-compile-time-tracker.com/compare.php?from=3638e00f5a7e6b57ff62ced62388a60011d1a8db&to=ac7c640db15bebf3e5826e51b15aeaa55ac03263&stat=instructions:u

>From ac7c640db15bebf3e5826e51b15aeaa55ac03263 Mon Sep 17 00:00:00 2001
From: Ramkumar Ramachandra <artagnon at tenstorrent.com>
Date: Wed, 26 Aug 2026 13:34:02 +0100
Subject: [PATCH] [SCEV] Find existing AR in cache in getAR (NFC)

For compile-time. Also use the extended findExistingSCEVInCache to clean
up code in proveNoWrapByVaryingStart.
---
 llvm/include/llvm/Analysis/ScalarEvolution.h |  8 +++---
 llvm/lib/Analysis/ScalarEvolution.cpp        | 26 ++++++++++++--------
 2 files changed, 21 insertions(+), 13 deletions(-)

diff --git a/llvm/include/llvm/Analysis/ScalarEvolution.h b/llvm/include/llvm/Analysis/ScalarEvolution.h
index 4f59bc99cb697..3cdd9d11fe640 100644
--- a/llvm/include/llvm/Analysis/ScalarEvolution.h
+++ b/llvm/include/llvm/Analysis/ScalarEvolution.h
@@ -2546,9 +2546,11 @@ class ScalarEvolution {
   /// an add rec on said loop.
   void getUsedLoops(const SCEV *S, SmallPtrSetImpl<const Loop *> &LoopsUsed);
 
-  /// Look for a SCEV expression with type `SCEVType` and operands `Ops` in
-  /// `UniqueSCEVs`.  Return if found, else nullptr.
-  SCEV *findExistingSCEVInCache(SCEVTypes SCEVType, ArrayRef<SCEVUse> Ops);
+  /// Look for a SCEV expression with type \p SCEVType and operands \p Ops in
+  /// UniqueSCEVs. If \p SCEVType is scAddRecExpr, the loop \p L must be passed.
+  /// Return if found, else nullptr.
+  SCEV *findExistingSCEVInCache(SCEVTypes SCEVType, ArrayRef<SCEVUse> Ops,
+                                const Loop *L = nullptr);
 
   /// Get reachable blocks in this function, making limited use of SCEV
   /// reasoning about conditions.
diff --git a/llvm/lib/Analysis/ScalarEvolution.cpp b/llvm/lib/Analysis/ScalarEvolution.cpp
index f6ad701ce819e..865920b450118 100644
--- a/llvm/lib/Analysis/ScalarEvolution.cpp
+++ b/llvm/lib/Analysis/ScalarEvolution.cpp
@@ -1439,15 +1439,8 @@ bool ScalarEvolution::proveNoWrapByVaryingStart(const SCEV *Start,
 
   for (unsigned Delta : {-2, -1, 1, 2}) {
     const SCEV *PreStart = getConstant(StartAI - Delta);
-
-    FoldingSetNodeID ID;
-    ID.AddInteger(scAddRecExpr);
-    ID.AddPointer(PreStart);
-    ID.AddPointer(Step);
-    ID.AddPointer(L);
-    void *IP = nullptr;
-    const auto *PreAR =
-      static_cast<SCEVAddRecExpr *>(UniqueSCEVs.FindNodeOrInsertPos(ID, IP));
+    const auto *PreAR = static_cast<SCEVAddRecExpr *>(
+        findExistingSCEVInCache(scAddRecExpr, {PreStart, Step}, L));
 
     // Give up if we don't already have the add recurrence we need because
     // actually constructing an add recurrence is relatively expensive.
@@ -3689,6 +3682,14 @@ const SCEV *ScalarEvolution::getAddRecExpr(SmallVectorImpl<SCEVUse> &Operands,
            "SCEVAddRecExpr operand is not available at loop entry!");
 #endif
 
+  if (auto *AR = static_cast<SCEVAddRecExpr *>(
+          findExistingSCEVInCache(scAddRecExpr, Operands, L))) {
+    if (AR->getNoWrapFlags(Flags) != Flags)
+      AR->setNoWrapFlags(
+          StrengthenNoWrapFlags(this, scAddRecExpr, Operands, Flags));
+    return AR;
+  }
+
   if (Operands.back()->isZero()) {
     Operands.pop_back();
     return getAddRecExpr(Operands, L, SCEV::FlagAnyWrap); // {X,+,0}  -->  X
@@ -3834,11 +3835,16 @@ const SCEV *ScalarEvolution::getGEPExpr(SCEVUse BaseExpr,
 }
 
 SCEV *ScalarEvolution::findExistingSCEVInCache(SCEVTypes SCEVType,
-                                               ArrayRef<SCEVUse> Ops) {
+                                               ArrayRef<SCEVUse> Ops,
+                                               const Loop *L) {
+  assert((SCEVType != scAddRecExpr || L) &&
+         "L must be passed to find existing AddRecs");
   FoldingSetNodeID ID;
   ID.AddInteger(SCEVType);
   for (SCEVUse Op : Ops)
     ID.AddPointer(Op.getOpaqueValue());
+  if (L)
+    ID.AddPointer(L);
   void *IP = nullptr;
   return UniqueSCEVs.FindNodeOrInsertPos(ID, IP);
 }