[llvm] [LoopVectorize] In LoopVectorize.cpp start using getSymbolicMaxBackedgeTakenCount (PR #108833)

[David Sherwood via llvm-commits]

https://github.com/david-arm updated https://github.com/llvm/llvm-project/pull/108833

>From 459661e07ecdc14b5edd68fed9209bf080b4422e Mon Sep 17 00:00:00 2001
From: David Sherwood <david.sherwood at arm.com>
Date: Mon, 16 Sep 2024 14:09:02 +0000
Subject: [PATCH] [LoopVectorize] In LoopVectorize start using
 getSymbolicMaxBackedgeTakenCount

LoopVectorizationLegality currently only treats a loop as legal
to vectorise if PredicatedScalarEvolution::getBackedgeTakenCount
returns a valid SCEV, or more precisely that the loop must have
an exact backedge taken count. Therefore, in LoopVectorize.cpp
we can safely replace all calls to getBackedgeTakenCount with
calls to getSymbolicMaxBackedgeTakenCount, since the result is
the same.

This also helps prepare the loop vectoriser for PR #88385.
---
 llvm/lib/Transforms/Vectorize/LoopVectorize.cpp |  7 ++++++-
 llvm/lib/Transforms/Vectorize/VPlan.cpp         | 13 +++++++++----
 2 files changed, 15 insertions(+), 5 deletions(-)

diff --git a/llvm/lib/Transforms/Vectorize/LoopVectorize.cpp b/llvm/lib/Transforms/Vectorize/LoopVectorize.cpp
index bd493fb2c1ba19..1921e3c2430ae8 100644
--- a/llvm/lib/Transforms/Vectorize/LoopVectorize.cpp
+++ b/llvm/lib/Transforms/Vectorize/LoopVectorize.cpp
@@ -4075,7 +4075,12 @@ LoopVectorizationCostModel::computeMaxVF(ElementCount UserVF, unsigned UserIC) {
     unsigned MaxVFtimesIC =
         UserIC ? *MaxPowerOf2RuntimeVF * UserIC : *MaxPowerOf2RuntimeVF;
     ScalarEvolution *SE = PSE.getSE();
-    const SCEV *BackedgeTakenCount = PSE.getBackedgeTakenCount();
+    // Currently only loops with countable exits are vectorized so it's safe to
+    // use getSymbolicMaxBackedgeTakenCount as it should give the same result
+    // as getBackedgeTakenCount.
+    const SCEV *BackedgeTakenCount = PSE.getSymbolicMaxBackedgeTakenCount();
+    assert(!isa<SCEVCouldNotCompute>(BackedgeTakenCount) &&
+           "Invalid loop count");
     const SCEV *ExitCount = SE->getAddExpr(
         BackedgeTakenCount, SE->getOne(BackedgeTakenCount->getType()));
     const SCEV *Rem = SE->getURemExpr(
diff --git a/llvm/lib/Transforms/Vectorize/VPlan.cpp b/llvm/lib/Transforms/Vectorize/VPlan.cpp
index 6ddbfcf0ecfe58..4316c7dd9b6430 100644
--- a/llvm/lib/Transforms/Vectorize/VPlan.cpp
+++ b/llvm/lib/Transforms/Vectorize/VPlan.cpp
@@ -880,11 +880,16 @@ VPlanPtr VPlan::createInitialVPlan(Type *InductionTy,
   auto Plan = std::make_unique<VPlan>(Entry, VecPreheader);
 
   // Create SCEV and VPValue for the trip count.
-  const SCEV *BackedgeTakenCount = PSE.getBackedgeTakenCount();
-  assert(!isa<SCEVCouldNotCompute>(BackedgeTakenCount) && "Invalid loop count");
+
+  // Using getSymbolicMaxBackedgeTakenCount instead of getBackedgeTakenCount,
+  // since they should be identical as we currently only vectorize loops when
+  // all exits are countable.
+  const SCEV *BackedgeTakenCountSCEV = PSE.getSymbolicMaxBackedgeTakenCount();
+  assert(!isa<SCEVCouldNotCompute>(BackedgeTakenCountSCEV) &&
+         "Invalid loop count");
   ScalarEvolution &SE = *PSE.getSE();
-  const SCEV *TripCount =
-      SE.getTripCountFromExitCount(BackedgeTakenCount, InductionTy, TheLoop);
+  const SCEV *TripCount = SE.getTripCountFromExitCount(BackedgeTakenCountSCEV,
+                                                       InductionTy, TheLoop);
   Plan->TripCount =
       vputils::getOrCreateVPValueForSCEVExpr(*Plan, TripCount, SE);