[llvm] [LV] Initial support for stores in early exit loops (PR #137774)

[David Sherwood via llvm-commits]

================
@@ -1732,13 +1794,53 @@ bool LoopVectorizationLegality::isVectorizableEarlyExitLoop() {
 
   // TODO: Handle loops that may fault.
   Predicates.clear();
-  if (!isDereferenceableReadOnlyLoop(TheLoop, PSE.getSE(), DT, AC,
-                                     &Predicates)) {
+
+  if (HasStore && EELoad.has_value()) {
+    LoadInst *LI = *EELoad;
+    if (isDereferenceableAndAlignedInLoop(LI, TheLoop, *PSE.getSE(), *DT, AC,
+                                          &Predicates)) {
+      ICFLoopSafetyInfo SafetyInfo;
+      SafetyInfo.computeLoopSafetyInfo(TheLoop);
+      // FIXME: We may have multiple levels of conditional loads, so will
+      //        need to improve on outright rejection at some point.
+      if (!SafetyInfo.isGuaranteedToExecute(*LI, DT, TheLoop)) {
+        LLVM_DEBUG(
+            dbgs() << "Early exit condition load not guaranteed to execute.\n");
+        reportVectorizationFailure(
+            "Early exit condition load not guaranteed to execute",
+            "Cannot vectorize early exit loop when condition load is not "
+            "guaranteed to execute",
+            "EarlyExitLoadNotGuaranteed", ORE, TheLoop);
+      }
+    } else {
+      LLVM_DEBUG(dbgs() << "Early exit condition load potentially unsafe.\n");
+      reportVectorizationFailure("Uncounted loop condition not known safe",
+                                 "Cannot vectorize early exit loop with "
+                                 "possibly unsafe condition load",
+                                 "PotentiallyFaultingEarlyExitLoop", ORE,
+                                 TheLoop);
+      return false;
+    }
+  } else if (HasStore) {
+    LLVM_DEBUG(dbgs() << "Found early exit store but no condition load.\n");
     reportVectorizationFailure(
-        "Loop may fault",
-        "Cannot vectorize potentially faulting early exit loop",
-        "PotentiallyFaultingEarlyExitLoop", ORE, TheLoop);
+        "Early exit loop with store but no condition load",
+        "Cannot vectorize early exit loop with store but no condition load",
+        "NoConditionLoadForEarlyExitLoop", ORE, TheLoop);
     return false;
+  } else {
+    // Read-only loop.
+    // FIXME: as with the loops with stores, only the loads contributing to
----------------
david-arm wrote:

At least prior to this patch, this statement is untrue because the way the vectorised loop is constructed in vplan the latch block is always executed regardless of whether we would exit early in the scalar loop, i.e.

```
for.body:
  %ld1 = load i8, ...
  ...

for.inc:
  %ld2 = load i8, ...
  ...
```

will get vectorised into a single block:

```
vector.body:
  ...
  %vld1 = load <16 x i8>, ...
  %vld2 = load <16 x i8>, ...
  ...
  %vcmp = icmp eq <16 x i8> %vld1, splat (i8 3)
  %early.exit.cmp = reduce or %vcmp
  %latch.exit.cmp = icmp ...
  %final.cmp = or i1 %early.exit.cmp, %latch.exit.cmp
  br i1 %final.cmp, label %middle.split, label %vector.body
```

Are you saying that with this patch we now create a masked load for `%vld2` instead based on the early exit condition?

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