[SCEV] Splip collecting and applying rewrite info from loop guards (NFC)

Split off collecting rewrite info from loop guards to
collectRewriteInfoFromLoopGuards. This allows users of applyLoopGuards to
collect rewrite info once in cases where the same loop guards are
applied multiple times.

This is used to collect rewrite info once in howFarToZero, which saves a
bit of compile-time:
stage1-O3: -0.04%
stage1-ReleaseThinLTO: -0.02%
stage1-ReleaseLTO-g: -0.04%
stage2-O3: -0.02%
https://llvm-compile-time-tracker.com/compare.php?from=117b53ae38428ca66eaa886fb432e6f09db88fe4&to=4ffb7b2e1c99081ccebe6f236c48a0be2f64b6ff&stat=instructions:u

Notably this improves mafft by -0.9% with -O3, -0.11% with LTO and
-0.12% with stage2-O3.

Author: fhahn

llvm/include/llvm/Analysis/ScalarEvolution.h

@@ -1299,8 +1299,17 @@ class ScalarEvolution {
   /// sharpen it.
   void setNoWrapFlags(SCEVAddRecExpr *AddRec, SCEV::NoWrapFlags Flags);
 
+  /// Collect rewrite map for loop guards for loop \p L, together with flags
+  /// indidcating if NUW and NSW can be preserved during rewriting.
+  std::tuple<DenseMap<const SCEV *, const SCEV *>, bool, bool>
+  collectRewriteInfoFromLoopGuards(const Loop *L);
+
   /// Try to apply information from loop guards for \p L to \p Expr.
   const SCEV *applyLoopGuards(const SCEV *Expr, const Loop *L);
+  const SCEV *
+  applyLoopGuards(const SCEV *Expr, const Loop *L,
+                  const DenseMap<const SCEV *, const SCEV *> &RewriteMap,
+                  bool PreserveNUW, bool PreserveNSW);
 
   /// Return true if the loop has no abnormal exits. That is, if the loop
   /// is not infinite, it must exit through an explicit edge in the CFG.

llvm/lib/Analysis/ScalarEvolution.cpp

@@ -10490,8 +10490,11 @@ ScalarEvolution::ExitLimit ScalarEvolution::howFarToZero(const SCEV *V,
   if (!isLoopInvariant(Step, L))
     return getCouldNotCompute();
 
+  const auto &[RewriteMap, PreserveNUW, PreserveNSW] =
+      collectRewriteInfoFromLoopGuards(L);
   // Specialize step for this loop so we get context sensitive facts below.
-  const SCEV *StepWLG = applyLoopGuards(Step, L);
+  const SCEV *StepWLG =
+      applyLoopGuards(Step, L, RewriteMap, PreserveNUW, PreserveNSW);
 
   // For positive steps (counting up until unsigned overflow):
   //   N = -Start/Step (as unsigned)
@@ -10508,7 +10511,8 @@ ScalarEvolution::ExitLimit ScalarEvolution::howFarToZero(const SCEV *V,
   //   N = Distance (as unsigned)
   if (StepC &&
       (StepC->getValue()->isOne() || StepC->getValue()->isMinusOne())) {
-    APInt MaxBECount = getUnsignedRangeMax(applyLoopGuards(Distance, L));
+    APInt MaxBECount = getUnsignedRangeMax(
+        applyLoopGuards(Distance, L, RewriteMap, PreserveNUW, PreserveNSW));
     MaxBECount = APIntOps::umin(MaxBECount, getUnsignedRangeMax(Distance));
 
     // When a loop like "for (int i = 0; i != n; ++i) { /* body */ }" is rotated,
@@ -10549,7 +10553,8 @@ ScalarEvolution::ExitLimit ScalarEvolution::howFarToZero(const SCEV *V,
         getUDivExpr(Distance, CountDown ? getNegativeSCEV(Step) : Step);
     const SCEV *ConstantMax = getCouldNotCompute();
     if (Exact != getCouldNotCompute()) {
-      APInt MaxInt = getUnsignedRangeMax(applyLoopGuards(Exact, L));
+      APInt MaxInt = getUnsignedRangeMax(
+          applyLoopGuards(Exact, L, RewriteMap, PreserveNUW, PreserveNSW));
       ConstantMax =
           getConstant(APIntOps::umin(MaxInt, getUnsignedRangeMax(Exact)));
     }
@@ -10566,7 +10571,8 @@ ScalarEvolution::ExitLimit ScalarEvolution::howFarToZero(const SCEV *V,
 
   const SCEV *M = E;
   if (E != getCouldNotCompute()) {
-    APInt MaxWithGuards = getUnsignedRangeMax(applyLoopGuards(E, L));
+    APInt MaxWithGuards = getUnsignedRangeMax(
+        applyLoopGuards(E, L, RewriteMap, PreserveNUW, PreserveNSW));
     M = getConstant(APIntOps::umin(MaxWithGuards, getUnsignedRangeMax(E)));
   }
   auto *S = isa<SCEVCouldNotCompute>(E) ? M : E;
@@ -15096,7 +15102,7 @@ class SCEVLoopGuardRewriter : public SCEVRewriteVisitor<SCEVLoopGuardRewriter> {
 
 public:
   SCEVLoopGuardRewriter(ScalarEvolution &SE,
-                        DenseMap<const SCEV *, const SCEV *> &M,
+                        const DenseMap<const SCEV *, const SCEV *> &M,
                         bool PreserveNUW, bool PreserveNSW)
       : SCEVRewriteVisitor(SE), Map(M) {
     if (PreserveNUW)
@@ -15191,7 +15197,8 @@ class SCEVLoopGuardRewriter : public SCEVRewriteVisitor<SCEVLoopGuardRewriter> {
   }
 };
 
-const SCEV *ScalarEvolution::applyLoopGuards(const SCEV *Expr, const Loop *L) {
+std::tuple<DenseMap<const SCEV *, const SCEV *>, bool, bool>
+ScalarEvolution::collectRewriteInfoFromLoopGuards(const Loop *L) {
   SmallVector<const SCEV *> ExprsToRewrite;
   auto CollectCondition = [&](ICmpInst::Predicate Predicate, const SCEV *LHS,
                               const SCEV *RHS,
@@ -15600,9 +15607,6 @@ const SCEV *ScalarEvolution::applyLoopGuards(const SCEV *Expr, const Loop *L) {
     }
   }
 
-  if (RewriteMap.empty())
-    return Expr;
-
   // Let the rewriter preserve NUW/NSW flags if the unsigned/signed ranges of
   // the replacement expressions are contained in the ranges of the replaced
   // expressions.
@@ -15626,6 +15630,22 @@ const SCEV *ScalarEvolution::applyLoopGuards(const SCEV *Expr, const Loop *L) {
       RewriteMap.insert({Expr, Rewriter.visit(RewriteTo)});
     }
   }
+  return {RewriteMap, PreserveNUW, PreserveNSW};
+}
+
+const SCEV *ScalarEvolution::applyLoopGuards(const SCEV *Expr, const Loop *L) {
+  const auto &[RewriteMap, PreserveNUW, PreserveNSW] =
+      collectRewriteInfoFromLoopGuards(L);
+  return applyLoopGuards(Expr, L, RewriteMap, PreserveNUW, PreserveNSW);
+}
+
+const SCEV *ScalarEvolution::applyLoopGuards(
+    const SCEV *Expr, const Loop *L,
+    const DenseMap<const SCEV *, const SCEV *> &RewriteMap, bool PreserveNUW,
+    bool PreserveNSW) {
+  if (RewriteMap.empty())
+    return Expr;
+
   SCEVLoopGuardRewriter Rewriter(*this, RewriteMap, PreserveNUW, PreserveNSW);
   return Rewriter.visit(Expr);
 }