[llvm] [LoopVectorize] Add support for vectorisation of more early exit loops (PR #88385)

[Graham Hunter via llvm-commits]

================
@@ -1439,6 +1496,103 @@ bool LoopVectorizationLegality::canVectorizeLoopNestCFG(
   return Result;
 }
 
+void LoopVectorizationLegality::recordExitingBlocks() {
+  SmallVector<BasicBlock *, 8> ExitingBlocks;
+  TheLoop->getExitingBlocks(ExitingBlocks);
+
+  SmallVector<BasicBlock *, 4> CountableExitingBBs;
+  PSE.getSE()->getCountableExitingBlocks(TheLoop, &CountableExitingBBs);
+
+  // There could be multiple exiting blocks with an exact exit-not-taken
+  // count. Find all of the uncountable early exit blocks, i.e. the ones with
+  // an unknown count.
+  SmallVector<BasicBlock *, 4> UncountableExitingBBs;
+  SmallVector<BasicBlock *, 4> UncountableExitBBs;
+  for (BasicBlock *BB1 : ExitingBlocks) {
+    if (!is_contained(CountableExitingBBs, BB1)) {
+      UncountableExitingBBs.push_back(BB1);
+
+      for (BasicBlock *BB2 : successors(BB1)) {
+        if (!TheLoop->contains(BB2)) {
+          UncountableExitBBs.push_back(BB2);
+          break;
+        }
+      }
+    }
+  }
+
+  CountableExitingBlocks = std::move(CountableExitingBBs);
+  UncountableExitingBlocks = std::move(UncountableExitingBBs);
+  UncountableExitBlocks = std::move(UncountableExitBBs);
+}
+
+bool LoopVectorizationLegality::isVectorizableEarlyExitLoop() {
+  // At least one of the exiting blocks must be the latch.
+  BasicBlock *LatchBB = TheLoop->getLoopLatch();
+  if (!LatchBB) {
+    reportVectorizationFailure("Loop does not have a latch",
+                               "Cannot vectorize early exit loop",
+                               "NoLatchEarlyExit", ORE, TheLoop);
+    LLVM_DEBUG(dbgs() << "LV: Loop does not have a latch.\n");
+    return false;
+  }
+
+  recordExitingBlocks();
+
+  // We only support one uncountable early exit.
+  if (getUncountableExitingBlocks().size() != 1) {
+    reportVectorizationFailure("Loop has too many uncountable exits",
+                               "Cannot vectorize early exit loop",
+                               "TooManyUncountableEarlyExits", ORE, TheLoop);
+    LLVM_DEBUG(
+        dbgs() << "LV: Loop does not have one uncountable exiting block.\n");
+    return false;
+  }
+
+  // The only supported early exit loops so far are ones where the early
+  // exiting block is a unique predecessor of the latch block.
+  BasicBlock *LatchPredBB = LatchBB->getUniquePredecessor();
+  if (!LatchPredBB || LatchPredBB != getUncountableExitingBlocks()[0]) {
+    reportVectorizationFailure("Early exit is not the latch predecessor",
+                               "Cannot vectorize early exit loop",
+                               "EarlyExitNotLatchPredecessor", ORE, TheLoop);
+    LLVM_DEBUG(
+        dbgs() << "LV: Early exit block is not unique predecessor of latch\n");
+    return false;
+  }
+
+  if (Reductions.size() || FixedOrderRecurrences.size()) {
+    reportVectorizationFailure(
+        "Found reductions or recurrences in early-exit loop",
+        "Cannot vectorize early exit loop", "RecurrencesInEarlyExitLoop", ORE,
+        TheLoop);
+    LLVM_DEBUG(
+        dbgs() << "LV: Found reductions/recurrences in early exit loop.\n");
+    return false;
+  }
+
+  LLVM_DEBUG(
+      dbgs()
+      << "LV: Found an early exit. Retrying with speculative exit count.\n");
+  if (!PSE.getSE()->hasExactPredicatedBackedgeTakenCount(TheLoop, LatchBB)) {
----------------
huntergr-arm wrote:

I don't think this ScalarEvolution method needs to be added either.
```suggestion
  if (isa<SCEVCouldNotCompute>(PSE.getSE()->getExitCount(TheLoop, LatchBB)))
```

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