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

[Florian Hahn via llvm-commits]

================
@@ -1070,6 +1070,51 @@ bool LoopVectorizationLegality::canVectorizeMemory() {
     return false;
   }
 
+  // For loops with uncountable early exiting blocks that are not the latch
+  // it's necessary to perform extra checks, since the vectoriser is currently
+  // only capable of handling simple search loops.
+  if (IsEarlyExitLoop) {
+    // We don't support calls or any memory accesses that write to memory.
+    if (LAI->getNumStores()) {
+      reportVectorizationFailure(
+          "Writes to memory unsupported in early exit loops",
+          "Cannot vectorize early exit loop with writes to memory",
+          "WritesInEarlyExitLoop", ORE, TheLoop);
+      return false;
+    }
+
+    // The vectoriser cannot handle loads that occur after the early exit block.
+    BasicBlock *LatchBB = TheLoop->getLoopLatch();
+    assert(LatchBB->getUniquePredecessor() ==
+               getUncountableExitingBlocks()[0] &&
+           "Expected latch predecessor to be the early exiting block");
+
+    for (Instruction &I : *LatchBB) {
----------------
fhahn wrote:

Do we have any restrictions that there is only the latch after the early exit? If not, all blocks between early exit and latch would need to be check?

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