[llvm] [LV] Add initial legality checks for ee loops with stores (PR #145663)

[David Sherwood via llvm-commits]

================
@@ -1771,13 +1821,65 @@ bool LoopVectorizationLegality::isVectorizableEarlyExitLoop() {
 
   // TODO: Handle loops that may fault.
   Predicates.clear();
-  if (!isDereferenceableReadOnlyLoop(TheLoop, PSE.getSE(), DT, AC,
-                                     &Predicates)) {
-    reportVectorizationFailure(
-        "Loop may fault",
-        "Cannot vectorize potentially faulting early exit loop",
-        "PotentiallyFaultingEarlyExitLoop", ORE, TheLoop);
-    return false;
+  if (HasStore) {
+    // Record load for analysis by isDereferenceableAndAlignedInLoop
+    // and later by dependence analysis.
+    if (BranchInst *Br = dyn_cast<BranchInst>(
+            SingleUncountableEdge->first->getTerminator())) {
+      // FIXME: Handle exit conditions with multiple users, more complex exit
+      //        conditions than br(icmp(load, loop_inv)).
+      ICmpInst *Cmp = dyn_cast<ICmpInst>(Br->getCondition());
+      if (Cmp && Cmp->hasOneUse() &&
+          TheLoop->isLoopInvariant(Cmp->getOperand(1))) {
+        LoadInst *Load = dyn_cast<LoadInst>(Cmp->getOperand(0));
+        if (Load && Load->hasOneUse() && !TheLoop->isLoopInvariant(Load)) {
+          if (isDereferenceableAndAlignedInLoop(Load, 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(*Load, DT, TheLoop)) {
+              EarlyExitLoad = Load;
+            } else {
+              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 {
+            reportVectorizationFailure(
+                "Uncounted loop condition not known safe",
+                "Cannot vectorize early exit loop with "
+                "possibly unsafe condition load",
+                "PotentiallyFaultingEarlyExitLoop", ORE, TheLoop);
+            return false;
+          }
+        }
+      }
+    }
+
+    if (!EarlyExitLoad) {
+      reportVectorizationFailure(
+          "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:

I'd still prefer this FIXME to be removed because I think it's misleading. The existing method of vectorisation for loops without stores is to execute the vector loop up to the point of the early exit, which definitely requires the loads to be dereferenceable. Your FIXME implies that with that style of vectorisation we can remove the checks, which isn't true. Also, if we know at runtime the loop contains no stores and the loads are dereferenceable anyway, then why artificially bail out before the early exit when we don't need to? The FIXME suggests that in future we will unconditionally move to this new style of vectorisation you're proposing, whereas I expect there to be a choice based on the trade-offs for this specific loop.

I'm happy for you to add a TODO that says something like "TODO: With an alternative style of vectorisation where we exit the vector loop prior to the early exit we can remove the restriction that loads must be dereferenceable."

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