[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 d8c2e12d224f775ad0d101190bf54c2f6f88f32b 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 1/2] [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 08e78cb49c69fc..353c35efea776d 100644
--- a/llvm/lib/Transforms/Vectorize/LoopVectorize.cpp
+++ b/llvm/lib/Transforms/Vectorize/LoopVectorize.cpp
@@ -4054,7 +4054,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);
 

>From d12dbae1750515389753d7b3d3c39f1c9d78ce06 Mon Sep 17 00:00:00 2001
From: David Sherwood <david.sherwood at arm.com>
Date: Mon, 30 Sep 2024 12:12:56 +0000
Subject: [PATCH 2/2] Address review comment

* Updated comments around calls to getSymbolicMaxBackedgeTakenCount
---
 llvm/lib/Transforms/Vectorize/LoopVectorize.cpp | 7 ++++---
 llvm/lib/Transforms/Vectorize/VPlan.cpp         | 7 ++++---
 2 files changed, 8 insertions(+), 6 deletions(-)

diff --git a/llvm/lib/Transforms/Vectorize/LoopVectorize.cpp b/llvm/lib/Transforms/Vectorize/LoopVectorize.cpp
index 353c35efea776d..b02ea32348107a 100644
--- a/llvm/lib/Transforms/Vectorize/LoopVectorize.cpp
+++ b/llvm/lib/Transforms/Vectorize/LoopVectorize.cpp
@@ -4054,9 +4054,10 @@ LoopVectorizationCostModel::computeMaxVF(ElementCount UserVF, unsigned UserIC) {
     unsigned MaxVFtimesIC =
         UserIC ? *MaxPowerOf2RuntimeVF * UserIC : *MaxPowerOf2RuntimeVF;
     ScalarEvolution *SE = PSE.getSE();
-    // Currently only loops with countable exits are vectorized so it's safe to
-    // use getSymbolicMaxBackedgeTakenCount as it should give the same result
-    // as getBackedgeTakenCount.
+    // Currently only loops with countable exits are vectorized, but calling
+    // getSymbolicMaxBackedgeTakenCount allows enablement work for loops with
+    // uncountable exits whilst also ensuring the symbolic maximum and known
+    // back-edge taken count remain identical for loops with countable exits.
     const SCEV *BackedgeTakenCount = PSE.getSymbolicMaxBackedgeTakenCount();
     assert(!isa<SCEVCouldNotCompute>(BackedgeTakenCount) &&
            "Invalid loop count");
diff --git a/llvm/lib/Transforms/Vectorize/VPlan.cpp b/llvm/lib/Transforms/Vectorize/VPlan.cpp
index 4316c7dd9b6430..096ad3ba4768cb 100644
--- a/llvm/lib/Transforms/Vectorize/VPlan.cpp
+++ b/llvm/lib/Transforms/Vectorize/VPlan.cpp
@@ -881,9 +881,10 @@ VPlanPtr VPlan::createInitialVPlan(Type *InductionTy,
 
   // Create SCEV and VPValue for the trip count.
 
-  // Using getSymbolicMaxBackedgeTakenCount instead of getBackedgeTakenCount,
-  // since they should be identical as we currently only vectorize loops when
-  // all exits are countable.
+  // Currently only loops with countable exits are vectorized, but calling
+  // getSymbolicMaxBackedgeTakenCount allows enablement work for loops with
+  // uncountable exits whilst also ensuring the symbolic maximum and known
+  // back-edge taken count remain identical for loops with countable exits.
   const SCEV *BackedgeTakenCountSCEV = PSE.getSymbolicMaxBackedgeTakenCount();
   assert(!isa<SCEVCouldNotCompute>(BackedgeTakenCountSCEV) &&
          "Invalid loop count");