[llvm] [LoopVectorize] Teach LoopVectorizationLegality about more early exits (PR #107004)

[Igor Kirillov via llvm-commits]

================
@@ -1442,6 +1486,95 @@ bool LoopVectorizationLegality::canVectorizeLoopNestCFG(
   return Result;
 }
 
+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);
+    return false;
+  }
+
+  SmallVector<BasicBlock *, 8> ExitingBlocks;
+  TheLoop->getExitingBlocks(ExitingBlocks);
+
+  // Keep a record of all the exiting blocks with exact exit counts, as well as
+  // those with inexact counts.
+  SmallVector<const SCEVPredicate *, 4> Predicates;
+  for (BasicBlock *BB1 : ExitingBlocks) {
+    const SCEV *EC =
+        PSE.getSE()->getPredicatedExitCount(TheLoop, BB1, &Predicates);
+    if (isa<SCEVCouldNotCompute>(EC)) {
+      UncountableExitingBlocks.push_back(BB1);
+
+      unsigned NumExitBlocks = 0;
+      for (BasicBlock *BB2 : successors(BB1)) {
+        if (!TheLoop->contains(BB2)) {
+          UncountableExitBlocks.push_back(BB2);
+          NumExitBlocks++;
+        }
+      }
+      if (NumExitBlocks > 1) {
----------------
igogo-x86 wrote:

If `BB1` ends with conditional BranchInst, it has two jumps. Both jumps can't go outside of the loop; otherwise, `BB1` wouldn't belong to `TheLoop`. So, the only way `BB1` can have multiple Exit Blocks is when it ends with SwitchInst. So, couldn't we check if BB1 ends with a switch?

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