[llvm] [SCEV] Avoid erase+insert in constant folding (PR #101642)

[Nikita Popov via llvm-commits]

https://github.com/nikic updated https://github.com/llvm/llvm-project/pull/101642

>From 99634c22d8e16400eb7be454f3f0fb95a342e480 Mon Sep 17 00:00:00 2001
From: Nikita Popov <npopov at redhat.com>
Date: Fri, 2 Aug 2024 10:09:15 +0200
Subject: [PATCH 1/2] [SCEV] Avoid erase+insert in constant folding (NFC)

Based on Philip's suggestion in #101473, this changes the constant
folding implementation to keep the SCEVConstant as the first
element in the Ops, instead of erasing it and then later reinserting.

This gives a small compile-time improvement:
http://llvm-compile-time-tracker.com/compare.php?from=85c5265feae82d8e26869adf4505b448b3a17534&to=67271fa74cb49ff958865bc56e3eea79a934bafc&stat=instructions:u
---
 llvm/lib/Analysis/ScalarEvolution.cpp | 45 +++++++++++++++------------
 1 file changed, 25 insertions(+), 20 deletions(-)

diff --git a/llvm/lib/Analysis/ScalarEvolution.cpp b/llvm/lib/Analysis/ScalarEvolution.cpp
index a19358dee8ef49..58dd536c882d0a 100644
--- a/llvm/lib/Analysis/ScalarEvolution.cpp
+++ b/llvm/lib/Analysis/ScalarEvolution.cpp
@@ -777,8 +777,8 @@ CompareSCEVComplexity(EquivalenceClasses<const SCEV *> &EqCacheSCEV,
 /// results from this routine.  In other words, we don't want the results of
 /// this to depend on where the addresses of various SCEV objects happened to
 /// land in memory.
-static void GroupByComplexity(SmallVectorImpl<const SCEV *> &Ops,
-                              LoopInfo *LI, DominatorTree &DT) {
+static void GroupByComplexity(MutableArrayRef<const SCEV *> Ops, LoopInfo *LI,
+                              DominatorTree &DT) {
   if (Ops.size() < 2) return;  // Noop
 
   EquivalenceClasses<const SCEV *> EqCacheSCEV;
@@ -844,34 +844,39 @@ static const SCEV *
 constantFoldAndGroupOps(ScalarEvolution &SE, LoopInfo &LI, DominatorTree &DT,
                         SmallVectorImpl<const SCEV *> &Ops, FoldT Fold,
                         IsIdentityT IsIdentity, IsAbsorberT IsAbsorber) {
-  const SCEVConstant *Folded = nullptr;
+  bool HasConst = false;
   for (unsigned Idx = 0; Idx < Ops.size();) {
     const SCEV *Op = Ops[Idx];
     if (const auto *C = dyn_cast<SCEVConstant>(Op)) {
-      if (!Folded)
-        Folded = C;
-      else
-        Folded = cast<SCEVConstant>(
-            SE.getConstant(Fold(Folded->getAPInt(), C->getAPInt())));
-      Ops.erase(Ops.begin() + Idx);
-      continue;
+      if (!HasConst) {
+        // Move the constant to the start.
+        std::swap(Ops[0], Ops[Idx]);
+        HasConst = true;
+      } else {
+        Ops[0] = SE.getConstant(
+            Fold(cast<SCEVConstant>(Ops[0])->getAPInt(), C->getAPInt()));
+        Ops.erase(Ops.begin() + Idx);
+        // Revisit this index.
+        continue;
+      }
     }
     ++Idx;
   }
 
-  if (Ops.empty()) {
-    assert(Folded && "Must have folded value");
-    return Folded;
-  }
+  if (Ops.size() == 1)
+    return Ops[0];
 
-  if (Folded && IsAbsorber(Folded->getAPInt()))
-    return Folded;
+  if (HasConst && IsIdentity(cast<SCEVConstant>(Ops[0])->getAPInt())) {
+    Ops.erase(Ops.begin());
+    HasConst = false;
+  }
 
-  GroupByComplexity(Ops, &LI, DT);
-  if (Folded && !IsIdentity(Folded->getAPInt()))
-    Ops.insert(Ops.begin(), Folded);
+  if (Ops.size() == 1 ||
+      (HasConst && IsAbsorber(cast<SCEVConstant>(Ops[0])->getAPInt())))
+    return Ops[0];
 
-  return Ops.size() == 1 ? Ops[0] : nullptr;
+  GroupByComplexity(MutableArrayRef(Ops).drop_front(HasConst), &LI, DT);
+  return nullptr;
 }
 
 //===----------------------------------------------------------------------===//

>From 74a9c24727c07e21dfa3a61b8ac802604e1144f4 Mon Sep 17 00:00:00 2001
From: Nikita Popov <npopov at redhat.com>
Date: Fri, 16 Aug 2024 15:18:00 +0200
Subject: [PATCH 2/2] address review comments

---
 llvm/lib/Analysis/ScalarEvolution.cpp | 26 ++++++++++++++------------
 1 file changed, 14 insertions(+), 12 deletions(-)

diff --git a/llvm/lib/Analysis/ScalarEvolution.cpp b/llvm/lib/Analysis/ScalarEvolution.cpp
index 58dd536c882d0a..80e957ce109685 100644
--- a/llvm/lib/Analysis/ScalarEvolution.cpp
+++ b/llvm/lib/Analysis/ScalarEvolution.cpp
@@ -848,17 +848,17 @@ constantFoldAndGroupOps(ScalarEvolution &SE, LoopInfo &LI, DominatorTree &DT,
   for (unsigned Idx = 0; Idx < Ops.size();) {
     const SCEV *Op = Ops[Idx];
     if (const auto *C = dyn_cast<SCEVConstant>(Op)) {
-      if (!HasConst) {
-        // Move the constant to the start.
-        std::swap(Ops[0], Ops[Idx]);
-        HasConst = true;
-      } else {
+      if (HasConst) {
         Ops[0] = SE.getConstant(
             Fold(cast<SCEVConstant>(Ops[0])->getAPInt(), C->getAPInt()));
         Ops.erase(Ops.begin() + Idx);
         // Revisit this index.
         continue;
       }
+
+      // Move the constant to the start, without changing order otherwise.
+      std::rotate(Ops.begin(), Ops.begin() + Idx, Ops.begin() + Idx + 1);
+      HasConst = true;
     }
     ++Idx;
   }
@@ -866,15 +866,17 @@ constantFoldAndGroupOps(ScalarEvolution &SE, LoopInfo &LI, DominatorTree &DT,
   if (Ops.size() == 1)
     return Ops[0];
 
-  if (HasConst && IsIdentity(cast<SCEVConstant>(Ops[0])->getAPInt())) {
-    Ops.erase(Ops.begin());
-    HasConst = false;
+  if (HasConst) {
+    if (IsIdentity(cast<SCEVConstant>(Ops[0])->getAPInt())) {
+      Ops.erase(Ops.begin());
+      HasConst = false;
+      if (Ops.size() == 1)
+        return Ops[0];
+    } else if (IsAbsorber(cast<SCEVConstant>(Ops[0])->getAPInt())) {
+      return Ops[0];
+    }
   }
 
-  if (Ops.size() == 1 ||
-      (HasConst && IsAbsorber(cast<SCEVConstant>(Ops[0])->getAPInt())))
-    return Ops[0];
-
   GroupByComplexity(MutableArrayRef(Ops).drop_front(HasConst), &LI, DT);
   return nullptr;
 }