[llvm] Adding the support for vectorization of loops with load based tripcount (PR #209390)

[Manish Srivastava via llvm-commits]

================
@@ -15684,28 +15761,45 @@ const SCEV *PredicatedScalarEvolution::getPredicatedSCEV(const SCEV *Expr) {
     Expr = Entry.second;
 
   const SCEV *NewSCEV = SE.rewriteUsingPredicate(Expr, &L, *Preds);
+
+  if (!TripCountInvariantMap.empty()) {
+    NewSCEV =
+        SCEVParameterRewriter::rewrite(NewSCEV, SE, TripCountInvariantMap);
+  }
+
   Entry = {Generation, NewSCEV};
 
   return NewSCEV;
 }
 
 const SCEV *PredicatedScalarEvolution::getBackedgeTakenCount() {
   if (!BackedgeCount) {
-    SmallVector<const SCEVPredicate *, 4> Preds;
-    BackedgeCount = SE.getPredicatedBackedgeTakenCount(&L, Preds);
-    for (const auto *P : Preds)
-      addPredicate(*P);
+    if (!TripCountInvariantPreds.empty()) {
+      BackedgeCount = SE.computeBackedgeTakenCountWithTripCountInvariants(
+          &L, TripCountInvariantPreds, /*SymbolicMax=*/false);
+    } else {
+      SmallVector<const SCEVPredicate *, 4> Preds;
+      BackedgeCount = SE.getPredicatedBackedgeTakenCount(&L, Preds);
+      for (const auto *P : Preds)
+        addPredicate(*P);
+    }
   }
   return BackedgeCount;
 }
 
 const SCEV *PredicatedScalarEvolution::getSymbolicMaxBackedgeTakenCount() {
   if (!SymbolicMaxBackedgeCount) {
-    SmallVector<const SCEVPredicate *, 4> Preds;
-    SymbolicMaxBackedgeCount =
-        SE.getPredicatedSymbolicMaxBackedgeTakenCount(&L, Preds);
-    for (const auto *P : Preds)
-      addPredicate(*P);
+    if (!TripCountInvariantPreds.empty()) {
+      SymbolicMaxBackedgeCount =
+          SE.computeBackedgeTakenCountWithTripCountInvariants(
+              &L, TripCountInvariantPreds, /*SymbolicMax=*/true);
+    } else {
+      SmallVector<const SCEVPredicate *, 4> Preds;
+      SymbolicMaxBackedgeCount =
+          SE.getPredicatedSymbolicMaxBackedgeTakenCount(&L, Preds);
+      for (const auto *P : Preds)
+        addPredicate(*P);
+    }
----------------
mk-srivastava wrote:

I went through the mentioned PR. 

Based on you suggestion I have pushed new changes. 

In the new changes the predicate creation is done in `LoopVectorizer` by using SCEVs `ComparePredicate` rather than introducing a new predicate. 

But about:

> Can we avoid making changes to SCEV?

There are slight differences between the main problem here. The whole change is dependent upon whether `BackedgeTakenCount` should not return `SCEVCouldNotCompute` as several legality checks as well as LAA's runtime checks are dependent upon it. So we need to apply the assumption while the `BackedgeTakenCount` function is being computed. 
The function entirely is inside SCEV so making change their is unavoidable, although I have done minimal changes to  SCEV. 

Please feel free to correct me if you feel I missed something important. 



https://github.com/llvm/llvm-project/pull/209390