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

[David Sherwood 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) {
----------------
david-arm wrote:

Yeah, if you look in `isVectorizableEarlyExitLoop` you'll see this code:

```
  // 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 != getUncountableExitingBlocks()[0]) {
    reportVectorizationFailure("Early exit is not the latch predecessor",
                               "Cannot vectorize early exit loop",
                               "EarlyExitNotLatchPredecessor", ORE, TheLoop);
    return false;
  }
```

which is why a few lines above I am able to add this assert:

```
    assert(LatchBB->getUniquePredecessor() ==
               getUncountableExitingBlocks()[0] &&
           "Expected latch predecessor to be the early exiting block");
```

Although if I move the new code from `canVectorizeMemory` to `isVectorizableEarlyExitLoop` I can get rid of the assert.

I added this restriction as a simple optimisation to avoid messier CFGs, such as conditionally executed blocks after the early exit:

```
loop:
...
  br i1 %cmp, label %early.exit, label %loop.bb1

loop.bb1:
...
  br i1 %cmp2, label %do.stuff, label %loop.latch

do.stuff:
...
  br label %loop.latch

loop.latch:
...
  br i1 %cmp3, label %normal.exit, label %loop
```

In practice we can't yet permit any useful work after the early exit anyway until we can teach the vectoriser that subsequent blocks are effectively if-converted and require a mask, and the operations would need to occur before the early exit.

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