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

[Paul Walker via llvm-commits]

================
@@ -1070,6 +1070,50 @@ 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;
+    }
+
+    if (LAI->getNumCalls()) {
+      reportVectorizationFailure(
+          "Calls unsupported in early exit loops",
+          "Cannot vectorize early exit loop with function calls",
+          "CallsInEarlyExitLoop", ORE, TheLoop);
+      return false;
+    }
+
+    // The vectoriser cannot handle loads that occur after the early exit block.
+    BasicBlock *LatchBB = TheLoop->getLoopLatch();
+    for (Instruction &I : *LatchBB) {
+      if (I.mayReadFromMemory()) {
+        reportVectorizationFailure(
+            "Loads not permitted after early exit",
+            "Cannot vectorize early exit loop with loads after early exit",
+            "LoadsAfterEarlyExit", ORE, TheLoop);
+        return false;
+      }
+    }
+
+    // The vectoriser does not yet handle loops that may fault, but this will
+    // be improved in a follow-on patch.
----------------
paulwalker-arm wrote:

Up to you but my suggestion was to replace the original comment with the TODO.

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